diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 48147d5d3234..62145f00e6ed 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -394,6 +394,8 @@ /tests/unittest/_torch/modeling_v2 @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju /tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju /tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju +/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v41.py @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju +/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v41_replay.py @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju # The two catalog categories whose contracts state kernel behaviour the # attention and MoE owners are the authority on; co-owned rather than reassigned. /tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju @xxi-nv @yuxianq diff --git a/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu b/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu index cd33862afb4c..a1b8bd2e7af6 100644 --- a/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu +++ b/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu @@ -108,7 +108,7 @@ __global__ void indexerKCacheGatherUnifiedKernel(uint8_t const* __restrict__ k_c // Convert flat index to memory offset using strides (k cache pool from cpp kv cache manager is non-contiguous) int64_t src_offset = flatIndexToMemoryOffset(flat_idx, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3); - int64_t dst_offset = token_idx * head_dim + head_dim_idx; + int64_t dst_offset = static_cast(token_idx) * head_dim + head_dim_idx; // 4 bytes read from non-contiguous cache, write to contiguous output *reinterpret_cast(&out_fp8[dst_offset]) = *reinterpret_cast(&k_cache[src_offset]); @@ -118,7 +118,7 @@ __global__ void indexerKCacheGatherUnifiedKernel(uint8_t const* __restrict__ k_c { int64_t src_offset_scale = flatIndexToMemoryOffset(flat_idx_scale_base, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3); - int64_t dst_offset_scale = token_idx * scale_size; // scale_size = 4 + int64_t dst_offset_scale = static_cast(token_idx) * scale_size; // scale_size = 4 // 4 bytes read for scale *reinterpret_cast(&out_scale[dst_offset_scale]) diff --git a/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu b/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu index f0c7c27a235c..9894cd3bc15e 100644 --- a/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu +++ b/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu @@ -103,7 +103,7 @@ __global__ void indexerKCacheScatterUnifiedKernel(uint8_t const* __restrict__ k_ // Convert flat index to memory offset using strides (k cache pool from cpp kv cache manager is non-contiguous) int64_t dst_offset = flatIndexToMemoryOffset(flat_idx, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3); - int64_t src_offset = token_idx * head_dim + head_dim_idx; + int64_t src_offset = static_cast(token_idx) * head_dim + head_dim_idx; // 4 bytes write *reinterpret_cast(&k_cache[dst_offset]) = *reinterpret_cast(&k_fp8_bytes[src_offset]); @@ -113,7 +113,7 @@ __global__ void indexerKCacheScatterUnifiedKernel(uint8_t const* __restrict__ k_ { int64_t dst_offset_scale = flatIndexToMemoryOffset(flat_idx_scale_base, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3); - int64_t src_offset_scale = token_idx * scale_size; // scale_size = 4 + int64_t src_offset_scale = static_cast(token_idx) * scale_size; // scale_size = 4 // 4 bytes write for scale *reinterpret_cast(&k_cache[dst_offset_scale]) diff --git a/cpp/tensorrt_llm/kernels/moe/routing/customMoeRoutingKernels.cu b/cpp/tensorrt_llm/kernels/moe/routing/customMoeRoutingKernels.cu index 01b15bd4bbde..7e9fcdeffac5 100644 --- a/cpp/tensorrt_llm/kernels/moe/routing/customMoeRoutingKernels.cu +++ b/cpp/tensorrt_llm/kernels/moe/routing/customMoeRoutingKernels.cu @@ -310,8 +310,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: @@ -425,7 +423,7 @@ __global__ void gate_forward_kernel( // C++ wrapper function (output tensors passed as parameters) // All tensors are float32 -template +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, bool const* image_mask, float const* vision_bias, cudaStream_t stream) @@ -434,18 +432,35 @@ void launch_gate_forward_kernel(float* scores_in, float* bias, int* input_ids, i 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, + gate_forward_kernel<<>>(scores_in, bias, input_ids, tid2eid, out_weights, out_indices, batch_size, route_scale, image_mask, vision_bias); } +template +void dispatch_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, bool is_hash, bool const* image_mask, float const* vision_bias, + cudaStream_t stream) +{ + if (is_hash) + { + launch_gate_forward_kernel(scores_in, nullptr, input_ids, tid2eid, out_weights, + out_indices, batch_size, route_scale, nullptr, nullptr, stream); + } + else + { + launch_gate_forward_kernel(scores_in, bias, nullptr, nullptr, out_weights, out_indices, + batch_size, route_scale, image_mask, vision_bias, stream); + } +} + 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, bool const* image_mask, float const* vision_bias, - cudaStream_t stream) + int batch_size, int n_experts, int topk, float route_scale, bool is_hash, bool const* image_mask, + float const* vision_bias, cudaStream_t stream) { auto* scores = static_cast(scores_in); auto* bias_ptr = static_cast(bias); @@ -454,33 +469,27 @@ void gate_forward(void* scores_in, // [batch_size, nExperts] - pre-computed from auto* weights = static_cast(out_weights); auto* indices = static_cast(out_indices); - switch (n_experts) + // Supported (n_experts, topk) pairs: 256/6 and 384/6 (DeepSeek-V4 / V4.1 routed MoE), 128/3 (DeepSeek-V4.1 DSpark + // draft MoE). + if (n_experts == 256 && topk == 6) { - case 256: - if (is_hash) - { - launch_gate_forward_kernel<256, true>(scores, nullptr, input_ids_ptr, tid2eid_ptr, weights, indices, - batch_size, route_scale, nullptr, nullptr, stream); - } - else - { - launch_gate_forward_kernel<256, false>(scores, bias_ptr, nullptr, nullptr, weights, indices, batch_size, - route_scale, image_mask, vision_bias, 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, nullptr, nullptr, stream); - } - else - { - launch_gate_forward_kernel<384, false>(scores, bias_ptr, nullptr, nullptr, weights, indices, batch_size, - route_scale, image_mask, vision_bias, stream); - } - break; - default: TLLM_CHECK_WITH_INFO(false, "gate_forward only supports n_experts 256 or 384"); + dispatch_gate_forward_kernel<256, 6>(scores, bias_ptr, input_ids_ptr, tid2eid_ptr, weights, indices, batch_size, + route_scale, is_hash, image_mask, vision_bias, stream); + } + else if (n_experts == 384 && topk == 6) + { + dispatch_gate_forward_kernel<384, 6>(scores, bias_ptr, input_ids_ptr, tid2eid_ptr, weights, indices, batch_size, + route_scale, is_hash, image_mask, vision_bias, stream); + } + else if (n_experts == 128 && topk == 3) + { + dispatch_gate_forward_kernel<128, 3>(scores, bias_ptr, input_ids_ptr, tid2eid_ptr, weights, indices, batch_size, + route_scale, is_hash, image_mask, vision_bias, stream); + } + else + { + TLLM_CHECK_WITH_INFO(false, + "gate_forward supports (n_experts, topk) in {(256, 6), (384, 6), (128, 3)}, got (%d, %d)", n_experts, topk); } sync_check_cuda_error(stream); } diff --git a/cpp/tensorrt_llm/kernels/moe/routing/customMoeRoutingKernels.h b/cpp/tensorrt_llm/kernels/moe/routing/customMoeRoutingKernels.h index 14137f655a96..48e8be3ada2f 100644 --- a/cpp/tensorrt_llm/kernels/moe/routing/customMoeRoutingKernels.h +++ b/cpp/tensorrt_llm/kernels/moe/routing/customMoeRoutingKernels.h @@ -38,8 +38,8 @@ void gate_forward(void* scores_in, // [batch_size, nExperts] - pre-computed from 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, bool const* image_mask, float const* vision_bias, - cudaStream_t stream); + int batch_size, int n_experts, int topk, float route_scale, bool is_hash, bool const* image_mask, + float const* vision_bias, cudaStream_t stream); } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/moe/routing/moeGateOp.cpp b/cpp/tensorrt_llm/thop/moe/routing/moeGateOp.cpp index 5507cedb53b7..fc2199414090 100644 --- a/cpp/tensorrt_llm/thop/moe/routing/moeGateOp.cpp +++ b/cpp/tensorrt_llm/thop/moe/routing/moeGateOp.cpp @@ -26,8 +26,6 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -static constexpr int kTOPK = 6; - // PyTorch wrapper function that calls the kernel implementation void gate_forward(th::Tensor scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight) th::Tensor bias, // empty tensor if hash mode @@ -38,9 +36,9 @@ void gate_forward(th::Tensor scores_in, // [batch_size, nExperts] - pre-computed int64_t topk, double route_scale, bool is_hash, std::optional image_mask, std::optional vision_bias) { - TORCH_CHECK(topk == kTOPK, "topk must be ", kTOPK); auto const n_experts = scores_in.size(1); - TORCH_CHECK(n_experts == 256 || n_experts == 384, "n_experts must be 256 or 384"); + TORCH_CHECK((n_experts == 256 && topk == 6) || (n_experts == 384 && topk == 6) || (n_experts == 128 && topk == 3), + "(n_experts, topk) must be one of (256, 6), (384, 6), (128, 3); got (", n_experts, ", ", topk, ")"); TORCH_CHECK(scores_in.scalar_type() == torch::kFloat32, "scores_in must be float32"); TORCH_CHECK(out_weights.scalar_type() == torch::kFloat32, "out_weights must be float32"); TORCH_CHECK(out_indices.scalar_type() == torch::kInt32, "out_indices must be int32"); @@ -71,7 +69,8 @@ void gate_forward(th::Tensor scores_in, // [batch_size, nExperts] - pre-computed kernels::gate_forward(scores_in.data_ptr(), is_hash ? nullptr : bias.data_ptr(), is_hash ? input_ids.data_ptr() : nullptr, is_hash ? tid2eid.data_ptr() : nullptr, out_weights.data_ptr(), out_indices.data_ptr(), batch_size, static_cast(n_experts), - static_cast(route_scale), is_hash, image_mask.has_value() ? image_mask->data_ptr() : nullptr, + static_cast(topk), static_cast(route_scale), is_hash, + image_mask.has_value() ? image_mask->data_ptr() : nullptr, vision_bias.has_value() ? vision_bias->data_ptr() : nullptr, stream); } diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index a167ec0ae3e7..5b4c93035842 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -22,9 +22,9 @@ from tensorrt_llm.llmapi import (AttentionDpConfig, AutoDecodingConfig, CudaGraphConfig, DFlashDecodingConfig, DraftTargetDecodingConfig, - Eagle3DecodingConfig, KvCacheConfig, MoeConfig, - MTPDecodingConfig, NGramDecodingConfig, - TorchCompileConfig) + DSparkDecodingConfig, Eagle3DecodingConfig, + KvCacheConfig, MoeConfig, MTPDecodingConfig, + NGramDecodingConfig, TorchCompileConfig) example_prompts = [ "Hello, my name is", @@ -320,6 +320,13 @@ def setup_llm(args, **kwargs): spec_config = DFlashDecodingConfig( max_draft_len=args.spec_decode_max_draft_len, speculative_model=args.draft_model_dir) + elif spec_decode_algo == "DSPARK": + # Without --draft_model_dir the drafter is the target checkpoint's own + # mtp.* stages; TorchLlmArgs resolves that, and requires + # max_draft_len to equal the drafter's trained block size. + spec_config = DSparkDecodingConfig( + max_draft_len=args.spec_decode_max_draft_len, + speculative_model=args.draft_model_dir) elif spec_decode_algo == "DRAFT_TARGET": spec_config = DraftTargetDecodingConfig( max_draft_len=args.spec_decode_max_draft_len, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 61be49f5be0d..bf6193b74d1b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -20,7 +20,7 @@ export TRTLLM_MODELING_V2=require # off | auto | require | `TRTLLM_MODELING_V2` | Behaviour | |---|---| | unset or `off` (default) | The resolver returns immediately. Nothing in this package is imported and behaviour is byte-for-byte what it is today. | -| `auto` | Uses a target when one matches this exact configuration; falls back to the built-in implementation when none does. | +| `auto` | Uses a target when one matches this exact configuration; falls back to the built-in implementation when none does. A family whose targets bring their own input processor (DeepSeek-V4.1) raises instead of falling back, because the processor is chosen per checkpoint and the built-in model expects another one. | | `require` | Raises instead of falling back, naming the criterion that did not match. | **Use `require` for anything you will attribute to modeling_v2.** Under `auto`, @@ -36,8 +36,10 @@ its workers resolve the built-in is the silent split this package exists to prevent. An environment variable rather than an LLM-API field is a deliberate trade: it -keeps the entire concept inside this package, so the only change modeling_v2 -needs anywhere else is the `_resolve_class` hook. The cost is that the switch +keeps the entire concept inside this package, so the only changes modeling_v2 +needs anywhere else are the `_resolve_class` hook and, for a target family that +reads its own multimodal inputs, one call in `create_input_processor` +(`modeling_v2_input_processor`). The cost is that the switch does not appear in a run's recorded `llm.args` and cannot be set through `--extra_llm_api_options`. @@ -256,10 +258,11 @@ the names but not the submodule paths, so all five files importing them failed at import. They now use the canonical path. The lesson generalizes: with no pin, a target's contact with private engine -surface is checked only by running it. Each target declares that surface as -`REQUIRED_TRTLLM_OPS`, `test_modeling_v2_target_contract.py` asserts every name -in it exists, and the first-forward metadata field check catches the rest -- -which turns a drifting engine from a wrong answer into a loud failure. +surface is checked only by running it. The step contract check does that by +calling `_build_step_args` on the metadata it was handed, which reads every +field the target consumes -- so a rename upstream becomes a loud failure rather +than a wrong answer. Declaring the same surface a second time, as a list of +names beside the code that reads it, only adds something to forget to update. ## Two facts this migration surfaced about upstream diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py index d1b0f9599ba8..333bb300f01d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/__init__.py @@ -17,8 +17,9 @@ ``off`` (unset, the default) is byte-for-byte today's behaviour: the resolver returns immediately and nothing in this package is imported. ``auto`` uses a target when one matches and falls back to the built-in implementation when none -does. ``require`` raises instead of falling back -- see ``_router_index`` for -why that mode is not optional. +does, unless the checkpoint's targets bring their own input processor, which +the built-in model must not run behind. ``require`` raises instead of falling +back -- see ``_router_index`` for why that mode is not optional. This package sits *beside* ``_torch/models/`` rather than inside it, which is load-bearing twice over: ``is_builtin_zoo_module`` matches on the zoo's @@ -31,11 +32,14 @@ from ._router_index import ( MODELING_V2_ENV, MODELING_V2_ROUTERS, + MODELING_V2_VALIDATE_ENV, ROUTING_BACKENDS, ModelingV2Context, ModelingV2Mode, assert_backend_can_route, + modeling_v2_input_processor, modeling_v2_resolve, + step_contract_enabled, ) __all__ = [ @@ -44,6 +48,9 @@ "MODELING_V2_ROUTERS", "ModelingV2Context", "ModelingV2Mode", + "MODELING_V2_VALIDATE_ENV", "assert_backend_can_route", + "modeling_v2_input_processor", "modeling_v2_resolve", + "step_contract_enabled", ] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index d1b2a7840899..65b9bda30fce 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -43,7 +43,9 @@ #: The switch. An environment variable rather than an LLM-API field, so that #: nothing outside this package has to carry the concept: the only upstream -#: change modeling_v2 needs is the ``_resolve_class`` hook itself. +#: changes modeling_v2 needs are the ``_resolve_class`` hook itself and, for a +#: target family that reads its own multimodal inputs, the input-processor +#: hook in ``create_input_processor``. #: #: Assigning to ``os.environ`` from a script does not set it. Worker ranks #: receive the environment as it stood when MPI initialized, and long-lived @@ -55,10 +57,37 @@ #: which every rank re-applies to its own environment before it builds a model. MODELING_V2_ENV = "TRTLLM_MODELING_V2" +#: Turns the step contract check on. Off by default, so a served engine runs +#: the forward and nothing else: the checks describe facts fixed at engine +#: construction, and a check that can only ever hold is a check whose failure +#: mode is its own bug taking down a working deployment. +#: +#: It is an environment variable for the same reason the switch above is, and +#: with the same caveat about ranks. When modeling_v2 stops being experimental +#: this belongs in the LLM args instead, where a caller can set it per engine +#: rather than per process. +MODELING_V2_VALIDATE_ENV = "TRTLLM_MODELING_V2_VALIDATE" + + +def step_contract_enabled() -> bool: + """Whether targets should check their step contract on the first forward. + + Read once per model instance at construction rather than per forward: the + answer cannot change within a process, and the forward is the hot path. + """ + return os.environ.get(MODELING_V2_VALIDATE_ENV, "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + # architectures[0] -> routing module, relative to this package. MODELING_V2_ROUTERS = { "GptOssForCausalLM": "models.gpt_oss.routing", "DeepseekV3ForCausalLM": "models.deepseek_v3.routing", + "DeepseekV41ForConditionalGeneration": "models.deepseek_v41.routing", } #: Backends whose model construction reaches ``modeling_v2_resolve``. The @@ -252,6 +281,11 @@ def modeling_v2_resolve(config: "ModelConfig") -> Optional[str]: exists to prevent is silent: asking for a target that does not exist, getting the in-tree implementation, and reading the resulting curve as modeling_v2's. + + ``auto`` raises too when the checkpoint's targets bring their own input + processor (see ``modeling_v2_input_processor``): that processor was chosen + on the checkpoint alone, and falling back would pair it with a built-in + model that expects another placeholder layout. """ mode = ModelingV2Mode.from_env() if mode is ModelingV2Mode.OFF: @@ -271,6 +305,11 @@ def modeling_v2_resolve(config: "ModelConfig") -> Optional[str]: if target is None: if mode is ModelingV2Mode.REQUIRE: raise ValueError(explain_no_match(arch, routing, trace)) + if _input_processor_module(routing, pretrained_config) is not None: + # Replay the tree with a trace, only to say which criterion missed. + trace = Trace() + routing.route(ctx, trace) + raise ValueError(explain_no_auto_fallback(arch, routing, trace)) return None # The synthetic name is only a registry key: importing the target module @@ -280,18 +319,59 @@ def modeling_v2_resolve(config: "ModelConfig") -> Optional[str]: return target -def explain_no_match(arch: str, routing, trace: Trace) -> str: - """Say which criterion the configuration failed, not just that it did.""" - if routing is None: - known = ", ".join(sorted(MODELING_V2_ROUTERS)) or "(none)" - return ( - f"{MODELING_V2_ENV}=require, but no target exists for " - f"architecture {arch!r}; routed architectures: " - f"{known}" - ) +def modeling_v2_input_processor(pretrained_config: Any) -> Optional[type]: + """The input processor a modeling_v2 target family brings for this checkpoint, or None. + + ``create_input_processor`` takes the processor registered on the class under + ``architectures[0]``, and the multimodal placeholder registry is keyed by + ``model_type``; neither goes through ``modeling_v2_resolve``. A routing + module whose targets read their own multimodal inputs defines + ``input_processor(pretrained_config)``, naming a module (relative to this + package) that carries ``INPUT_PROCESSOR``, ``MODEL_TYPE`` and + ``PLACEHOLDER_METADATA``. This registers that placeholder metadata for + ``MODEL_TYPE`` and returns the processor class. None -- modeling_v2 off, no + routing module, or none of its targets brings a processor for this + checkpoint -- leaves the choice to the built-in registry. + + The choice is the checkpoint's alone, because the input processor is built + where the GPU and the parallel mapping are not known. So a configuration + no target matches raises when the model is built, under ``auto`` as under + ``require`` (``modeling_v2_resolve``): the built-in model never runs behind + this processor. + """ + if ModelingV2Mode.from_env() is ModelingV2Mode.OFF: + return None + architectures = getattr(pretrained_config, "architectures", None) + if not architectures: + return None + name = _input_processor_module(routing_module(architectures[0]), pretrained_config) + if name is None: + return None + + # Function-local: the inputs registry imports torch and the model zoo, + # and nothing on the resolve path should. + from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY + + module = import_module(f"{_PACKAGE}.{name}") + # Registered from outside the built-in zoo, so it holds the slot even if + # the zoo's own module for this model type is imported afterwards. + MULTIMODAL_PLACEHOLDER_REGISTRY.set_placeholder_metadata( + module.MODEL_TYPE, module.PLACEHOLDER_METADATA, registrant_module=module.__name__ + ) + module.INPUT_PROCESSOR._registered_model_type = module.MODEL_TYPE + return module.INPUT_PROCESSOR + + +def _input_processor_module(routing, pretrained_config: Any) -> Optional[str]: + """The input-processor module the routing module names for this checkpoint, or None.""" + pick = getattr(routing, "input_processor", None) + return None if pick is None else pick(pretrained_config) + + +def _trace_lines(routing, trace: Trace) -> List[str]: + """Where the routing module's decision tree stopped, one line per criterion.""" lines = [ - f"{MODELING_V2_ENV}=require, but no target matched architecture " - f"{arch!r}. The decision tree in " + f"The decision tree in " f"{routing.__name__.rpartition('.')[0].replace('.', '/')}/routing.py " f"got as far as:" ] @@ -300,4 +380,37 @@ def explain_no_match(arch: str, routing, trace: Trace) -> str: lines.append(f" {label:<10} {value!r:<48} {mark}") if not trace.steps: lines.append(" (the tree rejected the configuration before its first recorded criterion)") + return lines + + +def explain_no_auto_fallback(arch: str, routing, trace: Trace) -> str: + """Why ``auto`` refuses the built-in fallback for this checkpoint, and the ways out.""" + lines = [ + f"{MODELING_V2_ENV}=auto, but no modeling_v2 target matches this " + f"{arch!r} configuration, and its targets bring their own input " + f"processor and multimodal placeholder placement, which the built-in " + f"implementation does not expect. Falling back would lay out image " + f"prompts for a model that reads them differently, so this raises " + f"instead. Set {MODELING_V2_ENV}=off to run the built-in " + f"implementation with its own input processor, or run on an SM and " + f"with a parallel member that a target supports " + f"(`python -m {_PACKAGE}.explain` lists the criteria).", + ] + lines.extend(_trace_lines(routing, trace)) + return "\n".join(lines) + + +def explain_no_match(arch: str, routing, trace: Trace) -> str: + """Say which criterion the configuration failed, not just that it did.""" + if routing is None: + known = ", ".join(sorted(MODELING_V2_ROUTERS)) or "(none)" + return ( + f"{MODELING_V2_ENV}=require, but no target exists for " + f"architecture {arch!r}; routed architectures: " + f"{known}" + ) + lines = _trace_lines(routing, trace) + lines[0] = ( + f"{MODELING_V2_ENV}=require, but no target matched architecture {arch!r}. " + lines[0] + ) return "\n".join(lines) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_target.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_target.py new file mode 100644 index 000000000000..07adc9cce0e1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_target.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""What runs a step, separated from what a model statically is. + +A model's core carries the static structure: routing identity, weights, +parallel and quantization configuration. A `Target` +carries one way of running a forward over that structure. A core has two, +and the runtime picks between them per call. + +The split exists because the two phases are not one computation with different +arguments. On MLA the generation path works in latent space (absorption) while +the context path materializes K and V; `TrtllmAttention` refuses `mixed` for +MLA outright, which is what forces the choice up into Python. gpt_oss is not +MLA and so has no such branch -- this layer lands there first for the mechanism, +not for a speedup. + +Two rules the base class exists to hold: + + A target reads the core, it does not copy it. The core replaces `_rope` + when a contract check finds the table too short for the engine's admitted + `max_seq_len`, and parameters are still meta tensors when the shell is + constructed. Both have already cost this tree a bug; a target that snapshots + operands at construction reintroduces them. + + The phase is read once per routable unit and never cached. The trunk reads + it once per forward. A speculative draft loop rewrites `num_contexts` in + place between draft steps, so a drafter has to read it per step -- a value + computed on the first call is wrong on the rest, and under graph capture + would be frozen wrong at every position. +""" + +from __future__ import annotations + +import enum +from abc import ABC, abstractmethod +from typing import Any + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import step_contract_enabled + + +class Phase(enum.Enum): + """Which target a step routes to. + + Two members, not three: a mixed batch -- context and generation requests + in one step, which in-flight batching produces routinely -- goes to + PREFILL. That makes PREFILL the general case and DECODE a specialization + that may assume it has no context rows at all. + """ + + PREFILL = "prefill" + DECODE = "decode" + + +def phase_of(md: Any) -> Phase: + """The routing predicate, in full. + + One host-side integer read, no synchronization, and nothing else consulted: + a decode-only CUDA graph is a per-capture constant environment, and a + predicate that reached for a second field would be reading one this call + is not guaranteed to carry. + + `num_contexts == 0` is reachable two ways and both mean the same thing -- + `_apply_steady_gen_fast_prepare` writes it explicitly for an unchanged + generation-only batch, and an ordinary prepare arrives at it whenever the + scheduled batch holds no context request. + """ + return Phase.DECODE if md.num_contexts == 0 else Phase.PREFILL + + +class Target(ABC): + """One way of running a forward over a core's weights. + + Instantiated after weights load, never in `__init__` -- see the module + docstring. Holds the core by reference; every operand is resolved at read + time through `self.core`. + + Also owns the opt-in step-contract check (`_check_step_contract` below): + for a model split into targets, `step_args` is itself the probe -- the + one projection whose content differs by phase -- so the check belongs + beside it. A core with no target split (deepseek, today) carries its own + copy of the same check instead. + """ + + def __init__(self, core: Any) -> None: + self.core = core + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it: read once here + # rather than per forward, and False for the whole life of a served + # engine. "pending" rather than "enabled" because the check runs once + # -- everything it looks at is fixed at engine construction. Each + # target owns its own flag: the two targets' `step_args` projections + # differ, so probing one does not prove the other's fields exist, and + # a decode-only engine run must still clear decode's flag on its own + # first step rather than inherit prefill's. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md: Any) -> None: + """Opt-in first-forward fail-fast, run when TRTLLM_MODELING_V2_VALIDATE + asks for it: the metadata fields this target consumes must exist (they + are private trtllm surface). Everything checked is fixed at engine + construction -- once per target instance is sound, and off in a served + engine, where the only thing this could still do is fail. + + Called unconditionally every forward; the early return below is the + gate. That costs one Python call per forward, not per layer, against + a decode step measured in milliseconds -- accepted in exchange for + not scattering the `_contract_pending` check across every call site. + """ + if not self._contract_pending: + return + # Calling the projection is the check: it reads every metadata field + # this target consumes, so a rename or removal upstream surfaces here + # rather than mid-forward. Deriving it this way is the point -- a + # hand-kept list of the same names drifts silently the first time + # `step_args` gains a field and nobody updates the copy. + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + @abstractmethod + def step_args(self, md: Any) -> dict: + """The per-step batch state this target hands the attention op. + + Separate from `forward` because it is the one projection whose *content* + differs by phase even when the surrounding computation does not: a + decode target states the values its routing already guarantees rather + than reading them back off the metadata. Also the probe + `_check_step_contract` calls -- see above. + """ + + @abstractmethod + def forward(self, attn_metadata: Any, *args: Any, **kwargs: Any) -> Any: + """Run one step. Each target runs its own `_check_step_contract` first + -- see above -- rather than the core running one check for both.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_weights.py new file mode 100644 index 000000000000..ea4002374674 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_weights.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The weight table and the three passes over it, shared by every model. + +A model states its weights once, as a `W` per role, and subclasses +`ModelWeights` to say how the dimensions are derived. Declaration, checkpoint +manifest and load all unroll that one table, so a role cannot be declared +without a source or loaded into a shape nobody declared. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Callable + +import torch +import torch.nn as nn + +#: `layers=ALL_LAYERS` is the common case and the default. A sentinel rather +#: than a lambda so the table reads as data, and rather than `True` so the +#: field has one type of answer: which layers. +ALL_LAYERS = object() + + +@dataclass(frozen=True) +class W: + """One weight role: what it is, where its bytes come from, what happens on the way. + + `shape` takes the dimension bundle and returns the declared shape; every + shape in these models is derived from configuration rather than constant. + + `src` says where the bytes come from, in one of three forms: + + * a checkpoint key template, with `{p}` standing for the layer prefix; + * a sequence of `(template, slice_fn)` pairs, when one parameter is + assembled from several checkpoint tensors side by side -- `slice_fn` + takes the dimension bundle and returns an index into the destination; + * a callable `(d, i) -> [(key, index), ...]`, taking the dimension bundle + and the layer index. The first two forms are special cases of it; this + one exists for sources that are not a fixed template -- a per-rank + expert window, or a prefix that differs between layer families. A `key` + may itself be a tuple, meaning the transform is handed those tensors + together. + + `layers` says which layers carry this weight: `ALL_LAYERS` for every one, + `None` for a weight that is not per-layer at all, or a callable + `d -> iterable[int]` when only some layers have it. + + `transform` runs on the source tensor before the copy, on the + destination's device. `None` means the checkpoint tensor is already the + declared shape. + """ + + name: str + shape: Callable[[Any], tuple[int, ...]] + src: str | tuple[tuple[str, Callable], ...] | Callable + dtype: torch.dtype | None = None + transform: Callable | None = None + layers: object | None = ALL_LAYERS + + def layer_indices(self, d: SimpleNamespace) -> list[int] | None: + """The layers this weight exists on, or None if it is not per-layer.""" + if self.layers is None: + return None + if self.layers is ALL_LAYERS: + return list(range(d.num_layers)) + return list(self.layers(d)) + + def sources(self, d: SimpleNamespace, i: int | None, prefix: str) -> list[tuple[Any, Any]]: + """`[(checkpoint key or key tuple, index into the param | None), ...]`.""" + if callable(self.src): + return list(self.src(d, i)) + if isinstance(self.src, str): + return [(self.src.format(p=prefix), None)] + return [(template.format(p=prefix), slice_fn(d)) for template, slice_fn in self.src] + + +class ModelWeights: + """Shared base for a model's weight table. + + Subclasses supply `WEIGHTS` and `dims`; everything else here is one pass + over that table. A model whose load differs structurally overrides `load` + rather than growing a flag on this class. + """ + + WEIGHTS: tuple[W, ...] = () + + LAYER_PREFIX: str = "model.layers.{i}" + + # Param-key suffix after which to release transform scratch. The expert + # transforms allocate several GB per layer; without this, peak load + # memory is the whole model rather than one layer. + RELEASE_AFTER: str | None = None + + def dims(self, core) -> SimpleNamespace: + """The dimension bundle every pass is computed from. + + The single place a derived width or padded size is worked out. Reads + `core.model_config` only, so it can run before the core has set + anything on itself. + """ + raise NotImplementedError + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + """Allocate every weight the table declares. + + Meta-init intercepts `torch.empty` here -- real CUDA storage arrives + when the engine materializes the registry. + """ + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None, str]]: + """`[(param key, layer index | None, checkpoint layer prefix), ...]`. + + One place decides which parameters an entry declares, so `declare` and + `manifest` cannot disagree about whether a weight exists. + """ + layers = entry.layer_indices(d) + if layers is None: + return [(entry.name, None, "")] + return [(f"l{i}_{entry.name}", i, self.LAYER_PREFIX.format(i=i)) for i in layers] + + def manifest(self, core) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + """param key -> [(checkpoint key, index into the param | None, transform | None)].""" + d = self.dims(core) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i, prefix in self._targets(entry, d): + rows[key] = [ + (ckpt, index, entry.transform) for ckpt, index in entry.sources(d, i, prefix) + ] + return rows + + def expected_unconsumed(self, core, weights) -> set[str]: + """Checkpoint keys this rank is meant to leave untouched.""" + return set() + + def load(self, model, weights) -> None: + core = model.model + manifest = self.manifest(core) + consumed: set[str] = set() + + def fill(param: nn.Parameter, ckpt_key: str, index, transform) -> None: + assert ckpt_key in weights, f"checkpoint key missing: {ckpt_key}" + src = weights[ckpt_key][:] # materialize the lazy slice + dst = param.data if index is None else param.data[index] + if transform is not None: + src = transform(src.to(dst.device, non_blocking=True), core) + assert dst.shape == src.shape, (ckpt_key, tuple(dst.shape), tuple(src.shape)) + assert src.dtype == dst.dtype, (ckpt_key, src.dtype, dst.dtype) + dst.copy_(src, non_blocking=True) + consumed.add(ckpt_key) + + for param_key, sources in manifest.items(): + for ckpt_key, index, transform in sources: + fill(core.w[param_key], ckpt_key, index, transform) + if self.RELEASE_AFTER is not None and param_key.endswith(self.RELEASE_AFTER): + torch.cuda.empty_cache() + + # Shell-registered exception: the base class owns lm_head (untied). + fill(model.lm_head.weight, "lm_head.weight", None, None) + + torch.cuda.synchronize() + leftover = set(weights.keys()) - consumed - self.expected_unconsumed(core, weights) + assert not leftover, f"unconsumed checkpoint keys: {sorted(leftover)[:8]}" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py new file mode 100644 index 000000000000..f277bd5a2372 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/_op.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""What a catalog entry is, expressed as something that runs. + +An entry used to be a wrapper plus a prose contract. The prose said what the +op computes, which inputs it accepts, where it goes silently wrong, and which +architecture it had been measured on -- and nothing checked that any of it was +still true. A kernel change upstream could falsify a page of it without a +single test going red. + +This is the executable form of that contract: + + raw_call the op call a target makes + reference what the op is supposed to compute, in plain torch + is_valid the inputs the op takes and answers wrongly + compare how close the two have to be, and why + CELLS the input configurations this entry is certified over + ARCHS the architectures those cells were driven on + note what is left, and nothing that could have been one of the above + +`reference` and `CELLS` are the part that cannot go stale: the entry's test +drives every cell through both `raw_call` and `reference` and gates them with +`compare`, so a claim that stops holding stops the build. + +`CELLS` is also the answer to "can my model use this op?". A target author +reads the cell list and its reasons, not a document -- and what they read is +what CI proved on the last commit. Cells cover what the shipped targets +actually run; an input outside them is not known to work, which for most of +these ops means "nobody measured it" rather than "known broken". + +`is_valid` is never called on the way to the kernel. It states, in executable +form, the inputs the op accepts and quietly answers wrongly -- the class of +mistake a caller cannot detect for themselves. Running it per call would buy a +target nothing it does not already get from reading it once, and would put its +own bugs on the hot path of a served engine. The test drives it: `validating()` +in the test tree interposes it, so every cell also proves the guard admits what +it is supposed to admit, and the rejection cases prove it refuses what it is +supposed to refuse. +""" + +from __future__ import annotations + +import contextlib +import enum +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, Mapping + +import torch + + +@contextlib.contextmanager +def true_fp32_matmul(): + """Make torch's fp32 matmul actually fp32 for the duration. + + torch 2.12 defaults ``matmul.fp32_precision`` to ``tf32`` and + ``allow_tf32`` to True on this hardware, so a plain ``a.float() @ + b.float()`` is a *TF32* product -- ~1e-3 relative, which is 30x the error + of the ops being tested. Left alone the reference is the inaccurate side of + the comparison and the entry fails against a correct kernel. Measured: with + TF32 off the op is bit-identical to torch and both sit 1.9e-5 from a + float64 product; with TF32 on the reference alone moves by 0.035. + + Both switches have to move. Setting only ``allow_tf32`` leaves + ``fp32_precision`` at ``tf32`` and the product is still TF32 -- which is + exactly the failure this docstring is here to stop someone rediscovering. + """ + prev_allow = torch.backends.cuda.matmul.allow_tf32 + prev_prec = getattr(torch.backends.cuda.matmul, "fp32_precision", None) + torch.backends.cuda.matmul.allow_tf32 = False + if prev_prec is not None: + torch.backends.cuda.matmul.fp32_precision = "ieee" + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = prev_allow + if prev_prec is not None: + torch.backends.cuda.matmul.fp32_precision = prev_prec + + +#: One unit in the last place of each accumulating dtype's mantissa. +ULP = {torch.bfloat16: 2.0**-8, torch.float16: 2.0**-11, torch.float32: 2.0**-23} + + +def assert_within_ulp( + out: torch.Tensor, + ref: torch.Tensor, + element_ulp: float, + rms_ulp: float, + ulp: float | None = None, +) -> None: + """Gate a kernel against a reference that differs only in accumulation order. + + torch's default band cannot express this comparison. Its per-element `rtol` + is meaningless wherever cancellation drives `|ref|` toward zero -- a long + dot product over random operands puts a handful of outputs there, and a + 1e-4 absolute difference then reads as a relative difference of 10. And its + bf16 `atol` of 1e-5 sits orders of magnitude below one ulp of the output's + own scale, so it rejects differences no bf16 result could represent. + + So both gates are in units of the output's scale rather than its value: + `element_ulp` against each row's largest magnitude, `rms_ulp` against the + relative RMS in aggregate. A caller must say what it measured -- a + tolerance without a derivation is one that gets widened again next time. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + ulp = ULP[ref.dtype] if ulp is None else ulp + flat_out, flat_ref = ( + out.float().reshape(-1, out.shape[-1]), + ref.float().reshape(-1, ref.shape[-1]), + ) + row_scale = flat_ref.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + torch.testing.assert_close( + flat_out / row_scale, flat_ref / row_scale, rtol=0.0, atol=element_ulp * ulp + ) + rel_rms = ( + (flat_out - flat_ref).pow(2).mean().sqrt() / flat_ref.pow(2).mean().sqrt().clamp_min(1e-9) + ).item() + assert rel_rms <= rms_ulp * ulp, f"relative RMS {rel_rms:.3e} exceeds {rms_ulp} ulp" + + +#: Bumped once per engine step by the target that owns the entries. Only read +#: when validation is on: outside it, nothing compares generations and the +#: counter costs one integer increment a forward. +_STEP_GENERATION = 0 + + +def advance_step_generation() -> None: + """Called once per forward, before any `bind_const` or `bind_layered`.""" + global _STEP_GENERATION + _STEP_GENERATION += 1 + + +class Arch(str, enum.Enum): + """A GPU architecture an entry can be certified on. + + An enum rather than a string so that a typo is an AttributeError here + rather than an entry that silently claims nothing, and so that adding an + architecture is one edit that every entry's declaration is checked against. + """ + + SM_103 = "sm_103" + + +@dataclass(frozen=True) +class Cell: + """One input configuration an entry is certified over. + + `spec` is entry-specific: the entry's test knows how to turn it into + tensors, and nothing here interprets it. That keeps environment building -- + a KV-cache manager, a 4-rank communicator -- in the test where it belongs, + while the *set of configurations* stays with the op, which is what a target + author needs to read. + + `why` is not decoration. A cell that cannot say what it covers beyond the + ones already listed is a cell that costs CI time for nothing. + """ + + why: str + spec: Mapping[str, Any] = field(default_factory=dict) + + +class OpWrapper(ABC): + """Base for a catalog entry. + + An entry does not care *when* a value becomes known -- that was the old + design's axis (construction vs. per-step), and it was the caller's + business, not the interface's. What the entry actually needs to know is + whether a value varies by layer. Two binding methods state exactly that: + + bind_const(**values) a value, the same for every layer + bind_layered(layer_id, **values) one layer's values, read back by `layer=` + + Both are callable from anywhere a target chooses -- construction, the + first forward, every forward -- and the caller decides which. Both may be + called more than once; a later call updates only the keys it names and + leaves every other binding alone. `bind_layered` takes one layer at a + time rather than a whole table so that a caller whose layers are not + uniform -- some dense, some not -- can bind only the layers that have a + given operand, instead of padding the gaps itself. `__call__` merges + `_const`, then the row `_layered` holds for `layer=` (so a layered value + wins over a flat const for the same argument), then the call's own + kwargs last of all, and hands the result to `raw_call`. + + This is what gives gpt_oss's `attention_window_size` a home: per-layer, + but not knowable at construction (it derives from `attn_metadata`'s + resolved `max_seq_len`, which the KV cache manager does not produce until + after weights load). It is rebound with `bind_layered` every forward + rather than once at construction. + + The cost of collapsing the time axis: a reader at a call site can no + longer tell a volatile per-forward binding from a permanent one by the + method name alone -- both read `bind_const`. Make that obvious in + context instead, e.g. a comment naming what gets rebound each step -- + see gpt_oss's attention binding for the pattern. + + `raw_call` is the certified surface, and it takes everything explicitly -- + it must not read any bound state off `self`. That is what keeps `CELLS` + covering the entry's whole input space regardless of what a target chose to + bind, and it is why `reference` and `is_valid` mirror `raw_call` rather + than `__call__`. A gate in the claims suite enforces it. + + Instances are stateful: binding is state, and an op that manufactures its + own fixture holds a device tensor. Each target owns its own instances -- + there is no module-level singleton to contend over, which is what makes + per-target binding safe. Two models genuinely disagree about values like + `num_heads`; a shared instance could only hold one of them. + """ + + ARCHS: frozenset[Arch] = frozenset() + CELLS: tuple[Cell, ...] = () + note: str = "" + + def __init__(self) -> None: + """Set up empty binding state. + + No binding arguments: a target binds explicitly, from wherever a + value becomes knowable, with `bind_const` or `bind_layered` -- there + is no longer a stage reserved for construction. + """ + self._const: dict[str, Any] = {} + #: name -> {layer_id: value}, filled one layer at a time by + #: `bind_layered`. + self._layered: dict[str, dict[int, Any]] = {} + # Per key: the step generation in effect when that key was last + # bound. Read only by the validation harness -- see `bind_const`. + self._generation: dict[str, int] = {} + + def bind_const(self, **values: Any) -> None: + """Bind values that are the same for every layer. + + Updates rather than replaces: a later call adds or overwrites only + the keys it names, leaving every other binding -- including ones + made at construction -- in place. A target rebinding its per-step + projections once a forward is not required to also restate its + construction-time constants on every call. + + Also stamps each bound key with the step generation in effect right + now. A key bound before the first `advance_step_generation()` call + -- a construction-time constant -- carries generation 0 forever, + because nothing ever rebinds it. A key bound while a forward is + underway carries that forward's generation, and gets left behind the + moment the next forward advances the counter without rebinding it. + That drift is exactly the bug `_validating.validating` watches for: + the base class cannot tell a construction-time constant from a + volatile per-forward value by name anymore, but it can tell whether + a given key's last bind has fallen behind the current forward. + """ + self._const.update(values) + for key in values: + self._generation[key] = _STEP_GENERATION + + def bind_layered(self, layer_id: int, **values: Any) -> None: + """Bind one layer's values, read back by `layer=` at the call. + + One layer per call rather than a whole table: a caller binds + exactly the layers that have a given operand, which is what a model + whose layers are not uniform -- some dense, some not -- needs, and a + whole-table signature could not express without the caller padding + the gaps itself. The call site is unchanged either way -- still + `op(..., layer=i)`. + + Same update-not-replace and generation bookkeeping as `bind_const`, + for the same reason: a layered value is not inherently a + construction-time thing -- gpt_oss's `attention_window_size` is + rebound this way on every forward rather than once at construction, + and does so one layer at a time just like a construction-time table + would. + """ + for name, value in values.items(): + self._layered.setdefault(name, {})[layer_id] = value + self._generation[name] = _STEP_GENERATION + + def __call__(self, *args: Any, layer: int | None = None, **kwargs: Any) -> Any: + """Merge const, the layer row, and the call's own kwargs, in that order. + + Later stages win: the layer row overrides a flat const for the same + argument, and a call-time kwarg overrides both -- a target can still + pass a per-call override through the same call a layered value is + bound on. An unbound entry has empty `_const` and `_layered`, so a + call that never passes `layer=` is a plain passthrough -- which is + what keeps deepseek's call sites byte-for-byte the same call they + always were. + """ + merged = dict(self._const) + if layer is not None: + merged.update({name: rows[layer] for name, rows in self._layered.items()}) + merged.update(kwargs) + return self.raw_call(*args, **merged) + + @abstractmethod + def raw_call(self, *args: Any, **kwargs: Any) -> Any: + """Make the op call, taking every argument explicitly. The hot path. + + Must not read bound state off `self` -- see the class docstring. The + test tree interposes validation by patching this method on the + subclass, which is why it has to live there. + """ + + @abstractmethod + def reference(self, *args: Any, **kwargs: Any) -> Any: + """What the op is supposed to compute, in plain torch. + + Mirrors `raw_call`'s signature, parameter for parameter and name for + name, so the two can be driven from one cell's argument list. An + argument this computation does not need keeps its name anyway: dropping + it would hide that the op takes it, and renaming it would break every + call site that passes it by keyword. Say which ones are ignored in the + docstring; nothing in this repo's lint objects to an unused parameter. + + Collectives are the one documented exception to mirroring at all. A + collective's output is not a function of what the calling rank holds, so + a reference given only the local `input` could state nothing; those take + the group's inputs in rank order and say so in their own docstring. + + Accumulate in fp32 and round once at the end: the reference has to be + the more accurate side, or `compare`'s band measures the reference + rather than the op. + """ + + def is_valid(self, *args: Any, **kwargs: Any) -> None: + """Raise on an input the op accepts but computes wrongly. + + Not called on the way to the kernel -- see the module docstring. Default + is nothing to check; override only where the op is known to take an + input and return a plausible wrong answer. A condition the op already + rejects loudly does not belong here: restating it only swaps its error + for a worse one. + + Mirrors `raw_call`'s signature, like `reference`, because `validating` + forwards the call's arguments verbatim -- an override that accepts only + the arguments it inspects raises TypeError on every call site that + passes the others. Parameters keep the op's own names for the same + reason `reference`'s do. + """ + return None + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Gate `out` against `reference`. + + The default is torch's dtype-aware band, which is right for an op that + rounds once. An op that accumulates through several GEMMs will exceed it + for reasons that are not defects, and must widen the band *and say what + it measured* -- a tolerance without a derivation is a tolerance that + will be widened again next time it fails. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + torch.testing.assert_close(out, ref) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.md deleted file mode 100644 index 71e421866607..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.md +++ /dev/null @@ -1,83 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 5} ---- - -# flashinfer_silu_and_mul - -**Wraps** `torch.ops.trtllm.flashinfer_silu_and_mul` (one call). - -## Semantics - -SiLU-gated elementwise multiply (the SwiGLU MLP activation), fused in one -kernel. The last dim of the input holds the gate half followed by the up -half: - -``` -d = x.shape[-1] // 2 -out[..., i] = silu(x[..., i]) * x[..., d + i] for i in [0, d) -silu(v) = v / (1 + exp(-v)) -``` - -Both halves are loaded and the silu/multiply are computed in fp32 inside -the kernel; the result is cast back to the input dtype. - -Fusion boundary: the single call computes activation and gating multiply -only. The caller owns the gate/up projection that produces `x` (whether as -one fused matmul or two) and the down projection that consumes the output. -There is no quantization of the output (a separate -`torch.ops.trtllm.silu_and_mul` op exists with an optional quant scale) and -no gelu variant (a separate `torch.ops.trtllm.flashinfer_gelu_tanh_and_mul` -op exists). - -## Signature - -```python -def flashinfer_silu_and_mul(x: torch.Tensor) -> torch.Tensor -``` - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `x` | `[..., 2 * d]`, e.g. `[num_tokens, 2 * intermediate]` | fp16 / bf16 | contiguous | CUDA | -| returns | `[..., d]` | same as `x` | newly allocated, contiguous | CUDA (same device as `x`) | - -The output is a new tensor (`x` is not mutated). - -## Metadata consumed - -None. Stateless. - -## Preconditions - -- `x` is on a CUDA device and fully contiguous: the kernel indexes rows as - `token * 2d` on the raw pointer, so any non-contiguity silently corrupts - results. -- `x.shape[-1]` is even and a multiple of 16 elements for fp16/bf16 - (equivalently `d * itemsize % 16 == 0`). This is stricter than the - op's own runtime check — see Notes. -- `x.shape[-1] >= 16` for fp16/bf16 (`d` must hold at least one 16-byte - vector, or the computed block size is 0 and the launch is invalid). -- Dtype is fp16 or bf16. Other dtypes (including fp32) are rejected by the - kernel's dtype dispatch. -- No upper bound on `d` beyond memory: `d` larger than - `1024 * (16 / itemsize)` (e.g. > 8192 for bf16) takes a scalar remainder - loop, verified correct on this machine. - -## Notes - -- The op is registered only when flashinfer is importable - (`IS_FLASHINFER_AVAILABLE`). -- Alignment trap: the op's Python-side check only validates - `x.shape[-1] * itemsize % 16 == 0` (raising `ValueError`), but the - vectorized load of the up half starts at element offset `d`, so `d` - itself must be 16-byte aligned. Inputs that pass the check with - `d * itemsize % 16 != 0` (e.g. bf16 with last dim 24 or 2728) crash with - CUDA `misaligned address` — observed on sm_100. Keep the last dim a - multiple of 16 elements. -- The misaligned-address failure is sticky: it poisons the CUDA context - for the rest of the process. -- Programmatic dependent launch (PDL) is controlled by the env var - `TRTLLM_ENABLE_PDL` (default enabled) inside the trtllm custom op; it - affects scheduling only, not results. -- silu uses the fast-math `__expf`; observed error vs an fp32 torch - reference stays within default `assert_close` tolerances for fp16/bf16. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py index 7817e4effc85..95679fc4685e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py @@ -1,30 +1,106 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""SiLU-gated multiply (SwiGLU activation) via the flashinfer silu_and_mul kernel.""" +"""Gated SiLU activation via the flashinfer silu_and_mul kernel.""" + +from __future__ import annotations import torch +from torch.nn import functional as F import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from .._op import Arch, Cell, OpWrapper -def flashinfer_silu_and_mul(x: torch.Tensor) -> torch.Tensor: - """Return `silu(x[..., :d]) * x[..., d:]` with `d = x.shape[-1] // 2` as a new tensor.""" - # The op's own check is on the *row*, `x.shape[-1] * itemsize % 16 == 0`, - # which a final dimension of 24 passes. The kernel vectorizes over the two - # halves, so what has to be 16-byte aligned is the *half*: at 24 the - # second half starts mid-vector and the launch dies with `CUDA misaligned - # address`, poisoning the context rather than raising. This is the - # stricter precondition the contract states. - width = x.shape[-1] - assert width % 2 == 0, f"x.shape[-1] must be even to split in half; got {width}" - half_bytes = (width // 2) * x.element_size() - assert half_bytes % 16 == 0, ( - f"x.shape[-1] // 2 must be a whole number of 16-byte vectors; " - f"{width} halves to {half_bytes} bytes, which is not a multiple of 16 " - f"-- the kernel would fault on the misaligned second half" - ) - assert half_bytes >= 16, ( - f"x.shape[-1] // 2 must hold at least one 16-byte vector; {width} " - f"gives {half_bytes} bytes, for which the computed block size is 0" + +class FlashinferSiluAndMul(OpWrapper): + """`silu(x[..., :d]) * x[..., d:]` with `d = x.shape[-1] // 2`, new tensor out. + + The gate and up projections arrive concatenated along the last dim, which + is why the op takes one tensor and halves it rather than taking two. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # deepseek-r1 is the only shipped caller. It reaches this op from two + # widths -- the routed/shared expert intermediate and the dense layers' + # -- at decode and prefill token counts. + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode-shaped, 2 x moe_intermediate (2048)", + spec=dict(shape=(1, 4096), dtype=torch.bfloat16), + ), + Cell( + why="prefill-shaped: enough rows to cross the kernel's grid tiling", + spec=dict(shape=(2048, 4096), dtype=torch.bfloat16), + ), + Cell( + why="the dense layers' wider intermediate, 2 x 18432", + spec=dict(shape=(16, 36864), dtype=torch.bfloat16), + ), + # DeepSeek-V4.1's vision tower MLP: [gate ; up] at d = 2816, P patches + # per tower call, every P in 1..17037 certified by the prefix sweep in + # the test; these are the sweep's two ends. `data="vit"` is that test's + # input: N(0, 3) with planted zero, large-gate, large-up, tiny-gate and + # large-negative-gate rows. + Cell( + why="deepseek-v4.1 vision [P, 5632] P=1: the tower MLP, " + "one row, bottom of the prefix sweep 1..17037", + spec=dict(shape=(1, 5632), dtype=torch.bfloat16, data="vit"), + ), + Cell( + why="deepseek-v4.1 vision [P, 5632] P=17037: the tower MLP, " + "its largest call, top of the prefix sweep 1..17037", + spec=dict(shape=(17037, 5632), dtype=torch.bfloat16, data="vit"), + ), ) - return torch.ops.trtllm.flashinfer_silu_and_mul(x) + + note = """ + Registered only when flashinfer is importable (`IS_FLASHINFER_AVAILABLE`). + + fp16 and fp32 are accepted by the op. No shipped target uses them here, so + they are not cells. + + DeepSeek-V4.1: the vision tower's MLP calls it on cublas_mm's contiguous + bf16 `[P, 5632]` output (`[gate ; up]` in the checkpoint's w1 order, + d = 2816), unclamped, eagerly, with P any patch count 1..17037. Metadata + none. Every P is certified: the 17037-row call meets the reference, every + prefix of it is bitwise the op on that many rows, nothing is written past + the last row, and one kernel (block 352, grid P) runs at every count. The + kernel is built with fast math: gates in (-88.72, -87.33] and results + below 2^-126 in magnitude come back exactly 0 where the fp32 reference is a + tiny nonzero (below 1e-36 x max(1, |up|)). The default gate passes them; + the test's per-element 1-ulp gate admits exactly that set, measured over + every finite bf16 gate. Not certified for V4.1: CUDA-graph capture at this + width, fp16, more than 17037 rows, a non-contiguous or row-strided input, + NaN or inf input. + """ + + def raw_call(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.flashinfer_silu_and_mul(x) + + def reference(self, x: torch.Tensor) -> torch.Tensor: + """fp32-accumulated; the kernel computes in fp32 and rounds once.""" + gate, up = x.float().chunk(2, dim=-1) + return (F.silu(gate) * up).to(x.dtype) + + def is_valid(self, x: torch.Tensor) -> None: + # The op's own check is on the *row*, `x.shape[-1] * itemsize % 16 == 0`, + # which a final dimension of 24 passes. The kernel vectorizes over the + # two halves, so what has to be 16-byte aligned is the *half*: at 24 the + # second half starts mid-vector and the launch dies with `CUDA + # misaligned address`, poisoning the context rather than raising. + width = x.shape[-1] + assert width % 2 == 0, f"x.shape[-1] must be even to split in half; got {width}" + half_bytes = (width // 2) * x.element_size() + assert half_bytes % 16 == 0, ( + f"x.shape[-1] // 2 must be a whole number of 16-byte vectors; " + f"{width} halves to {half_bytes} bytes, which is not a multiple of 16 " + f"-- the kernel would fault on the misaligned second half" + ) + assert half_bytes >= 16, ( + f"x.shape[-1] // 2 must hold at least one 16-byte vector; {width} " + f"gives {half_bytes} bytes, for which the computed block size is 0" + ) + + +flashinfer_silu_and_mul = FlashinferSiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.md deleted file mode 100644 index 33a26545df88..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.md +++ /dev/null @@ -1,158 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 9} ---- - -# fused_qk_norm_rope - -**Wraps** `torch.ops.trtllm.fused_qk_norm_rope` (one call). - -## Semantics - -In-place, on a packed QKV activation `qkv` of shape -`[num_tokens, (num_heads_q + num_heads_k + num_heads_v) * head_dim]` whose -rows lay out `num_heads_q` query heads, then `num_heads_k` key heads, then -`num_heads_v` value heads, each `head_dim` wide. One call computes, per -token, for every **q and k head** `x` (fp32 internally, one final cast back -to bf16): - -``` -# 1. per-head RMS norm over the full head_dim (skipped if is_qk_norm=False) -h = x / sqrt(mean(x^2) + eps) * w # w = q_weight or k_weight - # use_gemma: * (1 + w) instead -# 2. RoPE on the first rotary_dim dims of h; h[rotary_dim:] passes through -inv_freq[j] = base^(-2j/rotary_dim) for j in [0, rotary_dim/2) -angle[j] = position_id * inv_freq[j] -cos, sin = cos(angle) * attention_factor, sin(angle) * attention_factor -# is_neox=True — half-split pairs (x1, x2) = (h[j], h[j + rotary_dim/2]) -# is_neox=False — interleaved pairs (x1, x2) = (h[2j], h[2j+1]) -(x1, x2) -> (x1*cos - x2*sin, x1*sin + x2*cos) -``` - -Value heads are passed through bit-exactly untouched. - -YaRN scaling (`factor`, `low`, `high`, `attention_factor`) replaces -`inv_freq` with the blend - -``` -ramp[j] = clamp((j - low) / (high - low), 0, 1) # high==low: high += 0.001 -inv_freq[j] = (1 - ramp[j]) * base^(-2j/rotary_dim) - + ramp[j] * base^(-2j/rotary_dim) / factor -``` - -which reduces to plain RoPE at `factor=1.0` (the `low`/`high` values are -then irrelevant). `attention_factor` scales cos/sin — but **not at -`factor == 1.0` exactly**, where any `attention_factor != 1.0` raises -`Assertion failed: attention_factor == 1.0f` -(`fusedQKNormRopeKernel.cu:322`). `factor = 1.0000001` accepts it. -The gate is on `factor` alone and is not otherwise documented. - -Interleaved mRoPE (`use_mrope=True`) takes 3 position rows -(temporal/height/width) per token. Frequency index `j` reads its position -from row 1 when `j % 3 == 1 and j < 3*mrope_section1`, from row 2 when -`j % 3 == 2 and j < 3*mrope_section2`, and from row 0 otherwise. Only this -interleaved layout is implemented — not the contiguous-section mRoPE -variant. - -Fusion boundary: q-norm + k-norm + RoPE, nothing else. The caller owns the -QKV projection before the call and everything after it (KV-cache append, -attention, output projection). A sibling op `fused_dit_qk_norm_rope` -exists for the DiT-style variant of this fusion. - -## Signature - -```python -def fused_qk_norm_rope( - qkv: torch.Tensor, - num_heads_q: int, - num_heads_k: int, - num_heads_v: int, - head_dim: int, - rotary_dim: int, - eps: float, - q_weight: torch.Tensor, - k_weight: torch.Tensor, - base: float, - is_neox: bool, - position_ids: torch.Tensor, - factor: float = 1.0, - low: float = 0.0, - high: float = 0.0, - attention_factor: float = 1.0, - is_qk_norm: bool = True, - use_gemma: bool = False, - use_mrope: bool = False, - mrope_section1: int = 0, - mrope_section2: int = 0, -) -> None -``` - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `qkv` | `[num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]` (2D only) | bf16 | contiguous | CUDA | -| `num_heads_q/k/v` | scalar | Python int | — | — | -| `head_dim` | scalar, one of 64 / 128 / 256 | Python int | — | — | -| `rotary_dim` | scalar, even, `<= head_dim` | Python int | — | — | -| `eps` | scalar | Python float | — | — | -| `q_weight`, `k_weight` | `[head_dim]` | bf16 | contiguous | CUDA (same device) | -| `base` | scalar (RoPE theta) | Python float | — | — | -| `is_neox` | scalar | Python bool | — | — | -| `position_ids` | `[num_tokens]`, or `[3, num_tokens]` row-major when `use_mrope` | int32 | contiguous | CUDA (same device) | -| `factor`, `low`, `high`, `attention_factor` | scalar (YaRN; defaults = plain RoPE) | Python float | — | — | -| `is_qk_norm` | scalar (False: RoPE only, weights ignored) | Python bool | — | — | -| `use_gemma` | scalar (True: scale by `1 + w`) | Python bool | — | — | -| `use_mrope`, `mrope_section1`, `mrope_section2` | scalar (interleaved mRoPE) | Python bool / int | — | — | -| returns | — | — | — | — | - -Returns `None`. `qkv` is mutated in place: q and k head slices are -overwritten with the normed + rotated values; v head slices are untouched. - -## Metadata consumed - -None. Stateless — all RoPE frequencies are derived on the fly from the -scalar arguments; no cos/sin cache and no attention metadata are read. -`position_ids` are the absolute positions of the `num_tokens` tokens, in -row order of `qkv` (batch layout is the caller's concern; the kernel is -purely per-token). - -## Preconditions - -- `qkv` is 2D, contiguous, bf16, on CUDA, with - `qkv.shape[1] == (num_heads_q + num_heads_k + num_heads_v) * head_dim`. - bf16 is the only accepted dtype (fp16/fp32 raise). -- `head_dim` is one of 64, 128, 256. Others (32, 80, 96, 192 verified) - raise `Unsupported head dimension for fusedQKNormRope`. -- `q_weight` and `k_weight` are contiguous bf16 `[head_dim]` tensors on - the same device; they must be passed (and valid) even when - `is_qk_norm=False`, though their values are then unused. -- `position_ids` is contiguous int32 on the same device: flat - `[num_tokens]` normally, `[3, num_tokens]` row-major when - `use_mrope=True`. int64 raises. A 3D upstream view must be reshaped and - made contiguous by the caller. -- `rotary_dim` is even and at most `head_dim`; dims beyond `rotary_dim` - are normed (if `is_qk_norm`) but not rotated. **Only the evenness half - is enforced** (`Assertion failed: rotary_dim must be even`); - `rotary_dim > head_dim` is accepted silently and rewrites the q and k - head columns — measured at 160 and 256 against `head_dim = 128`, - 2026-07-28. This is the caller's obligation, not the op's. -- Everything else above is validated by the op with a clear error, and in - every raising case `qkv` comes back bit-identical to its pre-call - contents. The two exceptions are the `rotary_dim` bound just named and - the `attention_factor` gate under *Semantics*. - -## Notes - -- The kernel checks dtype/shape/contiguity eagerly (thop-level asserts), - so contract violations fail loudly rather than corrupting `qkv`. -- Accuracy: internal math is fp32 with a single bf16 round at the end. - Because RoPE rotates a bf16-rounded pair, the per-element error is - absolute in the pair magnitude — with unit-scale inputs, max abs error - observed on sm_100 was 0.031, at most 78% of the test's - `rtol=1.6e-2, atol=2e-2` allowance (compare with `atol` sized to the q/k - magnitude, not the default bf16 `atol=1e-5`). The ceiling sits on the - neox-prefill and Gemma-norm cases; most of the file sits at 0.008-0.016. -- `low`/`high` are float half-dim indices in `[0, rotary_dim/2)` - (fractional values legal, matching YaRN's `truncate=False`). -- TRT-LLM's own caller derives `factor/low/high/attention_factor` from a - YaRN config and uses `rotary_dim = head_dim * partial_rotary_factor`; - this entry exposes them raw. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py deleted file mode 100644 index 3dfc7ae6b29b..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py +++ /dev/null @@ -1,56 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""In-place fused per-head QK RMS norm + RoPE on a packed QKV tensor.""" - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* - - -def fused_qk_norm_rope( - qkv: torch.Tensor, - num_heads_q: int, - num_heads_k: int, - num_heads_v: int, - head_dim: int, - rotary_dim: int, - eps: float, - q_weight: torch.Tensor, - k_weight: torch.Tensor, - base: float, - is_neox: bool, - position_ids: torch.Tensor, - factor: float = 1.0, - low: float = 0.0, - high: float = 0.0, - attention_factor: float = 1.0, - is_qk_norm: bool = True, - use_gemma: bool = False, - use_mrope: bool = False, - mrope_section1: int = 0, - mrope_section2: int = 0, -) -> None: - """In-place on qkv: per-head RMS norm of q/k heads, then RoPE. Returns None.""" - torch.ops.trtllm.fused_qk_norm_rope( - qkv, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - is_qk_norm, - use_gemma, - use_mrope, - mrope_section1, - mrope_section2, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/load_paged_kv_cache_for_mla.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/load_paged_kv_cache_for_mla.md deleted file mode 100644 index d6ed03181e5a..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/load_paged_kv_cache_for_mla.md +++ /dev/null @@ -1,316 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 10} ---- - -# load_paged_kv_cache_for_mla - -**Wraps** `torch.ops.trtllm.load_paged_kv_cache_for_mla` (one call). - -## Semantics - -The cache-read step of MLA context-phase attention with KV reuse: after a -context batch's full latent KV (cached prefix + this step's new tokens) is -resident in the paged MLA cache, one call gathers it back into dense -tensors so the caller can up-project and run context FMHA over the full -`[past + new]` range. - -For each context sequence `s` in `[0, num_contexts)`, with total latent-KV -length `L_s = cu_ctx_kv_lens[s+1] - cu_ctx_kv_lens[s]`, the call reads the -sequence's paged latent rows at positions `[0, L_s)` — each row is the -per-token `[compressed_kv | k_pe]` of width `C + R` -(`C = kv_lora_rank`, `R = qk_rope_head_dim`) — and writes them, -sequences concatenated in batch order, into two freshly allocated -contiguous tensors: - -``` -T = num_ctx_kv_tokens = sum(L_s) -compressed_kv[t, :] = cache_row(t)[:C] # [T, C] -k_pe[t, :] = cache_row(t)[C:] # [T, R] -``` - -The paged pool is never modified, and sequences at batch index -`>= num_contexts` (generation sequences) contribute nothing to the outputs. -Two pool configurations are certified, selected by `quant_mode`: - -**High-precision pool (`quant_mode = 0`).** The cache element type is -`out_dtype` and the gather is a bitwise copy. - -**fp8-e4m3 pool (`quant_mode` with the fp8-KV bit, `128`).** The cache -holds one `__nv_fp8_e4m3` byte per element and the gather dequantizes. -With `r = kv_scale_quant_orig[0]` (`1.0` when that tensor is omitted): - -``` -compressed_kv[t, :] = out_dtype( float(cache_row(t)[:C]) * r ) -k_pe[t, :] = out_dtype( float(cache_row(t)[C:]) * r ) -``` - -One fp32 multiply per element, by the exact fp32 number in -`kv_scale_quant_orig[0]`, rounded once to `out_dtype` — not a multiply in -`out_dtype`, which differs on 24.7% of the elements when `r` is not exactly -representable there (measured, `r = 1/1.5`, bf16). The byte→float step is -exact for all 256 e4m3 byte values, so at `r = 1` the gather is a lossless -widening. There is no per-block scale anywhere: the fp8 latent pool stores -nothing but e4m3 elements, and `r` is a single static number the caller -supplies. - -Fusion boundary: only this cache read happens inside the call. The caller -still owns writing the latent rows in the first place (the cached prefix -from earlier steps plus this step's new tokens — a sibling -context-preprocessing op appends the latter, quantizing them on the way in -under its own **write-side** scale), the `kv_b_proj` up-projection of -`compressed_kv`, and the subsequent context attention. **This call applies -`kv_scale_quant_orig` unconditionally and checks nothing about how the pool -came to hold what it holds** — see Preconditions. - -## Signature - -```python -def load_paged_kv_cache_for_mla( - out_dtype: torch.dtype, - num_contexts: int, - num_ctx_kv_tokens: int, - max_ctx_kv_len: int, - cu_ctx_kv_lens: torch.Tensor, - kv_cache_block_offsets: torch.Tensor, - host_kv_cache_pool_pointers: torch.Tensor, - host_kv_cache_pool_mapping: torch.Tensor, - kv_scale_quant_orig: Optional[torch.Tensor], - layer_idx: int, - kv_lora_rank: int, - qk_rope_head_dim: int, - tokens_per_block: int, - attention_window_size: int, - beam_width: int, - quant_mode: int, -) -> Tuple[torch.Tensor, torch.Tensor] -``` - -With `S` = total sequences in the batch (contexts first), `C` = -`kv_lora_rank`, `R` = `qk_rope_head_dim`, `T` = `num_ctx_kv_tokens`: - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `out_dtype` | — | the **output** element type, not the cache's: `torch.bfloat16` / `torch.float16` / `torch.float32`. On the fp8 path all three are certified over the same e4m3 pool; on the high-precision path it also fixes the cache's element width (see Preconditions) and bf16 / fp16 are certified there | — | — | -| `num_contexts` | leading context-phase sequence count | Python int | — | — | -| `num_ctx_kv_tokens` | `T = sum(L_s)` over context seqs, cached **plus** new tokens (sizes the outputs) | Python int | — | — | -| `max_ctx_kv_len` | `max(L_s)` over context seqs | Python int | — | — | -| `cu_ctx_kv_lens` | `[>= num_contexts+1]`; `[0, cumsum(L_s)]`; only the first `num_contexts+1` entries are read | int64 | contiguous | CUDA | -| `kv_cache_block_offsets` | `[1, >= S, 2, max_blocks_per_seq]`; raw block ids per seq (K and V rows identical, kv_factor=1 pool); contexts first | int32 | contiguous | CUDA | -| `host_kv_cache_pool_pointers` | `[num_pools, 2]`: (primary ptr, secondary ptr=0) | int64 | contiguous | CPU | -| `host_kv_cache_pool_mapping` | `[num_layers, 2]`: (pool index, layer-within-pool) row per layer | int32 | contiguous | CPU | -| `kv_scale_quant_orig` | the **read-side** dequantization scale; element `[0]` is the only one read. `None` behaves exactly as `1.0` and is what the engine's own call site passes. Read only when `quant_mode` has the fp8-KV bit — silently ignored otherwise | fp32 (others raise) | any 1-element-or-larger shape; a 0-dim scalar also works | CUDA (a CPU tensor is also accepted — see Notes) | -| `layer_idx` | row into `host_kv_cache_pool_mapping` | Python int | — | — | -| `kv_lora_rank` | `C` (512 certified) | Python int | — | — | -| `qk_rope_head_dim` | `R` (64 certified) | Python int | — | — | -| `tokens_per_block` | pool page size; 32 and 64 certified on both paths (32 is what a default `KvCacheConfig` produces) | Python int | — | — | -| `attention_window_size` | `>= max(L_s)`; production passes the manager's `max_seq_len` (smaller values imply cyclic-cache addressing, not certified) | Python int | — | — | -| `beam_width` | `1` | Python int | — | — | -| `quant_mode` | `0` = high-precision (bf16/fp16) KV cache; `128` (`QuantMode.FP8_KV_CACHE`) = e4m3 KV cache. Both certified. Other quantization bits set **alongside** `128` are not read — `1152` (`FP8_KV_CACHE \| FP8_1x128_128x128`, the combination trtllm's auto-deploy MLA path builds for an fp8 latent cache) and `384` (`\| FP8_QDQ`) return bit-identical results. `64` and `8192` are rejected (see Notes) | Python int | — | — | -| returns `compressed_kv` | `[T, C]`, freshly allocated | `out_dtype` | contiguous | CUDA | -| returns `k_pe` | `[T, R]`, freshly allocated | `out_dtype` | contiguous | CUDA | - -## Metadata consumed - -No thread-local or registered-layer state: all inputs are explicit -arguments. The addressing tensors are exactly what a -`TrtllmAttentionMetadata` prepared with -`enable_context_mla_with_cached_kv=True` and its `KVCacheManager` expose: -`ctx_kv_indptr` → `cu_ctx_kv_lens`, -`num_ctx_cached_tokens + num_ctx_tokens` → `num_ctx_kv_tokens`, -`max_ctx_kv_len` → `max_ctx_kv_len`, `metadata.kv_cache_block_offsets`, -`manager.kv_cache_pool_pointers`, `manager.kv_cache_pool_mapping`, -`manager.tokens_per_block`, `manager.max_seq_len` → -`attention_window_size`. Nothing about `kv_scale_quant_orig` comes from -metadata; it is the caller's number. - -## Preconditions - -- The paged pool addressed by `host_kv_cache_pool_pointers` + - `host_kv_cache_pool_mapping[layer_idx]` is an MLA latent cache: - kv_factor 1, one kv head, row width `C + R`, page size - `tokens_per_block`. Per page the layout is `[tokens_per_block, C + R]` - (page base = pool base + `block_id * tokens_per_block * (C + R)` - elements). -- **`quant_mode` alone fixes the pool's element width**, because the pool - arrives as a raw pointer and nothing else describes it: one byte per - element when the fp8-KV bit (`128`) is set, `sizeof(out_dtype)` when it - is not. So the caller owes a matching pool — e4m3 under `quant_mode=128`, - `out_dtype` under `quant_mode=0`. A mismatch in either direction is - **silent** (see Notes). -- `out_dtype` is one of fp16 / fp32 / bf16; anything else raises - `RuntimeError: out_dtype only support float16, float32, bfloat16` (tested - with e4m3, int8 and fp64). The output tensors are allocated before that - check, but no kernel runs. Note that under `quant_mode=128` `out_dtype` - is necessarily **not** the cache dtype: the cache is e4m3 and e4m3 is not - an accepted `out_dtype`. -- Every position `[0, L_s)` of every context sequence has already been - written to the pool (cached prefix and this step's new tokens alike). - The op reads raw pool memory, so slots never written simply propagate - whatever they currently hold. -- `kv_cache_block_offsets` covers `ceil(L_s / tokens_per_block)` pages for - every context sequence. -- `cu_ctx_kv_lens` is int64 (the kernel rejects other index dtypes), on - device, starts at 0, and is consistent with `num_ctx_kv_tokens` - (`cu_ctx_kv_lens[num_contexts] == num_ctx_kv_tokens`) and - `max_ctx_kv_len` (`>=` every `L_s`). -- `num_contexts > 0`, `num_ctx_kv_tokens > 0` and `max_ctx_kv_len > 0`: all - three are checked host-side and raise `RuntimeError` before any device - work. Production cannot reach a violation — the MLA context path runs only - under `num_contexts > 0`, and the cached-KV branch additionally requires - `num_ctx_cached_tokens > 0` — but a caller driving the op directly owns - them. -- `L_s <= attention_window_size` for every context sequence. -- `beam_width == 1`. - -**fp8 pool only (`quant_mode` fp8-KV bit set):** - -- `kv_scale_quant_orig` is an fp32 tensor whose element `[0]` is the - dequantization scale, or `None` for `1.0`. A `float16` or `float64` - tensor raises `RuntimeError: expected scalar type Float but found - Half`/`Double`; nothing else about it is validated (see Notes). -- **`kv_scale_quant_orig` must be the reciprocal of the scale the pool was - written at, and nothing anywhere checks that.** The write-side scale is - an argument of a different op (`kv_scale_orig_quant` on the sibling that - appends the new context rows) and this call never sees it. A pool written - at `w` and read at `r` yields exactly `float(e4m3(row * w)) * r`, so any - `r != 1/w` is a silently mis-scaled prefill with no error at any layer — - measured at `(w, r)` = (2, 2), (4, 1), (1, 3), each landing exactly - `w * r` away from the original rows. Production for - DeepSeek-R1-0528-FP4 is `w = r = 1.0`: all 124 per-layer `k_scale` / - `v_scale` tensors in that checkpoint are 0-dim fp32 scalars equal to 1.0 - (read from its safetensors shards), and the engine's own call sites - hard-code `None` on both the write and the read side. -- `out_dtype` must have the range to hold `max|e4m3 element| * r`. The - product is **not** clamped to `out_dtype`: with an e4m3 `±448` in the - pool and `r = 256`, fp16 output is `±inf` while bf16 and fp32 hold - `±114688` (measured). A NaN byte in the pool (`0x7f` / `0xff`) becomes a - NaN output. -- The pool must actually hold e4m3 bytes at `C + R` bytes per row. The op - reads them as raw bytes; no header, no per-block scale, no metadata. - -A caller violating none of these gets a correct result. - -## Notes - -- Schema argument names disagree with certified usage: the schema calls the - third and fourth arguments `num_ctx_cached_tokens` and - `max_ctx_cached_kv_len`, but the production call path passes cached - **plus** new token totals (`num_ctx_cached_tokens + num_ctx_tokens`, - `max_ctx_kv_len`). The totals are what size the outputs and bound the - read range; the wrapper parameter names reflect that. -- Precision, high-precision pool (sm_100): bf16→bf16 and fp16→fp16 gathers - are bitwise (`rtol=0, atol=0` holds), including rows crossing page - boundaries and pool-mapping rows other than 0 (`layer_idx=1` of a - two-layer pool certified). -- Precision, fp8 pool (sm_100): the gather is **bit-exact** against - `(float(e4m3 byte) * r)` rounded once to `out_dtype`, over - `r ∈ {omitted, 1.0, 1/1.5, 0.5, 0.25, 2.0}` × `out_dtype ∈ {bf16, fp16, - fp32}`, and separately over pool contents written at `w ∈ {0.25, 0.5, - 1.0, 1.5, 2.0, 4.0}`. Both output halves are covered by the same gate, so - the `C`/`R` split of the 1-byte row is pinned at byte `C`. -- fp8 domain (sm_100): all 256 e4m3 byte values — `±448`, `±0`, the - subnormals down to `2^-9`, and both NaN bytes — come back **exactly** at - `r = 1.0` in fp16, bf16 and fp32. The byte→float step neither saturates - nor flushes to zero; it is a plain conversion, and only the multiply that - follows can leave `out_dtype`'s range (see Preconditions). The *write* - side of the round trip behaves differently — the sibling append op's own - contract records that it clamps to `±448` where - `torch.Tensor.to(float8_e4m3fn)` yields the NaN byte — so a torch mirror - of the whole round trip is only valid while every `|value * w| <= 448`. - Every case certified here stays inside that. -- Page size (sm_100): certified at `tokens_per_block` 32 and 64 on both - paths. 32 is what a `KvCacheConfig` at its defaults hands the op; 64 is - what a tuned MLA target opts into. Both were run through the same bitwise - gate on real `KVCacheManager` state, over gathered ranges that end on a - page boundary and mid-page alike — high-precision pool, at 32: 37 / 96 - (three exact pages) / 400 (twelve pages plus sixteen rows) / 512 (sixteen - exact pages) in bf16 through pool-mapping row 1 of a two-layer pool, a - mixed batch with trailing generation sequences (128 = four exact pages, - and 63), and 64 / 608 / 7 in fp16; the deepest sequence walked nineteen - offsets-row entries. fp8 pool, at 32: gathered ranges of 37 / 96 / 400 / - 512 (1045 rows in one call) plus a mixed batch (128 and 63); at 64: - 1024 / 1024 / 135 (2183 rows in one call). The gather stayed bitwise in - every case, so - `page = t // tokens_per_block`, `slot = t % tokens_per_block` holds at - both sizes and both element widths. Other page sizes are untested rather - than known-bad. -- The outputs are new allocations sized `[T, C]`/`[T, R]`; a `T == 0` batch - is **rejected** by the op rather than merely unexercised (see - Preconditions), and production never reaches it anyway — the path is - guarded behind `num_ctx_cached_tokens > 0`. -- The pool is not written on either path: a whole-pool byte snapshot around - a 1045-row fp8 gather shows zero changed bytes, and a sentinel row placed - one slot past each context sequence's range never appears in the output. - Two identical calls return bitwise-identical tensors. -- A cache element-type or element-width mismatch (sm_100) is **silent in - every form measured**. The two mismatches an fp8 pool makes reachable — - reading it two bytes wide, and reading a two-byte pool one byte wide — - are gated in this entry's test; the other observations below are from - scratch probes. - - *Same width, wrong type* (bf16 pool read as fp16, or fp16 pool read as - bf16, `quant_mode=0` both times): addressing is unchanged and the cache - bits are copied out verbatim, so the output is bit-identical to the - correct gather and numerically wrong — max abs diff 2.0 one way and 512 - the other, nothing raised, nothing out of bounds. - - *Read wider than the pool* (fp32 `out_dtype` over a 2-byte pool, or - `quant_mode=0` over an e4m3 pool): the stride doubles, so token `t` of - block `b` is read from pool row - `2 * (b * tokens_per_block + t % tokens_per_block)` and consumes two - rows. Verified bit for bit both ways — 565 gathered rows in the - fp32-over-bf16 case, of which 310 came back holding a sentinel written - nowhere near that sequence, and the exact doubled row indices in the - e4m3 case (block 0 token 1 → row 2, block 1 token 0 → row 64). Past - half the pool it leaves the allocation entirely: at block 1584 of a - 2048-block pool compute-sanitizer reported invalid 16-byte `__global__` - reads 77.9 MB beyond the nearest allocation, while the call still - returned without raising. - - *Read narrower than the pool* (`quant_mode=128` over a bf16 pool): the - 1-byte slab for block `b` starts at byte - `b * tokens_per_block * (C + R)`, i.e. inside bf16 block `b // 2`, and - each token consumes `C + R` bytes. Verified bit for bit at blocks 0 and - 1. This direction stays inside the allocation — the read walks half the - bytes the pool holds — but every value is garbage, including NaNs - wherever a bf16 byte happens to be `0x7f`/`0xff` (107 of 36864 elements - in one scratch probe). -- `kv_scale_quant_orig` is read **only** when `quant_mode` has the fp8-KV - bit. With `quant_mode = 0` a scale of 4.0 produces byte-identical output - to `None`: a silently ignored argument, and the shape a caller who set up - the scale but forgot the `quant_mode` bit lands in. -- `kv_scale_quant_orig` validation is otherwise thin. Only the dtype is - checked, and loudly. A `[2]` fp32 CUDA tensor is accepted and only `[0]` - is used, and a 0-dim fp32 scalar tensor is accepted — both tested. A - **CPU** fp32 tensor is accepted too: the pointer is handed to the kernel - unchecked, and in a scratch probe on this machine it produced results - bit-identical to the equivalent CUDA tensor at two distinct values. That - works because this device reports - `CU_DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS = 1` (driver 595.58.03), so - it is a property of the machine rather than of the op — not something to - rely on, and left out of the test for that reason. Device placement is - not validated either way, so a wrong-device scale will not announce - itself. -- Rejected `quant_mode` values, both tested: - `64` (`QuantMode.INT8_KV_CACHE`) raises - `RuntimeError: [TensorRT-LLM][ERROR] Assertion failed: Only FP8 KV cache - is supported for now (../tensorrt_llm/thop/mlaPreprocessOp.cpp:134)`; - `8192` (`NVFP4_KV_CACHE`) fails earlier and for a different reason — - `Expected hostKvCachePoolPointers.dim() == 3`, the nvfp4 pool-pointer - layout, checked before the quant_mode branch — so it is not the same - guard. What the guard rejects is a *KV-cache* quantization bit other than - fp8; bits outside that group ride along unread, which is why `1152` - (`FP8_KV_CACHE | FP8_1x128_128x128`, the combination trtllm's - auto-deploy MLA path builds for an fp8 latent cache) and `384` - (`| FP8_QDQ`) are bit-identical to plain `128` (tested). -- The engine's own MLA backend passes `None` for `kv_scale_quant_orig` at - this call site while still passing its real `quant_mode`, so an fp8 KV - cache is dequantized at 1.0 there regardless of any calibrated per-layer - scale. Harmless for a checkpoint whose scales are 1.0; the `r != 1` - behaviour certified here is reachable only from a direct op caller. -- Sibling ops exist for adjacent MLA context-phase roles: - `trtllm::mla_rope_append_paged_kv_assign_q` (RoPE + append of the new - context tokens that this op then reads back, and the owner of the - write-side `kv_scale_orig_quant`), - `trtllm::load_chunked_kv_cache_for_mla` (slice-wise variant for chunked - prefill, with its own `kv_scale_quant_orig`), - `trtllm::merge_chunked_attention_for_mla`, and the context FMHA that - consumes the up-projected K/V. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/load_paged_kv_cache_for_mla.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/load_paged_kv_cache_for_mla.py deleted file mode 100644 index 2855bf2d1667..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/load_paged_kv_cache_for_mla.py +++ /dev/null @@ -1,50 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Gather each context sequence's full MLA latent KV from the paged cache.""" - -from typing import Optional, Tuple - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* - - -def load_paged_kv_cache_for_mla( - out_dtype: torch.dtype, - num_contexts: int, - num_ctx_kv_tokens: int, - max_ctx_kv_len: int, - cu_ctx_kv_lens: torch.Tensor, - kv_cache_block_offsets: torch.Tensor, - host_kv_cache_pool_pointers: torch.Tensor, - host_kv_cache_pool_mapping: torch.Tensor, - kv_scale_quant_orig: Optional[torch.Tensor], - layer_idx: int, - kv_lora_rank: int, - qk_rope_head_dim: int, - tokens_per_block: int, - attention_window_size: int, - beam_width: int, - quant_mode: int, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Copy positions [0, L_s) of every context sequence out of the paged MLA - latent cache into two new contiguous tensors (compressed_kv, k_pe).""" - compressed_kv, k_pe = torch.ops.trtllm.load_paged_kv_cache_for_mla( - out_dtype, - num_contexts, - num_ctx_kv_tokens, - max_ctx_kv_len, - cu_ctx_kv_lens, - kv_cache_block_offsets, - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping, - kv_scale_quant_orig, - layer_idx, - kv_lora_rank, - qk_rope_head_dim, - tokens_per_block, - attention_window_size, - beam_width, - quant_mode, - ) - return compressed_kv, k_pe diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_append_paged_kv_assign_q.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_append_paged_kv_assign_q.md deleted file mode 100644 index 768f3c875a12..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_append_paged_kv_assign_q.md +++ /dev/null @@ -1,372 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 6} ---- - -# mla_rope_append_paged_kv_assign_q - -**Wraps** `torch.ops.trtllm.mla_rope_append_paged_kv_assign_q` (one call). - -## Semantics - -The context-phase (prefill) preprocessing step of MLA attention when a -cached KV prefix exists (KV reuse or chunked prefill), over one prepared -batch. Two configurations are certified, `beam_width=1` and DeepSeek MLA -geometry (`lora_size=512`, `rope_size=64`, `nope_size=128`) in both: - -- **matching-dtype latent pool** (`quant_mode=0`, - `kv_scale_orig_quant=None`): bf16 or fp16 activations with a paged latent - KV cache of the same dtype. The appended row is a plain dtype-preserving - write. -- **fp8-e4m3 latent pool** (`quant_mode=128`): bf16 activations with a - one-byte-per-element e4m3 pool. Only the **appended pool row** is - quantized — `q` and `latent_cache` are still rotated in place as bf16. - -The rest of this section describes the matching-dtype pool; the fp8 -subsection below states exactly which of its three effects changes (only -the third) and which do not. - -The batch has `num_contexts` context-phase sequences first; trailing -generation sequences are ignored entirely. Rows of `q` and `latent_cache` -are the context sequences' **new (uncached) tokens**, concatenated in batch -order. For context sequence `s`, with - -``` -cached_s = cu_ctx_cached_kv_lens[s+1] - cu_ctx_cached_kv_lens[s] # prefix already in cache -kv_s = cu_seq_lens[s+1] - cu_seq_lens[s] # cached + new -new_s = kv_s - cached_s # rows this op consumes -``` - -one call computes, for each new token `i in [0, new_s)` (tensor row `t`, -absolute position `p = cached_s + i`), with `R = rope_size`, -`C = lora_size`, `N = nope_size`, `H = head_num`: - -``` -# GPT-J (interleaved-pair) RoPE, fp32 math, one rounding to the I/O dtype. -# (cos_d, sin_d) = cos_sin_cache row p, pair d, d in [0, R/2) -rope(x)[2d] = x[2d] * cos_d - x[2d+1] * sin_d -rope(x)[2d+1] = x[2d] * sin_d + x[2d+1] * cos_d - -# 1. q RoPE in place (per head h): the rope tail of each head is rotated, -# reading its own pre-rotation contents ("assign q"). -q[t, h*(N+R)+N : (h+1)*(N+R)] = rope(q[t, h*(N+R)+N : (h+1)*(N+R)]) -# 2. k_pe RoPE in place on the latent row: -latent_cache[t, C:] = rope(latent_cache[t, C:]) -# 3. latent-cache append, at page/slot of position p in seq s's blocks: -cache_row(s, p) = concat(latent_cache[t, :C], # bitwise copy - rope(k_pe)) # same rotated values as 2. -``` - -Fusion boundary: q RoPE, k_pe RoPE, and the paged-cache append happen -inside the call. The caller still owns the projections that produced `q` -and `latent_cache`, reading the full `[cached + new]` latent KV back out of -the pool, the `kv_b_proj` up-projection, and the context FMHA itself. -`q`'s nope regions, `latent_cache`'s compressed region, and every -already-cached pool row (positions `< cached_s`) are bitwise untouched. - -### fp8-e4m3 latent pool (`quant_mode=128`) - -`quant_mode=128` is `QuantMode`'s fp8-KV-cache bit. It switches the -**latent pool's** element type to fp8-e4m3, one byte per element, while -`q` and `latent_cache` stay bf16. Certified at the DeepSeek-R1-0528 cell — -`head_num = 128`, `tokens_per_block = 32`, `lora/nope/rope = 512/128/64`, -`beam_width = 1`, single-layer pool — with the KV scaling factor omitted -(the production call) and explicit 1.0 / 1/1.5 / 0.5 / 0.25 / 2.0 swept -beside it, over fresh-prefill (`cached_s = 0`) and cached-prefix -(`cached_s > 0`) context sequences in the same call. - -One fp32 scalar steers it: - -``` -w = kv_scale_orig_quant[0] # 1.0 when that tensor is None -e4m3(x) = round-to-nearest cast of an fp32 x to float8_e4m3fn, - saturating at +-448 (see Notes) -``` - -Effects 1 and 2 above are **unchanged** — `q`'s rope tails and -`latent_cache`'s k_pe are rotated in place and written back as bf16, -un-quantized. Effect 3 becomes, with `bf16(.)` the rounding to the -activation dtype that the in-place write already performs: - -``` -# 3. latent-cache append, quantized, at page/slot of position p in seq s: -cache_row(s, p) = e4m3(concat(latent_cache[t, :C], - bf16(rope(k_pe))) * w) -``` - -`w` is a **per-tensor static scale**, not a block or dynamic one: every -appended row is the same multiply-and-round, bit-exact against that mirror -at each scale tested. The multiply happens in fp32 *after* the rounding to -bf16, so `e4m3(bf16(x) * w)` and `e4m3(bf16(x * w))` are different tensors -for a non-power-of-two `w` and the kernel computes the first (measured, see -Notes). - -**`q` is not quantized here, and neither is `latent_cache`.** Both come -back as ordinary bf16 rotations, matching the same reference the -`quant_mode=0` path matches, and neither is an e4m3 round trip (6.2% of -their elements survive an e4m3 round trip unchanged — the same fraction as -the untouched nope slice of the same tensor, and far from the 100% a -quantized tensor would show). The FMHA that consumes `q` afterwards is a -separate call — `trtllm::attention` with -`attention_input_type = context_only` and `latent_cache=None` — and on the -fp8 MLA context path that op was measured to quantize its own `q`/`k`/`v` -to e4m3 internally. That is that op's fact, established separately and on -its fresh-prefill flavor rather than the cached-KV flavor this op feeds; -what matters here is that nothing this op writes into `q` is pre-quantized, -so the pair is not a double quantization. - -The `w` on the pool row is likewise **not** applied to `q`: two runs over -identical inputs at `w = 1.0` and `w = 0.25` returned bitwise-identical `q` -and `latent_cache` while the appended rows moved (tested). - -`w` is a write-side factor only. Whatever later reads the pool back has to -supply its own read-side factor `1/w`; this op takes no such argument and -checks no relation. The generation-phase sibling -`trtllm::mla_rope_generation`, which appends into the same pool for decode -steps, takes both (`kv_scale_orig_quant` and `kv_scale_quant_orig`) and -uses them independently; its own append was separately measured to be the -same `e4m3(row * w)` static rule, so a pool written by both phases is -uniformly scaled as long as the caller passes the same `w` to each. - -## Signature - -```python -def mla_rope_append_paged_kv_assign_q( - q: torch.Tensor, - latent_cache: torch.Tensor, - num_contexts: int, - cu_ctx_cached_kv_lens: torch.Tensor, - cu_seq_lens: torch.Tensor, - max_input_uncached_seq_len: int, - cos_sin_cache: torch.Tensor, - head_num: int, - nope_size: int, - rope_size: int, - lora_size: int, - kv_cache_block_offsets: torch.Tensor, - host_kv_cache_pool_pointers: torch.Tensor, - host_kv_cache_pool_mapping: torch.Tensor, - kv_scale_orig_quant: Optional[torch.Tensor], - layer_idx: int, - tokens_per_block: int, - attention_window_size: int, - beam_width: int, - quant_mode: int, -) -> None -``` - -With `nc` = `num_contexts`, `S` = total sequences in the batch (contexts -first), `T` = `sum(new_s)` over context sequences, and `C`/`R`/`N`/`H` as -above: - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `q` | `[T, H*(N+R)]`; per head only the trailing `R` written, in the activation dtype on both paths (never quantized) | bf16 / fp16 (bf16 only on the fp8 path) | 2-D enforced (other ranks rejected), contiguous | CUDA | -| `latent_cache` | `[T, C+R]` = per-token `[compressed_kv \| k_pe]`; only `[..., C:]` written, in the activation dtype on both paths | same as `q` | 2-D enforced, contiguous | CUDA | -| `num_contexts` | leading context-phase sequence count | Python int | — | — | -| `cu_ctx_cached_kv_lens` | `[>= nc+1]`; `[0, cumsum(cached_s)]`; only the first `nc+1` entries read | int64 | contiguous | CUDA | -| `cu_seq_lens` | `[>= nc+1]`; `[0, cumsum(kv_s)]`; only the first `nc+1` entries read | int64 | contiguous | CUDA | -| `max_input_uncached_seq_len` | `max(new_s)` over context seqs (grid bound; production passes the exact max) | Python int | — | — | -| `cos_sin_cache` | `[1, max_pos * R * 2]`: per position `R` fp32 `(cos, sin)` pairs, second `R/2` a duplicate of the first (duplicated/`duplicate_data=True` layout; only pairs `[0, R/2)` are read per position) | fp32 | contiguous | CUDA | -| `head_num` | query heads `H`; 16 (matching-dtype pool) and 128 (fp8 pool) certified | Python int | — | — | -| `nope_size` | `N` (128 certified) | Python int | — | — | -| `rope_size` | `R` (64 certified) | Python int | — | — | -| `lora_size` | `C` (512 certified) | Python int | — | — | -| `kv_cache_block_offsets` | `[1, >= S, 2, max_blocks_per_seq]`; raw block ids per seq (K and V rows identical, kv_factor=1 pool); contexts first | int32 | contiguous | CUDA | -| `host_kv_cache_pool_pointers` | `[num_pools, 2]`: (primary ptr, secondary ptr=0) | int64 | contiguous | CPU | -| `host_kv_cache_pool_mapping` | `[num_layers, 2]`: (pool index, layer-within-pool) row per layer | int32 | contiguous | CPU | -| `kv_scale_orig_quant` | `quant_mode=0`: `None`. `quant_mode=128`: `[1]` holding the **write-side** factor `w` — everything this call quantizes is multiplied by it. `None` = 1.0, which is what the engine's own MLA call site passes. Only element `[0]` is read (a longer tensor is accepted) | fp32 (fp16 rejected, see *Notes*) | contiguous | CUDA | -| `residual_dim` | — | `int` | Must be `0` or `rope_size`; the op rejects non-zero unless the KV pool is FP4. `0` on every path this entry certifies (bf16 and fp8-e4m3 pools), which is also what the in-tree caller passes | — | -| `layer_idx` | row into `host_kv_cache_pool_mapping` | Python int | — | — | -| `tokens_per_block` | pool page size; 32 and 64 certified on the matching-dtype pool, 32 on the fp8 pool (32 is what a default `KvCacheConfig` produces) | Python int | — | — | -| `attention_window_size` | `>= max(kv_s)`; production passes the manager's `max_seq_len` (smaller values imply cyclic-cache addressing, not certified) | Python int | — | — | -| `beam_width` | `1` | Python int | — | — | -| `quant_mode` | `0` (matching-dtype pool) and `128` = `QuantMode.FP8_KV_CACHE` (fp8-e4m3 pool) certified. A **non-fp8 quantized** KV cache is rejected — see *Notes*. `128 \| 256` (an fp8-weights bit alongside) was accepted rather than rejected; its output was not compared | Python int | — | — | -| returns | — (mutates `q`, `latent_cache`, and the paged pool) | — | — | — | - -## Metadata consumed - -No thread-local or registered-layer state: all inputs are explicit -arguments. The length/addressing tensors are exactly what a -`TrtllmAttentionMetadata` prepared with -`enable_context_mla_with_cached_kv=True` and its `KVCacheManager` expose: -`ctx_cached_token_indptr` → `cu_ctx_cached_kv_lens`, `ctx_kv_indptr` → -`cu_seq_lens`, `max_ctx_seq_len` → `max_input_uncached_seq_len`, -`metadata.kv_cache_block_offsets`, `manager.kv_cache_pool_pointers`, -`manager.kv_cache_pool_mapping`, `manager.tokens_per_block`, -`manager.max_seq_len` → `attention_window_size`. - -## Preconditions - -- The paged pool addressed by `host_kv_cache_pool_pointers` + - `host_kv_cache_pool_mapping[layer_idx]` is an MLA latent cache: - kv_factor 1, one kv head, row width `C + R`, page size - `tokens_per_block`. Per page the layout is `[tokens_per_block, C + R]` - and the page slab is `tokens_per_block * (C + R)` **elements**, whose - byte width the op takes from `quant_mode` and the activation dtype, not - from anything the cache manager tells it: the activation dtype's width at - `quant_mode=0` (2 bytes for the certified bf16/fp16), 1 byte (e4m3) at - `quant_mode=128`. The pool's real element type must match — a - `KVCacheManager` built with `DataType.BF16` / `DataType.HALF` and with - `DataType.FP8` respectively is what was certified. -- Every context sequence has block capacity for `kv_s` tokens **before** - the call (block offsets in `kv_cache_block_offsets` cover - `ceil(kv_s / tokens_per_block)` pages). -- `q` and `latent_cache` are 2-D, contiguous, of the same dtype, with - exactly `T` rows: the context sequences' new tokens in batch order - (production slices `q[:num_ctx_tokens]`; generation-token rows must not - be included). -- `q`'s rope tails and `latent_cache`'s k_pe hold **pre-rotation** values; - the call rotates them in place. A table with cos=1/sin=0 makes the RoPE - identity (pre-rotated inputs pass through). -- `cu_ctx_cached_kv_lens` / `cu_seq_lens` are int64 (the kernel rejects - other index dtypes), on device, start at 0, and satisfy - `cached_s <= kv_s` per sequence. -- `max_input_uncached_seq_len >= new_s` for every context sequence - (certified with the exact max), and `kv_s <= attention_window_size`, - `kv_s - 1 < max_pos` of the table. -- `cos_sin_cache` uses the duplicated `(cos, sin)`-pair layout above - (`RopeParams(..., duplicate_data=True).create_rope_const_params()` - produces it; MLA models set `duplicate_data=True`). -- **Matching-dtype pool only:** `quant_mode=0` and - `kv_scale_orig_quant=None`. -- **fp8 pool only** (`quant_mode=128`): - - `q` and `latent_cache` are bf16 (fp16 activations over an fp8 pool were - not run). - - `kv_scale_orig_quant` is a `[1]` fp32 CUDA tensor or `None` (= 1.0). - `None` is the production configuration and is what the engine's own MLA - call site passes. There is no read-side factor here: a consumer that - dequantizes the pool needs `1/w` from somewhere else, and nothing in - this call checks it. - - Input magnitudes must survive e4m3: `|value * w|` below `2**-10` - (~9.8e-4) rounds to zero, and anything above 448 saturates to `±448` - (measured — see *Notes*). The certified runs are unit-normal - activations at `w <= 2`. - -## Notes - -- Schema under-annotation: the registered schema marks **nothing** mutable — - it declares `(Tensor q, Tensor latent_cache, ...) -> ()`, with no - `Tensor(a!)` anywhere — yet the call rotates `q` and `latent_cache` in - place and appends to the paged pool. Do not rely on the schema's alias - info (e.g. under torch.compile functionalization). -- Precision of the in-place rotations (bf16 and fp16, sm_100, **both** - paths — the fp8 path rotates `q` and `latent_cache` exactly as the - matching-dtype path does): the `compressed_kv` copy is bitwise. Both RoPE - outputs match an fp32 reference rounded once to the I/O dtype and are - bitwise on all but 1.4e-5 to 1.7e-5 of elements: the kernel and a torch - fp32 reference evaluate `x*cos ∓ y*sin` in different orders, so where the - correctly-rounded fp32 result lands near a rounding midpoint of the I/O - dtype the two round to adjacent values. Measured in bf16 at `H = 16` on a - ~400-token context call: 7 such elements in 411 648 rotated `q` values at - page size 32 and 8 at page size 64, all within one bf16 ulp (0 in 25 728 - rotated `k_pe` values either way) — it is arithmetic, not addressing. At - `H = 128` on the same batch shape the tail reaches further: 46 elements in - 3 293 184 differ, one of them by 4 bf16 ulps, but the **absolute** - deviation stays ≤ 0.0078 because the elements that drift are the ones - where `x*cos - y*sin` nearly cancels. Default `assert_close` tolerances - absorb all of it with margin, which is what the test gates on; a caller - needing bit-exactness against a torch reference will not get it. -- Precision of the fp8 append (sm_100): **everything came back bit-exact** - against `e4m3(row * w)`, the `compressed_kv` half and the roped `k_pe` - half alike, at `w` ∈ {1.0 (both as `None` and as an explicit tensor), - 1/1.5, 0.5, 0.25, 2.0}. e4m3's 3-bit mantissa absorbs the evaluation-order - difference the bf16 rotation pays an ulp for. The test still gates the - roped half at one e4m3 ulp (`2**-3` relative) plus a bit-exact-fraction - floor rather than bitwise, because the tie it covers is arithmetic; - neither half has been approached. -- **The double rounding is real and a mirror must reproduce it.** The - quantized `k_pe` is `e4m3(bf16(rope(k_pe)) * w)` — and that bf16 value is - the one the call also writes back into `latent_cache` (both came back - bit-exact against the same torch reference), so a caller already holding - the post-call `latent_cache` can mirror the pool from it directly. - Quantizing the unrounded fp32 RoPE result instead disagrees with the - kernel on 2.4-3.3% of the appended `k_pe` bytes (measured at every scale), - which is enough to fail a bit-exact gate and small enough to pass a sloppy - one. Applying the scale before the bf16 rounding — - `e4m3(bf16(rope * w))` — disagrees on 1.05% of bytes at `w = 1/1.5` and is - identical at the power-of-two scales, so only a non-power-of-two scale - separates the two orders. -- A `quant_mode` asking for a **non-fp8 quantized** KV cache is rejected: - `quant_mode = 64` (`QuantMode.INT8_KV_CACHE`) raises - `RuntimeError: [TensorRT-LLM][ERROR] Assertion failed: Only FP8 KV cache - is supported for now (../tensorrt_llm/thop/mlaPreprocessOp.cpp:316)`. - Tested. `quant_mode = 8192` (`NVFP4_KV_CACHE`) also fails, but earlier and - for a different reason — `Expected hostKvCachePoolPointers.dim() == 3`, - the nvfp4 pool-pointer layout — so it is not the same guard. -- e4m3 range behaviour at the boundaries (scratch probe, `w = 1.0`): the - kernel **saturates**, writing `+448` for an input of 1000 or `+inf` and - `-448` for -1000, and flushes 1e-5 to `+0`. `torch.Tensor.to(float8_e4m3fn)` - does *not* saturate — it yields the NaN byte `0x7F` for 1000 and for - `inf` — so a torch mirror and the kernel agree only while every - `|value * w|` stays ≤ 448. Every certified run does. -- `kv_scale_orig_quant` validation is thin: a `[2]` fp32 CUDA tensor is - accepted and only `[0]` is used, and a **CPU** fp32 tensor is accepted - too — in one scratch probe on this machine it even produced the same bytes - as the equivalent CUDA tensor, so device placement is not checked and a - wrong-device scale will not announce itself. A `float16` tensor is - rejected loudly (`RuntimeError: expected scalar type Float but found - Half`). -- **Not idempotent**: the q and k_pe rotations read their own pre-call - contents in place, so calling twice on the same buffers composes two - rotations (unlike ops that read a separate source). Re-preparing inputs - is required before any retry. -- `cached_s = 0` (fresh prefill, nothing cached) is valid and certified on - both paths; positions then start at 0. `cached_s > 0` — a context call - over a prefix already in the pool, which is what block reuse makes the - engine do — is certified on both paths too, and the prefix rows come back - bitwise unchanged. -- Pool writes are exactly bounded on both paths: one `C+R`-element row per - new context token, at the `(page, slot)` its absolute position addresses, - and nothing else in the pool moves — checked by a whole-pool byte snapshot - on every case, including a 402-row single call spanning four sequences and - a mixed batch whose trailing generation sequences own their own pages. - That is also what pins the page-slab geometry described under - *Preconditions*. -- Page size (sm_100): certified at `tokens_per_block` 32 and 64 on the - matching-dtype pool, 32 on the fp8 pool. 32 is what - a `KvCacheConfig` at its defaults hands the op; 64 is what a tuned MLA - target opts into. At 32 the cached prefixes were chosen to put the first - new token at every alignment a 32-slot page has — 0 (fresh prefill), - 32 and 96 (slot 0 of a fresh page, after one and three full pages), - 1 / 33 / 100 (mid-page), 31 (a page's last slot, so the run crosses on its - very first token), 512 (sixteen full pages) and 511 (the last slot of page - 15, which the single new token closes) — with new-token runs of 1 to 300 - tokens, the longest walking ten pages in one call, in bf16 (including - through pool-mapping row 1 of a two-layer pool) and in fp16, plus a mixed - batch whose trailing generation sequences own their own pages. Every effect - held at the same gates as at page 64: RoPE within the one-ulp band above, - `compressed_kv` copy bitwise, cached prefix and the untouched slices - bitwise. The fp8 cases run page 32 only, over the same alignment set — - first new token at 0, 17, 31, 32, 96, 100 and 511, runs of 1 to 300 tokens, - and a 402-row four-sequence call. Other page sizes are untested rather than - known-bad. -- Both `q.dim() == 2` and `latent_cache.dim() == 2` are enforced with - runtime errors. A comment in the runtime describes `latent_cache` as - `[tokens, 1, C+R]`, but the op itself rejects 3-D — callers flatten to - rows first. -- Only the DeepSeek-V3 geometry (`C=512`, `R=64`, `N=128`) was exercised; - other sizes are unverified. -- Sibling ops exist for adjacent MLA roles: - `trtllm::load_paged_kv_cache_for_mla` (gathers the full `[cached + new]` - latent KV this op just completed, and takes its own read-side - `kv_scale_quant_orig`), `trtllm::load_chunked_kv_cache_for_mla` - and `trtllm::merge_chunked_attention_for_mla` (chunked-prefill variants), - `trtllm::mla_rope_generation` (the generation-phase counterpart), and the - context FMHA that consumes the up-projected K/V. -- Not certified on either path: `beam_width > 1`, - `attention_window_size` below the KV length (cyclic-cache addressing), and - MTP / speculative-decoding batches. Not certified on the fp8 path - specifically: - `tokens_per_block = 64`, head counts other than 128, fp16 activations, - multi-layer latent pools, `|value * w| > 448`, and a YaRN-scaled rope - table (the fp8 cases run the unscaled theta-10000 table; the table's - *content* is not an axis this op's gate exercises). - - -## Accepted KV-cache formats - -The op takes an fp8-e4m3 latent pool and an NVFP4 one; an int8 pool -(`quant_mode=64`) is rejected with `Only FP8 and NVFP4 KV caches are -supported for now`. **NVFP4 is accepted by the op but not certified here** — -no cell in this entry's test drives it, so it stays outside the envelope. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_append_paged_kv_assign_q.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_append_paged_kv_assign_q.py deleted file mode 100644 index e7f39b6182bd..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_append_paged_kv_assign_q.py +++ /dev/null @@ -1,63 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""MLA context-phase RoPE of q_pe/k_pe in place + latent paged-KV-cache append.""" - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* - - -def mla_rope_append_paged_kv_assign_q( - q: torch.Tensor, - latent_cache: torch.Tensor, - num_contexts: int, - cu_ctx_cached_kv_lens: torch.Tensor, - cu_seq_lens: torch.Tensor, - max_input_uncached_seq_len: int, - cos_sin_cache: torch.Tensor, - head_num: int, - nope_size: int, - rope_size: int, - lora_size: int, - kv_cache_block_offsets: torch.Tensor, - host_kv_cache_pool_pointers: torch.Tensor, - host_kv_cache_pool_mapping: torch.Tensor, - kv_scale_orig_quant: Optional[torch.Tensor], - # ``residual_dim`` must be 0 or ``rope_size``, - # and the op rejects non-zero unless the KV pool is FP4. Every caller - # here runs a bf16 or fp8-e4m3 pool, so 0 is the only legal value. - residual_dim: int, - layer_idx: int, - tokens_per_block: int, - attention_window_size: int, - beam_width: int, - quant_mode: int, -) -> None: - """RoPE each new context token's q_pe (in q) and k_pe (in latent_cache) - in place at its absolute position, and append the token's latent row - [compressed_kv | rope(k_pe)] to the paged MLA KV cache. Returns None.""" - torch.ops.trtllm.mla_rope_append_paged_kv_assign_q( - q, - latent_cache, - num_contexts, - cu_ctx_cached_kv_lens, - cu_seq_lens, - max_input_uncached_seq_len, - cos_sin_cache, - head_num, - nope_size, - rope_size, - lora_size, - kv_cache_block_offsets, - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping, - kv_scale_orig_quant, - residual_dim, - layer_idx, - tokens_per_block, - attention_window_size, - beam_width, - quant_mode, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_generation.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_generation.md deleted file mode 100644 index a28e160aa6c0..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_generation.md +++ /dev/null @@ -1,425 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 15} ---- - -# mla_rope_generation - -**Wraps** `torch.ops.trtllm.mla_rope_generation` (one call). - -## Semantics - -The generation-phase (decode) preprocessing step of MLA absorbed attention, -over one prepared batch. Two configurations are certified, and they differ -in what the call writes — read the one you are on: - -- **bf16 latent pool** (`quant_mode=0`, every fp8/scale tensor `None`): the - roped q lands in `fused_q`'s tail and the appended cache row is bf16. -- **fp8-e4m3 latent pool** (`quant_mode=128`): `fused_q` is **not written at - all**; the roped q lands quantized in `quant_q_buffer`, the appended cache - row lands quantized, and the two decode-FMHA scale buffers are filled. - -Both are certified with bf16 activations (`fused_q`, `q_pe`, -`latent_cache`), `rope_append=True`, `beam_width=1`, -`block_ids_per_seq=None`, `out_scale=None` and no helix parallelism, at -`predicted_tokens_per_seq` (`P` below) 1, 2, 3 and 4. - -The batch has `num_contexts` context-phase sequences first, then `G` -generation-phase sequences. The op processes **only the generation -sequences**. Each of them contributes `P` new query tokens — `P = 1` is an -ordinary decode step, `P > 1` the MTP shape, where `P = max_draft_len + 1`. -`fused_q`/`q_pe`/`latent_cache`/`quant_q_buffer` therefore have `G * P` -rows, **token-major within a sequence**: row `n` is generation sequence -`g = n // P`'s `t = n % P`-th new token. Per-sequence tensors -(`sequence_length`, `host_past_key_value_lengths`, `host_context_lengths`, -`kv_cache_block_offsets`) cover the whole batch in order and stay one entry -per sequence; the op indexes them at `num_contexts + g`. - -**Every row gets its own rope angle.** With -`L_g = sequence_length[num_contexts + g]` the total KV length of generation -sequence `g` *including all `P` of its new tokens*, row `(g, t)` sits at -0-based absolute position - -``` -pos(g, t) = L_g - P + t # t = 0 .. P-1, consecutive, ending at L_g - 1 -``` - -so the `P` tokens of one sequence occupy `P` consecutive positions and each -is rotated at its own. At `P = 1` this is the familiar `L_g - 1`. The -position comes from the **device** `sequence_length` tensor: running the -same prepared step with `host_past_key_value_lengths` perturbed to -`sequence_length - 1` left every position, every pool slot and -`cu_kv_seqlens` unchanged (observed on the fp8 path at `P = 3`). - -**On the bf16 path**, one call then computes, with `R = qk_rope_head_dim` -and `C = kv_lora_rank` (the fp8 path replaces effects 1 and 2 — see below): - -``` -# GPT-J (interleaved-pair) RoPE, fp32 math, one rounding to bf16. -# (cos_d, sin_d) = rotary_cos_sin row pos(g,t), pair d, d in [0, R/2) -rope(x, p)[2d] = x[2d] * cos_d - x[2d+1] * sin_d -rope(x, p)[2d+1] = x[2d] * sin_d + x[2d+1] * cos_d - -# for every row n = g*P + t: -# 1. q RoPE into the fused-q tail (per head h): -fused_q[n, h, C:] = rope(q_pe[n, h, :], pos(g,t)) -# 2. latent-cache append, at page/slot of position pos(g,t) in seq g's blocks: -cache_row(g, pos(g,t)) = concat(latent_cache[n, :C], # bitwise copy - rope(latent_cache[n, C:], pos(g,t))) # k_pe rotated -# 3. decode-FMHA scheduler buffers (trtllm-gen MQA layout): -cu_q_seqlens[0:G+1] = arange(G+1) * num_heads * P -cu_kv_seqlens[0:G+1] = [0, cumsum(L_g over generation seqs)] -fmha_scheduler_counter[0] = 0 -``` - -Fusion boundary: q RoPE, k_pe RoPE, paged-cache append and scheduler-buffer -fill happen inside the call. The caller still owns the up/down projections, -the absorbed-q BMM that fills `fused_q[..., :C]`, and the subsequent -latent-space attention. `q_pe` and `latent_cache` are inputs only: despite -the mutable schema annotation on `q_pe` neither is modified (observed). - -On the bf16 path `fused_q[..., :C]` is bitwise untouched and never read, so -the absorbed-q BMM and this call may run concurrently. **On the fp8 path -they may not** — see below. - -### fp8-e4m3 latent pool (`quant_mode=128`) - -`quant_mode=128` is `QuantMode`'s fp8-KV-cache bit. It switches the -**latent pool's** element type to fp8-e4m3, one byte per element, while -`fused_q`, `q_pe` and `latent_cache` all stay bf16. Certified at the -DeepSeek-R1-0528 cell — `H = 128`, `tokens_per_block = 32`, -`C/R/nope/v = 512/64/128/128`, `q_lora_rank = 1536`, single-layer pool — at -`P` 1, 2, 3 and 4, with the KV scaling factor omitted (the production call) -and explicit 1.0 / 1.5 / 2.0 swept beside it, at `q_scaling` 1.0 and -DeepSeek-R1's YaRN value `1/mscale^2 = 0.5336594`. - -Two independent fp32 scalars steer it, each with its own role and neither -derived from the other: - -``` -w = kv_scale_orig_quant[0] # 1.0 when that tensor is None -r = kv_scale_quant_orig[0] # 1.0 when that tensor is None -e4m3(x) = round-to-nearest cast of an fp32 x to float8_e4m3fn -``` - -One call then computes, with `rope` as above and `bf16(.)` the rounding to -the activation dtype that happens before quantization, for every row -`n = g*P + t` at `p = pos(g,t)`: - -``` -# 1. q RoPE, quantized. fused_q is READ, never written: -quant_q_buffer[n, h, :C] = e4m3(fused_q[n, h, :C] * w) -quant_q_buffer[n, h, C:] = e4m3(bf16(rope(q_pe[n, h, :], p)) * w) -# 2. latent-cache append, quantized, at page/slot of position p: -cache_row(g, p) = e4m3(concat(latent_cache[n, :C], - bf16(rope(latent_cache[n, C:], p))) * w) -# 3. decode-FMHA scales (a per-batch scalar pair — not per sequence, and -# P does not enter them): -x = r*r / (q_scaling * sqrt(qk_nope_head_dim + qk_rope_head_dim)) -mla_bmm1_scale = [x, x * log2(e)] -mla_bmm2_scale = [r] -# 4. scheduler buffers: exactly as on the bf16 path -``` - -`w` is a **per-tensor static scale**, not a block or dynamic one: every -appended row and every quantized q element is the same multiply-and-round, -bit-exact against that mirror at each scale tested. - -Fusion boundary on this path: the same RoPE / append / scheduler fill, plus -the q quantization and the two scale derivations. The caller still owns the -absorbed-q BMM — and now must have **finished** it before this call, since -`fused_q[..., :C]` is read here to build `quant_q_buffer`'s head. (The -engine's own MLA module enforces exactly that by dropping the auxiliary -stream when the KV cache is fp8.) - -**Why `r` appears in both scales.** The fp8 MLA generation FMHA — the -sibling op `trtllm::attention` called with -`attention_input_type = generation_only`, `is_mla_enable=True` — takes its -query from `quant_q_buffer` rather than from the bf16 fused q, takes its -BMM1 scale from `mla_bmm1_scale[1]` (the log2-domain copy; `[0]` is inert -there) and its BMM2 scale from `mla_bmm2_scale[0]`, and reads neither KV -scale tensor. Dequantization is therefore the producer's job: `r*r` undoes -the `w` on the query and on `K`, and `r` undoes the `w` on `V = K[:, :C]`, -leaving the plain softmax scale `1 / (q_scaling * sqrt(nope + R))`. That -cancellation is only correct when `r = 1/w`; this op does not check it, and -a non-reciprocal pair produces a silently mis-scaled decode (pinned here by -running `w = 0.25` against `r = 3.0` and observing each drive its own half). -The consuming call's behaviour above is that op's fact, established -separately — this entry's test gates the values written, not their use. - -## Signature - -```python -def mla_rope_generation( - fused_q: torch.Tensor, - q_pe: torch.Tensor, - latent_cache: torch.Tensor, - rotary_cos_sin: Optional[torch.Tensor], - cu_q_seqlens: torch.Tensor, - cu_kv_seqlens: torch.Tensor, - fmha_scheduler_counter: torch.Tensor, - mla_bmm1_scale: Optional[torch.Tensor], - mla_bmm2_scale: Optional[torch.Tensor], - quant_q_buffer: Optional[torch.Tensor], - sequence_length: torch.Tensor, - host_past_key_value_lengths: torch.Tensor, - host_context_lengths: torch.Tensor, - num_contexts: int, - kv_cache_block_offsets: Optional[torch.Tensor], - host_kv_cache_pool_pointers: Optional[torch.Tensor], - host_kv_cache_pool_mapping: Optional[torch.Tensor], - kv_scale_orig_quant: Optional[torch.Tensor], - kv_scale_quant_orig: Optional[torch.Tensor], - out_scale: Optional[torch.Tensor], - block_ids_per_seq: Optional[torch.Tensor], - helix_tensor_params: List[Optional[torch.Tensor]], - predicted_tokens_per_seq: int, - layer_idx: int, - num_heads: int, - num_kv_heads: int, - head_size: int, - tokens_per_block: int, - attention_window_size: int, - beam_width: int, - quant_mode: int, - q_scaling: float, - q_lora_rank: int, - kv_lora_rank: int, - qk_nope_head_dim: int, - qk_rope_head_dim: int, - v_head_dim: int, - rope_append: bool, -) -> None -``` - -With `G` = generation sequences, `P` = `predicted_tokens_per_seq`, -`S` = total sequences (`num_contexts + G`), `H` = `num_heads`, -`C` = `kv_lora_rank`, `R` = `qk_rope_head_dim`, `D = C + R` -(= `head_size`). The four per-token tensors have `G * P` rows, ordered -token-major within a sequence (row `n = g*P + t`); everything else is one -entry per sequence and does **not** scale with `P`: - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `fused_q` | `[G*P, H, D]`; bf16 pool: only `[..., C:]` written, `[..., :C]` untouched and unread. fp8 pool: **entirely read-only**, and `[..., :C]` must already hold the absorbed q | bf16 | contiguous | CUDA | -| `q_pe` | `[G*P, H, R]`; read only | bf16 | last dim contiguous; strided head-dim views (e.g. a slice of packed `[G*P, H, nope+R]` q) work | CUDA | -| `latent_cache` | `[G*P, D]` = per-token `[compressed_kv \| k_pe]`; read only | bf16 | contiguous | CUDA | -| `rotary_cos_sin` | `[1, max_pos * R * 2]`: per position `R` fp32 `(cos, sin)` pairs, second `R/2` a duplicate of the first (duplicated/`duplicate_data=True` layout; only pairs `[0, R/2)` are read per position) | fp32 | contiguous | CUDA | -| `cu_q_seqlens` | `[>= G+1]`; exactly `G+1` entries written whatever `P` is (entries past `G+1` are left alone — observed) | int32 | contiguous | CUDA | -| `cu_kv_seqlens` | `[>= G+1]`; same, `G+1` entries | int32 | contiguous | CUDA | -| `fmha_scheduler_counter` | `[1]`; written (zeroed) | uint32 | — | CUDA | -| `mla_bmm1_scale` | bf16 pool: `None`. fp8 pool: `[2]`, written with `[x, x*log2(e)]`; may be `None`, and is then simply not written | fp32 | contiguous | CUDA | -| `mla_bmm2_scale` | bf16 pool: `None`. fp8 pool: `[1]`, written with `[kv_scale_quant_orig]`; may be `None`, and is then simply not written | fp32 | contiguous | CUDA | -| `quant_q_buffer` | bf16 pool: `None`. fp8 pool: `[G*P, H, D]`, **required** — one e4m3 byte per element, fully overwritten | uint8 (a `float8_e4m3fn` view of the same bytes behaves identically) | contiguous | CUDA | -| `sequence_length` | `[S]`; per-seq total KV length incl. **all `P`** of this step's tokens. This is the tensor the `P` positions and `cu_kv_seqlens` are derived from | int32 | contiguous | CUDA | -| `host_past_key_value_lengths` | `[S]`; pass the same values as `sequence_length` (total KV length, **not** past-only, despite the name). Perturbing it to `sequence_length - 1` changed no observable effect at the certified fp8 cell, so it is `sequence_length` that steers this op; larger divergences were not measured | int32 | contiguous | CPU | -| `host_context_lengths` | `[S]`; per-seq prompt length | int32 | contiguous | CPU | -| `num_contexts` | leading context-phase sequence count | Python int | — | — | -| `kv_cache_block_offsets` | `[1, >= S, 2, max_blocks_per_seq]`; raw block ids per seq (K and V rows identical, kv_factor=1 pool) | int32 | contiguous | CUDA | -| `host_kv_cache_pool_pointers` | `[num_pools, 2]`: (primary ptr, secondary ptr=0) | int64 | contiguous | CPU | -| `host_kv_cache_pool_mapping` | `[num_layers, 2]`: (pool index, layer-within-pool) row per layer | int32 | contiguous | CPU | -| `kv_scale_orig_quant` | `[1]`, the **write-side** factor `w`: everything quantized by this call is multiplied by it. `None` = 1.0 (what the engine's own MLA call site passes). Ignored on the bf16 path | fp32 | contiguous | CUDA | -| `kv_scale_quant_orig` | `[1]`, the **read-side** factor `r`: folded into `mla_bmm1_scale` (as `r*r`) and copied into `mla_bmm2_scale`. `None` = 1.0. Not used for anything else, and not derived from `kv_scale_orig_quant`. Ignored on the bf16 path | fp32 | contiguous | CUDA | -| `out_scale` | `None`; fp8-output scale otherwise (not certified) | — | — | — | -| `block_ids_per_seq` | `None`; flash-MLA layout otherwise (not certified) | — | — | — | -| `helix_tensor_params` | `[None, None]`; helix position offsets / inactive-rank mask otherwise (not certified) | Python list | — | — | -| `predicted_tokens_per_seq` | `P`, the query tokens each generation sequence carries. `1` (ordinary decode) and `2`, `3`, `4` (the MTP path, `max_draft_len + 1`) certified. It is a single scalar for the whole batch — a ragged per-sequence draft length cannot be expressed here. Values above 4 untested | Python int | — | — | -| `layer_idx` | row into `host_kv_cache_pool_mapping` | Python int | — | — | -| `num_heads` | query heads `H`; 16 and 128 certified | Python int | — | — | -| `num_kv_heads` | `1` (MLA latent cache) | Python int | — | — | -| `head_size` | `D = kv_lora_rank + qk_rope_head_dim` | Python int | — | — | -| `tokens_per_block` | pool page size; 32 and 64 certified on the bf16 path, 32 on the fp8 path (32 is what a default `KvCacheConfig` produces) | Python int | — | — | -| `attention_window_size` | `>= max total KV length`; smaller values imply cyclic-cache addressing (not certified) | Python int | — | — | -| `beam_width` | `1` | Python int | — | — | -| `quant_mode` | `0` (bf16 latent pool) and `128` (fp8-e4m3 latent pool) certified | Python int | — | — | -| `q_scaling` | softmax-scale divisor baked into `mla_bmm1_scale`; live on the fp8 path (1.0 and 0.5336594 certified), inert on the bf16 path, which writes no scale at all (bitwise-identical outputs at 1.0 and 7.5, observed) | Python float | — | — | -| `q_lora_rank` | unused by both certified paths (`0` and `1536` both fine) | Python int | — | — | -| `kv_lora_rank` | `C` | Python int | — | — | -| `qk_nope_head_dim`, `qk_rope_head_dim` | model MLA dims (`R` must match `q_pe`/table; their sum is the `sqrt` in `mla_bmm1_scale`) | Python int | — | — | -| `v_head_dim` | model v head dim; not read on either certified path | Python int | — | — | -| `rope_append` | `True` certified (rotate k_pe + append to cache) | Python bool | — | — | -| returns | — (mutates `cu_q_seqlens`, `cu_kv_seqlens`, `fmha_scheduler_counter`, the paged pool, plus `fused_q`'s tail on the bf16 path or `quant_q_buffer`/`mla_bmm1_scale`/`mla_bmm2_scale` on the fp8 path) | — | — | — | - -## Metadata consumed - -No thread-local or registered-layer state: all inputs are explicit -arguments. The length/addressing tensors are exactly what a prepared -`TrtllmAttentionMetadata` and its `KVCacheManager` expose: -`kv_lens_cuda_runtime` → `sequence_length`, `kv_lens_runtime` → -`host_past_key_value_lengths`, `prompt_lens_cpu_runtime` → -`host_context_lengths`, `metadata.kv_cache_block_offsets`, -`manager.kv_cache_pool_pointers`, `manager.kv_cache_pool_mapping`. - -## Preconditions - -- The paged pool addressed by `host_kv_cache_pool_pointers` + - `host_kv_cache_pool_mapping[layer_idx]` is an MLA latent cache: - kv_factor 1, one kv head, row width `D`, page size `tokens_per_block`. - Per page the layout is `[tokens_per_block, D]` and the page slab is - `tokens_per_block * D` **elements**, whose byte width the op derives from - `quant_mode` alone: 2 bytes (bf16) at `quant_mode=0`, 1 byte (e4m3) at - `quant_mode=128`. The pool's real element type must match — a - `KVCacheManager` built with `DataType.BF16` and `DataType.FP8` - respectively is what was certified. -- Every generation sequence has block capacity for `L_g` tokens **before** - the call (block offsets in `kv_cache_block_offsets` cover - `ceil(L_g / tokens_per_block)` pages); with a `KVCacheManager`, call - `impl.add_token(request_id)` **`P` times** per sequence per step, then - prepare the metadata. The `P` rows of a sequence may straddle a page - boundary; the op follows the block-offset row across it. -- `sequence_length` / `host_past_key_value_lengths` hold `L_g` (past + all - `P` new tokens), and `L_g <= attention_window_size`, - `L_g - 1 < max_pos` of the table, `L_g >= P`. -- Batch order: context sequences first; `fused_q`/`q_pe`/`latent_cache` - contain generation tokens only, `G * P` rows in generation-sequence order - with a sequence's `P` tokens contiguous and in position order. -- The `G * P` cache rows written by one call must address **distinct - physical slots**. This holds automatically for a `KVCacheManager`-allocated - batch (each row goes to its own position, distinct sequences hold distinct - pages), and it is the reason `P > 1` needs no extra care from the caller. If - a caller-built block-offset table aliases them onto one slot, the writes - race and leave torn cache rows, nondeterministically and with no error — - reproduced here on purpose as a control, at `P = 2` with 8 sequences - aliased onto one page: 6 of 6 identical armed calls left different pool - images, while the un-aliased geometry repeated bitwise in the same - process. -- `rotary_cos_sin` uses the duplicated `(cos, sin)`-pair layout above - (`RopeParams(..., duplicate_data=True).create_rope_const_params()` - produces it; MLA models set `duplicate_data=True`). A table with - cos=1/sin=0 makes both RoPEs identity (pre-rotated inputs pass through). -- `cu_q_seqlens`/`cu_kv_seqlens`/`fmha_scheduler_counter` are - caller-allocated with the dtypes above; contents need not be initialized. -- **bf16 path only:** `quant_mode=0` and every optional fp8/scale tensor - `None`. -- **fp8 path only** (`quant_mode=128`): - - `quant_q_buffer` must be allocated `[G*P, H, D]` at one byte per element. - It is **not** presence-checked: passing `None` is not rejected, the - kernel launches anyway and the run dies with `CUDA error: an illegal - memory access was encountered` (observed in a scratch probe at this - configuration; the CUDA context is lost, the process aborts, and the KV - manager's own `release_pools` fails on the way out). Its previous - contents are irrelevant — the call overwrites every byte. - - `fused_q[..., :C]` must already hold the absorbed q **when the call is - issued**: this path reads it. Overlapping the absorbed-q BMM with this - call on a second stream is a race here, unlike on the bf16 path. - - `mla_bmm1_scale` `[2]` fp32 and `mla_bmm2_scale` `[1]` fp32 may be - `None` — accepted, and the buffer is then simply not written (observed; - the rest of the outputs are unaffected). A decode consuming those scales - still needs them, so in practice both are allocated. - - `kv_scale_orig_quant` and `kv_scale_quant_orig` are `[1]` fp32 CUDA - tensors or `None` (= 1.0). They are used independently; pass reciprocals - (`kv_scale_orig_quant = 1/s`, `kv_scale_quant_orig = s`) unless a - deliberately asymmetric round trip is what you want. `None` for both is - the production configuration and is what the engine's own MLA call site - passes. - - Input magnitudes must survive e4m3: `|value * w|` below `2**-10` - (~9.8e-4) rounds to zero, and 448 is the largest finite e4m3 value (what - the kernel does above it was not measured). The certified runs are - unit-normal activations at `w <= 1`. - -## Notes - -- Schema annotation is wrong in both directions: it marks `fused_q` and - `q_pe` mutable (`Tensor(a!)`), but `q_pe` is never modified — and on the - fp8 path neither is `fused_q` — while `cu_q_seqlens`, `cu_kv_seqlens`, - `fmha_scheduler_counter`, `quant_q_buffer`, `mla_bmm1_scale` and - `mla_bmm2_scale` — all filled by the call — carry no annotation at all. Do - not rely on the schema's alias info (e.g. under torch.compile - functionalization). -- Precision, bf16 pool (sm_100): the `compressed_kv` copy is bitwise. Both - RoPE outputs match an fp32 reference rounded once to bf16 to within **one - bf16 ulp**, and are bitwise on all but ~1.5e-5 of elements: the kernel and - a torch fp32 reference evaluate `x*cos ∓ y*sin` in different orders, so - where the correctly-rounded fp32 result lands exactly on a bf16 rounding - midpoint the two round to adjacent values. Measured on a 64-sequence - decode step: 1 such element in 65 536 rotated `q_pe` values and 0 in 4096 - rotated `k_pe` values, with identical counts at page sizes 32 and 64 — it - is arithmetic, not addressing. Default `assert_close` tolerances absorb it - with margin, which is what the test gates on. -- Precision, fp8 pool (sm_100): **everything came back bit-exact**, both the - pure copies (`quant_q_buffer`'s absorbed-q head, the appended - `compressed_kv` half) and the roped halves. e4m3's 3-bit mantissa absorbs - the evaluation-order difference the bf16 pool pays one ulp for, but the - double rounding is real and a mirror must reproduce it: the kernel rounds - the fp32 RoPE result to **bf16 first** and quantizes that. Quantizing the - fp32 RoPE result directly disagrees with the kernel on ~3% of the roped q - elements (and up to 9% of the much smaller per-row `k_pe` samples). The - test still gates the roped halves at one e4m3 ulp - (`2**-3` relative) plus a bit-exact-fraction floor rather than bitwise, - because the tie it covers is arithmetic; neither half has been approached. -- The fp8 scales are exact enough to gate tightly: `mla_bmm2_scale` is a - bitwise copy of `kv_scale_quant_orig`, not a computation — a scratch probe - passing 0.7 (whose fp32 value is 0.699999988079071) got that exact fp32 - number back, and the test gates the copy bitwise at 1.0 / 1.5 / 2.0 / 3.0. - `mla_bmm1_scale` matches a double-precision reference rounded once to fp32 - to within 7.2e-8 relative (0.6 fp32 ulp) across the swept grid, exactly on - several cases. -- Pool writes are exactly bounded: `P` `D`-byte rows per generation - sequence, at the slots its `P` positions address, and nothing else in the - pool moves — checked by a whole-pool byte snapshot on every fp8 case, at - `P` 1 through 4, 64-sequence batches included. -- Re-running the same prepared step overwrites the same `P` cache slots - (positions are derived from `sequence_length`), so a repeated call is - idempotent, not double-appending — and on the fp8 path the second run - reproduces `quant_q_buffer` and the pool bitwise. This holds at `P > 1` - too: the rows do not advance by `P` on the second call. -- The scheduler buffers feed the trtllm-gen decode MQA: `cu_q_seqlens` is - in units of q rows (`H` rows per generation *token*, so `H * P` per - generation sequence), `cu_kv_seqlens` in tokens (the cumulative - `sequence_length` over the generation sequences, which already includes - all `P` new tokens). They are filled identically on both paths, and both - fills were observed at `P` 1, 2, 3 and 4 — `cu_q_seqlens[i] = i * H * P` - exactly, and no entry past index `G` is written. -- Page size (sm_100): certified at `tokens_per_block` 32 and 64 on the bf16 - path, 32 on the fp8 path. 32 is what a `KvCacheConfig` at its defaults - hands the op; 64 is what a tuned MLA target opts into. At 32 the appended - rows were placed at every alignment a 32-slot page has — position 31 (a - page's last slot), 32 and 96 (slot 0 of a fresh page after one and three - full pages), 64 (after two), and mid-page — including 64-sequence batches - run for two consecutive steps in which sequences cross into a new page - between the steps and others open a fresh page on the first, and mixed - batches whose leading context sequence is skipped. At `P > 1` the same - sweep covers a sequence whose `P` rows **straddle** the boundary, at every - split from 1 row in the old page to `P-1` (`P` = 2, 3 and 4, fp8 and bf16 - pools) — a shape unreachable at `P = 1`, where one call writes one slot - per sequence. Every effect held at the same gates as at page 64. Other - page sizes are untested rather than known-bad. -- `q_pe` in production is a strided slice of the packed q tensor - (`[G*P, H, nope+R]` split); the kernel handles that stride — certified for - both contiguous and packed-slice layouts, on both paths, at `P` 1-4. -- The `rotary_cos_sin` *content* is not a certified axis: what the op is - gated on is that it applies the table's `(cos, sin)` pairs, and both - certified configurations ran the unscaled theta-10000 table. A YaRN-scaled - table of the same layout has not been run through this op here. -- Sibling ops exist for adjacent MLA roles: - `trtllm::mla_rope_append_paged_kv_assign_q` (context-phase RoPE + cache - append), `trtllm::attention` / `trtllm::mla_custom_op_inplace` (the - attention core that consumes `fused_q` or `quant_q_buffer` and these - scheduler buffers), and `trtllm::merge_chunked_attention_for_mla`. -- `predicted_tokens_per_seq` is the *only* way a multi-token generation - query reaches this op: its signature carries no spec-decoding mask, tree - or draft-token argument of any kind, so there is nothing else for an MTP - caller to set here. (What the resulting query block is then allowed to - attend to is the attention op's fact, not this one's.) -- Not certified on either path: `predicted_tokens_per_seq` above 4, - `block_ids_per_seq` (the flash-MLA layout), helix, `out_scale` (fp8 - attention output), `attention_window_size` below the KV length, and - multi-layer pools. Not certified on the fp8 path specifically: - `tokens_per_block = 64`, head counts other than 128, and activation dtypes - other than bf16. - - -## Parameters this entry pins to their defaults - -Every value certified below leaves the op on its default behaviour, and -matches what the in-tree caller passes on the same path. - -| Parameter | Certified value | Why | -|---|---|---| -| `kv_cache_scale_orig_quant` | `None` | With no value the op uses `kv_scale_orig_quant` for the cache scale (`dsv3RopeOp.cpp`), which is exactly what it did before this parameter existed. | -| `residual_dim` | `0` | Must be `0` or `rope_size`; non-zero requires an FP4 KV pool, and this entry certifies bf16 and fp8-e4m3 pools. | -| `kv_norm_weight` / `kv_norm_eps` | `None` / `1e-6` | A non-`None` weight folds the `kv_a_layernorm` into this kernel, which then reads `latent_cache` **raw**. A caller that already normalized would normalize twice. Only DeepSeek-V4's sparse module folds it upstream. | -| `precomputed_cu_seqlens` | `False` | The kernel fills the cu-seqlens buffers itself; `True` only when the Q half is skipped. | -| `precomputed_fmha_scheduler` | `False` | The scheduler counter and the two bmm scales come from this call, not from a sparse index kernel. | -| `kv_only` / `kv_done_elsewhere` | `False` / `False` | Both halves run in this one call. | -| `quant_scale_qkv` | `None` | `q_nope` in `quant_q_buffer` is not pre-quantized. | diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_generation.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_generation.py deleted file mode 100644 index ba5d77856847..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/mla_rope_generation.py +++ /dev/null @@ -1,122 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""MLA generation-phase RoPE + latent KV-cache append + FMHA scheduler-buffer fill.""" - -from typing import List, Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* - - -def mla_rope_generation( - fused_q: torch.Tensor, - q_pe: torch.Tensor, - latent_cache: torch.Tensor, - rotary_cos_sin: Optional[torch.Tensor], - cu_q_seqlens: torch.Tensor, - cu_kv_seqlens: torch.Tensor, - fmha_scheduler_counter: torch.Tensor, - mla_bmm1_scale: Optional[torch.Tensor], - mla_bmm2_scale: Optional[torch.Tensor], - quant_q_buffer: Optional[torch.Tensor], - sequence_length: torch.Tensor, - host_past_key_value_lengths: torch.Tensor, - host_context_lengths: torch.Tensor, - num_contexts: int, - kv_cache_block_offsets: Optional[torch.Tensor], - host_kv_cache_pool_pointers: Optional[torch.Tensor], - host_kv_cache_pool_mapping: Optional[torch.Tensor], - kv_scale_orig_quant: Optional[torch.Tensor], - kv_scale_quant_orig: Optional[torch.Tensor], - # When None the op falls back to kv_scale_orig_quant - # (dsv3RopeOp.cpp:280). - kv_cache_scale_orig_quant: Optional[torch.Tensor], - out_scale: Optional[torch.Tensor], - block_ids_per_seq: Optional[torch.Tensor], - helix_tensor_params: List[Optional[torch.Tensor]], - predicted_tokens_per_seq: int, - layer_idx: int, - num_heads: int, - num_kv_heads: int, - head_size: int, - # 0 or rope_size, and non-zero requires an FP4 KV pool. - residual_dim: int, - tokens_per_block: int, - attention_window_size: int, - beam_width: int, - quant_mode: int, - q_scaling: float, - q_lora_rank: int, - kv_lora_rank: int, - qk_nope_head_dim: int, - qk_rope_head_dim: int, - v_head_dim: int, - rope_append: bool, - # Every default below leaves the op on its default behaviour. - # kv_norm_weight non-None would fold the kv_a_layernorm into - # this kernel, which then reads latent_cache RAW -- a caller that already - # normalized would be normalizing twice. - kv_norm_weight: Optional[torch.Tensor] = None, - kv_norm_eps: float = 1e-6, - precomputed_cu_seqlens: bool = False, - precomputed_fmha_scheduler: bool = False, - kv_only: bool = False, - kv_done_elsewhere: bool = False, - quant_scale_qkv: Optional[torch.Tensor] = None, -) -> None: - """RoPE q_pe, append latent_cache rows to the paged MLA KV cache, and fill - the decode-FMHA scheduler buffers. Returns None. - - Where the roped q_pe lands depends on the pool: into fused_q's tail over a - bf16 pool, into quant_q_buffer over an fp8 one (where fused_q is read, not - written). See the contract.""" - torch.ops.trtllm.mla_rope_generation( - fused_q, - q_pe, - latent_cache, - rotary_cos_sin, - cu_q_seqlens, - cu_kv_seqlens, - fmha_scheduler_counter, - mla_bmm1_scale, - mla_bmm2_scale, - quant_q_buffer, - sequence_length, - host_past_key_value_lengths, - host_context_lengths, - num_contexts, - kv_cache_block_offsets, - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping, - kv_scale_orig_quant, - kv_scale_quant_orig, - kv_cache_scale_orig_quant, - out_scale, - block_ids_per_seq, - helix_tensor_params, - predicted_tokens_per_seq, - layer_idx, - num_heads, - num_kv_heads, - head_size, - residual_dim, - tokens_per_block, - attention_window_size, - beam_width, - quant_mode, - q_scaling, - q_lora_rank, - kv_lora_rank, - qk_nope_head_dim, - qk_rope_head_dim, - v_head_dim, - rope_append, - kv_norm_weight, - kv_norm_eps, - precomputed_cu_seqlens, - precomputed_fmha_scheduler, - kv_only, - kv_done_elsewhere, - quant_scale_qkv, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.md deleted file mode 100644 index 94f95df61644..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.md +++ /dev/null @@ -1,1944 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 47} ---- - -# thop_attention - -**Wraps** `tensorrt_llm.bindings.internal.thop.attention` (one call). - -This is a pybind binding, not a `torch.ops` op; its inclusion in the catalog -is an approved policy exception to the `torch.ops.trtllm.*` entry shape. - -## Semantics - -The fused attention core of one decoder layer over one prepared batch, with -**fully explicit state**: unlike the registered-layer attention entry -points, this binding reads no Python-side registry, no thread-local -metadata, and no layer objects. Every piece of batch state, cache -addressing, and layer config is an argument. - -The certified surface is bf16 activations over a caller-owned paged pool — -bf16 (`quant_mode=0`) or fp8-e4m3 (`quant_mode=128`, see *Standard -configuration over an fp8-e4m3 KV pool* and *MLA over an fp8-e4m3 latent -pool*; the two configurations quantize differently and are certified -separately) — with no other quantization and no spec-dec / sparse / cross / -helix / beam features (their arguments held at the inert values listed in -*Signature*), in two configurations: - -- **Standard** (`is_mla_enable=False`): packed-QKV GQA/MHA self-attention, - paged KV-cache append + masked FMHA, mixed batches in one call, over a - pool holding one layer or several layers side by side (see *Paged KV - cache addressing*), on either context execution path - (`use_paged_context_fmha` `False` or `True` — the latter is what lets a - context call carry a cached prefix, see *Paged-context FMHA*), optionally - with per-query-head **attention sinks** (bf16 pool, causal mask — see - *Attention sinks*) and/or a **sliding window** (`attention_window_size` - smaller than the sequence — see *Sliding window*). -- **MLA** (`is_mla_enable=True`): DeepSeek-style multi-head latent - attention over a paged **latent** cache (kv_factor 1: one latent row of - width `C + R` per token, `C = kv_lora_rank`, `R = qk_rope_head_dim`), - bf16 or fp8-e4m3. - A batch is served by **one call per phase** — `attention_input_type=1` - (context_only) over the context rows, then `=2` (generation_only) over - the generation rows — sharing the same full-batch state tensors. Mixed - `attention_input_type=0` is rejected for MLA. The context phase has two - certified flavors selected by `latent_cache`: fresh prefill - (`latent_cache` given: in-kernel RoPE + latent append) and no-append - (`latent_cache=None`: pure FMHA over explicit K/V — the cached-KV and - chunked-prefill context flows, including `softmax_stats_tensor` output). - Both flavors, the generation call and the mixed-batch pairing of two of - them are certified over **both** pool element types; the chunked - partial-pass pattern of the no-append flavor is bf16-only. The generation - call is additionally certified at `predicted_tokens_per_seq` 1, 2, 3 and 4 - — one speculative-decoding step, where a generation sequence contributes a - `P`-token draft chain instead of a single token. - Mixed `attention_input_type=0` is **not** rejected — nothing in the - binding checks it (the "MLA cannot be mixed" `ValueError` lives one layer - up, in the Python attention backend, which this op bypasses). On a - single-phase batch it is bitwise identical to that phase's own value; on a - genuinely mixed batch it runs **both** phase paths off one set of scalars, - and `head_size`, the `q`/`output` row widths and - `cu_q_seqlens`/`cu_kv_seqlens` can serve only one of them — so at most one - row group comes back right, the other is silently wrong or never written, - and the generation write overruns a context-sized `output`. Pass the - phase's own value. - -### Standard configuration - -One call, for a batch of `ns` sequences — context-phase sequences first -(`host_request_types[i] == 0`, `i < num_contexts`), generation-phase after -(`== 1`). Rows of `q` are: all context tokens sequence by sequence -(`num_ctx_tokens` rows total), then one row per generation sequence -(`predicted_tokens_per_seq == 1`). For each sequence `s` with new-token -count `l_s` and total KV length `kv_s = sequence_length[s]`: - -``` -# 1. slice K and V out of the packed q rows and append them to the paged -# cache at token positions kv_s - l_s .. kv_s - 1 -# (a generation sequence appends one token at position kv_s - 1; a context -# sequence appends l_s of them, l_s < kv_s only on the paged-context path -# — see *Paged-context FMHA*) -# 2. for each new query token i (absolute position p = kv_s - l_s + i), -# each query head h, over kv head h // (Hq/Hkv): -keys = [0, p] # mask_type == 1 (causal) -keys = [p - W + 1, p] # mask_type == 1, W = attention_window_size - # < p + 1 (see *Sliding window*) -keys = [0, kv_s - 1] # mask_type == 0 (padding: no mask) -logits = q[i, h] @ K[keys].T / (q_scaling * sqrt(D)) -out[i, h] = softmax(logits) @ V[keys] # fp32 accum -# with attention_sinks given (see *Attention sinks*): -out[i, h] = softmax(concat(logits, sink[h]))[:-1] @ V[keys] -``` - -Results are written into `output[:num_tokens]` in place; the op returns -nothing. Generation sequences always attend over all `kv_s` cached+current -tokens (a decode step is `l_s = 1`), or over the newest `W` of them when a -sliding window is active. - -Fusion boundary: KV-cache append, masking, softmax, and GQA attention -happen inside the call (plus RoPE when `position_embedding_type` selects -it — not certified for this configuration). The caller owns the QKV -projection, any q/k norm, external RoPE, and the output projection. - -#### Paged-context FMHA - -`use_paged_context_fmha` selects the **context** execution path; it does -not change the generation path. Both values are certified. - -- `False` — the context FMHA reads K and V out of the **packed `q` rows** - and never touches the pool, so a context sequence can only attend to its - own new tokens: `kv_s` must equal `l_s` (the sequence starts empty). -- `True` — the context FMHA reads K and V out of the **paged pool**, after - this call's own append has written the new tokens into it. `kv_s > l_s` - is then legal: the sequence's first `kv_s - l_s` tokens are a cached - prefix that earlier calls wrote, query row `i` sits at absolute position - `p = kv_s - l_s + i`, and the causal mask is bottom-right aligned exactly - as *Standard configuration* states. This is what KV-cache reuse and - chunked prefill need, and an engine turns it on for every call of the - batch as soon as either is enabled (trtllm's defaults enable reuse). - -Two consequences pinned by test: - -- **With nothing cached the flag is observationally inert.** A fresh - prefill and a decode step run from identical state at both values come - back **bitwise identical** in `output` and in pool content (certified on - 32/8/128, 32/4/128 and 64/8/64). The K/V source does move even then — - aliasing every page of a fresh-prefill sequence onto one garbage page - leaves the `False` output bitwise unchanged but moves the `True` output - 31x outside the tolerance band — it is the append running first that - leaves the pool holding exactly the packed rows. -- **With a cached prefix the flag is load-bearing and its absence is - silent.** The same cached-prefix context call at `False` returns without - raising, 51x outside the band around the correct answer. Nothing in the - op checks the combination. - -The append is unchanged by the flag: the new tokens land bit-exactly at -their absolute positions and nothing outside the token ranges is written, -so a repeated identical call is idempotent on both paths. - -Certified surface for `True`: the standard bf16-pool configuration, -`mask_type=1` (causal), `is_fused_qkv=True`, `q_scaling=1.0`, -`tokens_per_block` 32, geometries 32/8/128, 32/4/128 and 64/8/64, on a -single-layer pool and on a real 4-layer `KVCacheManager` pool, with and -without `attention_sinks` and a sliding window. `True` with `mask_type=0` -(padding), with the fp8-e4m3 pool, or with MLA is not certified (MLA -context takes separate q/k/v and its own cached-KV flavor instead). - -#### Head geometry - -Head **counts are a free axis**, not an enumerated list: `num_heads` (`Hq`) -and `num_kv_heads` (`Hkv`) are plain runtime parameters of the same -kernels, and every pair satisfying `Hq % Hkv == 0` is served — MHA -(`Hq == Hkv`), MQA (`Hkv == 1`), and any GQA ratio, power of two or not. -Certified on the bf16 pool: `Hq`/`Hkv` of 32/8 and 32/4 (the two shipped -targets), 64/8, 16/1, 28/4, 12/3, 8/2, 4/4 — ratios 1, 4, 7, 8 and 16, with -non-power-of-2 head counts (28, 12, 3) among them. All of them run the same -kernels at the same accuracy (see *Notes*); nothing distinguishes the -"round" geometries. - -`head_size` (`D`) is the **enumerated** axis: it selects the FMHA kernel, -and only the shipped head dims exist. Certified: 64, 128, 256 (all on the -bf16 pool; 128 also on the fp8-e4m3 pool). A scratch probe additionally saw -80 work (its generation kernel is JIT-generated), while 32, 96 and 192 are -rejected — hard, see *Preconditions*. Treat any `D` outside {64, 128, 256} -as untested rather than merely unsupported. Certified `Hq`/`Hkv`/`D` -triples on the bf16 pool: 32/8/128, 32/4/128, 64/8/64, 16/1/128, 28/4/128, -12/3/128, 8/2/128, 8/2/256, 4/4/64. - -`Hq % Hkv != 0` is the one geometry error the op does not report. A -context-only call at 6/4/128 returned normally, having computed only the -first `(Hq // Hkv) * Hkv = 4` head columns of `output` and left the other -two all-zero; the same geometry does raise once a generation-phase call -selects a decode kernel (`numQHeads should be multiple of numKVHeads`, from -the XQA dispatcher), and `Hq = 7`, `Hkv = 2` aborts the process outright -inside kernel selection. Because the context path fails silently, the -wrapper asserts `num_heads % num_kv_heads == 0` before the call — one of the -entry's two guard asserts (the other guards `attention_sinks`), holding for -every certified configuration including MLA (8/8, 16/16, 32/32 and 128/128 -context, 8/1, 16/1, 32/1 and 128/1 generation). - -#### Attention sinks - -`attention_sinks` is an optional fp32 `[Hq]` CUDA tensor of per-query-head -**sink logits**. When given, each head's softmax gains one extra logit -column that is dropped from the weighted sum: the sink lands in the -denominator only, so the attention weights of every row sum to less than 1 -and the output shrinks accordingly. For query row `i`, head `h`, over the -masked logit range of *Standard configuration*: - -``` -out[i, h] = softmax(concat(logits, sink[h]))[:-1] @ V[:n_keys] - = mass[i, h] * (softmax(logits) @ V[:n_keys]) -mass[i, h] = sigmoid(logsumexp(logits) - sink[h]) # < 1, per (row, head) -``` - -Facts pinned by test, not by inspection: - -- The sink is a logit in the **scaled**-score domain — it is compared - against `q @ K.T / (q_scaling * sqrt(D))`, *after* the softmax scale, not - against the raw dot products. (Certified at `head_size = 64`, - `q_scaling = 1.0`, so a scale of 1/8; the rival "sink joins the unscaled - row" hypothesis sits 12.8x to 41x outside the tolerance band on every - certified case.) -- It is honoured in **both** phases, which take different FMHA kernels: - the context kernel, the short-history decode kernel, and the - multi-CTA-KV decode kernel that folds partial softmax states across CTAs - (`...MultiCtasKvCga...`, selected at 600 and 2000 cached tokens). A - single `attention_input_type=0` call mixing a context and a generation - sequence applies it to both row groups. -- Indexing is by **query** head: sink `h` belongs to `output`'s head column - `h`, not to the kv head `h // (Hq/Hkv)`. -- The KV-cache append is unaffected — the appended pages stay bit-exact and - nothing outside the token ranges is written. -- `attention_sinks=None` and a sink that cannot contribute (`-100.0` or - `-inf` per head) produce **bitwise identical** output in both phases: the - argument's only effect is the one added denominator term. - -Certified surface for sinks: the standard configuration over a bf16 pool -(`quant_mode=0`), `mask_type=1` (causal), `is_fused_qkv=True`, -`q_scaling=1.0`, geometry 64/8/64 — the gpt-oss-120b tp1 layer shape — -both without a window (`attention_window_size = max_seq_len`) and with the -sliding window of *Sliding window* below, and on both context execution -paths including a context call over a cached prefix (see *Paged-context -FMHA*). Sinks combined with `mask_type=0`, with the fp8-e4m3 pool, with -`q_scaling != 1.0`, or with MLA are not certified. - -#### Sliding window - -`attention_window_size` (`W`) is a **per-call token count**. When a -sequence's KV length exceeds it, query row at absolute position `p` attends - -``` -keys [max(0, p - W + 1), p] # exactly min(p + 1, W) keys -``` - -— `W` keys, never `W + 1`; a sequence still shorter than `W` is plain -causal. `W >= every sequence's total KV` is the no-window behavior. The -window applies identically to context rows and generation rows, and with -`attention_sinks` the sink logit joins the denominator of that **windowed** -logit set (each key of a uniform-logit row of `n = min(p + 1, W)` keys -weighs `1 / (n + exp(sink[h]))`). - -Facts pinned by test, not by inspection: - -- The boundary is exact. A one-hot-V probe (all-zero K, so every unmasked - logit is 0 and the softmax is uniform) reads the per-key weights straight - out of `output`: keys below `p - W + 1` come back **bitwise zero**, keys - from `p - W + 1` up to `p` all carry the same `1 / (n + exp(sink[h]))` - with `n = min(p + 1, W)`. Certified at `W = 128` in both phases and at - `W = 33` and `100` — the window is a token count, not a page count, and - may be smaller than `tokens_per_block`. -- **The window is a mask only: cache addressing is unchanged.** The op - appends every new token at its **absolute** position — page - `t // tokens_per_block` of the sequence's `kv_cache_block_offsets` row, - slot `t % tokens_per_block` — and wraps nothing, neither at `W` nor at - `max_seq_len`. Whatever cyclic reuse of pool memory happens is the - caller's page mapping, not the op's (see *Paged KV cache addressing under - a sliding window*). -- A single context prefill longer than `W` is legal in one call: each row - gets its own window. At `use_paged_context_fmha=False` the context FMHA - reads the packed `q` rows, never the pool — aliasing every page of the - sequence onto one garbage page left the context output bitwise identical - (the append still runs, so the pool itself is then wrong). At `True` the - same aliasing moves the output far off: that path reads the pool (see - *Paged-context FMHA*). -- `sequence_length` must stay the **global** cached+new count: it is both - the origin the window is measured back from and the append position. - Capping it to `W` moves the write to slot `W - 1` and truncates the - attended range — silently, no error. -- `attention_window_size` is per call, so layers with different windows - share one pool, one layer→pool mapping and one block-offset table inside - the same batch (certified on a real 4-layer `KVCacheManager` pool with - windows alternating 128 / full — the gpt-oss layer alternation). The - shared pool must then still be sized for the full-attention layers; a - smaller pool per window group needs multi-pool addressing, which is not - certified. - -Certified surface for the window: the standard bf16-pool configuration, -`mask_type=1`, `is_fused_qkv=True`, `q_scaling=1.0`, geometry 64/8/64, -`tokens_per_block` 32, `W` of 128 / 100 / 33, with and without -`attention_sinks`, over prefill (including prefills 1.5-2.3x the window), -decode, mixed batches, histories up to 2000 cached tokens, and 250 decode -steps over a wrapping page ring; at `W = 128` also over a context call -whose cached prefix (200 tokens) runs well past the window, on the -paged-context path. `W` with `mask_type=0`, with the fp8-e4m3 pool, or with -MLA is not certified. - -### Standard configuration over an fp8-e4m3 KV pool - -`quant_mode=128` (the FP8-KV-cache bit alone — a bf16 checkpoint serving -with a quantized cache) switches the pool element type to fp8-e4m3 while -`q` and `output` stay bf16. Both fp32 `[1]` CUDA scale tensors are then -live: `kv_scale_orig_quant` holding `1/s` (applied on write) and -`kv_scale_quant_orig` holding `s` (applied on read), `s` = the KV-cache -scaling factor — `s = 1.0` for uncalibrated bf16-checkpoint serving, the -production default. Passing `None` for either is silently equivalent to -`1.0` rather than an error (see *Preconditions*), so at any other `s` the -tensors are the caller's only defence. The call's three roles were pinned -against per-role references: - -- **Append** (both phases): each new K/V token row lands in the pool as - `e4m3(k * kv_scale_orig_quant)` — fp32 multiply, round-to-nearest cast. - Bit-exact against that mirror for power-of-2 scales (1.0 and 2.0 - certified); at `s = 1.5` ~0.6% of elements differ from the mirror by - exactly one e4m3 ulp (the kernel's intermediate precision differs - slightly from a pure fp32 multiply), never by more. Pages outside the - written ranges stay bitwise untouched. -- **Context FMHA** reads the **bf16 packed `q` rows, not the pool**: - context output matches the bf16-KV fp32 reference at the same - tolerances as the bf16-pool surface (an fp8-round-trip reference is - ~1.4e-1 off). Prefill accuracy is unaffected by cache quantization; - only what later reads the cache sees e4m3. -- **Generation FMHA** attends over dequantized pool rows - `k_hat = v_hat = pool_row * kv_scale_quant_orig`, with `q` additionally - put through the same e4m3 round-trip in-kernel (quantize by - `kv_scale_orig_quant`, dequantize by `kv_scale_quant_orig` — pinned by - a non-power-of-2 scale sweep: only that scale collapses the error). - Decode output matched an fp32 reference over `(q_hat, K_hat, V_hat)` - within `2^-4 * (1 + |ref|)` — observed max abs err 3.1e-2 on - magnitude-~1 outputs, at most 44% of that allowance; the residual is - the kernel's internal e4m3 handling of the softmax probabilities (the - decode kernel computes every MMA operand in e4m3). - -Certified fp8-pool surface: GQA 32/8/128, `tokens_per_block` 32, causal -mask; on a single-layer pool: context-only / generation-only / mixed -batches, scales 1.0, 1.5, 2.0, decode histories up to 200 tokens with -page-boundary crossings; on a 4-layer shared pool (real `DataType.FP8` -manager state, `s = 1.0`): per-layer context-only prefill and -generation-only page-crossing decode with bit-exact per-layer appends and -bitwise sibling-layer isolation (see the multi-layer note in *Notes*; the -non-1.0 scale sweep and mixed batches were not re-run multi-layer). The -bf16-pool-only extras (padding mask, the paged-context path, and every head -geometry other than 32/8/128) were not re-run under fp8. MLA over an fp8 -pool is certified separately, and behaves differently in every role but the -append — see *MLA over an fp8-e4m3 latent pool*. Nothing in this section -transfers to it: in particular "context FMHA reads the bf16 `q` rows, not -the pool" is a fact about **this** configuration only. - -### MLA configuration - -Both phases scale `QK^T` by -`1 / (q_scaling * sqrt(qk_nope_head_dim + qk_rope_head_dim))` — the -generation phase too, despite its `head_size` being `C + R`. `q_scaling` is -a live argument, not a fixed 1.0: it was swept over four values in both -phases and moves only that scale (see the `q_scaling` row of *Geometry and -masking scalars*, and *Notes*). - -#### MLA head count - -`H` = `num_heads` is a plain runtime parameter of the context call, which -runs as MHA (`num_kv_heads = H`, `head_size = nope + R`); the generation -call runs as latent MQA (`num_kv_heads = 1`, `head_size = C + R`) and there -`H` does reach kernel selection (see the JIT paragraph below). Certified: -**`H` = 8, 16, 32 and 128** — 32 is the deepseek-v3-lite tp1 layer -shape, 16 a TP-slice-like count, 8 the tep4 slice of that same checkpoint -(32 attention heads split over 4 tensor-parallel ranks), and 128 the -DeepSeek-R1-0528 layer shape, which an attention-DP rank runs whole because -DP partitions requests rather than heads — all at the same -`C`/`R`/`nope`/`v_head_dim` (512/64/128/128). The first three were run -through the -identical four-case set at `tokens_per_block = 64` — fresh prefill including -a page-crossing sequence, cached-KV no-append context, generation decode over -a page boundary, and a mixed batch served by two phase calls — against the -same fp32 references and the same tolerances, and they land in the same -accuracy envelope: every `H = 8` case sits between 36.8% and 45.9% of the -tolerance allowance, at or below the `H = 32` run of the same case in all -eight comparisons and below the `H = 16` run in all but one (the page-64 -mixed batch, 38.5% against 37.9%) — per-case fractions in *Notes*. `H = 8` -runs that whole four-case set at -`tokens_per_block = 32` as well, on the page-32 geometries of *MLA page -size*, so it is certified at both page sizes across all four cases — the -coverage `H = 32` has. Those four cases all pass `q_lora_rank = 1536`, but -`q_lora_rank = 0` at `H = 8` no longer rests on the `H = 32` measurement: -the inertness sweep of *MLA q-LoRA rank* is driven at `H = 8` as well, which -puts it on the `...VarSeqQ8...` decode kernel this head count compiles -instead of on a kernel a tep4 target never runs. - -**`H = 128`** runs that same four-case set at `tokens_per_block = 32` only, -on the page-32 geometries of *MLA page size*, and runs it **twice**: once on -the unscaled rope table with `q_scaling = 1.0`, so the head count is the only -thing moving against the `H = 8` and `H = 32` runs of the same cases, and -once at the complete DeepSeek-R1-0528 cell — the checkpoint's YaRN-scaled -rope table plus `q_scaling = 1/mscale² ≈ 0.53366` (see *In-kernel RoPE -arguments* and the `q_scaling` row of *Geometry and masking scalars*), which -is the combination a DeepSeek-R1 layer actually passes. It lands in the same -accuracy envelope as the smaller counts: 37.2-63.3% of the tolerance -allowance on the baseline set, 37.2-60.7% at the R1 cell — per-case fractions -in *Notes*. `H = 128` is **not** certified at `tokens_per_block = 64`, and -none of the four cases was re-run at `q_lora_rank = 0` (they all pass 1536, -the R1 value); the rank's inertness rests on the `H = 32` and `H = 8` sweeps. - -`H = 128` is also the **only** head count certified over an fp8-e4m3 latent -pool, where the same four cases run again — baseline cell and R1 cell alike — -against fp8 references at the wider fp8 tolerance (see *MLA over an fp8-e4m3 -latent pool*). - -Everything scales with `H` exactly as the shape tables below state: `q`/`k` -row width `H * (nope+R)`, context `output` width `H * v_head_dim`, the `v` -split view's row stride `H * (nope + v_head_dim)` elements, generation `q` -width `H * (C+R)`, generation `output` width `H * C`, `q_pe` `[G*P, H, R]`, -and `cu_q_seqlens = arange(G+1) * H * P` (`P` = `predicted_tokens_per_seq`, -1 unless a speculative-decoding step). The paged latent pool is untouched by -the head count (its row width is `C + R`). - -The **context** phase JIT-compiles nothing at any certified head count. The -**generation** phase compiles a decode kernel whose name carries a q-tile -that the head count moves at the low end only: `H = 8` takes -`...PagedKvDenseP{32,64}VarSeqQ8Kv128StaticSwapsAbForGen` where `H` = 16, 32 -**and 128** all take `...VarSeqQ16...`. So `H = 8` decode is a genuinely -different compiled kernel, while 16, 32 and 128 are separate compile-cache -entries behind one kernel name — each still pays its own ~5.5 s compile (see -*Notes*). Head counts other than 8, 16, 32 and 128 are untested rather than -known-bad, and a new one may pay a decode JIT compile of its own. - -Two parts of the MLA surface stay certified at `H = 16` only, because no -other head count was run through them: the -`softmax_stats_tensor` output, and the chunked partial-pass pattern of the -no-append flavor (padding mask, per-pass stats, zero-KV sequences, -downstream merge). Both are independent of the head count in the shape -tables — the stats tensor is `[>= Tc, >= H, 2]` — but neither has a run at -any other head count behind it. - -#### MLA page size - -`tokens_per_block` is certified at **32 and 64** for MLA. 32 is what a -`KvCacheConfig` at its defaults produces, so it is the page size an MLA -target hits before anyone tunes it; 64 is the tuned value. Like the head -count this is a kernel-selection axis for the generation phase only: it -JIT-compiles a decode kernel whose name carries the page size — -`...PagedKvDenseP32VarSeqQ16Kv128StaticSwapsAbForGen` at 32 versus -`...P64...` at 64 (see *Notes*) — while the context phase triggers no JIT -compile at either size. - -Certified at page 32: all three MLA call flavors at `H = 32` (the -deepseek-v3-lite tp1 layer shape), at `H = 8` (its tep4 slice) and at -`H = 128` (the DeepSeek-R1-0528 layer shape, on the baseline rope/scale and -again at the full R1 cell), including -the mixed-batch pairing of two of them at all three counts, plus the -generation flavor also at `H = 16` because that is the one flavor where the -page size changes the compiled kernel and the decode compile cache is keyed -by head count. The page-32 geometry is chosen -for 32-slot pages rather than reused from the page-64 cases — a 64-token -prefill fills a 64-slot page exactly but spans two 32-slot pages, so the -crossing arithmetic differs: - -- **fresh prefill**: sequences of 96 (three exact pages — a 64-slot page - never ends there) and 33 (a one-token spill into a second page), append - checked page by page; -- **generation decode**: a 64-token history filling two pages exactly, so - the first decode token opens page 2, alongside a 31-token history whose - first decode token takes page 0's last slot and whose second opens page - 1 — a decode read range crossing a boundary mid-case, two steps, all - four head counts; -- **mixed batch**: a 64-token history decoded next to a fresh 33-token - context sequence, the two phase calls sharing one page-32 offsets table; -- **cached-KV no append**: prefixes of 96 / 31 / 0 reaching KV lengths of - 128 (four exact pages), 40 and 25, with the pages reserved as production - reserves them and the pool verified bitwise untouched. - -Page 32 is also the only size certified over an fp8-e4m3 latent pool, where -all four geometries above run again at `H = 128` (see *MLA over an fp8-e4m3 -latent pool*). - -Accuracy at page 32 lands on the page-64 magnitudes (per-case fractions in -*Notes*, including the one page-32 case that sets the file's headroom -ceiling — a flavor that never touches the pool, so its number tracks its KV -geometry rather than the page size), the paged latent append holds to the -same gate, and the generation call still writes nothing. Page sizes other -than 32 and 64 are untested for MLA rather than known-bad; the standard -configuration is certified at 32 only. - -#### MLA q-LoRA rank - -`q_lora_rank` names the rank of the checkpoint's `q_a_proj`: DeepSeek-V3 -down-projects the query through a rank-1536 LoRA pair, while a checkpoint -whose config carries `"q_lora_rank": null` (deepseek-v3-lite) has no q-LoRA -at all and projects the query directly. Certified: **1536 and 0**. - -The argument is **inert**: no observable of any certified MLA call depends -on it. That matches where the query projection sits — entirely outside this -call, which receives finished `q` rows (absorbed or not) whether or not a -q-LoRA produced them. Measured on sm_100 by sweeping -`1536 → 0 → 1536 → 4096` (the last larger than `C`, so arithmetic that -touched the value would move something) over identical inputs, in all three -MLA call flavors at `tokens_per_block = 32`, on the page-32 geometries of -*MLA page size*, at **two** of the certified head counts — `H = 32` (the -deepseek-v3-lite tp1 layer shape) and `H = 8` (its tep4 slice), two -independent sweeps with their own seeds; `H = 128` and `H = 16` were not -swept. Both swept counts are measured rather than one inferred from the other -because the generation phase compiles a different decode kernel at each -(`...P32VarSeqQ16...` at `H = 32`, `...P32VarSeqQ8...` at `H = 8` — see -*MLA head count*), so an `H = 32` sweep alone is evidence about a kernel a -tep4 target never runs. Not swept: the mixed batch (a pairing of two of -those flavors), `H = 16`, `tokens_per_block = 64` (both sweeps are -page-32), and the no-append flavor's chunked partial-pass pattern with -`softmax_stats_tensor` — the sweep takes its one-shot cached-KV pattern. -Every observable came back **bitwise identical** at both head counts: the -`output` rows, the whole paged latent pool, the in-place-roped `q`/`k` of -the fresh-prefill flavor, and the byte count the op resized `workspace_` -to. The repeated `1536` is the run-to-run determinism control that makes -those comparisons mean something. Kernel selection does not see it either: -a process running only the two sweeps compiled exactly one decode kernel -per head count — one `...VarSeqQ16...` serving the four `H = 32` rank -values and one `...VarSeqQ8...` serving the four `H = 8` ones — where a -real selection axis (the head count) compiles once per value; see *Notes*. - -`None` is **not** an accepted value on the MLA path even though the -signature admits it: the C++ unwraps the optional unconditionally and the -call raises `RuntimeError: bad optional access`. Pass the checkpoint's rank, -or `0` when it has none. - -**Context phase — fresh prefill** (`attention_input_type=1`, separate -`q`/`k`/`v`, `is_fused_qkv=False`, `latent_cache` given): certified context -sequences of this flavor have no cached prefix, so token `i` of a sequence -sits at absolute position `i`. Per context sequence of length `L`, with -per-token latent row `latent_cache[t] = [ckv_t | k_pe_t]`: - -``` -# GPT-J (interleaved-pair) RoPE at position i, fp32 math, from the -# duplicated-layout rotary_cos_sin table: -# rope(x)[2d] = x[2d]*cos_d - x[2d+1]*sin_d -# rope(x)[2d+1] = x[2d]*sin_d + x[2d+1]*cos_d -q'[i, h] = [ q[i, h, :nope] | rope_i(q[i, h, nope:]) ] # written back into q -k'[i, h] = [ k[i, h, :nope] | rope_i(k_pe_i) ] # k_pe broadcast to all - # heads, written back into k -cache_row(seq, i) = [ ckv_i | rope_i(k_pe_i) ] # paged latent append -out[i, h] = softmax_causal(q'[i, h] @ K'.T * scale) @ V # headDim nope+R, - # headDimV v_head_dim -``` - -`output[:num_ctx_tokens]` is written. The rope slice of the incoming `k` -(`k[..., nope:]`) is ignored (the kernel fills it from `latent_cache`); -the nope slices of `q`/`k` are left bitwise intact, and `latent_cache` is -read-only. Fusion boundary (context): q_pe/k_pe RoPE, latent-cache append, -and causal FMHA happen inside; the caller owns the projections that build -`q`, `k_nope`, `v`, and `latent_cache`, and the output projection. - -**Context phase — no append** (`attention_input_type=1`, separate -`q`/`k`/`v`, `is_fused_qkv=False`, `latent_cache=None`): the in-kernel -RoPE and cache append are skipped entirely — the call is a pure masked -FMHA over caller-supplied K/V, used for context over a cached KV prefix -and for chunked-prefill partial passes. Per context sequence `s` with -`n = host_context_lengths[s]` q rows and `m = sequence_length[s]` K/V rows -(this call's KV range; `m` is independent of `n`): - -``` -out[i, h] = softmax(q[i, h] @ K[:j_i].T * scale) @ V[:j_i] # fp32 accum - j_i = (m - n) + i + 1 # mask_type == 1: causal, bottom-right aligned - # (query i sits at absolute position m - n + i) - j_i = m # mask_type == 0: padding (whole KV range) -``` - -and, when `softmax_stats_tensor` is given, per (q row `t`, head `h`) over -the same masked logit range (natural-log domain, fp32): - -``` -softmax_stats[t, h] = (max_j logit_j, sum_j exp(logit_j - max)) -``` - -Nothing is rotated, appended, or mutated except `output` (and the stats -tensor): `q` is consumed as-is (production pre-rotates its rope tails -upstream — a sibling op, `trtllm::mla_rope_append_paged_kv_assign_q`, -exists for that step and the latent append), `k` carries rope(k_pe) in its -per-head tails, and `q`, `k`, `v`, and the paged pool were verified -bitwise untouched by the call. This flavor is certified on pure-context -batches (`ns == num_contexts`). Two usage patterns are certified: - -- **One-shot cached-KV context** (causal): K/V cover each sequence's full - `[cached + new]` range (`sequence_length = cached + new >` q length), - one call computes the final context output. -- **Chunked partial passes**: per chunk loop, K/V cover one chunk slice of - the cached range per sequence — `sequence_length[s]` = that sequence's - chunk length, `0` allowed (the sequence sits the loop out; its output - and stats rows are then **undefined** and must be skipped downstream) — - with `mask_type=0` and stats emitted; then one final causal pass over - the new tokens only (`sequence_length == context_lengths`), stats - emitted. Each pass's output/stats pair is folded into the running - full-range result by a sibling merge op - (`trtllm::merge_chunked_attention_for_mla`) under the production - copy/merge/skip plan — verified end to end: the merged output and stats - equal single-pass attention over the full `[cached + new]` range. - -**Generation phase** (`attention_input_type=2`, `q` = absorbed "fused q", -`is_fused_qkv=True`, `k`/`v` `None`): a pure latent-MQA **read** of the -paged cache. With `P = predicted_tokens_per_seq`, each generation sequence -contributes `P` query rows — one for an ordinary decode step, `P > 1` for a -speculative-decoding step whose `P` rows are that sequence's draft chain. -The rows are **token-major within a sequence**: row `n` is sequence -`g = n // P`'s `t = n % P`-th token, and it sits at absolute position -`L_g - P + t`, where `L_g = sequence_length[num_contexts + g]` is the total -KV length **including all `P` of this step's tokens**. Per row: - -``` -K = cache rows [0, L_g - P + t] # [L_g-P+t+1, C+R] -V = K[:, :C] -out[n, h] = softmax(q[n, h] @ K.T * scale) @ V # [H, C] per token -``` - -`output[:G*P]` is written (`G` = generation-sequence count). At `P = 1` the -key range is the whole cache, which is the ordinary decode rule. At `P > 1` -the **within-block mask is causal and bottom-right aligned against the -sequence's own KV length**: draft row `t` sees the cache up to and including -its own position and **not** rows `t+1 .. P-1` of its own block, whose tokens -do not exist yet. - -That is measured rather than inferred, and by a readout rather than a -tolerance: with each cached latent row's `ckv` half set to a one-hot vector, -the decode's `V` is that one-hot and the output row becomes the attention -weight vector itself, so the attended key set can be read off column by -column. At `P` 1-4, `H = 128`, page 32, over both pool element types, every -row weighted exactly `[0, L_g - P + t]` and every other column came back -**bitwise zero** — including the 1 to 3 future-sibling columns a full -within-block mask would have weighted at 0.023-0.039 each. `mask_type` does -not change this (see its row in *Geometry and masking scalars*), and on -sm_100 no mask tensor is involved at all: for a linear-tree draft the Python -backend forces `is_spec_decoding_enabled` off there, so -`spec_decoding_packed_mask` and its siblings are `None` and -`predicted_tokens_per_seq` is the only thing producing the mask. - -No RoPE is applied and nothing is written to the cache: the generation-phase -preprocessing — the RoPE of `q_pe` at each row's own position, the latent -append of all `P` rows at positions `L_g - P .. L_g - 1`, and the -decode-scheduler-buffer fill — happens **before** this call (a sibling op, -`trtllm::mla_rope_generation`, exists for it). - -Where the roped `q_pe` has to **land** is pool-dependent, and the line above -describes only the bf16 case. Over a **bf16** latent pool this call reads its -query from `q`, so the preprocessing must leave the roped `q_pe` in fused q's -tail (`q[..., C:]`) — that is the layout certified here. Over an -**fp8-e4m3** pool this call does not read `q` at all: the query arrives -already quantized in `quant_q_buffer` (see *MLA over an fp8-e4m3 latent -pool*), so whether a producer ever materializes a roped tail in `q` is -outside what this op observes. What it requires on that path is only that -`quant_q_buffer` hold the right bytes. - -The `latent_cache` and `q_pe` arguments must be non-None (presence-checked) -but their data is **not consumed** — verified by passing garbage in both -while output, pool, and `q` stayed correct/untouched, at `P` 1 through 4; they -are not mutated either. The caller owns the absorbed-q BMM that fills -`q[..., :C]` and the `v_b_proj` BMM applied to the `[G*P, H*C]` output -afterwards. - -#### MLA over an fp8-e4m3 latent pool - -`quant_mode=128` switches the **latent pool's** element type to fp8-e4m3 -(one byte per element; the pool is `[num_blocks, tokens_per_block, C+R]` -e4m3) while `q`, `k`, `v`, `latent_cache`, `q_pe` and `output` all stay -bf16. Certified at the DeepSeek-R1-0528 layer geometry — `H = 128`, -`tokens_per_block = 32`, `C/R/nope/v = 512/64/128/128`, -`q_lora_rank = 1536`, causal mask, single-layer pool, one call per phase — -for **all three MLA call flavors** (fresh-prefill context, no-append context, -generation) and for the **mixed batch** that pairs two of them, at KV-cache -scaling factor `s = 1.0` (the production value: a -`kv_cache_quant_algo: FP8` checkpoint whose per-layer `k_scale`/`v_scale` -are 1.0). `s` enters as the same -two fp32 `[1]` CUDA tensors the standard configuration uses — -`kv_scale_orig_quant` = `1/s`, `kv_scale_quant_orig` = `s` — and passing -`None` for both is again silently the `s = 1.0` path (bitwise identical -output and pool, verified). `s` = 1.5 and 2.0 are swept beside 1.0 on the -fresh-prefill, generation and no-append flavors; the mixed batch runs -`s = 1.0` only. - -Two rope/scale cells are covered, and not uniformly. The **fresh-prefill** -and **generation** flavors run both: the unscaled theta-10000 table at -`q_scaling = 1.0`, and the complete DeepSeek-R1-0528 cell — that checkpoint's -YaRN-scaled `rotary_cos_sin` table together with -`q_scaling = 1/mscale² ≈ 0.53366` (see *In-kernel RoPE arguments*). The -**no-append** flavor, the **mixed batch** and the `quant_mode`-bit checks run -the R1 cell only. Both axes are certified over the bf16 pool as axes of their -own; the fp8 runs pin them **in combination with the pool**, which is what a -DeepSeek-R1 layer actually passes and what separate certifications do not -cover. Both reach the -fp8 path and both are gated: an fp8 R1-cell prefill sits 12.7x outside a -reference built at `q_scaling = 1.0` and 3.7x outside one built from the -unscaled table (the matching reference uses 0.40 of the same band); its -decode sits 4.2x outside the `q_scaling = 1.0` reference and the no-append -flavor 10.8x. The rope table is -pinned far harder by the append than by either output gate: the appended -`k_pe` half is bit-exact against a table-aware mirror while a mirror built -from the unscaled table differs in 2336 of the 6144 e4m3 bytes of a 96-token -request (38.0%, against a gate that allows 6). - -**Only the append resembles the standard configuration's fp8 surface.** The -two read paths are different code and different arguments; read them here, -not there. - -- **Append** (context call, both halves of the row): each new latent row - lands as `e4m3(row * kv_scale_orig_quant)` — fp32 multiply, round-to- - nearest cast. Bit-exact against that mirror at every scale run (1.0, 1.5, - 2.0), for the bitwise-copied `ckv` half **and** the in-kernel-roped `k_pe` - half alike: e4m3's 3-bit mantissa absorbs the fp32 evaluation-order - difference that costs the bf16 pool its ~8-in-500 000 one-ulp elements. - Pages outside the written rows stay bitwise zero, including when a - request's pages are scattered and out of order. -- **Context FMHA runs in fp8, not bf16 — in *both* context flavors.** The op - quantizes `q`, `k` and `v` to e4m3 itself and computes the masked FMHA on - those operands, so **prefill accuracy is affected by cache quantization** — - not because the context reads the pool back (the `s != 1.0` runs match a - model built from the caller's own bf16 operands quantized at 1.0, which is - not what the pool rows hold), but because the op quantizes its inputs. - This is not inherited from one flavor to the other: it was measured - separately on the fresh-prefill flavor (`latent_cache` given) and on the - no-append flavor (`latent_cache=None`), which appends nothing and whose K/V - arrive **already dequantized** from an fp8 pool — it quantizes them again - anyway. - Pinned bitwise on each: with the softmax - collapsed onto one key, the output row comes back as `e4m3(V_row)` - exactly (max abs deviation 0, 99.98% of the raw bytes equal — the rest - signed zeros), where a bf16 path would return `V_row`, 60-63 bf16 ulps - away. On - realistic inputs each flavor matches an fp32 reference over e4m3-rounded - q/k/v within `2^-3 + 2^-4 * |ref|` (fresh prefill 40-49% of it, no-append - 41%); the fresh-prefill run sits 26x (baseline cell) to 49.8x (R1 cell) - outside the bf16 band around the un-quantized reference — the mixed batch's - context call 58.1x — while the no-append run sits 2.2x - outside the *fp8* band around it. - The quantization uses **scale 1.0** — `kv_scale_orig_quant` is *not* - applied to `q`/`k`/`v` — while the FMHA nevertheless applies the - read-side dequantization factors that scale implies: `s^2` on the softmax - scale and `s` on the output. The two only cancel at `s = 1.0`, so - **both context flavors are correct at `s = 1.0` alone**. At `s` = 1.5 and - 2.0 the output is what the `s = 1.0` math would give with those two factors - applied (measured within 0.58x / 0.84x of the same band on the fresh flavor - and 0.86x on the no-append one), i.e. 21x-48x - away from the correct result, silently; the no-append peaked readout returns - exactly `s * e4m3(V_row)`, bit for bit, which is where the output factor is - read off directly. There is no caller-side knob that - fixes it: `quant_scale_qkv`, the argument that looks like it would set the - quantization scale, is inert here — a scratch probe passing `[1/s]` in it - at `s = 2.0` returned bitwise-identical output. -- **The no-append flavor still touches no pool.** Under `quant_mode=128` as - under `0`, that call reads and writes nothing but `output`: `q`, `k`, `v` - and the whole e4m3 pool came back bitwise intact from a call whose pages - held real cached content. Its operands in the certified run were built the - way a target builds them — the cached prefix written into the pool as - `e4m3(row * kv_scale_orig_quant)`, read back as - `bf16(float(cache_byte) * kv_scale_quant_orig)` and pushed through a - `kv_b_proj`-shaped matmul, with only the new tokens fresh — so the flavor - is certified on operands that have been through the pool, not just on - operands of the right shape. -- **Generation FMHA takes its query from `quant_q_buffer`, and its two - scales from `mla_bmm1_scale` / `mla_bmm2_scale`.** `q` is not read at all - on this path (replacing it with garbage left the output bitwise - identical), and neither kv scale tensor is read either (dropping both is - bitwise identical at `s = 2.0`, where a read-side dequant would be - glaring). Mechanically the kernel computes, over the **raw** e4m3 numbers - in `quant_q_buffer` and in the pool: - - ``` - out[g, h] = softmax2(qq[g, h] @ Kraw.T * mla_bmm1_scale[1]) - @ Kraw[:, :C] * mla_bmm2_scale[0] - # softmax2 = exp2-domain softmax, so mla_bmm1_scale[1] is the natural-log - # scale times log2(e); mla_bmm1_scale[0] is never read (zeroing it changes - # nothing, zeroing [1] flattens the softmax) - ``` - - The caller therefore owns the dequantization. Filling the three buffers - like this — - - ``` - quant_q_buffer = e4m3(fused_q * kv_scale_orig_quant) # [G*P, H, C+R] - x = 1 / (q_scaling * sqrt(nope + R)) * s**2 - mla_bmm1_scale = [x, x * log2(e)] # fp32 [2] - mla_bmm2_scale = [s] # fp32 [1] - ``` - - — makes the call compute exactly latent MQA over the e4m3 round trip of - both operands (`s^2` undoing the `1/s` on the query and on `K`, `s` the - one on `V = K[:, :C]`), at the plain softmax scale: - - ``` - q_hat = e4m3(q * 1/s) * s K_hat = pool_row * s V_hat = K_hat[:, :C] - out[n, h] = softmax(q_hat[n, h] @ K_hat.T / (q_scaling*sqrt(nope+R))) @ V_hat - ``` - - with `K_hat` cut to row `n`'s own key range when `predicted_tokens_per_seq` - is above 1 (see *Semantics*), and the two bmm scales unchanged by it — they - are per-batch scalars, so an fp8 MTP decode needs no per-draft-row - rescaling. This whole recipe was re-run at `P` = 1, 2, 3 and 4. - - and that holds at **every** scale certified — 1.0, 1.5 and 2.0 all matched - this reference within 18% of the fp8 band, the non-power-of-two scale - included. Leaving `s` out of the two bmm - scales — the mistake the standard configuration's "dequantized on read" - wording invites — lands 2.9x (`s = 1.5`) and 3.6x (`s = 2.0`) outside it. - All three buffers are presence-checked: `None` raises - `Assertion failed: quant_q_buf is nullptr. (attentionOp.cpp:1050)`, - `bmm1_scale is nullptr. (:1051)`, `bmm2_scale is nullptr. (:1052)`. - Their *contents* are unchecked, and `quant_q_buffer` is read through a raw - pointer (an identical `uint8` view of the same bytes behaves identically). - The generation call still writes nothing: the pool comes back bitwise - unchanged. - - That block is a recipe for **this** op's caller — the certified runs passed - exactly those bytes, built in plain torch — and not a description of how a - producer arrives at them. In particular `quant_q_buffer` is written in terms - of the finished bf16 fused q this call would have read on the bf16 path; - a producer may instead assemble it from the absorbed `q[..., :C]` plus a - freshly roped `q_pe`, and never materialize a roped tail in `q` at all - (which is what the sibling preprocessing op does here). Both routes have to - land on the same bytes; only the bytes are this op's business. - -**Quantization bits outside the KV-cache group ride along unread.** A target -derives `quant_mode` from its checkpoint's quant config and will not get a -bare `128`: an nvfp4 checkpoint served with an fp8 KV cache produces `1152` -(`FP8_KV_CACHE | FP8_1x128_128x128`), and fp8-QDQ weights add `FP8_QDQ` for -`384`. Both were run against a bare `128` on all three MLA flavors at the R1 -cell and are **bit-identical** — output and pool alike. Nothing validates the -extra bits in either direction; this is a measurement, not a guarantee about -bits that were not tried (`64` = INT8 KV and `8192` = NVFP4 KV were not tried -on this op). - -The generation flavor is additionally certified over this pool at -`predicted_tokens_per_seq` = 1, 2, 3 and 4 — one speculative-decoding step per -call — including the mixed batch that pairs it with a context call. See the -`predicted_tokens_per_seq` and `mask_type` rows of *Geometry and masking -scalars* and the generation block of *Semantics*; the context flavors run at -1 only. - -Not certified over the fp8 latent pool: the no-append flavor's **chunked -partial-pass pattern** (`mask_type=0`, `softmax_stats_tensor`, zero-KV -sequences — only its one-shot cached-KV pattern is certified), -`tokens_per_block = 64`, head counts other than 128, multi-layer latent -pools, `predicted_tokens_per_seq` above 4, and `s != 1.0` in either context -flavor (measured-wrong, see above, rather than untested). - -## Signature - -```python -def thop_attention( - q, k, v, output, output_sf, workspace_, - sequence_length, host_past_key_value_lengths, host_total_kv_lens, - context_lengths, host_context_lengths, host_request_types, - max_context_q_len_override, - kv_cache_block_offsets, host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping, cache_indirection, - kv_scale_orig_quant, kv_scale_quant_orig, out_scale, - rotary_inv_freq, rotary_cos_sin, latent_cache, q_pe, - block_ids_per_seq, attention_sinks, - is_fused_qkv, update_kv_cache, predicted_tokens_per_seq, - local_layer_idx, num_heads, num_kv_heads, head_size, tokens_per_block, - max_num_requests, max_context_length, max_seq_len, - attention_window_size, beam_width, mask_type, quant_mode, q_scaling, - position_embedding_type, rope_dim, rope_base, rope_scale_type, - rope_scale, rope_short_m_scale, rope_long_m_scale, rope_max_positions, - rope_original_max_positions, use_paged_context_fmha, - attention_input_type, is_mla_enable, chunked_prefill_buffer_batch_size, - q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, - v_head_dim, rope_append, mrope_rotary_cos_sin, mrope_position_deltas, - helix_position_offsets, helix_is_inactive_rank, attention_chunk_size, - softmax_stats_tensor, is_spec_decoding_enabled, use_spec_decoding, - is_spec_dec_tree, spec_decoding_generation_lengths, - spec_decoding_position_offsets_for_cpp, spec_decoding_packed_mask, - spec_decoding_bl_tree_mask_offset, spec_decoding_bl_tree_mask, - spec_bl_tree_first_sparse_mask_offset_kv, - sparse_kv_indices, sparse_kv_offsets, sparse_attn_indices, - sparse_attn_offsets, sparse_attn_indices_block_size, - # keyword tail, defaults mirror the binding: - num_sparse_topk=None, sparse_attn_kv_lens=None, - skip_softmax_threshold_scale_factor_prefill=None, - skip_softmax_threshold_scale_factor_decode=None, skip_softmax_stat=None, - cu_q_seqlens=None, cu_kv_seqlens=None, fmha_scheduler_counter=None, - mla_bmm1_scale=None, mla_bmm2_scale=None, quant_q_buffer=None, - flash_mla_tile_scheduler_metadata=None, flash_mla_num_splits=None, - sage_attn_num_elts_per_blk_q=0, sage_attn_num_elts_per_blk_k=0, - sage_attn_num_elts_per_blk_v=0, sage_attn_qk_int8=False, - num_contexts=0, num_ctx_tokens=0, trtllm_gen_jit_warmup=False, - aux_kv_cache_pool_ptr=None, is_cross=False, cross_kv=None, - relative_attention_bias=None, relative_attention_max_distance=0, - spec_decoding_target_max_draft_tokens=None, quant_scale_qkv=None, - dsv4_inv_rope_cos_sin_cache=None, enable_dsv4_epilogue_fusion=False, -) -> None -``` - -Exact per-argument types are in the wrapper; the binding is keyword-callable -and the wrapper forwards every argument by keyword. - -With `Hq` = `num_heads`, `Hkv` = `num_kv_heads`, `D` = `head_size`, -`ns` = batch size, `T` = total new tokens in the batch -(`num_ctx_tokens + (ns - num_contexts) * predicted_tokens_per_seq`). `T` is -the standard configuration's `q`/`output` row count; an **MLA** batch is -served by one call per phase, so its context call has `num_ctx_tokens` rows -and its generation call `(ns - num_contexts) * predicted_tokens_per_seq`. -Certified MLA dims (DeepSeek-V3 / R1 geometry): `C = kv_lora_rank = 512`, -`R = qk_rope_head_dim = 64`, `nope = qk_nope_head_dim = 128`, -`v_head_dim = 128`, at `H` = 8, 16, 32 **or** 128 query heads — a -plain runtime parameter for the context call as in the standard -configuration, and a decode-kernel selection axis for the generation call -(see *MLA head count*); every MLA shape below is written in terms of `H`. - -### Data tensors — standard configuration - -| Argument | Shape / value | Dtype | Device | -|---|---|---|---| -| `q` | `[T, (Hq + 2*Hkv) * D]`, packed per-token `q \| k \| v`, contiguous (both pool dtypes: activations stay bf16) | bf16 | CUDA | -| `k`, `v` | `None` — K/V ride inside packed `q` (`is_fused_qkv=True`; separate-KV mode not certified for this configuration) | — | — | -| `output` | `[T, Hq * D]`, contiguous; rows `[:T]` overwritten in place (both pool dtypes) | bf16 (same as `q`) | CUDA | -| `kv_scale_orig_quant` | `quant_mode=0`: `None`; `quant_mode=128`: `[1]` holding `1/s` — **required for any `s != 1.0`** (`None` is silently read as `1.0`, never rejected) | fp32 | CUDA | -| `kv_scale_quant_orig` | `quant_mode=0`: `None`; `quant_mode=128`: `[1]` holding `s` — **required for any `s != 1.0`** (same silent `1.0` fallback) | fp32 | CUDA | -| `attention_sinks` | `None` (no sink), or exactly `Hq` values, **contiguous** — one sink logit per query head (see *Attention sinks*; certified on the bf16 pool with `mask_type=1`). Shape beyond the element count is ignored (`[Hq]`, `[1, Hq]`, `[Hq, 1]` all behave identically) | fp32 | CUDA | -| `latent_cache`, `q_pe` | `None` | — | — | -| `cu_q_seqlens`, `cu_kv_seqlens`, `fmha_scheduler_counter` | `None` (the op fills internal ones when absent) | — | — | - -### Data tensors — MLA context call, fresh prefill (`latent_cache` given) - -`Tc` = `num_ctx_tokens` (this call's `q` rows), `Hq == Hkv == H`, -`head_size = nope + R` (192 certified). - -| Argument | Shape / value | Dtype | Device | -|---|---|---|---| -| `q` | `[Tc, H * (nope+R)]`, contiguous; rope slice of each head **overwritten in place** with rope(q_pe) | bf16 | CUDA | -| `k` | `[Tc, H * (nope+R)]`, contiguous; incoming rope slice ignored (may be uninitialized) and **overwritten in place** with rope(k_pe); nope slice consumed as K | bf16 | CUDA | -| `v` | `[Tc, H * v_head_dim]` row view whose **row stride is `H * (nope + v_head_dim)` elements** — the `[.., H*nope:]` split of a packed `[Tc, H*(nope+v)]` kv-projection buffer (certified); the context FMHA hard-codes this stride, so a fully contiguous `[Tc, H*v_head_dim]` tensor is misread (see *Notes*) | bf16 | CUDA | -| `output` | `[Tc, H * v_head_dim]`, contiguous; rows `[:Tc]` overwritten | bf16 | CUDA | -| `latent_cache` | `[Tc, C + R]` per-token `[ckv \| k_pe]`, contiguous; read-only | bf16 | CUDA | -| `q_pe` | `None` | — | — | -| `kv_scale_orig_quant`, `kv_scale_quant_orig` | `quant_mode=0`: `None`; `quant_mode=128`: `[1]` holding `1/s` and `s`. Only the append consumes `1/s`; the context FMHA applies `s^2`/`s` to a computation it quantized at 1.0, so **`s = 1.0` (or `None`) is the only correct value here** — see *MLA over an fp8-e4m3 latent pool* | fp32 | CUDA | -| `mla_bmm1_scale`, `mla_bmm2_scale`, `quant_q_buffer` | `None` (the context call neither checks nor reads them at either `quant_mode`) | — | — | -| `softmax_stats_tensor` | `None` (stats certified only for the no-append flavor) | — | — | -| `cu_q_seqlens`, `cu_kv_seqlens`, `fmha_scheduler_counter` | `None` (the op fills internal ones when absent) | — | — | - -### Data tensors — MLA context call, no append (`latent_cache=None`) - -`Tc` = `num_ctx_tokens` (this call's `q` rows); `Tkv` = sum of -`sequence_length[s]` over the context sequences = total K/V rows, the -per-sequence KV ranges concatenated in batch order. - -| Argument | Shape / value | Dtype | Device | -|---|---|---|---| -| `q` | `[Tc, H * (nope+R)]`, contiguous; consumed as-is (rope tails prepared upstream) and **not mutated** | bf16 | CUDA | -| `k` | `[Tkv, H * (nope+R)]`, contiguous, read-only; per-head tails hold rope(k_pe) broadcast to all heads (caller-built) | bf16 | CUDA | -| `v` | `[Tkv, H * v_head_dim]` row view, read-only; **row stride `H * (nope + v_head_dim)` elements required** (same hard-coded stride as the fresh flavor — certified as the split of packed `[Tkv, H*(nope+v)]` buffers; contiguous rows are misread) | bf16 | CUDA | -| `output` | `[Tc, H * v_head_dim]`, contiguous; rows `[:Tc]` overwritten (rows of zero-KV sequences: undefined) | bf16 | CUDA | -| `latent_cache` | `None` — selects this flavor (skips in-kernel RoPE + append; the paged pool is neither read nor written, at either `quant_mode`) | — | — | -| `q_pe` | `None` | — | — | -| `kv_scale_orig_quant`, `kv_scale_quant_orig` | `quant_mode=0`: `None`; `quant_mode=128`: `[1]` holding `1/s` and `s`. This flavor appends nothing, so `1/s` reaches nothing at all, while the FMHA still applies `s^2`/`s` to a computation it quantized at 1.0 — **`s = 1.0` (or `None`) is the only correct value here**, exactly as for the fresh-prefill flavor. See *MLA over an fp8-e4m3 latent pool* | fp32 | CUDA | -| `softmax_stats_tensor` | `None`, or `[>= Tc, >= H, 2]` contiguous, rows `[:Tc]` overwritten per (token, head) with `(max, sum)` as in *Semantics* (rows of zero-KV sequences: undefined). Presence-checked by the op: fp32, dim 0 `>=` the call's q rows, dim 1 `>= num_heads`, dim 2 `== 2`. Certified on the bf16 pool only | fp32 | CUDA | -| `mla_bmm1_scale`, `mla_bmm2_scale`, `quant_q_buffer` | `None` (this flavor neither checks nor reads them at either `quant_mode`) | — | — | -| `cu_q_seqlens`, `cu_kv_seqlens`, `fmha_scheduler_counter` | `None` (the op fills internal ones when absent) | — | — | - -### Data tensors — MLA generation call - -`G` = generation-sequence count, `P` = `predicted_tokens_per_seq` (query rows -per generation sequence: 1 for an ordinary decode step, `P > 1` for a -speculative-decoding step — 1, 2, 3 and 4 certified), `head_size = C + R` -(576 certified), `Hkv = 1`. Every per-row tensor is `G*P` rows tall and -**token-major within a sequence**: row `n` is sequence `n // P`'s -`n % P`-th token. - -| Argument | Shape / value | Dtype | Device | -|---|---|---|---| -| `q` | `[G*P, H * (C+R)]` fused q (absorbed q_nope in `[..., :C]`, rope(q_pe) in `[..., C:]`, both prepared before the call — at `P > 1` each row must have been roped at **its own** absolute position `L_g - P + n % P`; this call applies no rope and cannot tell), contiguous; not mutated. Under `quant_mode=128` it is **not read either** — the query comes from `quant_q_buffer` below — but it is still the tensor whose shape and dtype the call is configured around | bf16 | CUDA | -| `k`, `v` | `None` (`is_fused_qkv=True`) | — | — | -| `output` | `[G*P, H * C]`, contiguous; rows `[:G*P]` overwritten (per-head latent output — `v_b_proj` still to apply). A taller buffer's rows past `G*P` are left bitwise untouched | bf16 | CUDA | -| `latent_cache` | non-None required; **data not consumed and not mutated** (this step's rows must already be in the pool). Certified passing `[G*P, C + R]`, the production shape; the row count is not checked (`[G, …]` and `[1, …]` were also accepted, bitwise identically) | bf16 | CUDA | -| `q_pe` | non-None required; **data not consumed and not mutated**. Certified passing `[G*P, H, R]`, same unchecked-row-count caveat | bf16 | CUDA | -| `cu_q_seqlens` | `[G+1]` **required** (presence only): `arange(G+1) * H * P` (q-row prefix sums — this is what the sibling preprocessing op writes at every `P`) | int32 | CUDA | -| `cu_kv_seqlens` | `[G+1]`, `[0, cumsum(L_g)]` over generation sequences — **neither checked nor read on this path** | int32 | CUDA | -| `fmha_scheduler_counter` | `[1]` **required** (presence only), zeroed before the call | uint32 | CUDA | -| `quant_q_buffer` | `quant_mode=0`: `None`; `quant_mode=128`: `[G*P, H, C+R]` **required and read as the query** — `e4m3(fused_q * kv_scale_orig_quant)`, contiguous; `q` itself is then unread. Read through a raw pointer, so an equal-bytes `uint8` view behaves identically — and so a buffer shorter than `G*P` rows is read past its end rather than rejected | e4m3 | CUDA | -| `mla_bmm1_scale` | `quant_mode=0`: `None`; `quant_mode=128`: `[2]` **required**, `[x, x*log2(e)]` with `x = s^2 / (q_scaling * sqrt(nope+R))`. Element `[1]` is the one read. Per batch, **not** per `P` — the same pair serves every draft row | fp32 | CUDA | -| `mla_bmm2_scale` | `quant_mode=0`: `None`; `quant_mode=128`: `[1]` **required**, `[s]` — multiplies the output | fp32 | CUDA | -| `kv_scale_orig_quant`, `kv_scale_quant_orig` | **not read on this path at either `quant_mode`** (bitwise-identical output with both `None` at `s = 2.0`); pass them anyway if the same tensors serve the context call | fp32 | CUDA | - -`None` fails for two of the three — `cu_q_seqlens` (`seqQOffset is nullptr`, -`attentionOp.cpp:1045`) and `fmha_scheduler_counter` (`fmha_tile_counter is -nullptr`, `:1047`) — and both are presence-checked only, their contents -bitwise inert here. `cu_kv_seqlens` is accepted as `None`, and `None`, -all-zeros, non-monotonic and 1000x-inflated values all returned bitwise -identical output at `G = 4` and `G = 16` with pairwise-distinct `L_g`: the -paged decode kernel takes each sequence's KV length from `sequence_length` -(a one-token change there does move the output, from 0.016 to 0.499 max abs -error). Fill all three as the sibling generation-preprocessing op does — -that is the production shape and the certified one. - -**The scheduler buffers stay inert at `P > 1`** — the taller query block is -exactly the change that could have made them start mattering, so all of it -was re-measured at `P = 3`, `G = 2`, over the fp8 R1 cell. `cu_q_seqlens` -zeroed, left in the `P = 1` form `i*H`, given `i*P` without the head factor, -inflated 1000x, and made non-monotonic all returned **bitwise-identical** -output; so did `cu_kv_seqlens` zeroed, inflated and non-monotonic. Pass the -production values anyway: they are what the sibling op writes, and nothing -here promises the next version reads them the same way. - -### Shared data tensors - -| Argument | Shape / value | Dtype | Device | -|---|---|---|---| -| `output_sf` | `None` (NVFP4-output path not certified) | — | — | -| `workspace_` | persistent scratch tensor, any length (0 ok); the op grows it **in place** via `resize_()` when too small — 33 MB to 109 MB across the certified shapes. The MLA context requirement scales with tokens x heads and still sets the peak: 113 924 096 B for an `H = 128` context call over 545 tokens, against 52 579 072 B for the same call over 129 tokens and 100 024 320 B for the standard-configuration head-geometry sweep. The MLA **generation** call is well below that and, measured, does not grow with `predicted_tokens_per_seq`: 39 100 416 B at `H = 128`, page 32 for every `(G, P, L)` tried — `G` 2 and 4, `P` 1 and 4, `L` 50 and 500 — so a `P`-times-taller query block moves nothing here. Pass a plain resizable tensor, not a view; reuse it across calls to avoid re-allocation | int8 | CUDA | - -### Batch state (the "prepared metadata" of this op — all int32) - -| Argument | Shape | Content | Device | -|---|---|---|---| -| `sequence_length` | `[ns]` | per-sequence KV length **of this call**: total `cached + new` everywhere except an MLA chunked partial pass, where a context row holds that sequence's chunk length instead (0 = sits this loop out). For a generation row at `predicted_tokens_per_seq = P` the `new` part is all `P` of this step's tokens, and this tensor — not its host twin — is what the MLA generation call takes its per-sequence key range from (shifting it by one token moved every draft row's attended set by one, read out directly) | CUDA | -| `host_past_key_value_lengths` | `[ns]` | same values as `sequence_length`. On the MLA generation path the per-sequence values are otherwise inert, but the vector must not be **all** zero — see *Preconditions* | CPU | -| `host_total_kv_lens` | `[2]` | `[0]` = sum of `sequence_length` over context sequences (= the K/V row count for a no-append MLA context call), `[1]` = sum over generation sequences | CPU | -| `context_lengths` | `[ns]` | context sequences: this call's new-token (q-row) count — equals the prompt length only when nothing is cached; generation sequences: the original prompt length. On a context row it is load-bearing and unchecked (see *Preconditions*); on a generation row it is inert under a sliding window (0, 1 and the prompt length gave **bitwise identical** output) and inert on the MLA generation path at `predicted_tokens_per_seq` above 1 as well (0 and the full KV length both bitwise identical at `P = 3`) | CUDA | -| `host_context_lengths` | `[ns]` | same values as `context_lengths` | CPU | -| `host_request_types` | `[ns]` | 0 = context, 1 = generation; all 0s must precede all 1s | CPU | -| `num_contexts` (int) | — | number of context sequences. Binding default 0 is only correct for generation-only batches — always pass the true count | — | -| `num_ctx_tokens` (int) | — | sum of new-token counts over context sequences; same caveat | — | -| `max_context_q_len_override` | — | `None` (encoder CUDA-graph override; not certified) | — | - -CPU tensors may be pageable or pinned (pinning is a perf optimization). -For a mixed MLA batch, the two per-phase calls share these full-batch -(`ns`-sized) tensors and the same `num_contexts`/`num_ctx_tokens`; only -`q`/`output` and the phase-specific arguments differ between the calls. -The generation call indexes sequences starting at `num_contexts`. - -### Paged KV cache addressing - -The cache is a caller-owned pool addressed in **slab** units: one slab = -one layer's K (or V) block, `tokens_per_block * Hkv * D` elements -(standard) or `tokens_per_block * (C + R)` (MLA). A pool page holds the -slabs of every layer sharing the pool, layers side by side. Certified: -one pool, holding 1 layer (both configurations) or 4 layers (standard -configuration, bf16 and fp8-e4m3 pools both; multi-layer MLA pools not -certified). The certified multi-layer state — pool, mapping, block -offsets — was produced by a real 4-layer `KVCacheManager` (one per pool -dtype) and consumed unmodified, so the shapes below are exactly what a -multi-layer manager emits. - -Standard configuration (kv_factor 2, HND), `L` layers in the pool: pool -tensor `[num_blocks, L, 2, Hkv, tokens_per_block, D]` CUDA contiguous — -element type bf16 (`quant_mode=0`) or fp8-e4m3 (`quant_mode=128`) — -page `b` holds layer `l`'s K slab at -`[b, l, 0]` and its V slab at `[b, l, 1]` (single layer is the `L = 1` -case, `[num_blocks, 2, Hkv, tokens_per_block, D]`). The op sizes slabs -from `quant_mode`, never from the tensor (it sees only a raw pointer): a -pool whose dtype disagrees with `quant_mode` is silently corrupted. - -MLA (kv_factor 1, single layer): pool tensor -`[num_blocks, tokens_per_block, C + R]` CUDA contiguous — element type bf16 -(`quant_mode=0`) or fp8-e4m3 (`quant_mode=128`) — one latent row per token, -no separate K/V slabs. The slab width is `tokens_per_block * (C + R)` -**elements** regardless of the per-call `num_kv_heads`/`head_size` (context -passes `H`/192, generation `1`/576; both address the same pool), so the slab -is that many bytes under `quant_mode=128` and twice that under `quant_mode=0` -— the op derives the element width from `quant_mode` alone. Certified at -both element types with the requests' pages scattered out of order, every -page outside the reserved sets left bitwise zero. How many *bytes per token* -a KV-cache manager reserves for such a pool is the manager's business, not -this op's: the op is handed a base pointer plus slab offsets and addresses -exactly the geometry above. - -| Argument | Shape / value | Dtype | Device | -|---|---|---|---| -| `kv_cache_block_offsets` | `[1, max_num_requests, 2, max_blocks_per_seq]`; row `[0, s, 0, j]` = K-slab offset of sequence `s`'s `j`-th page, `[0, s, 1, j]` = V-slab offset. Offsets count slabs and are **layer-agnostic** — every layer of the pool shares one offsets row per sequence; layer selection comes from the pool mapping, not the offsets. Standard `L`-layer pool → page `p` has K `p * 2L`, V `p * 2L + 1` (single layer: K `2p`, V `2p + 1`); MLA pool → raw block id `p` in **both** rows. Only the first `ceil(kv_s / tokens_per_block)` entries per row are read | int32 | CUDA | -| `host_kv_cache_pool_pointers` | `[num_pools, 2]`: row `p` is pool `p`'s (base address = `pool.data_ptr()`, secondary-pool address; 0 = none). Certified at `num_pools` 1 and 2 | int64 | CPU | -| `host_kv_cache_pool_mapping` | `[num_layers, 2]`, one row per layer: row `local_layer_idx` = (pool index, layer index within pool). The row's **layer column drives the pool-base shift** — the call addresses slabs starting `layer_in_pool * kv_factor` slabs from the pool base (see *Notes*). The **pool column selects which base** to shift from. Certified: the identity rows a real manager produces — `[[0, 0] .. [0, 3]]` for the 4-layer standard pool, `[1, 2]` zeros for single-layer pools, and `[[0, 0], [0, 1], [1, 0], [1, 1]]` for two pools of two layers each, which is the shape KVCacheManagerV2 produces when a checkpoint's layers fall into more than one attention-window class | int32 | CPU | -| `local_layer_idx` (int) | row of `host_kv_cache_pool_mapping` for this layer (production: the layer's index among the rank's local layers); 0-3 certified (standard, one pool and two), 0 (MLA) | — | — | -| `tokens_per_block` (int) | pool page size; must be a power of two; 32 certified (standard), 32 and 64 certified (MLA — see *MLA page size*) | — | — | -| `update_kv_cache` (bool) | `True` (also for the MLA generation and no-append context calls, which nevertheless write nothing) | — | — | -| `cache_indirection` | `None` (beam search only) | — | — | -| `block_ids_per_seq` | `None` (FlashMLA path only) | — | — | - -A sequence's page set is shared by every layer of *its own* pool: the -per-layer calls of one batch pass identical offsets and pool pointers and -differ only in `local_layer_idx`. - -Two pools are certified. A layer's pool comes from the mapping row's pool -column, which selects both the pointer row and the offsets slot; layers in -different pools share nothing. The certified case drives every layer of a -two-pool, two-layers-each layout through prefill and decode and asserts that -each call leaves the *other* pool bitwise untouched -- collapsing the pool -selection to 0 still produces plausible outputs, since both pools hold validly -shaped pages, so only that comparison detects it. `num_pools > 2` is -uncertified, as are pools that differ from each other in geometry. - -### Paged KV cache addressing under a sliding window - -`attention_window_size` changes **nothing** about addressing: token `t` of a -sequence always lives at entry `t // tokens_per_block` of that sequence's -offsets row, slot `t % tokens_per_block`, and the append writes there. Every -memory saving is the caller's page mapping. What the op requires: - -- The offsets row is indexed by **absolute** page index, so it must have at - least `ceil(total_kv / tokens_per_block)` entries — that count keeps - growing with the sequence, window or no window, and bounds - `kv_cache_block_offsets.shape[3]`. -- Only entries holding at least one **in-window** token need to point at - valid live pages. An entry whose tokens are all older than - `total_kv - W` is never read: pointing it at a page filled with garbage - left a decode's output **bitwise identical** (checked entry by entry over - a 7-page sequence with `W = 128`; the first entry holding an in-window - token does change the output). Aged-out pages may therefore be recycled, - and stale ids may stay in the row — as long as they remain in-pool slab - offsets (out-of-pool ids were not tried and must not be). -- The same holds for a **context** call on the paged-context path, with the - in-window range measured from that call's rows: query row `i` of a - context sequence reads keys `[max(0, p - W + 1), p]` at - `p = kv_s - l_s + i`, so the call's read span is - `[max(0, kv_s - l_s + 1 - W), kv_s - 1]`. Measured entry by entry over an - 8-page sequence (`W = 128`, 200 cached tokens, 40 new): every entry - holding an in-window **cached** token changed the output (318x outside - the band with a decoy there), while entries below the window and the one - entry holding only this call's own new tokens left it **bitwise - identical** — that last one because the append lands wherever the offsets - point and the FMHA reads it straight back. The ring bound that follows - from that span (not separately certified) is - `P * tokens_per_block >= W + l_s - 1` for a cached-prefix context call, - on top of the decode bound below. -- A **bounded ring** is the resulting production shape: give the sequence - `P` physical pages and map absolute page index `j` to `ring[j % P]`. The - op's absolute-position append then lands token `t` at physical slot - `t mod (P * tokens_per_block)`, so the pool holds a moving window of the - sequence and genuinely wraps, overwriting exactly the slot of the token - that has just aged out. Requires - `P * tokens_per_block >= attention_window_size`; certified at - `P * tokens_per_block` of 128 (= `W`) and 192 over 250 wrapping decode - steps. One page short (96 < 128) is **silently wrong** — no exception, - output far off the reference. -- The same call's new tokens must map to **distinct** physical slots: - `P * tokens_per_block >= that call's new-token count`. A prefill longer - than the ring makes several tokens of one append target one slot; those - writes race (observed: two identical runs left different pool bytes, - torn rows). So a sliding-window sequence still needs - `ceil(prompt_len / tokens_per_block)` distinct pages for its prefill call - — the window shrinks the steady-state decode footprint, not the prefill's. -- The append remains surgical: one decode step changes exactly one - `(K, V)` row pair in the whole pool, the evicted slot's, and touches - nothing else. No compaction, zeroing, or relocation ever happens. - -### Geometry and masking scalars - -| Argument | Standard (certified) | MLA context | MLA generation | -|---|---|---|---| -| `is_fused_qkv` | `True` (packed QKV in `q`) | `False` | `True` | -| `attention_input_type` | 0 = mixed (either subset may be empty; 1/2 not certified here) | 1 = context_only | 2 = generation_only | -| `num_heads`, `num_kv_heads`, `head_size` | must match packed `q` and pool. Head counts are a free axis subject to `num_heads % num_kv_heads == 0`; `head_size` is enumerated by the available FMHA kernels (see *Head geometry*). Certified on the bf16 pool: 32/8/128, 32/4/128, 64/8/64, 16/1/128, 28/4/128, 12/3/128, 8/2/128, 8/2/256, 4/4/64; on the fp8-e4m3 pool: 32/8/128 | `H`/`H`/192 (`H`/`H`/`nope+R`); 8/8/192, 16/16/192, 32/32/192 and 128/128/192 certified on the bf16 latent pool, 128/128/192 only on the fp8-e4m3 one | `H`/1/576 (`H`/1/`C+R`); 8/1/576, 16/1/576, 32/1/576 and 128/1/576 certified on the bf16 latent pool, 128/1/576 only on the fp8-e4m3 one | -| `mask_type` | 1 = causal, 0 = padding — both certified on the bf16 pool (padding on context-only batches); fp8 pool certified with 1 only | 1 = causal (bottom-right aligned when KV > q); 0 = padding certified for the no-append chunked partial passes, on the bf16 latent pool only | **Inert — this path takes its mask from `predicted_tokens_per_seq` alone.** Pass 1. At `P = 1` a decode row attends to all `L_g` cached tokens; at `P > 1` draft row `t` attends to `[0, L_g - P + t]`, bottom-right-aligned causal within the block. `mask_type = 0` returns **bitwise-identical** output at `P` = 2 and 4 (fp8, `H = 128`, page 32) — the plausible guess that padding drops the within-block mask and lets a draft row see its later siblings is wrong | -| `q_lora_rank` | `None` | the checkpoint's `q_a_proj` rank — `1536` and `0` (no q-LoRA) certified and **inert**: no MLA path reads it (see *MLA q-LoRA rank*). `None` raises here, unlike in the standard configuration | same value as context | -| `kv_lora_rank`, `qk_nope_head_dim`, `qk_rope_head_dim` | all `None` | 512 / 128 / 64 certified | same values as context | -| `v_head_dim` | `None` | 128 (context V/output head width) | 512 (= `C`, latent output width) | -| `rope_append` | `None` | `True` | `True` (`False` widens the output to `C+R` per source — not certified) | -| `predicted_tokens_per_seq` | 1 | 1 | **1, 2, 3 and 4 certified** (`P`): the number of query rows each generation sequence contributes. 1 is an ordinary decode step; `P > 1` is one speculative-decoding step verifying a `P`-token draft chain, and `P` is then the *only* thing that produces the within-block mask (see `mask_type`). One scalar for the whole batch — a ragged per-sequence draft length cannot be expressed here. Certified over the fp8-e4m3 latent pool at the complete DeepSeek-R1-0528 cell and over a bf16 latent pool, both at `H = 128`, page 32; above 4, and at other head counts or page sizes, untested. `P` also reaches decode-kernel selection as the kernel's `maxSeqLenQ` — see *Notes* | -| `q_scaling` | softmax scale `1 / (q_scaling * sqrt(head_size))`; 1.0 certified | free positive scalar; scale uses `sqrt(nope+R)`. Certified 1.0, 0.53366 (`= 1/mscale²`, DeepSeek-R1's YaRN temperature — see *In-kernel RoPE arguments*), 2.0 and 0.25 on the bf16 latent pool at `H = 128`; 1.0 and 0.53366 on the fp8-e4m3 one, in both context flavors | same values; the scale still uses `sqrt(nope+R)`, **not** `sqrt(head_size)`. Over an fp8 pool it reaches the kernel only through the caller's `mla_bmm1_scale` (see *MLA over an fp8-e4m3 latent pool*) | -| `quant_mode` | 0 = bf16 pool; 128 = fp8-e4m3 pool (see *Standard configuration over an fp8-e4m3 KV pool*; other bits not certified) | 0, or 128 for an fp8-e4m3 latent pool. `1152` (`\| FP8_1x128_128x128`) and `384` (`\| FP8_QDQ`) are bit-identical to `128` here — only the KV-cache bit is read (see *MLA over an fp8-e4m3 latent pool*) | same value as context | -| `beam_width` | 1 | 1 | 1 | -| `max_num_requests` | `>= ns`; equals `kv_cache_block_offsets.shape[1]` | same | same | -| `max_context_length` | host launch/workspace bound `>=` max per-sequence q-row count of the call (certified `= max_seq_len`) | same | same | -| `max_seq_len` | `>=` max `sequence_length`; per-sequence page capacity bound. Does not select what a query attends to: with a window active, decodes at `max_seq_len` of `W`, `total + 1` and `4 * (total + 1)` returned **bitwise identical** output (it does steer kernel selection — see *Notes*) | same | same | -| `attention_window_size` | sliding-window token count `W`; `>= max_seq_len` = no window. Smaller = each query attends only the newest `W` keys (see *Sliding window*); 128 / 100 / 33 certified on the bf16 pool with `mask_type=1`. Per-call scalar: layers sharing a pool may differ in it | same (no window certified) | same (no window certified) | -| `use_paged_context_fmha` | context execution path; both values certified. `False` = context FMHA over the packed `q` rows, so every context sequence must start empty (`sequence_length == host_context_lengths`). `True` = context FMHA over the paged pool, which is what a context call over a cached prefix requires (`sequence_length > host_context_lengths`) — see *Paged-context FMHA*. With nothing cached the two are bitwise identical | `False` (MLA context always takes separate q/k/v) | `False` | -| `chunked_prefill_buffer_batch_size` | 1 | 1 | 1 | -| `attention_chunk_size` | `None` | `None` | `None` | -| `is_cross` / `cross_kv` | `False` / `None` | same | same | -| `trtllm_gen_jit_warmup` | `False` | `False` | `False` | - -### In-kernel RoPE arguments - -Standard configuration: RoPE disabled — `position_embedding_type=0` -(learned-absolute; no rotation applied), `rotary_inv_freq=None`, -`rotary_cos_sin=None`, `rope_dim=0`, `rope_base=10000.0`, -`rope_scale_type=0`, `rope_scale=1.0`, `rope_short_m_scale=1.0`, -`rope_long_m_scale=1.0`, `rope_max_positions=1024`, -`rope_original_max_positions=1024`, `mrope_rotary_cos_sin=None`, -`mrope_position_deltas=None`. Non-zero `rope_dim` with a rope-type -`position_embedding_type` applies RoPE to q/k inside the call — not -certified for this configuration. - -MLA (all calls receive the same values; only the fresh-prefill context -call rotates anything): `position_embedding_type=8` (yarn, as MLA models -set), `rope_dim=R`, and `rotary_cos_sin` = fp32 -`[1, max_positions * R * 2]` duplicated-layout table — per position, `R` -`(cos, sin)` pairs at flat offsets `p*2R + 2d` and `p*2R + 2d + 1`, whose -second `R/2` pairs duplicate the first; only pairs `[0, R/2)` are read per -position. - -**The table's content is the only rope input the op reads.** Two contents -are certified, both at `H = 128`, page 32, on the fresh-prefill context call -(the only MLA flavor that rotates anything): - -- the **unscaled** table, `angle(p, d) = p / theta^(2d/R)` at - `theta = 10000` — what a config without rope scaling produces, and what - every MLA case at `H` = 8, 16 and 32 above uses; -- the **YaRN-scaled** table a DeepSeek-R1-0528 config produces: `theta` - 10000, `factor` 40, `original_max_position_embeddings` 4096, - `beta_fast` 32, `beta_slow` 1, `mscale` 1.0, `mscale_all_dim` 1.0. - -The YaRN content, written out so a caller can build it without reaching into -TensorRT-LLM (fp32 throughout; `d` indexes the half-dimension `[0, R/2)` and -`j` the table's `R` pairs per position, `[0, R)`): - -``` -freq(d) = theta ** (2d / R) -low = max(0, floor(R * ln(orig_max_pos / (beta_fast * 2*pi)) - / (2 * ln(theta)))) # 10 for R1 -high = min(R - 1, ceil (R * ln(orig_max_pos / (beta_slow * 2*pi)) - / (2 * ln(theta)))) # 23 for R1 -ramp(d) = clamp((d - low) / max(high - low, 0.001), 0, 1) -inv_freq(d) = ramp(d) / (factor * freq(d)) + (1 - ramp(d)) / freq(d) -m(x) = 1.0 if factor <= 1 else 0.1 * x * ln(factor) + 1.0 -amplitude = m(mscale) / m(mscale_all_dim) # 1.0 for R1 -angle(p, j) = p * inv_freq(j mod (R/2)) # the duplication lives here -table[p*2R + 2j] = cos(angle(p, j)) * amplitude -table[p*2R + 2j + 1] = sin(angle(p, j)) * amplitude -``` - -Setting `factor = 1` collapses this to the unscaled table (verified: the two -constructions agree to 3.8e-6 max abs, fp32 evaluation-order noise). Two -consequences -worth naming: YaRN rescales only the **low-frequency** half of the spectrum -(`ramp` is 0 below `low`), so the two tables agree closely at small positions -and diverge with distance; and the table's `amplitude` is **1.0 whenever the -config's two mscales are equal**, which is R1's case — the model's YaRN -attention temperature `mscale = 0.1*ln(40) + 1 = 1.36885` is folded into the -softmax scale through `q_scaling = 1/mscale² ≈ 0.53366` instead, never into -the table. - -**Everything else about the rope configuration is inert on the MLA path** — -measured, not inferred. With the table held fixed, the seven scalars -`rope_base`, `rope_scale_type`, `rope_scale`, `rope_short_m_scale`, -`rope_long_m_scale`, `rope_max_positions`, `rope_original_max_positions` and -the `rotary_inv_freq` tensor moved **no** observable of a fresh-prefill -context call: output rows, the whole paged latent pool, the in-place-roped -`q`/`k`, all bitwise identical across R1's production scalar set, the -unscaled set above, a deliberately out-of-range set (`rope_base` 500000, -`rope_scale_type` 3, `rope_scale` 7.5, m-scales 3.0/9.0, position windows 77 -and 13 — both shorter than the sequences in flight), and `rotary_inv_freq` -zeroed or passed as `None`. Swapping the table in the same comparison does -move all of them, which is what makes the bitwise result mean something. -So: pass whatever scalars your config carries, and get the table right. - -Two obligations the table brings, neither of them checked: - -- It must hold **one row per absolute position the call ropes**. Nothing - bounds-checks it and `rope_max_positions` is not consulted: a 64-row table - under a 96-token prefill returned without raising and was wrong from output - row 64 on (max abs diff 3.9 against magnitude-4 outputs) — the kernel reads - past the end of the tensor. Certified with 1024 rows over sequences - reaching position 511. -- Row content does not depend on the row count, so a table truncated to - `max_seq_len` rows is **bit-identical** to the leading rows of the full - 163840-row table an engine builds for R1 (verified element for element). - Truncating is therefore safe; short-changing the positions in flight is - not. - -`rotary_inv_freq` = the `[R/2]` fp32 tensor from the same construction. It -is accepted, unread on this path (above), and still worth passing as -production does. - -### Feature groups not certified — pass these inert values - -| Group | Arguments → inert value | -|---|---| -| Quantized output / quantized QKV input | `out_scale=None`, `quant_scale_qkv=None` (fp8/fp4-quantized attention *output* and the fused DSv4 QKV-quantization path are not certified — on the MLA fp8 path `quant_scale_qkv` was additionally observed to be inert, a `[1]` fp32 `1/s` there leaving the context output bitwise identical. The fp8-e4m3 **KV pool** is certified in both configurations — see those sections, whose two `kv_scale_*` tensors then carry the scale, `None` being read as `1.0` rather than rejected. On the bf16 pool both stay `None`, and so does `quant_q_buffer`.) | -| Folded KV RMSNorm (MLA) | `kv_norm_weight=None`, `kv_norm_eps=1e-6`. Non-`None` weight folds the `kv_a_layernorm` into the KV kernel, which then reads `latent_cache` **raw** — a caller that already normalized would normalize twice. Only DeepSeek-V4's sparse module uses it upstream | -| Skip-correction (MLA) | `skip_correction_threshold=0.0`. A lossy trtllm-gen MLA option, SM100/SM103 only, off by default upstream (`enable_mla_skip_correction`). The engine forces it to 0.0 on a non-MLA layer regardless | -| Spec-dec tree mask | `force_prepare_spec_dec_tree_mask` — engine-set, `True` only for a dynamic tree; a linear-tree draft leaves it `False` | -| Sequence-count sizing | `max_num_sequences` — engine-set, defaults to `max_num_requests` | -| MLA DSv4 / FlashMLA / sparse-MLA sub-features | `flash_mla_tile_scheduler_metadata=None`, `flash_mla_num_splits=None`, `dsv4_inv_rope_cos_sin_cache=None`, `enable_dsv4_epilogue_fusion=False`, `aux_kv_cache_pool_ptr=None`, `sparse_attn_kv_lens=None` (`mla_bmm1_scale` / `mla_bmm2_scale` / `quant_q_buffer` are **certified** for the MLA generation call over an fp8-e4m3 latent pool, where all three are required — see that section; `None` everywhere else) | -| Speculative decoding — the **mask/offset tensor** group | `is_spec_decoding_enabled=False`, `use_spec_decoding=False`, `is_spec_dec_tree=False`, all `spec_decoding_*=None`, `spec_bl_tree_first_sparse_mask_offset_kv=None`, `spec_decoding_target_max_draft_tokens=None`. This does **not** mean speculative decoding is unavailable: `predicted_tokens_per_seq > 1` is certified for the MLA generation call with exactly these inert values, and that is also what the engine passes on sm_100 for a linear-tree draft, where `is_spec_decoding_enabled` is forced off (the Python backend computes it as `is_spec_decoding_enabled and (not trtllm_gen_arch or is_spec_dec_dynamic_tree)`), leaving the mask, position-offset and generation-length tensors all `None`. The within-block mask then comes from `predicted_tokens_per_seq` alone — see its row above. A **tree** draft, which is what these tensors describe, is not certified | -| Sparse / skip-softmax | `sparse_kv_indices=None`, `sparse_kv_offsets=None`, `sparse_attn_indices=None`, `sparse_attn_offsets=None`, `sparse_attn_indices_block_size=0`, `num_sparse_topk=None`, `skip_softmax_threshold_scale_factor_prefill=None`, `skip_softmax_threshold_scale_factor_decode=None`, `skip_softmax_stat=None` (`softmax_stats_tensor` is certified for the no-append MLA context call — see that table; `None` everywhere else) | -| SageAttention | `sage_attn_num_elts_per_blk_q=0`, `sage_attn_num_elts_per_blk_k=0`, `sage_attn_num_elts_per_blk_v=0`, `sage_attn_qk_int8=False` | -| Helix CP | `helix_position_offsets=None`, `helix_is_inactive_rank=None` | -| Relative bias (T5) | `relative_attention_bias=None`, `relative_attention_max_distance=0` | - -`attention_sinks` is **certified** for the standard bf16-pool causal -configuration — see *Attention sinks*; pass `None` everywhere else -(MLA calls, the fp8-e4m3 pool, `mask_type=0`). - -## Metadata consumed - -None — the op reads no thread-local model attrs, no registered layers, and -no global Python state. Every input above is an explicit argument. (The -batch-state tensor block plays the role that prepared attention metadata -plays for the registered-layer entry points.) - -## Preconditions - -- Batch order: context sequences before generation sequences; `q` rows in - that order. `num_contexts` / `num_ctx_tokens` passed explicitly and - consistent with `host_request_types`. -- `num_heads % num_kv_heads == 0` (the wrapper asserts it: a context-only - call at a non-multiple geometry returns a silently wrong `output`). -- `head_size` is one of the kernel's head dims. On sm_100 an unsupported - `head_size` **aborts the process**: the kernel-selection failure is - thrown where it cannot unwind (`std::terminate`), so it is neither - catchable nor loggable by the caller — observed at 32 (`Invalid TilePvM - as MMA only supports 64 or 128`), 96 and 192 (`Unsupported HeadDim for - BMM2-N`). Same for a `num_heads` that kernel selection cannot group into - CTAs (observed at `Hq = 7`, `Hkv = 2`: `Internal error numHeadsQ=7, - numHeadsPerCta=3, numCtasForAllHeads=2`). Verify a new `head_size` in a - throwaway process before wiring it into a target. -- Standard-configuration context sequences must start empty - (`sequence_length == host_context_lengths ==` new-token count) **unless** - `use_paged_context_fmha=True`; a cached prefix at `False` is silently - wrong (see *Paged-context FMHA*). The MLA fresh-prefill flavor requires - empty starts unconditionally; MLA context over cached tokens is certified - only via the no-append flavor (`latent_cache=None`, explicit full-range - or chunked K/V), and `use_paged_context_fmha` stays `False` for MLA. -- On the standard paged-context path (`use_paged_context_fmha=True`), a - context row's `context_lengths` / `host_context_lengths` is **this - call's** new-token (q-row) count, not the sequence's prompt or KV length. - Nothing checks it: passing the full KV length landed 101x outside the - tolerance band and `0` landed 53x outside, both without raising, and both - also corrupt the pool. `sequence_length` stays the global cached+new - count, as everywhere else. -- The paged-context path additionally requires every page holding a key the - context call reads to be valid at call time, and those pages to be - distinct from one another — the whole `[0, sequence_length)` range - without a window, the narrower range given in *Paged KV cache addressing - under a sliding window* with one. The packed path (`False`) reads no page - at all during a context call, so this obligation is new and nothing - checks it: aliasing a sequence's pages onto a single page returns a - result 31x outside the tolerance band without raising. -- No-append MLA context calls additionally require: `k`/`v` rows are - exactly the per-sequence KV ranges described by `sequence_length`, - concatenated in batch order, with `host_total_kv_lens[0]` equal to that - row count; `v`'s row stride is `H * (nope + v_head_dim)` elements; under - `mask_type=1` every context sequence satisfies `sequence_length >= - context_lengths` (the causal mask offset is their difference); under - `mask_type=0` any KV length including 0 is allowed, but the output and - stats rows of zero-KV sequences are undefined and must never be read - (production merge plans skip them). `q` arrives fully prepared — the - call applies no RoPE even though the rope-table arguments are passed. - Under `quant_mode=128` all of the above is unchanged, and the pool stays - untouched (it is not read to build `k`/`v` — assembling those from the - cached latent rows is entirely the caller's step); what changes is that the - op quantizes the `k`/`v` it is handed to e4m3 before the FMHA, so operands - the caller has already dequantized off an fp8 pool go through e4m3 a second - time. -- A caller-provided `softmax_stats_tensor` (no-append MLA context only) - is fp32 `[>= Tc, >= H, 2]`, contiguous, on the same device (the op - rejects other dtypes/sizes); certified at the exact size `[Tc, H, 2]`. -- For every generation sequence, the cached tokens at positions - `[0, kv_s - 1)` were actually written by prior calls addressing the same - pool pages, and `sequence_length` equals the true cached count + 1. -- Per-sequence pages cover `ceil(total_kv / tokens_per_block)` entries in - its `kv_cache_block_offsets` row, where `total_kv` is the sequence's - full `[cached + new]` KV length (also for chunked partial passes, whose - `sequence_length` entries are smaller); distinct sequences use distinct - pages; all offsets lie inside the pool. `total_kv <= max_seq_len` for - every sequence. -- With a sliding window the entry count above is unchanged (the row is - indexed by absolute page index), but only the entries holding in-window - tokens must be valid, and the pages behind them must be distinct over - each call's new-token range plus the whole window — see *Paged KV cache - addressing under a sliding window* for the ring rule - (`P * tokens_per_block >= max(attention_window_size, this call's - new tokens)`), which is not checked anywhere and fails silently. -- `attention_sinks`, when given, is fp32 (the op raises `Expected - attention_sinks to have float dtype` for anything else) and holds exactly - `num_heads` values in contiguous memory on the current CUDA device. Only - the dtype is checked by the op: the buffer is read as `num_heads` raw - fp32 values from `data_ptr()`, so every other violation is **silent** — - a shorter tensor is read past its end, a longer one has its tail ignored, - a strided view is read as its underlying memory (a stride-2 view holding - the right values returned results ~2.1 off), and an empty tensor - disables the sink entirely. The wrapper therefore asserts contiguity and - `numel() == num_heads`; nothing else guards this argument. A pageable - **CPU** fp32 tensor was also accepted and gave the correct result (the - values do reach the kernel), but CUDA is the certified device. -- All tensors contiguous unless noted (MLA context `v` is the strided - split view described in *Signature*), on the devices listed above; every - CUDA tensor on the current device; dtypes exactly as listed (host length - tensors are int32 — int64 is not certified). -- `quant_mode` and the pool element type must agree (0 ↔ bf16, - 128 ↔ fp8-e4m3): the op sizes and interprets slabs from `quant_mode` - alone and cannot detect a mismatched pool allocation. -- With `quant_mode=128`, `kv_scale_orig_quant` and `kv_scale_quant_orig` - are fp32 `[1]` CUDA tensors holding `1/s` and `s` for one scaling factor - `s` (mutual consistency is the caller's job). **`None` in either slot or - in both is not rejected — it is silently read as `1.0`**, bitwise - identical to passing 1.0 tensors in both the pool bytes and the output, - in both configurations, in every phase, and on both the single-layer and - the 4-layer manager pool. - A forgotten scale is therefore a wrong-number bug, not a crash: inert at - the production `s = 1.0`, and 1.19 off on magnitude-1.5 decode outputs - over a standard-configuration pool written at `s = 2.0`. - Certified `s`: 1.0 (the bf16-checkpoint production value), 1.5, 2.0. -- On the **MLA** path under `quant_mode=128` the same two tensors reach far - less of the computation, and `s != 1.0` is not survivable in the context - phase (full measurements in *MLA over an fp8-e4m3 latent pool*): - - the fresh-prefill append consumes `kv_scale_orig_quant` and is bit-exact - at any `s`; - - **both** context flavors quantize `q`/`k`/`v` at 1.0 but apply `s^2` and - `s` as if they had not, so a context call at `s != 1.0` returns a silently - wrong result — measured on the fresh-prefill flavor (21x / 48x outside the - band at 1.5 / 2.0) and on the no-append flavor (48x at 2.0), which appends - nothing and so has no use for `kv_scale_orig_quant` at all. Pass - `s = 1.0`, or `None`; - - the generation call reads **neither** tensor. Its query comes from - `quant_q_buffer` and its scales from `mla_bmm1_scale`/`mla_bmm2_scale`, - all three required (`None` raises `quant_q_buf is nullptr.` / - `bmm1_scale is nullptr.` / `bmm2_scale is nullptr.`, from - `attentionOp.cpp:1050-1052`), and the caller must fold `s` into them as - shown in that section. Nothing checks their contents. -- Bits of `quant_mode` outside the KV-cache group are **not** a reason to - strip the value a checkpoint's quant config produces: `1152` and `384` were - measured bit-identical to a bare `128` on every MLA flavor (see *MLA over an - fp8-e4m3 latent pool*). Bits that were not tried are unknown, not inert. -- MLA batches take one call per phase; both calls of a mixed batch use the - same full-batch state tensors. `attention_input_type=0` is **accepted, not - rejected**: inert (bitwise identical) on a single-phase batch, silently - unsatisfiable on a mixed one — see *Semantics*. -- Every MLA call needs an int `q_lora_rank` — `0` for a checkpoint without - a q-LoRA — even though the value is never read: `None` raises - `RuntimeError: bad optional access` (see *MLA q-LoRA rank*). -- MLA generation additionally requires, before the call: every generation - sequence's full latent history `[0, L_g)` resident in the pool — the - context call's append covers the prefill rows; this step's `P` rows at - positions `L_g - P .. L_g - 1` come from the caller's generation - preprocessing — and the caller-allocated `cu_q_seqlens` / `cu_kv_seqlens` / - `fmha_scheduler_counter` filled as listed in *Signature*. Over an - fp8-e4m3 pool the same preprocessing must also have filled - `quant_q_buffer` with the e4m3 query and both `mla_bmm*_scale` buffers - with the scales above; a decode row appended by the caller must be - quantized the same way the op's own append quantizes - (`e4m3(row * kv_scale_orig_quant)`), since the kernel applies no - per-row scale of its own. -- With `predicted_tokens_per_seq = P` above 1, MLA generation additionally - requires `L_g >= P` for every generation sequence — draft row 0 attends to - `[0, L_g - P]`, so a shorter `L_g` leaves it no keys. `L_g == P` is legal - and certified (row 0 then attends to exactly one key, its own). `q` / - `output` / `quant_q_buffer` must be `G*P` rows tall and token-major within - a sequence; `sequence_length[num_contexts + g]` must count **all `P`** of - this step's tokens. Nothing checks the row count — `q` and `quant_q_buffer` - are addressed from `P` and the sequence index, so a `G`-row buffer is read - past its end rather than rejected. -- `host_past_key_value_lengths` must be non-zero for at least one sequence. - Its per-sequence values are otherwise inert on the MLA generation path - (`sequence_length - 1`, `- P` and all-ones are each bitwise identical to - the true values at `P = 3`), but an **all-zero** vector makes the call - return without writing a single element of `output` — a sentinel-filled - buffer comes back untouched, at `P` = 1 and above alike, with no error. - Passing the same values as `sequence_length`, as everywhere else in this - contract, satisfies it. -- The MLA fresh-prefill context call clobbers the rope slice of each - `q`/`k` head in place: treat both tensors as consumed (their nope slices - stay intact, but do not rely on pre-call rope-slice contents - afterwards). The no-append context call mutates neither. -- `rotary_cos_sin` must hold a row for every absolute position the - fresh-prefill context call ropes, in the duplicated layout and with the - content its rope configuration implies (see *In-kernel RoPE arguments*). - Nothing checks either: a short table is read past its end and comes back - silently wrong from the first over-long position on, and the scalar rope - arguments beside it — including `rope_max_positions` — are not consulted at - all, so they cannot substitute for getting the table right. -- First generation-phase call per process JIT-compiles its FMHA kernel via - NVRTC (~5-10 s; one kernel per configuration — standard GQA and MLA - decode compile separately): the C++ locates the shipped kernel headers - (`tensorrt_llm/include/trtllm_gen_kernels/fmha/`) by running - `pip show tensorrt_llm` in a subprocess, so a `pip` that resolves the - installed package must be on PATH (this repo pins `pip` as a project - dependency for exactly this reason); otherwise the call raises - `NVRTC_ERROR_COMPILATION` ("could not open source file cuda.h"). - Context-phase kernels are precompiled into `libtensorrt_llm.so` and need - no JIT. -- The op is synchronous-schema but asynchronous: it launches on the current - CUDA stream and returns; synchronize before reading `output` on the host. - -## Notes - -- One Python-level invocation launches several device kernels (QKV/KV-cache - preprocessing, FMHA, postprocessing) inside the single C++ op — the same - attention core the TRTLLM backend entry points route into, certified here - as a single catalog call. -- On sm_100 the op selects trtllm-gen FMHA kernels. Generation kernels are - JIT-compiled per process (in-memory kernel cache; the log warns "Possible - JIT Cache Missing ... generateAndCompileKernel took ~5-10 s" on first - use of each kernel variant — the MLA decode variant is - `...HQk576HV512...PagedKv...`; the fp8-KV GQA decode variant is - `fmhaSm100aKernel_QkvE4m3OBfloat16H128PagedKvDense...` — its name pins - that q, K, and V are all e4m3 in the decode MMAs). Re-runs within the - process are fast. Context kernels — including the padding-mask and - softmax-stats variants of the no-append flavor, and the bf16 context - kernel that keeps serving the context phase under `quant_mode=128` — are - precompiled (no JIT observed). -- The generation-kernel variant is keyed by `head_size` and the q-tile size - the head group maps to, rather than by the head counts directly — several - distinct geometries share one name: the bf16 - variants observed are `...H128PagedKvDenseP32VarSeqQ8Kv128Static...` for - `Hq/Hkv` of 32/8, 32/4 and 28/4, `...Q16...` for 16/1, and `...H256...` / - `...H80...` for those head sizes. So a first call at a new head geometry - costs at most one extra ~6 s JIT compile, and often none. The mapping is - not constant in the head count either, though: MLA generation at - `H` = 16, 32 **and 128** all report - `fmhaSm100aKernel_QkvBfloat16OBfloat16HQk576HV512HVPerCta128PagedKvDenseP64VarSeqQ16Kv128StaticSwapsAbForGen` - (the `P32` sibling at page 32), while `H = 8` reports - `...VarSeqQ8Kv128StaticSwapsAbForGen` — same kernel - family, different q-tile, hence a different compiled kernel. The q-tile - therefore stops growing well below the head count: 16x the heads of the - `H = 8` case still map to `Q16`. Where head - counts do share a name the **compile cache is still keyed more finely**: a - process that runs 16, 32 and 128 pays the ~5.5 s `...Q16...` compile three - times, once per head count, while any number of calls at one head count — - across batch sizes and histories — compile once. The **page size is part of - the name** too: the same MLA generation call at `tokens_per_block = 32` - compiles - `fmhaSm100aKernel_QkvBfloat16OBfloat16HQk576HV512HVPerCta128PagedKvDenseP32VarSeqQ16Kv128StaticSwapsAbForGen` - (~5.4-5.9 s, observed once per head count that maps to it), a genuinely - different kernel from the `...P64...` one — so page 32 and page 64 MLA - decode are separate compiled variants, both certified. Full-test-file counts on sm_100 measured - from the run behind the current receipt, all at - `predicted_tokens_per_seq = 1`: `...P64VarSeqQ16...` twice (the - page-64 16/32 pair), `...P32VarSeqQ16...` **three times** (16/32/128), - `...P64VarSeqQ8...` once and - `...P32VarSeqQ8...` once (`H = 8`, each serving both that page size's decode - and mixed-batch cases, the page-32 one additionally serving the `H = 8` - q-LoRA-rank sweep's decode calls) — all seven ~5.4-5.7 s, a wall-clock - figure that moves ~0.15 s run to run. One `H = 128` compile serves all ten - of the file's `H = 128` decode calls at `P = 1` — 3 from the baseline decode - and mixed-batch cases, 3 from their R1-cell twins, 4 from the `q_scaling` - sweep. The MTP section adds two more bf16 compiles on top, one per `P` it - runs there (`...HVPerCta128...P32VarSeqQ16...` at `P = 2` and - `...HVPerCta256...P32VarSeqQ16...` at `P = 4`), for nine bf16 MLA decode - compiles in the file. The counts are the stable part: identical across - full-file runs. - Every JIT compile logged across the file is a `...ForGen` decode variant — - no MLA context call at any certified head count, page size or pool dtype - triggered one. `quant_mode` is part of the decode kernel name too: MLA - generation over an fp8-e4m3 latent pool compiles - `fmhaSm100aKernel_QkvE4m3OBfloat16HQk576HV512HVPerCta128PagedKvDenseP32VarSeqQ16Kv128StaticSwapsAbForGen` - — the `QkvE4m3` sibling of the bf16 `...P32VarSeqQ16...` name, so q, K and V - are all e4m3 in the decode MMAs — once for every fp8 MLA decode call at - `predicted_tokens_per_seq = 1`: the baseline-cell decode, the R1-cell - decode, the mixed-batch decode and the `quant_mode`-bits decodes all share - it. So on the fp8 path neither `q_scaling`, nor the rope table, nor the extra - `quant_mode` bits reach decode kernel selection, and no fp8 MLA **context** - call — either flavor — triggers a JIT compile at all. -- **`predicted_tokens_per_seq` does reach decode kernel selection**: it becomes - the kernel's `maxSeqLenQ`, which appears in the JIT log line and moves the - `HVPerCta` split in the name. Observed at `H = 128`, page 32, both pool - dtypes: `P` = 1 and 2 take `...HV512HVPerCta128PagedKvDenseP32VarSeqQ16...` - and `P` = 3 and 4 take the `...HVPerCta256...` sibling, with `P` = 1 and - `P` = 2 paying **separate** ~5.4-5.7 s compiles despite sharing a name — the - same finer-than-the-name cache keying the head count shows. The file's fp8 - MLA decode compiles are therefore three, at `maxSeqLenQ` 1, 2 and 3, with - `P = 4` reusing the `maxSeqLenQ = 3` entry. Batch size can move it as well: - a probe at `G = 4`, `P = 4` compiled a further variant whose name carries no - `HVPerCta` segment at all. So a target sweeping `max_draft_len` should - expect one decode JIT compile per `P` it runs, not one for the model. -- First call logs "Attention workspace size is not enough" and resizes - `workspace_` in place — expected when starting from an empty tensor. The - size tracks the call's tokens x heads: ~33 MB standard, ~39 MB MLA and - ~35 MB no-append MLA context at the 8/16/32-head shapes, but ~50 MB for an - `H = 128` MLA context call over 129 tokens and ~109 MB over 545, which is - the file's peak (the standard-configuration head-geometry sweep reaches - ~95 MB). `quant_mode=128` raises the MLA context requirement — the fp8 - path stages quantized q/k/v in the workspace: 61 033 728 B for the same - 129-token `H = 128` call that needs 52 579 072 B in bf16, and 142 612 480 B - over 512 tokens (measured outside the shipped cases, which stay below the - file's bf16 peak). The **no-append** flavor pays the same surcharge — it - quantizes the K/V it is handed, so it stages them too: the shipped R1-cell - case (66 q rows over 193 K/V rows at `H = 128`) asks for 52 816 640 B under - `quant_mode=128` against 43 288 832 B at the same shape in bf16. An engine - at production `max_num_tokens` will be well above all of - these — budget from the shape, not from these figures. -- The paged-cache append is bit-exact wherever it is a copy: after a - standard call, page `p` of a - sequence holds exactly the bf16 K/V slices packed into `q`; after an MLA - fresh-prefill context call it holds `[ckv | rope(k_pe)]` with `ckv` - bitwise-copied (verified for both configurations, including page-boundary - crossings, and for MLA at page sizes 32 and 64). The MLA RoPE half is - bit-identical to an fp32 reference rounded once to bf16 **except at - fp32→bf16 rounding ties**: kernel and torch reference evaluate - `x*cos ∓ y*sin` in different fp32 orders, and where the correctly-rounded - fp32 result lands exactly on a bf16 midpoint the two round to adjacent - bf16 values. Measured on sm_100 over 60 prefill runs spanning both page - sizes (495 360 roped elements): 8 such elements — 1.6e-5 of the total, - never more than one bf16 ulp, and the same count at page 32 as at page 64, - so it is arithmetic rather than addressing. The entry's test gates the - RoPE half at one bf16 ulp (`rtol = 2^-7`, `atol = 0`) plus a hard "at most - one element in a thousand differs at all" floor, and keeps the `ckv` half - strictly bitwise; that gate leaves the real failure modes far outside it — - measured on one 96-token page-32 request, an append shifted one slot lands - at 6.1e5 ulp with 99.9% of elements differing, the pool read with page - size 64 instead of 32 at 2.6e5 ulp / 33%, RoPE applied at position `i+1` - at 5.1e5 ulp / 69%, `k_pe` left un-roped at 4.4e5 ulp / 93%. The - standard-configuration fp8-pool - append is bit-exact against - `e4m3(k * kv_scale_orig_quant)` (fp32 multiply, RN cast) for power-of-2 - scales; at `s=1.5` ~0.6% of elements sit one e4m3 ulp off that mirror - (never more) — see the fp8 section. The **MLA** fp8 append is bit-exact - against the same mirror at every scale run (1.0, 1.5, 2.0) and on both - halves of the latent row, roped `k_pe` included: the e4m3 grid is coarse - enough to swallow the rounding tie that costs the bf16 latent append its - one-ulp elements. Its gate is one e4m3 ulp plus the same - one-in-a-thousand inexact floor, and no run has approached it — across the - appended-row checks the fp8 MLA cases make (three scales, both rope tables, - all four call shapes), not one byte differs from the mirror, in either half. - Two things make that bit-exactness load-bearing rather than merely tidy. - First, it is the sharpest place the **rope table** shows up: a mirror roped - by the unscaled theta-10000 table instead of R1's YaRN one differs in 2336 - of the 6144 e4m3 bytes of a 96-token request (38.0%) against a gate that - allows 6 — ~389x — so the same comparison that passes on the certified - content demonstrably fails on the wrong one. Second, the comparison can see - a **racing** append: replaying this entry's own documented arming sequence - (a call whose new tokens do not map to distinct physical slots — here a - 96-token prefill with all three of its absolute pages aliased onto one - physical page) made an identical seeded call leave different pool bytes on - every repeat — five of five in a scratch probe, 11 177 to 12 513 bytes - apart — while the certified geometry reproduced bitwise. The shipped - control asserts exactly that (four armed repeats must not all agree, three - certified ones must), and it fired in five of five fresh processes. Both - controls run in the test file, in the same processes as the certified - cases. -- Multi-layer pool addressing (standard configuration, sm_100): certified - against state produced by a real 4-layer `KVCacheManager` (bf16, - GQA 8/2/128, `tokens_per_block` 32) — identity mapping rows, - layer-agnostic offsets (K `8p`, V `8p + 1`), prefill plus page-crossing - decode at `local_layer_idx` 0-3. Each layer's append landed bit-exactly - in its own slab group with sibling layers bitwise untouched, and decode - read back each layer's own history. A doctored-mapping call - (`local_layer_idx=1`, row 1 rewritten to layer-in-pool 3) landed its - append in layer 3's slabs while layer 1 stayed bitwise intact: the - pool-base shift is `mapping[local_layer_idx][1] * kv_factor` slabs — - the row's layer column, `local_layer_idx` only selecting the row. The - same 4-layer prefill + page-crossing-decode certification was re-run - over an fp8-e4m3 pool from a `DataType.FP8` manager (`quant_mode=128`, - GQA 32/8/128, `tokens_per_block` 32, `s = 1.0` — identical mapping and - offset layout): every layer's context and generation append landed - bit-exactly against the `e4m3(k * kv_scale_orig_quant)` mirror in its - own slab group, sibling layers stayed bitwise untouched by every call, - and each layer's decode read back its own history — pinning that the - layer-base shift is computed in e4m3-sized slabs under `quant_mode=128`. - The doctored-mapping probe was not repeated under fp8. -- Accuracy (bf16, sm_100) vs fp32 references, compared with `atol=5e-3, - rtol=1.6e-2` (not the default bf16 `atol=1e-5`): standard max abs err - 1.6e-2 on magnitude-~1 outputs (4-layer shared pool the same: 1.6e-2 - prefill, 7.8e-3 decode); MLA fresh-prefill context max abs err - 1.6e-2; MLA generation max abs err 7.8e-3; no-append MLA context max abs - err 1.6e-2 (one-shot cached-KV), 7.8e-3 (chunked partial passes) rising to - 1.6e-2 on the final causal pass, 1.6e-2 (chunk passes merged vs a - single-pass full-range reference). The largest fraction of the combined - `atol + rtol*|ref|` allowance any case in the test file uses is 66% — the - maximum sits on the `q_scaling` sweep's `q_scaling = 1.0` context run at - `H = 128` (66.1%), just ahead of the paged-context 4-layer shared-pool case - (65.8%), the page-32 cached-KV no-append MLA context case at `H = 32` - (63.9%), the page-32 `H = 128` fresh prefill (63.3%), the fp8-pool - multi-layer context (61.5%) and the YaRN-table case (61.3%); - this bullet's standard GQA/MHA cases sit at 36-52% and its MLA - single-pass cases at 36-64% (the sink, sliding-window and paged-context - bullets below carry their own numbers; the latent append's one-ulp gate - is a separate, tighter one, and one case does approach it — the page-32 - `H = 32` mixed batch reaches 84% of the one-ulp gate on its roped - elements). Every figure in this bullet and this ceiling - were re-measured across the whole file in the state the current receipt - certifies, by an instrumented sibling run of the same seeded cases. - Head geometry does not move these numbers: the two shipped-target - geometries (32/8/128, 32/4/128) and the MQA / odd-ratio / odd-count / - `head_size` 256 sweep all land at max abs err 1.6e-2 and at most 53% of - the allowance, the same envelope the 8/2/128 and 4/4/64 cases sit in. - A scratch probe of wider ratios (64/8/128, 64/1/128, 128/1/128) stayed at - the same 1.6e-2 max abs err but pushed decode to ~75% of the allowance — - headroom narrows as the head group grows, so re-measure before adopting a - ratio above 16. The 64/8/64 geometry without sinks sits in the same - envelope: max abs err 1.6e-2, at most 50% of the allowance. -- MLA head count does not move the numbers either (bf16, sm_100, same fp32 - references and `atol=5e-3, rtol=1.6e-2`). At `H = 32`, page 64: - fresh-prefill - context max abs err 1.6e-2 (53% of the allowance), cached-KV no-append - context 7.8e-3 (50%), generation decode 7.8e-3 (37%), mixed batch 1.6e-2 - context (51%) / 7.8e-3 decode (36%) — the same magnitudes the `H = 16` - runs of those four cases produce in the same process (1.6e-2 at 49%, - 1.6e-2 at 40%, 7.8e-3 at 39%, 1.6e-2 at 38% / 7.8e-3 at 36%), so doubling - the head count costs no headroom. - Halving it below 16 costs none either, and the `H = 8` cell has the - mildest worst case of the three (45.9% against 48.5% at `H = 16` and 63.9% - at `H = 32`): at page 64, fresh-prefill context 1.6e-2 (46%), - cached-KV no-append context 1.6e-2 (40%), generation decode 7.8e-3 (37% - then 36%) over two steps, mixed batch 7.8e-3 context (38%) / 7.8e-3 decode - (36%); at page 32, fresh-prefill 7.8e-3 (39%), cached-KV no-append 7.8e-3 - (40%), decode 7.8e-3 (37% then 34%), mixed batch 7.8e-3 context (37%) / - 7.8e-3 decode (36%). Case by case, every `H = 8` figure is at or below the - `H = 32` one for the same case; against `H = 16` the only one above is the - page-64 mixed batch's context row, 38.5% where `H = 16` reads 37.9%. - Multiplying it by 4 to `H = 128` costs no headroom either, at page 32: - on the baseline rope/scale, fresh-prefill context 1.6e-2 (63%), generation - decode 7.8e-3 (40%) then 1.6e-2 (40%) over two steps, mixed batch 1.6e-2 - context (52%) / 7.8e-3 decode (37%), cached-KV no-append context 1.6e-2 - (48%); at the full R1 cell (YaRN table, `q_scaling ≈ 0.53366`) the same - four cases read 1.6e-2 (61%), 1.6e-2 (49% then 43%), 1.6e-2 context (58%) / - 7.8e-3 decode (37%), and 1.6e-2 (54%). Same max abs errors as every smaller - count; the fractions sit inside the 34-64% band the others span, and the - head count is not what orders them (the `H = 128` fresh prefill's 63% is - the highest of its set, where at `H = 8` and `H = 32` the no-append case - is). - The paged latent append held to its gate (see the append bullet above) at - all four head counts, including page-boundary crossings, and every count - leaves `latent_cache`, `q_pe`, `q`/`k`/`v` and the pool untouched exactly - where the per-flavor tables say they do. At `H = 8` the appended rows came - back bitwise exact on both halves in all six append-checking cases (no - fp32→bf16 rounding tie among their roped elements), so the one-ulp gate was - not approached there; the same held at `H = 128`, whose 7 append-checking - cases make 16 appended-row checks between them — every one bitwise exact on - both halves, on both table contents. -- MLA page size does not move the numbers either (bf16, sm_100, same fp32 - references and `atol=5e-3, rtol=1.6e-2`). At `tokens_per_block = 32`, - `H = 32`: fresh-prefill context max abs err 1.6e-2 (50% of the allowance), - generation decode 1.6e-2 (42%) then 7.8e-3 (37%) over two steps, mixed - batch 1.6e-2 context (42%) / 7.8e-3 decode (37%), cached-KV no-append - context 1.6e-2 (64%); the page-32 decode case at `H = 16` sits at 7.8e-3 / - 37% on both steps, and the four page-32 cases at `H = 8` (numbers in the - head-count bullet above) span 34-40%. Those are the page-64 - magnitudes and, the no-append case aside, the page-64 band. The no-append flavor never reads or writes - the pool, so its 64% tracks that case's KV geometry (prefixes 96/31/0 - reaching 128/40/25 keys) rather than the page size — it was the file's - ceiling until the `H = 128` `q_scaling` sweep edged past it at 66%. - The paged latent append held to the same gate at page 32, the generation - call still wrote nothing, and the no-append call still left `q`/`k`/`v` - and the pool bitwise intact. -- MLA q-LoRA rank moves nothing whatsoever, at either shipped head count - (bf16, sm_100, `tokens_per_block = 32`; `H = 32` is the deepseek-v3-lite - tp1 cell, `H = 8` its tep4 slice; same fp32 references and `atol=5e-3, - rtol=1.6e-2`). At `q_lora_rank = 0`, `H = 32`: fresh-prefill context max - abs err 1.6e-2 (47% of the allowance), generation decode 7.8e-3 (36% then - 37%) over two steps, cached-KV no-append context 1.6e-2 (44%). At - `q_lora_rank = 0`, `H = 8`: fresh-prefill context 1.6e-2 (44%), - generation decode 7.8e-3 (37% then 36%), cached-KV no-append context - 1.6e-2 (37%). All six sit between 35.9% and 46.7% of the allowance — the - band the page-32 cases above sit in — and the two head counts interleave - there rather than separating (the `H = 8` decode's first step, 36.5%, is - the one figure above its `H = 32` counterpart, 35.9%). Each fresh-prefill - case's appended latent rows came back bitwise exact on both halves (no - fp32→bf16 rounding tie among either run's 8256 roped elements, so the - one-ulp gate was not even approached), and the no-append calls again left - `q`/`k` and the pool bitwise intact. Since the four-value rank sweep is - bitwise identical throughout (see *MLA q-LoRA rank*), these are element - for element the numbers the `q_lora_rank = 1536` runs of the same six - cases produce. The rank reaches neither kernel selection nor the compile - cache at either count: a process running only the two four-value sweeps - generated - `...HQk576HV512HVPerCta128PagedKvDenseP32VarSeqQ16Kv128StaticSwapsAbForGen` - exactly **once** (the `H = 32` sweep, 8 decode calls) and - `...HQk576HV512HVPerCta128PagedKvDenseP32VarSeqQ8Kv128StaticSwapsAbForGen` - exactly **once** (the `H = 8` sweep, 8 decode calls); and in the entry's - full test run, measured with and without the `H = 8` sweep present, the - counts are identical either way — the first variant twice, once each for - the two head counts that share it (`H = 16` and `H = 32`), the second once - (`H = 8`) — so those 8 extra decode calls add no compile at all, against - one compile per value on a real selection axis. -- MLA `q_scaling` is a live axis of both phases, and only of the softmax - scale (bf16, sm_100, `H = 128`, `tokens_per_block = 32`, YaRN table). Swept - over 1.0, 0.53366, 2.0 and 0.25 on identical inputs, each run compared - against fp32 references built at all four scales: the matching reference is - the only one within tolerance, at 30-66% of the allowance, while the full - 4x4 cross matrix spans 20.1x-178x the allowance in both phases — the - tightest cell is 1.0 against 2.0 in the context phase (20.1x), and a - reference at 1.0 for a run at any other value (the "argument silently - ignored" hypothesis) spans 20.1x-70.5x. The axis is finer than the sweep - resolves: 0.5 against 0.53366, a 6.3% change of scale, separates by only - 2.95x-3.61x, so the shipped values are kept apart rather than the gate - loosened. Across the sweep the paged latent pool comes back **bitwise - identical** — `q_scaling` reaches the softmax scale and nothing else, not - the RoPE and not the append. It does not reach kernel selection either: the - sweep's four decode calls compile nothing beyond the one `H = 128` decode - kernel the file already pays for. -- The MLA rope table is read for its content only (bf16, sm_100, `H = 128`, - `tokens_per_block = 32`). A fresh-prefill context call driven by the - DeepSeek-R1-0528 YaRN table reproduces a table-aware fp32 reference at 61% - of the allowance and sits 27x outside a reference built from the unscaled - theta-10000 table; the converse run sits 27x outside the YaRN reference. - That separation is position-dependent — 8.6x over a 96-token sequence, - 16.9x at 256, 27.0x at 512, 31.7x at 960 — because YaRN rescales only the - low-frequency half of the spectrum, so a short probe would not have - resolved the two tables at all. Meanwhile the seven scalar rope arguments - and `rotary_inv_freq` are bitwise inert beside the table (see *In-kernel - RoPE arguments*), and a plain-torch rebuild of the YaRN formula matches - TensorRT-LLM's own table to 1.9e-6 max abs on cos/sin values in [-1, 1] — - a few fp32 ulp from a different evaluation order of the same blend, where - formula slips land 4-5 orders of magnitude away (dropping the - interpolation, or moving `beta_fast` 32 → 64, both differ by 2.0; taking - the amplitude as `mscale` rather than the `mscale`/`mscale_all_dim` ratio - differs by 3.7e-1). -- Accuracy with attention sinks (bf16, sm_100, 64/8/64 causal) vs fp32 - sink-aware references at the same `atol=5e-3, rtol=1.6e-2`: max abs err - 1.6e-2, at most 50% of the allowance, across context prefill, two decode - steps, a mixed batch, decode over 600- and 2000-token histories, and the - constant-V probe that reads the softmax row mass directly. Every case - also had to sit **outside** the tolerance band of two rival references — - "sink silently ignored" and "sink pre-scaled" — and the tightest of those - separations was 12.8x the allowance (decode over 2000 cached tokens, - where a small sink is genuinely negligible against 2000 keys, so the test - draws the sink near `log(kv_len)` to keep the mechanism observable). -- Accuracy with the sliding window (bf16, sm_100, 64/8/64 causal, `W` 128 / - 100 / 33) vs fp32 windowed references, sinks on and off, at the same - `atol=5e-3, rtol=1.6e-2`: max abs err 3.9e-3 (at most 30% of the - allowance) across prefill past the window, decode, a mixed batch, a - 2000-token history, a 300-token prefill, and 250 wrapping decode steps; - the per-layer alternating-window case sits at 1.6e-2 / 49%, its - full-attention layers being the wider ones. Each case is additionally - gated **outside** the tolerance band of a "window ignored" (full causal) - reference on top of the two sink rivals. The exact per-key weights read - out by the one-hot probe match `1 / (n + exp(sink[h]))` to 3.6e-3 - relative — 59% of the 6e-3 (~1.5 bf16 ulp) gate — while out-of-window - keys are bitwise zero, which is what pins the boundary at `W` rather than - `W + 1` keys. -- Accuracy on the paged-context path (bf16, sm_100, `use_paged_context_fmha - = True`) vs the same fp32 references at the same `atol=5e-3, rtol=1.6e-2`: - max abs err 1.6e-2 across every case — fresh prefill and decode at the - three shipped geometries (53% of the allowance), cached prefixes at - 32/8/128 and 32/4/128 including page-boundary crossings and a mixed - cached/fresh/generation batch (54%), the six-prefix page-grid batch - (59%), two chunked-prefill loops (42%), and the 4-layer shared-pool - cached-context cycle (66%, the file's ceiling). The gpt-oss cell sits - lower: 7.8e-3 / 40% over sinks with and without the 128-token window, and - 1.6e-2 / 49% on the 200-token windowed prefill that sets up the read-set - probe; each sink case is gated outside the sink-ignored (11.9x-56x), - sink-pre-scaled (11.3x-53x) and — with the window — window-ignored - (32x-43x) rivals. The one-hot probe on a cached-prefix context row reads - the per-key weights back to 2.2e-3 relative, 36% of the same 6e-3 gate, - with out-of-window keys bitwise zero. Wrong batch states land 31x-318x - outside the band - (aliased pages 31x, `use_paged_context_fmha=False` over a cached prefix - 51x, `context_lengths` = full KV length 101x or 0 53x, a decoy in place - of an in-window cached page 248x-318x); the test gates them at 20x. -- Sliding-window kernel variants (sm_100): the FMHA family is chosen from - the batch's KV length against the window, **not** from `max_seq_len`. A - call moves from `...PagedKvDense...` / `...PackedQkvCausal...` to - `...SlidingOrChunkedCausal...` exactly when the batch's longest KV run - *exceeds* `attention_window_size` — at window 128, KV 127 and 128 take the - dense kernel and KV 129 the sliding one. `max_seq_len` does not enter the - choice: at window 128 and KV 41 the values 128, 256 and 1024 all give - `...PagedKvDense...`, and at KV 200 both 128 and 256 give - `...SlidingOrChunkedCausal...`; the dispatcher's `maxSeqLenKv` is the - batch's real KV length in every case (41, 128, 129, 200), never clamped to - the window. One windowed sequence therefore pays two decode JIT compiles - over its lifetime: `...PagedKvDense...` while its history still fits the - window, `...SlidingOrChunkedCausal...` from the first step past it. Decode - JIT-compiles - `fmhaSm100aKernel_QkvBfloat16OBfloat16H64PagedKvSlidingOrChunkedCausalP32VarSeqQ8Kv128StaticSwapsAbForGen` - once (~6 s) and `...SlidingOrChunkedCausalP32MultiCtasKvCga...` for long - histories — selected on the *total* cached length, so a 128-token window - over 2001 cached tokens still takes it; the context variants are - precompiled (no JIT observed). Where both families can compute the same - row (a short sequence alone, then batched behind one past the window) they - agree bit for bit. -- Paged-context kernels are precompiled too: across a full run of this - entry's test file every JIT compile logged is a `...ForGen` decode - variant, none a context one, so switching `use_paged_context_fmha` on - costs no compile. It does not steer the decode kernel either — a decode - step at both flag values is bitwise identical. -- Attention-sink quirks (sm_100): the sink pointer is read as `num_heads` - fp32 values with no size, stride, or device check (see *Preconditions*). - Enabling sinks costs no extra JIT compile: at 64/8/64 the decode kernel - `fmhaSm100aKernel_QkvBfloat16OBfloat16H64PagedKvDenseP32VarSeqQ8Kv128StaticSwapsAbForGen` - was generated once and then served both the sink and the no-sink calls in - the same process — adding or removing the argument triggered no further - generation. Long histories select the - `...H64PagedKvDenseP32MultiCtasKvCga...` variant, which folds partial - softmax states across CTAs and applies the sink in that reduction. A - `-inf` or `-100.0` sink reproduces the no-sink output bit for bit — the - sink adds exactly one term to the softmax denominator and changes nothing - else in the code path. -- Accuracy (fp8-e4m3 pool, standard configuration, sm_100): context matches - the bf16 numbers above - (same kernel, same inputs). Decode vs the fp32 quantization-aware - reference of the fp8 section: max abs err 3.1e-2 across a 10-seed sweep, - at most 44% of the `2^-4 * (1 + |ref|)` allowance; a 200-token-history - decode sat at 9.3e-3. A reference that skips the read-side dequant - diverges to ~4.4e-1 — the gate separates the failure mode by ~7x. The - 4-layer shared pool sits at the same levels per layer: context max abs - err 1.6e-2 (61% of the bf16-tolerance allowance — the highest fraction - observed on any certified case), decode max abs err 3.5e-2 (46% of the - fp8 allowance). -- Accuracy (fp8-e4m3 **latent** pool, MLA, sm_100, `H = 128`, page 32, - `s = 1.0`) vs fp32 references over e4m3-rounded operands, compared with - `atol = 2^-3, rtol = 2^-4`: context prefill (96 + 33 tokens) max abs err - 7.0e-2, 49% of the allowance; generation decode 2.1e-2 / 2.3e-2 over two - steps, 17-18%; the `s` = 1.5 and 2.0 decodes 2.1e-2, 16%. The floor is 2^-3 - rather than the standard configuration's 2^-4 because the residual is the - kernel's e4m3 handling of the softmax probabilities, whose error scales - with the `|V| ~ 1` rows rather than with the output element: at a 2^-4 - floor the context case runs 0.72-1.08 of the allowance over six seeds — - one seed *over* — because a `129 x 128 x 128` context output samples that - noise 16x more often than a two-row decode does. Rival separations from - the same runs: the bf16-KV reference (cache quantization ignored) 26x - outside the **bf16** band and 1.3-1.6x outside this one; the `s = 1.0` - math against `s` = 1.5 / 2.0 context runs 21x / 48x; a zeroed - `mla_bmm1_scale[1]` 4.2x; a decode whose bmm scales omit the kv scale - 2.9x at `s = 1.5` and 3.6x at `s = 2.0`. What the loosened gate cannot resolve is the query's - own e4m3 rounding in decode (0.23x, against the correct model's 0.17x), - so no claim rests on it; the "is the context math fp8" question is settled - bitwise instead by the peaked-softmax V readout described in that - section. -- Accuracy over the fp8 latent pool at the **complete DeepSeek-R1-0528 cell** - (YaRN table + `q_scaling = 1/mscale²`, same `atol = 2^-3, rtol = 2^-4`): - fresh-prefill context (96 + 33) 40.1% of the allowance, generation decode - 19.8% / 21.4% over two steps, mixed batch 36.4% (context) and 14.7% - (decode), no-append context (prefixes 96/31/0 reaching 128/40/25) 41.4%, - and the no-append realistic case of the operand/scale test 34.6% — the same - 14-49% envelope the baseline-cell fp8 cases sit in, so neither the rope - table nor the softmax scale costs headroom. Rival separations measured on - these runs: a reference at `q_scaling = 1.0` sits 12.7x outside the band in - fresh prefill, 10.8x in the no-append flavor and 4.2x in decode; an - unscaled-rope-table reference 3.7x in fresh prefill at 96 tokens (10.5x at - 512, since YaRN rescales only the low-frequency half of the spectrum and the - two tables diverge with distance); the bf16-KV reference 49.8x outside the - **bf16** band in fresh prefill, 58.1x in the mixed batch, and 2.2x outside - the fp8 band in the no-append flavor. The rope table is pinned far harder - by the append than by any of these: see the append bullet above. -- Softmax stats (no-append MLA context, sm_100): the emitted `(max, sum)` - match an fp32 reference over the same bf16 inputs to 2e-6 abs (max stat) - / 1.2e-6 rel (sum stat) — natural-log domain, scaled logits. Feeding each - pass's output/stats into `trtllm::merge_chunked_attention_for_mla` under - the production copy/merge/skip plan reproduced single-pass full-range - attention (output within the tolerances above, merged stats at the same - 1e-6-level agreement), certifying the stats as directly consumable by - that sibling op. -- V-stride hard-coding: for bf16 separate-QKV context MLA, the trtllm-gen - kernel computes V's row stride as `num_kv_heads * (head_size - 64 + - v_head_dim)` elements — the packed kv-projection width — rather than - reading it from the tensor (per source; the binding's `v_stride_in_bytes` - plumbing is unused in this version). A scratch probe passing a fully - contiguous `[Tkv, H*v_head_dim]` V returned wrong results (max abs err - ~4) and such a call reads out of bounds past the buffer. Always pass the - strided split view, in both context flavors. -- `chunked_prefill_buffer_batch_size` is consumed only for fp8-context-MLA - workspace sizing per source; inert in the certified bf16 configurations - (certified at 1, including the chunked partial passes). Every fp8 MLA - context case — both flavors — passes 1 as well and none was swept over it, - so the argument is certified at that one value on the path where source says - it is live. The chunked partial-pass *pattern* of the no-append flavor - (`mask_type=0` + `softmax_stats_tensor`) is not certified over an fp8 pool; - its one-shot cached-KV pattern is. -- Re-running the same prepared batch overwrites the same cache slots - (append position derives from `sequence_length` minus new-token count), - so a repeated call is idempotent, not double-appending. The MLA - generation and no-append MLA context calls write nothing at all (pool - bitwise unchanged). -- MLA generation head geometry: trtllm-gen decode kernels exist for - `(C+R, C)` of (576, 512) and (320, 256) per source; only (576, 512) is - certified, at 8, 16, 32 and 128 query heads (the compiled kernel reports - `HVPerCta128` at all four — the head count moves the q-tile, not the - per-CTA head-value width, and it stops moving it above 16). -- Sibling ops exist for adjacent roles: - `torch.ops.trtllm.attn_custom_op_inplace` and - `torch.ops.trtllm.mla_custom_op_inplace` (registered-layer wrappers over - the same attention core, state via model extra attrs), - `torch.ops.trtllm.create_attn_outputs` / `create_mla_outputs` - (output-buffer allocation), `torch.ops.trtllm.mla_rope_generation` - (MLA generation preprocessing: q_pe RoPE + latent append + the scheduler - buffers this op consumes), - `torch.ops.trtllm.mla_rope_append_paged_kv_assign_q`, - `load_paged_kv_cache_for_mla`, `load_chunked_kv_cache_for_mla` and - `merge_chunked_attention_for_mla` (the MLA cached-KV / chunked-prefill - context flow), and the split-phase trtllm-gen bindings - `thop.trtllm_gen_context_preprocess` / `thop.trtllm_gen_generation_preprocess` - / `thop.trtllm_gen_context_postprocess`. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py index 17815fdf2c96..3e1dfe14a272 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py @@ -3,278 +3,655 @@ """Fused attention core with fully explicit state: paged KV-cache append + masked FMHA. Wraps the pybind binding ``tensorrt_llm.bindings.internal.thop.attention`` -(approved policy exception to the ``torch.ops.trtllm.*`` entry shape): the -same C++ attention op behind the TRTLLM backend, but every piece of batch -state and layer config arrives as an explicit argument — no registered -layers, no thread-local metadata. +(approved policy exception to the ``torch.ops.trtllm.*`` entry shape): the same +C++ attention op behind the TRTLLM backend, but every piece of batch state and +layer config arrives as an explicit argument -- no registered layers, no +thread-local metadata. """ -from typing import Optional +from __future__ import annotations + +from typing import Optional, Sequence import torch from tensorrt_llm.bindings.internal import thop +from .._op import ULP, Arch, Cell, OpWrapper, assert_within_ulp + +#: attention_input_type. 0 is the standard append-and-attend path; 1 and 2 are +#: MLA's context and generation phases. +STANDARD, MLA_CONTEXT, MLA_GENERATION = 0, 1, 2 + +#: quant_mode bit for an fp8 KV cache. +QM_FP8_KV_CACHE = 128 + +#: mask_type 0 (AttentionMaskType.padding): with no pool, full attention +#: within each sequence. +MASK_BIDIRECTIONAL = 0 + +#: The deepseek-v4.1 vision tower's head geometry, the only one run without a +#: KV pool. +VIT_HEADS, VIT_HEAD_DIM = 16, 64 + +#: (why, per-slot lengths, max_context_q_len_override, q scales) of the vision +#: cells. Slot lengths sum to the call's row count; an override of None is the +#: eager one-sequence call, an override B the captured two-slot call over a +#: bucket of B rows. The q scales set the logit std (about 1 and 2). +_VIT_CELLS = ( + ("deepseek-v4.1 vision eager P=1296: one image, no pool, 16/16/64", (1296,), None, (1.0, 2.0)), + ("deepseek-v4.1 vision eager P=9189: the per-image ceiling", (9189,), None, (1.0, 2.0)), + ("deepseek-v4.1 vision graph B=1024 P=7: slots (P, B-P), override B", (7, 1017), 1024, (2.0,)), + ("deepseek-v4.1 vision graph B=9216 P=9189: the largest image", (9189, 27), 9216, (2.0,)), + ("deepseek-v4.1 vision graph B=9216 P=9216: an empty pad slot", (9216, 0), 9216, (2.0,)), +) + + +class ThopAttention(OpWrapper): + """Append this step's K/V to the paged cache, then attend over it. + + The most load-bearing entry in the catalog and the only one wrapping a + pybind binding rather than a registered torch op. Both shipped targets go + through it, by different paths: gpt-oss on the standard path with attention + sinks and an alternating sliding window, deepseek on MLA's two phases. + + Writes rows `[:num_tokens]` of `output` and returns `output` -- the same + buffer, mutated and handed back so a call site can show what it touched + without allocating. + """ -def thop_attention( - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - output: torch.Tensor, - output_sf: Optional[torch.Tensor], - workspace_: Optional[torch.Tensor], - sequence_length: torch.Tensor, - host_past_key_value_lengths: torch.Tensor, - host_total_kv_lens: torch.Tensor, - context_lengths: torch.Tensor, - host_context_lengths: torch.Tensor, - host_request_types: torch.Tensor, - max_context_q_len_override: Optional[int], - kv_cache_block_offsets: Optional[torch.Tensor], - host_kv_cache_pool_pointers: Optional[torch.Tensor], - host_kv_cache_pool_mapping: Optional[torch.Tensor], - cache_indirection: Optional[torch.Tensor], - kv_scale_orig_quant: Optional[torch.Tensor], - kv_scale_quant_orig: Optional[torch.Tensor], - out_scale: Optional[torch.Tensor], - rotary_inv_freq: Optional[torch.Tensor], - rotary_cos_sin: Optional[torch.Tensor], - latent_cache: Optional[torch.Tensor], - q_pe: Optional[torch.Tensor], - block_ids_per_seq: Optional[torch.Tensor], - attention_sinks: Optional[torch.Tensor], - is_fused_qkv: bool, - update_kv_cache: bool, - predicted_tokens_per_seq: int, - local_layer_idx: int, - num_heads: int, - num_kv_heads: int, - head_size: int, - tokens_per_block: Optional[int], - max_num_requests: int, - max_context_length: int, - max_seq_len: int, - attention_window_size: int, - beam_width: int, - mask_type: int, - quant_mode: int, - q_scaling: float, - position_embedding_type: int, - rope_dim: int, - rope_base: float, - rope_scale_type: int, - rope_scale: float, - rope_short_m_scale: float, - rope_long_m_scale: float, - rope_max_positions: int, - rope_original_max_positions: int, - use_paged_context_fmha: bool, - attention_input_type: Optional[int], - is_mla_enable: bool, - chunked_prefill_buffer_batch_size: Optional[int], - q_lora_rank: Optional[int], - kv_lora_rank: Optional[int], - qk_nope_head_dim: Optional[int], - qk_rope_head_dim: Optional[int], - v_head_dim: Optional[int], - rope_append: Optional[bool], - mrope_rotary_cos_sin: Optional[torch.Tensor], - mrope_position_deltas: Optional[torch.Tensor], - helix_position_offsets: Optional[torch.Tensor], - helix_is_inactive_rank: Optional[torch.Tensor], - attention_chunk_size: Optional[int], - softmax_stats_tensor: Optional[torch.Tensor], - is_spec_decoding_enabled: bool, - use_spec_decoding: bool, - is_spec_dec_tree: bool, - spec_decoding_generation_lengths: Optional[torch.Tensor], - spec_decoding_position_offsets_for_cpp: Optional[torch.Tensor], - spec_decoding_packed_mask: Optional[torch.Tensor], - spec_decoding_bl_tree_mask_offset: Optional[torch.Tensor], - spec_decoding_bl_tree_mask: Optional[torch.Tensor], - spec_bl_tree_first_sparse_mask_offset_kv: Optional[torch.Tensor], - sparse_kv_indices: Optional[torch.Tensor], - sparse_kv_offsets: Optional[torch.Tensor], - sparse_attn_indices: Optional[torch.Tensor], - sparse_attn_offsets: Optional[torch.Tensor], - sparse_attn_indices_block_size: int, - num_sparse_topk: Optional[int] = None, - sparse_attn_kv_lens: Optional[torch.Tensor] = None, - skip_softmax_threshold_scale_factor_prefill: Optional[float] = None, - skip_softmax_threshold_scale_factor_decode: Optional[float] = None, - skip_softmax_stat: Optional[torch.Tensor] = None, - cu_q_seqlens: Optional[torch.Tensor] = None, - cu_kv_seqlens: Optional[torch.Tensor] = None, - fmha_scheduler_counter: Optional[torch.Tensor] = None, - mla_bmm1_scale: Optional[torch.Tensor] = None, - mla_bmm2_scale: Optional[torch.Tensor] = None, - quant_q_buffer: Optional[torch.Tensor] = None, - flash_mla_tile_scheduler_metadata: Optional[torch.Tensor] = None, - flash_mla_num_splits: Optional[torch.Tensor] = None, - sage_attn_num_elts_per_blk_q: int = 0, - sage_attn_num_elts_per_blk_k: int = 0, - sage_attn_num_elts_per_blk_v: int = 0, - sage_attn_qk_int8: bool = False, - num_contexts: int = 0, - num_ctx_tokens: int = 0, - trtllm_gen_jit_warmup: bool = False, - aux_kv_cache_pool_ptr: Optional[int] = None, - is_cross: bool = False, - cross_kv: Optional[torch.Tensor] = None, - relative_attention_bias: Optional[torch.Tensor] = None, - relative_attention_max_distance: int = 0, - spec_decoding_target_max_draft_tokens: Optional[int] = None, - quant_scale_qkv: Optional[torch.Tensor] = None, - dsv4_inv_rope_cos_sin_cache: Optional[torch.Tensor] = None, - enable_dsv4_epilogue_fusion: bool = False, - # Defaults match what the in-tree caller - # (attention/backends/fmha/fallback.py) passes on a dense, non-sparse, - # non-folded path -- which is the path both migrated targets are on. - max_num_sequences: Optional[int] = None, - force_prepare_spec_dec_tree_mask: bool = False, - kv_norm_weight: Optional[torch.Tensor] = None, - kv_norm_eps: float = 1e-6, - skip_correction_threshold: float = 0.0, -) -> None: - """Run the fused attention core over one prepared batch — paged KV-cache - append + masked FMHA (standard), or the MLA context/generation phase - selected by attention_input_type; writes rows [:num_tokens] of output. - Returns None.""" - # num_heads must be an integer multiple of num_kv_heads. A context-only - # call with a non-multiple returns without raising on sm_100 (observed at - # 6q/4kv d128: only the first (num_heads // num_kv_heads) * num_kv_heads - # head columns of output are computed, the rest come back all-zero), so - # the geometry cannot be left for the op to reject. - assert num_heads % num_kv_heads == 0, ( - f"num_heads ({num_heads}) must be a multiple of num_kv_heads ({num_kv_heads})" + ARCHS = frozenset({Arch.SM_103}) + + # The four shapes the shipped targets reach, which is four different code + # paths rather than four sizes. gpt-oss alternates its two window classes + # inside one engine, so both of its cells are live every step; deepseek runs + # context and generation against the same pool. + CELLS: tuple[Cell, ...] = ( + Cell( + why="gpt-oss sliding layer: GQA 64q/8kv d64, sinks, window 128", + spec=dict( + input_type=STANDARD, + num_heads=64, + num_kv_heads=8, + head_size=64, + v_head_dim=64, + sinks=True, + window=128, + quant_mode=0, + ), + ), + Cell( + why="gpt-oss full layer: the same heads with no window", + spec=dict( + input_type=STANDARD, + num_heads=64, + num_kv_heads=8, + head_size=64, + v_head_dim=64, + sinks=True, + window=None, + quant_mode=0, + ), + ), + Cell( + why="deepseek MLA context: explicit k/v, qk 576 -> v 128, fp8 pool", + spec=dict( + input_type=MLA_CONTEXT, + num_heads=128, + num_kv_heads=128, + head_size=576, + v_head_dim=128, + sinks=False, + window=None, + quant_mode=QM_FP8_KV_CACHE, + ), + ), + Cell( + why="deepseek MLA generation: one kv head, latent 576 -> 512, fp8 pool", + spec=dict( + input_type=MLA_GENERATION, + num_heads=128, + num_kv_heads=1, + head_size=576, + v_head_dim=512, + sinks=False, + window=None, + quant_mode=QM_FP8_KV_CACHE, + ), + ), + # The deepseek-v4.1 vision tower: the no-KV-cache mode, one image per + # call, eagerly over one sequence or captured over two slots (image, + # pad) with the launch fixed by `max_context_q_len_override`. Cells at + # the ends of what was measured; the sweeps between them are tests. + *( + Cell( + why=why, + spec=dict( + input_type=STANDARD, + num_heads=VIT_HEADS, + num_kv_heads=VIT_HEADS, + head_size=VIT_HEAD_DIM, + v_head_dim=VIT_HEAD_DIM, + sinks=False, + window=None, + quant_mode=0, + kv_cache=False, + mask_type=MASK_BIDIRECTIONAL, + seq_lens=seq_lens, + q_len_override=q_len_override, + cuda_graph=q_len_override is not None, + q_scales=q_scales, + ), + ) + for why, seq_lens, q_len_override, q_scales in _VIT_CELLS + ), ) - # attention_sinks is consumed as a raw buffer of exactly num_heads fp32 - # values read from data_ptr(); only its dtype is validated by the op. On - # sm_100 a stride-2 view whose values were correct produced silently wrong - # output (the underlying memory is what gets read), a 32-element tensor was - # read 64 elements deep past its end, a 128-element one had its tail - # silently ignored, and an empty tensor silently disabled the sink. - if attention_sinks is not None: - assert attention_sinks.is_contiguous() and attention_sinks.numel() == num_heads, ( - f"attention_sinks must be contiguous with exactly num_heads " - f"({num_heads}) elements, got shape " - f"{tuple(attention_sinks.shape)} contiguous=" - f"{attention_sinks.is_contiguous()}" + + note = """ + KV addressing over one pool and two is certified. gpt-oss alternates + sliding-window and full attention, and KVCacheManagerV2 gives each window + class its own layer group, so its mapping carries two pool ids; the op reads + the pool column itself and this entry passes the mapping through untouched. + A third pool is outside what was measured. + + `reference` covers the masked-softmax half. The append half -- that this + step's K/V land in the right pool slots -- is what the paged-cache entries + beside this one certify, and what the matrix here checks by reading the pool + back rather than by a reference. + + The softmax statistics the MLA chunked-prefill path writes are a second + output; `softmax_stats` states them separately so `reference` stays one + thing. + + DeepSeek-V4.1: the vision tower calls this op in the no-KV-cache mode -- + every pool argument and `tokens_per_block` None, so the op runs the + packed-QKV context FMHA with no append and no persistent state. bf16, + 16/16/64 heads, `is_fused_qkv` on the packed `[T, 3072]` rows, `mask_type` + 0 (bidirectional within each sequence, sequences never see each other), + `q_scaling` 1, `position_embedding_type` 0 (the 2D RoPE is the caller's), + `update_kv_cache=True` (False raises), every request a context one + (`num_contexts` = slots, `num_ctx_tokens` = T). One image per call: eagerly + over one sequence of P <= 9189 rows with `cu_q_seqlens` = (0, P); or + captured once per padded bucket B in 1024..9216, over two slots (image P, + pad B - P) with `max_context_q_len_override = B` and the host lengths + holding all B rows as one sequence (`[B, 0]`), after which a replay follows + whatever split the device `sequence_length` / `context_lengths` / + `cu_q_seqlens` hold -- bitwise what an eager call over that split returns. + Metadata consumed: the per-slot lengths in `sequence_length` and + `context_lengths` (device), `host_past_key_value_lengths` and + `host_context_lengths` (host), `host_total_kv_lens = [T, 0]`, + `cu_q_seqlens` (optional; None is bitwise the same as the true offsets). + The cells sit at the ends; the measured ranges are the tests': sequences + of 1..9216 rows and up to 16 sequences / 17037 rows per call against the + full reference, buckets 1024..9216 with P from 7 to B (an empty pad slot + included) and four-slot splits with empty slots, and on sampled rows up to + 702960 rows per call and 262144 per sequence. Gated by `compare` and, in + the tests, by the bf16 atol 5e-3 / rtol 1.6e-2 band the measurements used + (at most 0.46 of it at logit std 1 and 2), with a causal mask and + attention across images each required at least 10x outside it. Silently + wrong, and not rejected: `context_lengths` differing from + `sequence_length` (neither split comes back); a given `cu_q_seqlens` is + the split the FMHA uses, over `context_lengths`; `mask_type` 1 is causal + within each sequence; a capture without the override is fixed to its + capture-time split, and a replay over another split comes back wrong or + non-finite. These are tensor contents, so `is_valid` cannot see them. + Inert: `tokens_per_block` 32 vs None, `max_seq_len` / + `max_context_length` / `attention_window_size` (even below the longest + sequence), `attention_input_type` 0 vs 1; rows `[T:]` of `output` are never + written. `workspace_` grows in place, so a graph needs an eager call to + size it before capture. Not certified for DeepSeek-V4.1: another head + geometry, GQA/MQA, sinks, a window, paged context, `q_scaling` != 1, + in-kernel RoPE, fp8, generation rows, an override below the longest + sequence or on a cache-backed call, and a softmax sharper than logit std 2 + (one seed reached 1.03 of the band at std 3). + """ + + def raw_call( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + output: torch.Tensor, + output_sf: Optional[torch.Tensor], + workspace_: Optional[torch.Tensor], + sequence_length: torch.Tensor, + host_past_key_value_lengths: torch.Tensor, + host_total_kv_lens: torch.Tensor, + context_lengths: torch.Tensor, + host_context_lengths: torch.Tensor, + host_request_types: torch.Tensor, + max_context_q_len_override: Optional[int], + kv_cache_block_offsets: Optional[torch.Tensor], + host_kv_cache_pool_pointers: Optional[torch.Tensor], + host_kv_cache_pool_mapping: Optional[torch.Tensor], + cache_indirection: Optional[torch.Tensor], + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + out_scale: Optional[torch.Tensor], + rotary_inv_freq: Optional[torch.Tensor], + rotary_cos_sin: Optional[torch.Tensor], + latent_cache: Optional[torch.Tensor], + q_pe: Optional[torch.Tensor], + block_ids_per_seq: Optional[torch.Tensor], + attention_sinks: Optional[torch.Tensor], + is_fused_qkv: bool, + update_kv_cache: bool, + predicted_tokens_per_seq: int, + local_layer_idx: int, + num_heads: int, + num_kv_heads: int, + head_size: int, + tokens_per_block: Optional[int], + max_num_requests: int, + max_context_length: int, + max_seq_len: int, + attention_window_size: int, + beam_width: int, + mask_type: int, + quant_mode: int, + q_scaling: float, + position_embedding_type: int, + rope_dim: int, + rope_base: float, + rope_scale_type: int, + rope_scale: float, + rope_short_m_scale: float, + rope_long_m_scale: float, + rope_max_positions: int, + rope_original_max_positions: int, + use_paged_context_fmha: bool, + attention_input_type: Optional[int], + is_mla_enable: bool, + chunked_prefill_buffer_batch_size: Optional[int], + q_lora_rank: Optional[int], + kv_lora_rank: Optional[int], + qk_nope_head_dim: Optional[int], + qk_rope_head_dim: Optional[int], + v_head_dim: Optional[int], + rope_append: Optional[bool], + mrope_rotary_cos_sin: Optional[torch.Tensor], + mrope_position_deltas: Optional[torch.Tensor], + helix_position_offsets: Optional[torch.Tensor], + helix_is_inactive_rank: Optional[torch.Tensor], + attention_chunk_size: Optional[int], + softmax_stats_tensor: Optional[torch.Tensor], + is_spec_decoding_enabled: bool, + use_spec_decoding: bool, + is_spec_dec_tree: bool, + spec_decoding_generation_lengths: Optional[torch.Tensor], + spec_decoding_position_offsets_for_cpp: Optional[torch.Tensor], + spec_decoding_packed_mask: Optional[torch.Tensor], + spec_decoding_bl_tree_mask_offset: Optional[torch.Tensor], + spec_decoding_bl_tree_mask: Optional[torch.Tensor], + spec_bl_tree_first_sparse_mask_offset_kv: Optional[torch.Tensor], + sparse_kv_indices: Optional[torch.Tensor], + sparse_kv_offsets: Optional[torch.Tensor], + sparse_attn_indices: Optional[torch.Tensor], + sparse_attn_offsets: Optional[torch.Tensor], + sparse_attn_indices_block_size: int, + num_sparse_topk: Optional[int] = None, + sparse_attn_kv_lens: Optional[torch.Tensor] = None, + skip_softmax_threshold_scale_factor_prefill: Optional[float] = None, + skip_softmax_threshold_scale_factor_decode: Optional[float] = None, + skip_softmax_stat: Optional[torch.Tensor] = None, + cu_q_seqlens: Optional[torch.Tensor] = None, + cu_kv_seqlens: Optional[torch.Tensor] = None, + fmha_scheduler_counter: Optional[torch.Tensor] = None, + mla_bmm1_scale: Optional[torch.Tensor] = None, + mla_bmm2_scale: Optional[torch.Tensor] = None, + quant_q_buffer: Optional[torch.Tensor] = None, + flash_mla_tile_scheduler_metadata: Optional[torch.Tensor] = None, + flash_mla_num_splits: Optional[torch.Tensor] = None, + sage_attn_num_elts_per_blk_q: int = 0, + sage_attn_num_elts_per_blk_k: int = 0, + sage_attn_num_elts_per_blk_v: int = 0, + sage_attn_qk_int8: bool = False, + num_contexts: int = 0, + num_ctx_tokens: int = 0, + trtllm_gen_jit_warmup: bool = False, + aux_kv_cache_pool_ptr: Optional[int] = None, + is_cross: bool = False, + cross_kv: Optional[torch.Tensor] = None, + relative_attention_bias: Optional[torch.Tensor] = None, + relative_attention_max_distance: int = 0, + spec_decoding_target_max_draft_tokens: Optional[int] = None, + quant_scale_qkv: Optional[torch.Tensor] = None, + dsv4_inv_rope_cos_sin_cache: Optional[torch.Tensor] = None, + enable_dsv4_epilogue_fusion: bool = False, + # Defaults match what the in-tree caller + # (attention/backends/fmha/fallback.py) passes on a dense, non-sparse, + # non-folded path -- which is the path both migrated targets are on. + max_num_sequences: Optional[int] = None, + force_prepare_spec_dec_tree_mask: bool = False, + kv_norm_weight: Optional[torch.Tensor] = None, + kv_norm_eps: float = 1e-6, + skip_correction_threshold: float = 0.0, + ) -> torch.Tensor: + """Writes `output` in place and returns it -- the op destroys its + prior contents in rows `[:num_tokens]`, so a caller that wants them + has to have kept its own copy beforehand.""" + thop.attention( + q=q, + k=k, + v=v, + output=output, + output_sf=output_sf, + workspace_=workspace_, + sequence_length=sequence_length, + host_past_key_value_lengths=host_past_key_value_lengths, + host_total_kv_lens=host_total_kv_lens, + context_lengths=context_lengths, + host_context_lengths=host_context_lengths, + host_request_types=host_request_types, + max_context_q_len_override=max_context_q_len_override, + kv_cache_block_offsets=kv_cache_block_offsets, + host_kv_cache_pool_pointers=host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=host_kv_cache_pool_mapping, + cache_indirection=cache_indirection, + kv_scale_orig_quant=kv_scale_orig_quant, + kv_scale_quant_orig=kv_scale_quant_orig, + out_scale=out_scale, + rotary_inv_freq=rotary_inv_freq, + rotary_cos_sin=rotary_cos_sin, + latent_cache=latent_cache, + q_pe=q_pe, + block_ids_per_seq=block_ids_per_seq, + attention_sinks=attention_sinks, + is_fused_qkv=is_fused_qkv, + update_kv_cache=update_kv_cache, + predicted_tokens_per_seq=predicted_tokens_per_seq, + local_layer_idx=local_layer_idx, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_size=head_size, + tokens_per_block=tokens_per_block, + max_num_requests=max_num_requests, + max_context_length=max_context_length, + max_seq_len=max_seq_len, + attention_window_size=attention_window_size, + beam_width=beam_width, + mask_type=mask_type, + quant_mode=quant_mode, + q_scaling=q_scaling, + position_embedding_type=position_embedding_type, + rope_dim=rope_dim, + rope_base=rope_base, + rope_scale_type=rope_scale_type, + rope_scale=rope_scale, + rope_short_m_scale=rope_short_m_scale, + rope_long_m_scale=rope_long_m_scale, + rope_max_positions=rope_max_positions, + rope_original_max_positions=rope_original_max_positions, + use_paged_context_fmha=use_paged_context_fmha, + attention_input_type=attention_input_type, + is_mla_enable=is_mla_enable, + chunked_prefill_buffer_batch_size=chunked_prefill_buffer_batch_size, + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + rope_append=rope_append, + mrope_rotary_cos_sin=mrope_rotary_cos_sin, + mrope_position_deltas=mrope_position_deltas, + helix_position_offsets=helix_position_offsets, + helix_is_inactive_rank=helix_is_inactive_rank, + attention_chunk_size=attention_chunk_size, + softmax_stats_tensor=softmax_stats_tensor, + is_spec_decoding_enabled=is_spec_decoding_enabled, + use_spec_decoding=use_spec_decoding, + is_spec_dec_tree=is_spec_dec_tree, + spec_decoding_generation_lengths=spec_decoding_generation_lengths, + spec_decoding_position_offsets_for_cpp=spec_decoding_position_offsets_for_cpp, + spec_decoding_packed_mask=spec_decoding_packed_mask, + spec_decoding_bl_tree_mask_offset=spec_decoding_bl_tree_mask_offset, + spec_decoding_bl_tree_mask=spec_decoding_bl_tree_mask, + spec_bl_tree_first_sparse_mask_offset_kv=spec_bl_tree_first_sparse_mask_offset_kv, + sparse_kv_indices=sparse_kv_indices, + sparse_kv_offsets=sparse_kv_offsets, + sparse_attn_indices=sparse_attn_indices, + sparse_attn_offsets=sparse_attn_offsets, + sparse_attn_indices_block_size=sparse_attn_indices_block_size, + num_sparse_topk=num_sparse_topk, + sparse_attn_kv_lens=sparse_attn_kv_lens, + skip_softmax_threshold_scale_factor_prefill=skip_softmax_threshold_scale_factor_prefill, + skip_softmax_threshold_scale_factor_decode=skip_softmax_threshold_scale_factor_decode, + skip_softmax_stat=skip_softmax_stat, + cu_q_seqlens=cu_q_seqlens, + cu_kv_seqlens=cu_kv_seqlens, + fmha_scheduler_counter=fmha_scheduler_counter, + mla_bmm1_scale=mla_bmm1_scale, + mla_bmm2_scale=mla_bmm2_scale, + quant_q_buffer=quant_q_buffer, + flash_mla_tile_scheduler_metadata=flash_mla_tile_scheduler_metadata, + flash_mla_num_splits=flash_mla_num_splits, + sage_attn_num_elts_per_blk_q=sage_attn_num_elts_per_blk_q, + sage_attn_num_elts_per_blk_k=sage_attn_num_elts_per_blk_k, + sage_attn_num_elts_per_blk_v=sage_attn_num_elts_per_blk_v, + sage_attn_qk_int8=sage_attn_qk_int8, + num_contexts=num_contexts, + num_ctx_tokens=num_ctx_tokens, + trtllm_gen_jit_warmup=trtllm_gen_jit_warmup, + aux_kv_cache_pool_ptr=aux_kv_cache_pool_ptr, + is_cross=is_cross, + cross_kv=cross_kv, + relative_attention_bias=relative_attention_bias, + relative_attention_max_distance=relative_attention_max_distance, + spec_decoding_target_max_draft_tokens=spec_decoding_target_max_draft_tokens, + quant_scale_qkv=quant_scale_qkv, + dsv4_inv_rope_cos_sin_cache=dsv4_inv_rope_cos_sin_cache, + enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, + max_num_sequences=max_num_sequences, + force_prepare_spec_dec_tree_mask=force_prepare_spec_dec_tree_mask, + kv_norm_weight=kv_norm_weight, + kv_norm_eps=kv_norm_eps, + skip_correction_threshold=skip_correction_threshold, ) - thop.attention( - q=q, - k=k, - v=v, - output=output, - output_sf=output_sf, - workspace_=workspace_, - sequence_length=sequence_length, - host_past_key_value_lengths=host_past_key_value_lengths, - host_total_kv_lens=host_total_kv_lens, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - host_request_types=host_request_types, - max_context_q_len_override=max_context_q_len_override, - kv_cache_block_offsets=kv_cache_block_offsets, - host_kv_cache_pool_pointers=host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=host_kv_cache_pool_mapping, - cache_indirection=cache_indirection, - kv_scale_orig_quant=kv_scale_orig_quant, - kv_scale_quant_orig=kv_scale_quant_orig, - out_scale=out_scale, - rotary_inv_freq=rotary_inv_freq, - rotary_cos_sin=rotary_cos_sin, - latent_cache=latent_cache, - q_pe=q_pe, - block_ids_per_seq=block_ids_per_seq, - attention_sinks=attention_sinks, - is_fused_qkv=is_fused_qkv, - update_kv_cache=update_kv_cache, - predicted_tokens_per_seq=predicted_tokens_per_seq, - local_layer_idx=local_layer_idx, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - head_size=head_size, - tokens_per_block=tokens_per_block, - max_num_requests=max_num_requests, - max_context_length=max_context_length, - max_seq_len=max_seq_len, - attention_window_size=attention_window_size, - beam_width=beam_width, - mask_type=mask_type, - quant_mode=quant_mode, - q_scaling=q_scaling, - position_embedding_type=position_embedding_type, - rope_dim=rope_dim, - rope_base=rope_base, - rope_scale_type=rope_scale_type, - rope_scale=rope_scale, - rope_short_m_scale=rope_short_m_scale, - rope_long_m_scale=rope_long_m_scale, - rope_max_positions=rope_max_positions, - rope_original_max_positions=rope_original_max_positions, - use_paged_context_fmha=use_paged_context_fmha, - attention_input_type=attention_input_type, - is_mla_enable=is_mla_enable, - chunked_prefill_buffer_batch_size=chunked_prefill_buffer_batch_size, - q_lora_rank=q_lora_rank, - kv_lora_rank=kv_lora_rank, - qk_nope_head_dim=qk_nope_head_dim, - qk_rope_head_dim=qk_rope_head_dim, - v_head_dim=v_head_dim, - rope_append=rope_append, - mrope_rotary_cos_sin=mrope_rotary_cos_sin, - mrope_position_deltas=mrope_position_deltas, - helix_position_offsets=helix_position_offsets, - helix_is_inactive_rank=helix_is_inactive_rank, - attention_chunk_size=attention_chunk_size, - softmax_stats_tensor=softmax_stats_tensor, - is_spec_decoding_enabled=is_spec_decoding_enabled, - use_spec_decoding=use_spec_decoding, - is_spec_dec_tree=is_spec_dec_tree, - spec_decoding_generation_lengths=spec_decoding_generation_lengths, - spec_decoding_position_offsets_for_cpp=spec_decoding_position_offsets_for_cpp, - spec_decoding_packed_mask=spec_decoding_packed_mask, - spec_decoding_bl_tree_mask_offset=spec_decoding_bl_tree_mask_offset, - spec_decoding_bl_tree_mask=spec_decoding_bl_tree_mask, - spec_bl_tree_first_sparse_mask_offset_kv=spec_bl_tree_first_sparse_mask_offset_kv, - sparse_kv_indices=sparse_kv_indices, - sparse_kv_offsets=sparse_kv_offsets, - sparse_attn_indices=sparse_attn_indices, - sparse_attn_offsets=sparse_attn_offsets, - sparse_attn_indices_block_size=sparse_attn_indices_block_size, - num_sparse_topk=num_sparse_topk, - sparse_attn_kv_lens=sparse_attn_kv_lens, - skip_softmax_threshold_scale_factor_prefill=skip_softmax_threshold_scale_factor_prefill, - skip_softmax_threshold_scale_factor_decode=skip_softmax_threshold_scale_factor_decode, - skip_softmax_stat=skip_softmax_stat, - cu_q_seqlens=cu_q_seqlens, - cu_kv_seqlens=cu_kv_seqlens, - fmha_scheduler_counter=fmha_scheduler_counter, - mla_bmm1_scale=mla_bmm1_scale, - mla_bmm2_scale=mla_bmm2_scale, - quant_q_buffer=quant_q_buffer, - flash_mla_tile_scheduler_metadata=flash_mla_tile_scheduler_metadata, - flash_mla_num_splits=flash_mla_num_splits, - sage_attn_num_elts_per_blk_q=sage_attn_num_elts_per_blk_q, - sage_attn_num_elts_per_blk_k=sage_attn_num_elts_per_blk_k, - sage_attn_num_elts_per_blk_v=sage_attn_num_elts_per_blk_v, - sage_attn_qk_int8=sage_attn_qk_int8, - num_contexts=num_contexts, - num_ctx_tokens=num_ctx_tokens, - trtllm_gen_jit_warmup=trtllm_gen_jit_warmup, - aux_kv_cache_pool_ptr=aux_kv_cache_pool_ptr, - is_cross=is_cross, - cross_kv=cross_kv, - relative_attention_bias=relative_attention_bias, - relative_attention_max_distance=relative_attention_max_distance, - spec_decoding_target_max_draft_tokens=spec_decoding_target_max_draft_tokens, - quant_scale_qkv=quant_scale_qkv, - dsv4_inv_rope_cos_sin_cache=dsv4_inv_rope_cos_sin_cache, - enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, - max_num_sequences=max_num_sequences, - force_prepare_spec_dec_tree_mask=force_prepare_spec_dec_tree_mask, - kv_norm_weight=kv_norm_weight, - kv_norm_eps=kv_norm_eps, - skip_correction_threshold=skip_correction_threshold, - ) + return output + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_lens: Sequence[int], + kv_lens: Sequence[int], + num_heads: int, + num_kv_heads: int, + qk_head_dim: int, + v_head_dim: int, + softmax_scale: float, + causal: bool = True, + attention_sinks: Optional[torch.Tensor] = None, + window: Optional[int] = None, + e4m3_inputs: bool = False, + out_scale: float = 1.0, + ) -> torch.Tensor: + """Masked softmax attention over explicit per-sequence K/V, in fp32. + + The paged-cache exception to mirroring `raw_call`, and the largest one: + this op's output is a function of what the pool holds, and its 115 + arguments mostly say where to look. `k` and `v` are the gathered + histories the call should have attended over; the test reads them out of + the environment, which is the only place they exist. + + Causal is bottom-right aligned -- query `i` of a sequence sits at + absolute position `kv_len - q_len + i`, which is what makes a decode + step's single query attend to everything before it. + + `attention_sinks` adds one per-head logit to the softmax *denominator* + and drops it from the numerator, so the attention weights sum to less + than one. It is not scaled by `softmax_scale`: the rival hypothesis that + it is was measured and rejected, and the test still drives both so the + distinction cannot quietly stop holding. + + `window` keeps only the newest `window` keys per query row. It is a pure + mask -- the pool is shared with the full-attention layers and nothing is + evicted -- which is why one cell can differ from another by this alone. + + `e4m3_inputs` rounds q, k and v through e4m3, which is what the fp8 path + quantizes them to; `out_scale` is the second of that path's two + dequantization factors, the first riding in `softmax_scale`. + """ + rep = num_heads // num_kv_heads + total_q = sum(q_lens) + out = torch.full( + (total_q, num_heads * v_head_dim), float("nan"), dtype=torch.float32, device=q.device + ) + q_start = kv_start = 0 + for q_len, kv_len in zip(q_lens, kv_lens): + if kv_len == 0: + # The op leaves these rows undefined and the merge plans skip + # them; saying so here keeps a caller from reading them. + q_start += q_len + continue + q_seq = q[q_start : q_start + q_len].reshape(q_len, num_heads, qk_head_dim).float() + k_seq = k[kv_start : kv_start + kv_len].reshape(kv_len, num_kv_heads, qk_head_dim) + v_seq = v[kv_start : kv_start + kv_len].reshape(kv_len, num_kv_heads, v_head_dim) + k_seq, v_seq = k_seq.float(), v_seq.float() + if rep > 1: + k_seq = k_seq.repeat_interleave(rep, dim=1) + v_seq = v_seq.repeat_interleave(rep, dim=1) + if e4m3_inputs: + q_seq = q_seq.to(torch.float8_e4m3fn).float() + k_seq = k_seq.to(torch.float8_e4m3fn).float() + v_seq = v_seq.to(torch.float8_e4m3fn).float() + + scores = torch.einsum("ihd,jhd->hij", q_seq, k_seq) * softmax_scale + keep = self._mask(q_len, kv_len, causal, window, q.device) + scores = scores.masked_fill(~keep.unsqueeze(0), float("-inf")) + + if attention_sinks is None: + probs = torch.softmax(scores, dim=-1) + else: + sink = attention_sinks.float().view(num_heads, 1) + peak = torch.maximum(scores.max(dim=-1).values, sink) + weights = torch.exp(scores - peak.unsqueeze(-1)) + probs = weights / (weights.sum(-1) + torch.exp(sink - peak)).unsqueeze(-1) + + out[q_start : q_start + q_len] = ( + torch.einsum("hij,jhd->ihd", probs, v_seq).reshape(q_len, -1) * out_scale + ) + q_start += q_len + kv_start += kv_len + return out.to(q.dtype) + + def softmax_stats( + self, + q: torch.Tensor, + k: torch.Tensor, + q_lens: Sequence[int], + kv_lens: Sequence[int], + num_heads: int, + num_kv_heads: int, + qk_head_dim: int, + softmax_scale: float, + causal: bool = True, + window: Optional[int] = None, + ) -> torch.Tensor: + """The second output the MLA chunked-prefill path writes, `[Tq, H, 2]`. + + Per (token, head): the row max of the scaled logits and the sum of + `exp(logit - max)` over *this pass's* KV range. Separate from + `reference` because it is one path's extra product, not part of what + attention computes, and folding it in would make every other cell carry + a return value it has no claim about. + + Rows of zero-KV sequences stay NaN, as the op leaves them. + """ + rep = num_heads // num_kv_heads + total_q = sum(q_lens) + stats = torch.full( + (total_q, num_heads, 2), float("nan"), dtype=torch.float32, device=q.device + ) + q_start = kv_start = 0 + for q_len, kv_len in zip(q_lens, kv_lens): + if kv_len == 0: + q_start += q_len + continue + q_seq = q[q_start : q_start + q_len].reshape(q_len, num_heads, qk_head_dim).float() + k_seq = k[kv_start : kv_start + kv_len].reshape(kv_len, num_kv_heads, qk_head_dim) + k_seq = k_seq.float() + if rep > 1: + k_seq = k_seq.repeat_interleave(rep, dim=1) + scores = torch.einsum("ihd,jhd->hij", q_seq, k_seq) * softmax_scale + keep = self._mask(q_len, kv_len, causal, window, q.device) + scores = scores.masked_fill(~keep.unsqueeze(0), float("-inf")) + row_max = scores.max(dim=-1).values + row_sum = torch.exp(scores - row_max.unsqueeze(-1)).sum(dim=-1) + stats[q_start : q_start + q_len, :, 0] = row_max.transpose(0, 1) + stats[q_start : q_start + q_len, :, 1] = row_sum.transpose(0, 1) + q_start += q_len + kv_start += kv_len + return stats + + @staticmethod + def _mask(q_len: int, kv_len: int, causal: bool, window: Optional[int], device) -> torch.Tensor: + """`[q_len, kv_len]` keep-mask, bottom-right aligned.""" + if not causal and window is None: + return torch.ones(q_len, kv_len, dtype=torch.bool, device=device) + keep = torch.zeros(q_len, kv_len, dtype=torch.bool, device=device) + offset = kv_len - q_len + for i in range(q_len): + high = offset + i + 1 if causal else kv_len + low = 0 if window is None else max(0, high - window) + keep[i, low:high] = True + return keep + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """A softmax chain, so the band is the longest in this catalog. + + Attention accumulates twice -- once over the KV range inside the + softmax, once over the same range in the value product -- and the + kernel splits both differently from a single fp32 pass. 8 ulp + element-wise against the row scale and 4 of relative RMS, the same band + a two-GEMM MoE chain gets, and for the same reason. + + Scored in bf16 ulp whatever the output dtype: on the fp8 paths the + operands are e4m3 and the accumulation is bf16-limited, so fp32's ulp + would measure the container. + """ + assert_within_ulp(out, ref, element_ulp=8.0, rms_ulp=4.0, ulp=ULP[torch.bfloat16]) + + def is_valid( + self, + *, + num_heads: int, + num_kv_heads: int, + attention_sinks: Optional[torch.Tensor] = None, + **unused_kwargs: object, + ) -> None: + """Keyword-only, because every call site passes keywords. + + The one entry that does not mirror `raw_call`. This op takes 115 + arguments; naming all of them here would restate the signature without + adding a claim, and `reference` cannot mirror it either -- it is the + paged-cache exception and takes gathered K/V -- so there is no shared + argument list left for the mirroring to preserve. `**unused_kwargs` + absorbs the rest. + """ + # A context-only call with a non-multiple returns without raising: + # measured at 6q/4kv d128, only the first + # `(num_heads // num_kv_heads) * num_kv_heads` head columns of output are + # computed and the rest come back all-zero. The geometry cannot be left + # for the op to reject. + assert num_heads % num_kv_heads == 0, ( + f"num_heads ({num_heads}) must be a multiple of num_kv_heads ({num_kv_heads})" + ) + # attention_sinks is consumed as a raw buffer of exactly num_heads fp32 + # values read from data_ptr(); only its dtype is validated by the op. + # Measured: a stride-2 view whose values were correct produced silently + # wrong output (the underlying memory is what gets read), a 32-element + # tensor was read 64 elements deep past its end, a 128-element one had + # its tail silently ignored, and an empty tensor silently disabled the + # sink. + if attention_sinks is not None: + assert attention_sinks.is_contiguous() and attention_sinks.numel() == num_heads, ( + f"attention_sinks must be contiguous with exactly num_heads " + f"({num_heads}) elements, got shape " + f"{tuple(attention_sinks.shape)} contiguous=" + f"{attention_sinks.is_contiguous()}" + ) + + +thop_attention = ThopAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather.md deleted file mode 100644 index 82072d3843a6..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather.md +++ /dev/null @@ -1,378 +0,0 @@ ---- -receipts: - sm_103: {status: passed, world_size: 4} ---- - -# allgather - -**Wraps** `torch.ops.trtllm.allgather` (one call). - -## Semantics - -Every rank in `group` calls with its own slice of a tensor; every rank gets -back the whole thing, concatenated along **dim 0 in ascending rank order**: - -``` -out = cat([input_r0, input_r1, ..., input_r(G-1)], dim=0) -``` - -It is **not** in place — the result is a freshly allocated tensor, and -`input` is left untouched, so clobbering `input` right after the call -cannot reach the result. - -Only dim 0 is a gather axis. Every other dimension is carried through -untouched: `[n, 2560]` gathers to `[G*n, 2560]`, `[n, 2, 1280]` to -`[G*n, 2, 1280]`, and a 1-D `[n]` to `[G*n]`. Nothing is computed — the -call moves bytes, so the result is bit-for-bit the inputs' concatenation -in every dtype below. - -The `sizes` argument picks between two forms: - -**`sizes = None`** — every rank must hold the same number of rows `n`, and -`out.shape[0] = n * len(group)`. - -**`sizes = [n_0, ..., n_(G-1)]`** — rank `i` contributes `n_i` rows and -`out.shape[0] = sum(sizes)`. The list is in ascending rank order, must be -identical on every rank, and `sizes[my_rank]` must equal this rank's -`input.shape[0]`. Entries may be `0`: a rank with no rows contributes -nothing and the gather is still correct. This is the form attention data -parallelism needs, where per-rank token counts differ by construction. - -Both forms are certified, and they are certified separately because they are -not the same call underneath — an all-gather primitive takes one element -count. Read off NCCL's own trace of this entry's test: `sizes = None` -issues **one `ncclAllGather`** whose count is this rank's element count, -and the ragged form issues **one grouped `ncclBroadcast` per rank**, each -rooted at that rank with that rank's element count (for -`sizes = [1, 5, 9, 13]` at hidden 2560: counts 2560, 12800, 23040, 33280 at -roots 0, 1, 2, 3). So one ragged call enqueues `len(group)` NCCL operations -where a uniform call enqueues one — worth knowing when reasoning about the -call-order precondition, though the two forms were never mixed across ranks -in a probe, and the `sizes` list already has to be identical everywhere. -The two cost the same (see *Notes*), so the choice between them is about -what the caller can guarantee, not about speed. - -**Fusion boundary.** Inside the call: the collective and the output -allocation, nothing else. Outside: everything that produced `input` -(under attention DP, this rank's own attention output), any padding of -row counts to a uniform value, any quantization, and any reshaping — -including the reshape a caller needs to gather along an axis other than -dim 0, which this op cannot do. - -**Group membership.** `group` names **ranks in trtllm's MPI session -communicator** (`MPI_COMM_WORLD` under `mpirun`), not device ordinals or -`torch.distributed` ranks. A subset is legal: with `group = [0, 1]` at -world size 4, ranks 0 and 1 gather only between themselves and ranks 2 -and 3 must not call at all. The list is treated as a **set** — passing -`[3, 2, 1, 0]` produces the identical result to `[0, 1, 2, 3]`, still in -ascending rank order — so `group` cannot be used to permute the output. - -## Signature - -```python -def allgather( - input: torch.Tensor, - sizes: Optional[List[int]], - group: List[int], -) -> torch.Tensor -``` - -### Certified arguments - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `input` | `[rows, ...]`, at least 1-D; `[T, 2560]`, `[T, 2, 1280]` and `[T]` tested | bf16; also fp16, fp32, int32, uint8, float8_e4m3fn | **contiguous** | CUDA | -| `sizes` | `None`, or a list of `len(group)` non-negative ints in ascending rank order | Python `int` | — | host | -| `group` | list of MPI session ranks | Python `int` | — | host | -| returns | `[n * len(group), *input.shape[1:]]` for `sizes = None`, `[sum(sizes), *input.shape[1:]]` otherwise | = `input.dtype` | contiguous, newly allocated | = `input.device` | - -There are no other arguments — no strategy, no workspace, no output -buffer, no autotuner. Nothing selects a transport, so unlike the sibling -all-reduce there is no strategy axis to certify. - -Every dtype in the table was certified in both forms (uniform and ragged) -as a bitwise move: bf16 for hidden states, uint8 for packed NVFP4 bytes -and their scale factors, float8_e4m3fn for fp8 activations, int32 for -expert ids, fp32/fp16 for routing weights and non-bf16 activations. int64 -and bool were also observed to move correctly in probing and are not -certified. There is no arithmetic here for a dtype to be wrong about — -this op does not reinterpret bytes the way a summing collective does. - -## Metadata consumed - -No attention metadata, no KV cache, no registered layer, no workspace. -Two pieces of process state stand behind the call: - -1. **trtllm's MPI session communicator.** The op resolves `group` against - it. Under `mpirun` the session is `MPI_COMM_WORLD`; a trtllm engine's - worker ranks are already inside one, so a target's forward needs no - setup. World size 1 was not exercised; this entry's receipt is world - size 4. The sibling op `torch.ops.trtllm.allgather_pg` takes a - `c10d.ProcessGroup` instead, for the `TLLM_DISABLE_MPI=1` path; - `torch.ops.trtllm.allgather_list` gathers a list of tensors in one - call. Both are different ops and neither is this entry. -2. **The NCCL communicator cache**, keyed by the rank set. Built on first - use for a given `group`, which makes the first call for a group much - slower than the rest — and means that first call **cannot be inside a - CUDA-graph capture** (see *Preconditions*). Inside a trtllm serving - engine this communicator is the **caller's alone**: the runtime issues - nothing on it, so the only calls whose order has to line up are the - ones the caller makes. Measured — see *Notes*, "Inside a live serving - engine under attention data parallelism". - -## Preconditions - -- **Every rank in `group` calls, and their arguments agree**: same dtype, - same trailing shape, same `sizes` list, and — when `sizes is None` — - the same row count. Disagreement is **not survivable**: it hangs rather - than raising. Measured at world size 4, three ways, each in its own - process: ranks passing different row counts with `sizes = None`; a rank - holding more rows than its `sizes` entry; a rank holding fewer. In each - case some ranks returned a plausible-looking tensor and the rest never - came back, and the job had to be killed. This is why this entry's test - carries an external deadline. Do not read the hang as a guarantee, - though: every one of those three cases changes how many bytes a rank - moves, and a disagreement that leaves the byte counts equal on every rank - is silently wrong instead — see the call-order bullet below, where that - is measured. -- `input` is **contiguous**. The op reads `input.numel()` elements from - `input.data_ptr()` and ignores strides, so a strided view is gathered - from the wrong bytes and returns a plausible-looking wrong answer with - no error. The wrapper asserts this. -- `input.dim() >= 1`. A 0-d tensor **segfaults** inside - `AllgatherOp::run_list` and kills every rank in the job with no - diagnostic beyond the fault handler's stack. The wrapper asserts this. -- `len(sizes) == len(group)` when `sizes` is given. The op sizes its - output from `sizes` alone and does not cross-check it against `group`: - a list one entry short silently drops the highest rank's rows (measured - — the result is a correct gather of the remaining ranks, at the shorter - length), and a list one entry too long returns - `sum(sizes)` rows whose tail is uninitialized memory (measured once, in - probing; indexing past the group is undefined behaviour and not - something to rely on). The wrapper asserts this. -- `sizes[my_rank] == input.shape[0]`, in ascending rank order. The op - does not check it and cannot be made to — it takes no rank argument — - and a mismatch hangs, as above. The wrapper cannot check it either, for - the same reason: this one is the caller's. -- `input.device` is the device this rank set with - `torch.cuda.set_device`, and one rank owns one device. -- **A `group`'s first call must not be made inside a CUDA-graph capture.** - It is where the group's NCCL communicator gets built, and that build - raises under capture — see the *CUDA graphs* note for the error text and - the one-call fix. -- **Every rank issues the same sequence of calls on `group`, in the same - order**, graph replays included. Position on the communicator is what - pairs one rank's call with another's, so order is the caller's whole - responsibility — and disagreeing on it is *worse* than disagreeing on the - arguments, because it usually does not hang. Measured at world size 4 in - the certified configuration: - - | how the ranks disagreed | what happened | - |---|---| - | two same-shaped gathers issued in swapped order on one rank | **silently wrong on every rank.** No error, no hang. The swapping rank had 3/4 of its elements wrong in both results, every other rank exactly 1/4 — the swapping rank's block. A swapped *pair* puts the positions back, so a plain gather straight afterwards is bitwise correct. In this entry's test. | - | one rank issuing one gather more than the others, each iteration | **silently wrong on every rank, permanently**: the positions never realign, so every later call stays paired one off. Probing, outside the test, for that reason. | - | a gather swapped against a `torch.ops.trtllm.reducescatter` of the **same** byte count on one rank | **silently wrong on every rank** — the gather came back holding reduce-scatter payload and vice versa. Probing, outside the test. | - | the same swap where the two calls carry **different** byte counts | **wedged.** No rank returned; the job had to be killed. Probing, outside the test, for that reason. | - - So an ordering divergence surfaces as an accuracy loss whenever the - mispaired calls happen to move the same number of bytes, and as a hang - only when they do not. A caller must not wait for a hang to tell it that - its ranks have diverged. -- **The stream is the engine's to choose and the ranks need not agree on - it.** The call runs on `torch.cuda.current_stream()`, and a serving - engine moves that stream under the model — the same forward runs on - torch's graph-capture stream while a decode graph is being captured and - on the serving stream otherwise, so a target cannot pin it. Stream - identity plays no part in pairing the calls: certified with every rank on - a side stream, with one rank on a side stream while the others stayed on - the default, and with ranks alternating in opposite patterns so they - disagreed at every call index — all bitwise correct. Certified with the - side stream joined to the current one on both ends, which is what the - engine and a target's forward both do; two calls on one `group` running - *concurrently* on two streams of the same rank was not exercised. -- `group` entries must be ranks of the MPI session. With - `group = [0, 1, 2, 3, 4]` at world size 4, **rank 0 raised**: - - ``` - RuntimeError: [TensorRT-LLM][ERROR] Assertion failed: Failed: MPI error - ../tensorrt_llm/runtime/utils/mpiUtils.cpp:277 '6' - (../tensorrt_llm/runtime/utils/mpiUtils.cpp:277) - ``` - - and ranks 1, 2 and 3 never returned from the call — the job had to be - killed, twice out of twice. So this is a raise on one rank and a wedge - everywhere else, not a survivable error: measured in probing, outside - the test, for exactly that reason. -- `rows = 0` on **every** rank (an empty `[0, H]` input with - `sizes = None`) is accepted and returns an empty `[0, H]` tensor - (observed in probing, outside the test). A single rank at zero rows is - certified through the `sizes` form. - -## Notes - -The certified path: `mpirun`-launched ranks whose session communicator is -`MPI_COMM_WORLD`, one rank per B200 (sm_100), world size 4, group -`[0, 1, 2, 3]` and `[0, 1]`, bf16 hidden 2560 unless the dtype table says -otherwise. The op has no strategy, workspace or autotuner state, so there -is no second execution path here for a precondition to be quietly true -of — the only axis with two paths is `sizes`, and both are certified. - -**Inside a live serving engine under attention data parallelism.** The -question this answers is whether the runtime is a second party on the -communicator — whether its own cross-rank work can interleave with a -model's calls and pair against them. It is not. Measured by running a -4-rank trtllm serving engine (attention DP on, MoE expert-parallel over the -same four ranks) whose model issues 29 of these gathers and 29 -reduce-scatters per forward, with NCCL's own collective trace on for the -whole process — engine warm-up, 68 decode-graph captures and the served -requests: - -- **The runtime issues nothing on this communicator.** NCCL built exactly - **one** communicator per rank in the whole process, and all 14,036 - collectives logged on it were the model's own — 7,018 `AllGather` and - 7,018 `ReduceScatter` per rank, over 242 forwards. The engine contributed - none. -- **What the engine does instead is host-side MPI on a different - communicator.** Under attention DP it agrees the per-rank token counts - once per step — that is where `attn_metadata.all_rank_num_tokens` comes - from — with `MPIDist.tp_allgather`, an `MPI_Allgather` / `MPI_Allgatherv` - on a sub-communicator it builds with `MPI_Comm_create_group`: congruent - to the session communicator, never identical to it, and never NCCL. - Graph-capture consensus, batch-size consensus, request broadcast and - response gathering all go the same way. So a model's gather cannot be - paired against an engine collective — there is none on this communicator - to pair it against. Certified in this entry's test by driving the engine's - own `MPIDist` object between batches of gathers. -- **All four ranks' call sequences agreed exactly**: the 14,036 trace - entries matched across ranks in op kind, element count and stream index, - position by position. -- **The engine issues a model's collective on two different streams.** The - trace splits into 137 runs: one of 580 calls on the serving stream, then - 68 pairs of [58 calls on torch's graph-capture stream][116 on the serving - stream] — one forward captured, two run eagerly as the runner's warm-up. - Which stream a call lands on is therefore the engine's choice, and the - *Preconditions* entry above records that this is harmless: the ranks - agreed on it here, and they do not have to. - -What those runs do **not** establish. Ten runs of that engine were made; -eight finished and two stopped with every rank spinning in -`cuLaunchKernelEx` (`sched_yield`) — the launch queue full because the -device was not retiring work. One stopped in engine warm-up, all four ranks -in `cudaDeviceSynchronize` after the same warm-up forward; one stopped in a -served forward, all four ranks past the gather and inside the -non-collective work behind it. Nothing measured attributes either to this -call: the ranks' enqueued sequences agreed everywhere they were traced, and -the same collective pattern driven without an engine — the engine's real -host-side step sync plus 29 gathers and 29 reduce-scatters per step at -ragged attention-DP row counts up to 2048, nothing synchronising — ran -2000 steps (58,000 gathers per rank, more collectives than a whole -benchmark run of that engine) clean on those same shared GPUs. Both -stalls landed on the four GPUs that were sharing SMs roughly 50/50 with an -unrelated process; the four exclusive GPUs took four of the ten runs and -stalled on none. That is a correlation on a small sample, not a mechanism. -Note also that stock trtllm does **not** take this path at that -configuration: its MoE communication factory selected `DeepEPLowLatency` -there (measured), so a stock run is not a control for this op. - -**CUDA graphs — the surface a decode step actually needs.** Certified at -world size 4, group `[0, 1, 2, 3]`, bf16, under torch's default -`capture_error_mode="global"`, with every rank capturing and replaying the -same graph: - -| form | under capture | -|---|---| -| `sizes = None` (uniform) | certified at **all 35 of `1..32, 64, 128, 256` rows** | -| `sizes = [...]` (ragged) | certified — but the split is frozen at capture, see below | - -A trtllm engine instantiates one decode graph per configured batch size — -with `cuda_graph_config.max_batch_size = 256` and no explicit -`batch_sizes` list, 35 of them at `1..32, 64, 128, 256` — keeps all of -them alive in one memory pool, and replays them interleaved as the served -batch size moves. Under attention DP a decode call's per-rank row count -*is* that batch size, so the whole set is certified: all 35 captured into -one shared pool, none released, then replayed in four orders — ascending, -descending, shuffled, and largest-jump-first (`256, 1, 128, 2, 64, 3, ...`) -— each order twice, once one replay at a time with a synchronize and a -full check between graphs, and once with all 35 issued back to back and -nothing synchronizing between them. Every replay's payload is one no -earlier call used, and every result is bitwise equal to the concatenation -reference. The same is certified with **29 independent gathers inside each -of the 35 graphs** — 1015 captured collectives in one pool, the shape this -checkpoint's decode graph has (30 layers, layer 0 dense, so 29 -expert-parallel MoE calls each preceded by one gather). - -A replay re-runs the collective over whatever the input buffer holds at -replay time and writes into the same tensor the capture returned — keep it -and read it after each `replay()`. Replays also stay correct with eager -traffic in between: an eager uniform gather at 1500, 2048 or 8192 rows and -an eager ragged gather between every pair of replays, all bitwise correct, -which is the shape a server has when a prefill runs between two decode -steps. - -**`sizes` is frozen into a graph; row counts must be padded to capture -one.** `sizes` is a host-side argument, so a capture bakes in the split it -was given and a replay re-runs *that* split. Certified by capturing two -graphs with different sizes vectors into one pool and replaying them in -both orders across two payload rounds: each keeps its own row split and -its own output length. The consequence for an attention-DP target is -concrete — a graph-captured decode step must pad every rank to the same -row count and pass `sizes = None`, because the ragged split cannot vary -per replay; the ragged form belongs to the eager (non-captured) steps. - -**A `group`'s first-ever call cannot be captured**, because it is where the -NCCL communicator is built. Four ranks capturing their first call each -raised, out of `tensorrt_llm::_v1::getComm`: - -``` -RuntimeError: Failed, NCCL error ../tensorrt_llm/common/opUtils.cpp:175 -'unhandled cuda error (run with NCCL_DEBUG=INFO for details)' -``` - -which invalidates the capture, so the `with torch.cuda.graph(...)` block -itself then raises - -``` -AcceleratorError: CUDA error: operation failed due to a previous error -during capture -``` - -(`cudaErrorStreamCaptureInvalidated`). It is survivable, and the fix is one -line: make one eager call per `group` before any capture. After the failure -the same group gathers correctly eagerly, and a capture taken after that -replays correctly — both certified. The same failure was seen in probing -for a subgroup whose communicator did not yet exist in a process that -already had the full group's. One caveat for anyone catching it: torch's -graph context manager ends the capture before it restores the stream, so a -failed capture leaves its own stream current — put the resting one back -with `torch.cuda.set_stream`. - -**Numerics.** There are none. The op copies bytes, so every assertion in -this entry's test is bitwise (`rtol=0, atol=0`) rather than -tolerance-based, in every dtype, in both forms, eager and captured. A -caller does not have to reason about accumulation order the way the -sibling all-reduce forces it to. - -**Gathering along another axis is the caller's problem.** This op only -concatenates dim 0. trtllm's module-level helper reaches other axes by -reshaping around this call — `view` to 2-D before, `chunk`/`split` and -`cat` after — which is Python-level composition and would have to be -written from catalog entries, not hidden inside a wrapper. For the -attention-DP use (gather token rows before an expert-parallel MoE call) -dim 0 is already the token axis and no reshape is needed. Note also that -"dim 0" is the buffer's first dimension, not necessarily the token axis: -a scale-factor buffer in a swizzled layout has no token rows to gather. - -**Cost: the two forms are indistinguishable.** Measured on the certified -path (world size 4, bf16, hidden 2560, 200 timed iterations after 20 warm-up -calls, CUDA events, eager): `sizes = None` against an explicit uniform -`sizes` vector of the same total size came out at 12.5 vs 12.4 us at 1 row, -13.0 vs 13.2 at 8, 20.2 vs 20.1 at 128 and 80.1 vs 80.2 at 2048 — ratio -1.00-1.01 throughout. So a caller that already has the per-rank counts loses -nothing by passing them, and the reason a captured decode step pads to -uniform rows is the frozen split above, not the cost of the ragged form. -Below ~128 rows the call is launch-bound (12-13 us for a message 20x apart -in size), which is the regime a decode step runs in. - -**The inverse op exists.** `torch.ops.trtllm.reducescatter` has the same -`(input, sizes, group)` signature and undoes this one — it is a different -op and is not this entry. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather.py index aa3ef337825e..f6243391b0b6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather.py @@ -1,39 +1,184 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Concatenate every rank's rows of a tensor into the group-wide full set.""" +"""Gather every rank's rows into one tensor, via the trtllm allgather op.""" -from typing import List, Optional +from __future__ import annotations + +from typing import List, Optional, Sequence import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from .._op import Arch, Cell, OpWrapper + -def allgather( - input: torch.Tensor, - sizes: Optional[List[int]], - group: List[int], -) -> torch.Tensor: +class Allgather(OpWrapper): """Gather `input` from every rank in `group`, concatenated along dim 0. `sizes` gives each rank's `input.shape[0]` in ascending rank order, or is `None` when every rank holds the same number of rows. The result is a freshly allocated tensor on every rank; `input` is left untouched. """ - # Pure-metadata guards, each for a violation measured on this machine to - # end in silent corruption or a process kill rather than an error. - assert input.dim() >= 1, ( - "input must have at least one dimension; a 0-d tensor segfaults inside " - "AllgatherOp::run_list" - ) - assert input.is_contiguous(), ( - "input must be contiguous; the op reads it as packed memory and a " - "strided view is silently gathered from the wrong elements" - ) - assert sizes is None or len(sizes) == len(group), ( - f"len(sizes)={len(sizes) if sizes is not None else None} must equal " - f"len(group)={len(group)}; the op sizes its output from `sizes` alone, " - "so a short list silently drops the trailing ranks and a long one " - "appends uninitialized rows" + + ARCHS = frozenset({Arch.SM_103}) + + # deepseek-r1 at dep4 always passes `sizes=None`: it pads every rank to a + # common height first, so the ragged form the op also supports is not + # something a shipped target reaches. Row counts span what an attention-DP + # rank feeds its expert-parallel MoE call, from one decode token to a full + # prefill chunk. deepseek-v4.1 pads the same way; its cells follow. + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode: one row per rank, hidden 7168", + spec=dict(rows=1, trailing=(7168,), dtype=torch.bfloat16, sizes=None), + ), + Cell( + why="a small padded batch, the common case between decode and prefill", + spec=dict(rows=32, trailing=(7168,), dtype=torch.bfloat16, sizes=None), + ), + Cell( + why="prefill chunk height", + spec=dict(rows=2048, trailing=(7168,), dtype=torch.bfloat16, sizes=None), + ), + # deepseek-v4.1 under attention DP: per-rank rows n at the ends of the + # swept range 1..16384 (16384 is the largest max_num_tokens served) and + # at 1536, the largest n a decode graph captures (batch 256 x next_n 6). + # The matrix sweeps every n in between for each form. + *( + Cell( + why=f"deepseek-v4.1 {form} n={n}: {role}", + spec=dict(rows=n, trailing=trailing, dtype=dtype, sizes=None), + ) + for form, trailing, dtype in ( + ("moe bf16 [n, 5120]", (5120,), torch.bfloat16), + ("engram int32 [n, 48]", (48,), torch.int32), + ("engram int64 [n, 24]", (24,), torch.int64), + ) + for n, role in ( + (1, "decode end of the swept range"), + (1536, "top of the decode capture set"), + (16384, "max_num_tokens, prefill end of the swept range"), + ) + ), + *( + Cell( + why=f"deepseek-v4.1 image mask uint8 [n] n={n}: {role}", + spec=dict(rows=n, trailing=(), dtype=torch.uint8, sizes=None), + ) + for n, role in ( + (1, "eager only, decode end of the swept range"), + (16384, "eager only, prefill end of the swept range"), + ) + ), ) - return torch.ops.trtllm.allgather(input, sizes, group) + + note = """ + The ragged form (`sizes` as an explicit per-rank vector) and dtypes other + than bf16 (beyond the DeepSeek-V4.1 forms below) work and are exercised by + the op matrix, but no shipped target reaches them, so they are not cells. + + CUDA graph behaviour -- capture at every engine batch size, replays + surviving eager calls of other shapes, the captured sizes vector being held + -- is not a cell axis either: it is about the op across capture and replay + rather than about an input configuration, and the matrix carries it as its + own named checks. + + DeepSeek-V4.1: under attention DP every rank is padded to + max(all_rank_num_tokens) and gathers with `sizes=None`. The MoE input, + bf16 [n, 5120], is gathered on every layer; on the engram layers (only with + the device-resident engram table) the int64 [n, 24] engram ids are gathered + as their int32 [n, 48] view, and int64 [n, 24] is certified beside it. The + multimodal target, in every eager step and never under capture, views its + bool [T] image-span mask as uint8, zero-pads it at the end to n, gathers it + and reads the result back with `.view(torch.bool)`; bool itself is not + certified. No metadata. The matrix sweeps every per-rank n in 1..16384 for + the three core forms and for the mask (with the int32 id gather that follows + it), ascending and then shuffled, bitwise; it captures one decode graph per + per-rank n of the capture set (batch G in 1..32, 40..128 step 8, 192, 256; + n = G or 6G for next_n 1 or 6, plus 35 and 160: 84 values up to 1536), each + holding the forward's 44 gathers, and replays them bitwise against the + reference and against eager. A captured step must pad every rank to the same + n, because `sizes` is frozen into the graph. The ragged form at V4.1 sizes + is certified only at four dumped splits, eagerly. Not certified: n above + 16384; capture at an n outside that set; the mask under capture. The other + group layouts the target runs -- world 2 (one pp2 x dep2 stage), world 4 cut + into two groups of 2, world 8 across two trays, world 8 cut into two groups + of 4 -- were measured on the MPI-session symbol under an mpirun launcher; + this matrix runs one group of 4 on the ProcessGroup symbol and does not + certify them. + """ + + def raw_call( + self, + input: torch.Tensor, + sizes: Optional[List[int]], + group: List[int], + process_group: Optional[object] = None, + ) -> torch.Tensor: + """`process_group` selects the communicator flavor, mirroring the + product router in `_torch/distributed/ops.py`: None runs the + MPI-session symbol -- the one a dep4 serving world calls -- and a boxed + c10d ProcessGroup runs the `_pg` symbol the non-MPI backends use. + """ + if process_group is None: + return torch.ops.trtllm.allgather(input, sizes, group) + return torch.ops.trtllm.allgather_pg(input, sizes, group, process_group) + + def reference( + self, + per_rank_inputs: Sequence[torch.Tensor], + sizes: Optional[List[int]] = None, + group: Optional[List[int]] = None, + ) -> torch.Tensor: + """The concatenation the collective is supposed to make. + + The documented exception to mirroring `raw_call`: a collective's output + is not a function of what the calling rank holds, so a reference given + only `input` could not state anything at all. It takes the group's + inputs in ascending rank order instead. The matrix builds that list on + one rank from the seed convention every rank shares, which is what keeps + the reference arithmetic rather than a second collective. + """ + return torch.cat(list(per_rank_inputs), dim=0) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """A gate of exactly zero: the op moves bytes and computes nothing. + + Tightened from the base default rather than loosened. Any difference at + all is a wrong gather. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + if out.dtype is torch.float8_e4m3fn: + # torch.testing cannot compare float8; for a pure data move the byte + # pattern is the honest gate anyway. + out, ref = out.view(torch.uint8), ref.view(torch.uint8) + torch.testing.assert_close(out, ref, rtol=0, atol=0) + + def is_valid( + self, + input: torch.Tensor, + sizes: Optional[List[int]], + group: List[int], + process_group: Optional[object] = None, + ) -> None: + # Each guard is for a violation measured to end in silent corruption or + # a process kill rather than an error. + assert input.dim() >= 1, ( + "input must have at least one dimension; a 0-d tensor segfaults inside " + "AllgatherOp::run_list" + ) + assert input.is_contiguous(), ( + "input must be contiguous; the op reads it as packed memory and a " + "strided view is silently gathered from the wrong elements" + ) + assert sizes is None or len(sizes) == len(group), ( + f"len(sizes)={len(sizes) if sizes is not None else None} must equal " + f"len(group)={len(group)}; the op sizes its output from `sizes` alone, " + "so a short list silently drops the trailing ranks and a long one " + "appends uninitialized rows" + ) + + +allgather = Allgather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/reducescatter.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/reducescatter.md deleted file mode 100644 index c2512fb0a0b4..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/reducescatter.md +++ /dev/null @@ -1,423 +0,0 @@ ---- -receipts: - sm_103: {status: passed, world_size: 4} ---- - -# reducescatter - -**Wraps** `torch.ops.trtllm.reducescatter` (one call). - -## Semantics - -Every rank in `group` calls with a tensor of the same full shape. The op sums -them elementwise, splits the sum along **dim 0**, and gives each rank the -slice belonging to **its position in `group`, in ascending rank order**: - -``` -total = input_r0 + input_r1 + ... + input_r(G-1) # elementwise -out_i = total[offset_i : offset_i + n_i] # rank at group position i -``` - -It is **not** in place — the result is a freshly allocated tensor, and -`input` is left untouched, so clobbering `input` right after the call cannot -reach the result. - -Only dim 0 is split. Every other dimension is carried through untouched: -`[G*n, 2560]` scatters to `[n, 2560]`, `[G*n, 2, 1280]` to `[n, 2, 1280]`, -and a 1-D `[G*n]` to `[n]`. - -The `sizes` argument picks between two splits: - -**`sizes = None`** — the split is even. `input.shape[0]` must be a multiple of -`len(group)`; rank at position `i` gets rows `[i*n, (i+1)*n)` where -`n = input.shape[0] / len(group)`. - -**`sizes = [n_0, ..., n_(G-1)]`** — rank at position `i` gets `n_i` rows -starting at `sum(sizes[:i])`. The list is in ascending rank order, must be -identical on every rank, and `sum(sizes)` must equal `input.shape[0]`. -Entries may be `0`: a rank that keeps no rows still has to call, and the -reduction is still correct. This is the form attention data parallelism needs, -where per-rank token counts differ by construction. - -Both forms are certified. They are **not the same code path**, and NCCL's own -trace shows why: `sizes = None` issues **one `ncclReduceScatter`** whose count -is the elements one rank keeps (rows kept times the trailing dims), while a -`sizes` vector issues **one grouped `ncclReduce` per rank**, rooted at that -rank (for `sizes = [1, 5, 9, 13]` at hidden 2560: counts 2560 / 12800 / 23040 / -33280 rooted at 0 / 1 / 2 / 3, in that one call). An even explicit `sizes` -vector returns bits identical to `sizes = None` and costs the same, while a -genuinely uneven split costs about 1.5x more at 128 rows and above (see -*Notes*). Either way the whole call is bracketed by one -`ncclGroupStart`/`ncclGroupEnd` pair, so one call is **one** position on the -communicator whichever form it takes, and the ragged form's `G` reduces pair -across ranks atomically rather than one root at a time — measured, and matching -the disassembly of `ReducescatterOp::run_list` in `libth_common.so`. That is -what makes the call-order rule in *Preconditions* form-independent. - -**It is the inverse of the sibling all-gather, in layout.** The rows this op -hands back to rank `i` are exactly the rows rank `i` would have contributed to -a gather with the same `sizes` and `group` — certified end to end, including -`reducescatter(allgather(x)) == G * x`. What it is *not* is a byte-exact -inverse: the gather moves bytes, this one computes a sum, and the sum is taken -in the input dtype rather than in fp32. Read *Numerics* before relying on the -result of a chain. - -**Fusion boundary.** Inside the call: the reduction, the split, and the output -allocation, nothing else. Outside: everything that produced `input` (under -attention DP, this rank's expert-window output over the whole gathered token -set), any residual add, any scaling by routing weights, any quantization, any -padding of row counts to a uniform value, and any reshaping — including the -reshape a caller needs to split along an axis other than dim 0, which this op -cannot do. - -**Group membership.** `group` names **ranks in trtllm's MPI session -communicator** (`MPI_COMM_WORLD` under `mpirun`), not device ordinals or -`torch.distributed` ranks. A subset is legal: with `group = [2, 3]` at world -size 4, ranks 2 and 3 reduce only between themselves — rank 2 gets slice 0 and -rank 3 gets slice 1, i.e. the **position in the group**, not the MPI rank -(certified; a subset excluding rank 0 is what pins this). Ranks 0 and 1 must -not call. The list is treated as a **set** — passing `[3, 2, 1, 0]` produces -the identical result to `[0, 1, 2, 3]` — so `group` cannot be used to permute -the split. - -## Signature - -```python -def reducescatter( - input: torch.Tensor, - sizes: Optional[List[int]], - group: List[int], -) -> torch.Tensor -``` - -### Certified arguments - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `input` | `[rows, ...]`, at least 1-D; `[T, 2560]`, `[T, 2, 1280]` and `[T]` tested | bf16; also fp16, fp32, int32, int64, uint8, int8 | **contiguous** | CUDA | -| `sizes` | `None`, or a list of `len(group)` non-negative ints in ascending rank order summing to `input.shape[0]` | Python `int` | — | host | -| `group` | list of MPI session ranks | Python `int` | — | host | -| returns | `[input.shape[0] // len(group), *input.shape[1:]]` for `sizes = None`, `[sizes[my position], *input.shape[1:]]` otherwise | = `input.dtype` | contiguous, newly allocated | = `input.device` | - -There are no other arguments — no reduction operator (the reduction is always -a sum), no strategy, no workspace, no output buffer, no autotuner. Unlike the -sibling all-reduce there is no transport to select. - -Every float dtype in the table reduces as an arithmetic sum in **that dtype** -(see *Numerics*). The integer dtypes reduce as integer sums and wrap on -overflow in the narrow ones: four ranks of `int8` 100 came back as `-112` -(measured, outside the test, which keeps its integer payloads inside range). - -Dtypes that are **not** in the table divide into two kinds, and both are -traps: - -- `torch.float8_e4m3fn` is **accepted and silently wrong**. The op sums the - e4m3 *bit patterns* as unsigned bytes, wrapping at 256: four ranks each - holding `1.0` (byte `0x38`) return byte `0xE0`, i.e. `-32.0`, not `4.0`. - The wrapper rejects the dtype for this reason. Note that the sibling - all-gather moves fp8 correctly — it only copies — so a post-quantization - dispatch that works on the way out does **not** work on the way back. - `torch.bool` is the same shape of trap: the underlying bytes are summed - (four ranks of `True` leave byte value 4), which torch then reads as `True`, - so it behaves like an OR. Neither is certified. -- `torch.float64`, `torch.float8_e5m2`, `torch.float8_e4m3fnuz`, - `torch.float8_e5m2fnuz` and `torch.float8_e8m0fnu` **raise** - `RuntimeError: unsupported data type (../tensorrt_llm/runtime/torchUtils.h:113)` - — and that raise is terminal for the whole process. See *Preconditions*. - -## Metadata consumed - -No attention metadata, no KV cache, no registered layer, no workspace. -Two pieces of process state stand behind the call: - -1. **trtllm's MPI session communicator.** The op resolves `group` against it - and finds this rank's position in it. Under `mpirun` the session is - `MPI_COMM_WORLD`; a trtllm engine's worker ranks are already inside one, so - a target's forward needs no setup. World size 1 was not exercised; this - entry's receipt is world size 4. The sibling op - `torch.ops.trtllm.reducescatter_pg` takes a `c10d.ProcessGroup` instead, - for the `TLLM_DISABLE_MPI=1` path; `torch.ops.trtllm.reducescatter_list` - reduce-scatters a list of tensors in one call. Both are different ops and - neither is this entry. -2. **The NCCL communicator cache**, keyed by the rank set and shared with the - other collectives. Built on first use for a given `group`, which makes the - first call for a group much slower than the rest — and means that first - call **cannot be inside a CUDA-graph capture** (see *Preconditions*). - -## Preconditions - -- **Every rank in `group` calls, and their arguments agree**: same dtype, same - trailing shape, same `sizes` list, and the same `input.shape[0]`. - Disagreement is **not survivable**: it hangs rather than raising. Measured at - world size 4, each in its own process — ranks passing different row counts - with `sizes = None` (three ranks returned a plausible-looking tensor, one - never came back), and ranks passing different `sizes` vectors (no rank came - back). Both jobs had to be killed. This is why this entry's test carries an - external deadline. Do not read the hang as a guarantee, though: both of those - disagreements change how many bytes a rank moves, and a disagreement that - leaves the byte counts equal on every rank is silently wrong instead — see - the call-order bullet below, where that is measured. -- **A rank outside `group` must not call.** Measured: with `group = [0, 1]` at - world size 4, ranks 0 and 1 returned correctly and ranks 2 and 3 **never - returned** — they block inside the communicator bootstrap without ever - reaching the op's own membership check. The job had to be killed. -- `input` is **contiguous**. The op reads `input.numel()` elements from - `input.data_ptr()` and ignores strides, so a strided view is reduced from the - wrong bytes and returns a plausible-looking wrong answer with no error. The - wrapper asserts this. -- `input.dim() >= 1`. A 0-d tensor **segfaults** inside - `ReducescatterOp::run_list` and kills every rank in the job with no - diagnostic beyond the fault handler's stack. The wrapper asserts this. -- **The split must cover the input exactly**, and nothing checks it but the - wrapper: - - `sizes is None` requires `input.shape[0] % len(group) == 0`. Otherwise the - op takes `shape[0] // len(group)` rows per rank and the remainder is - **silently dropped** — 9 rows over 4 ranks returns 2 rows per rank, and the - 9th row is never reduced (measured; the 8 rows that are returned are - correct). - - `sizes` given requires `sum(sizes) == input.shape[0]`. A sum smaller than - the input silently ignores the tail rows (measured: the rows the split does - cover come back correct). A sum larger than the input makes the op read - **past the end of the buffer** — observed once, during this entry's test - bring-up, to return a correctly shaped tensor full of `NaN`. - - `len(sizes) == len(group)`. One entry short raises - `IndexError: vector::_M_range_check: __n (which is 3) >= this->size() - (which is 3)` on the highest rank while the others sit in the collective — - a raise on one rank and a wedge everywhere else, measured, not recoverable - inside the job. One entry too long is the out-of-bounds read above. -- `input.dtype` must be one the op supports, and this is a check the caller - has to make **before** calling rather than by catching. The unsupported - dtypes raise `RuntimeError: unsupported data type`, but the raise happens - after the op has opened an NCCL group and before it closes it, so the - group is left unbalanced and **every collective the process issues - afterwards silently returns garbage** — measured for this op (99.6% of - elements wrong on the very next call), and measured for - `torch.ops.trtllm.allgather` and `torch.ops.trtllm.allreduce` on the same - group too. It is the process that is finished, not just this entry's - communicator. Every other failure in this section that returns at all leaves - the group usable — the test checks that explicitly after the float8, - non-contiguous and uncovered-split calls. -- `input.device` is the device this rank set with `torch.cuda.set_device`, and - one rank owns one device. -- **A `group`'s first call must not be made inside a CUDA-graph capture.** It - is where the group's NCCL communicator gets built, and that build raises - under capture — see the *CUDA graphs* note for the error text and the - one-call fix. The cache is shared across collectives, so a group whose - communicator another op already built is warm for this one. -- **Every rank issues the same sequence of calls on `group`, in the same - order**, graph replays included. What pairs one rank's call with another's is - its **position** on the communicator — not the arguments, not the op, not the - stream — so call order is the caller's whole responsibility, and disagreeing - about it is *worse* than disagreeing about the arguments, because it does not - hang. Certified at world size 4 in the configuration of *Notes*: - - | how the ranks disagreed | what happened | - |---|---| - | two same-shaped calls issued in swapped order on one rank | **silently wrong on every rank.** No error, no hang, right shape, finite values. Each position returned the sum of whatever payloads met there, matched **bitwise** against a locally computed mix. Both `sizes` forms behave identically. A swapped *pair* puts the positions back, so a plain call straight afterwards is bitwise correct. In this entry's test. | - | one rank issuing one call more than the others | **silently wrong from the extra call onward, and it does not heal.** Every position after it stays paired one off; there is no resynchronization point inside a stream of collectives. Certified over five consecutive positions, every one of them bitwise equal to its own mix, and alignment came back only once the call counts were equalized. In this entry's test — where the other ranks issue one catch-up call at the end, which is what lets it terminate. What a rank left *permanently* ahead does was not probed: its last call has no partner, so the test would not have returned. | - | this op swapped against `torch.ops.trtllm.allgather` of the **same** element count on one rank | **silently wrong on every rank**, both calls returning at the shapes their own arguments imply. An even reduce-scatter of `[G*n, H]` and an all-gather of `[n, H]` both move `n*H` elements as NCCL counts them. The pair realigns. In this entry's test. | - | the same swap where the two calls carry **different** element counts | **wedged.** All four ranks entered the pair, none came out of it, and the job had to be ended from outside. Certified by the test's own second job, which is the only place it can live: a job that wedges never reports. | - | ranks disagreeing about the **stream** | **correct.** See the next bullet. | - - **How wrong** follows from the fact that this op computes rather than moves. - A collective that only copies spoils just the block the disagreeing rank - contributed; here that rank's payload is an addend of **every** element of - **every** rank's slice, so a single rank out of step makes every rank wrong - nearly everywhere — measured 0.982-0.993 of elements differing from the - intended result, on every rank, for every divergence above that returned. - The residue is elements where the two payloads happened to agree. - - **A mispaired result is deterministic, not noise**, which is the trap. It - carries no `NaN`, no `Inf` and no shape anomaly, it is bitwise identical - across repeats, and it is bitwise the **ring-order sum of the addends that - met** — the accumulation order of *Numerics*, applied to the wrong operands. - Certified on payloads with no exactness property, where a changed - accumulation order would show, and against a reference computed locally from - seeds, which is what makes it reproducible in any process that mispairs the - same way rather than only in this one (the cross-op mispairing was also - checked directly, returning identical bits in two separate processes). So a - divergence cannot be found by re-running the step and looking for - instability, and under attention data parallelism — where every rank pads to - `max(all_rank_num_tokens)` and every call therefore moves the same number of - bytes — it surfaces as a stable accuracy loss and nothing else. **Do not wait - for a hang to tell you the ranks have diverged.** -- **The stream is the engine's to choose and the ranks need not agree on it.** - The call runs on `torch.cuda.current_stream()`, and a serving engine moves - that stream under the model — the same forward runs on torch's graph-capture - stream while a decode graph is being captured and on the serving stream - otherwise, so a target cannot pin it. Stream identity plays no part in - pairing the calls, and none in the reduction order either: certified with - every rank on a side stream, with one rank on a side stream while the others - stayed on the default, and with ranks alternating in opposite patterns so - they disagreed at every call index — every result bitwise correct, including - against the ring-order chain of *Numerics* on payloads where a changed - accumulation order would show. Certified with the side stream joined to the - current one on both ends, which is what the engine and a target's forward - both do; two calls on one `group` running *concurrently* on two streams of - the same rank was not exercised. - -## Notes - -The certified path: `mpirun`-launched ranks whose session communicator is -`MPI_COMM_WORLD`, one rank per B200 (sm_100), world size 4, group -`[0, 1, 2, 3]` (subsets `[0, 1]` and `[2, 3]` for the membership claims), -bf16 hidden 2560 unless the dtype table says otherwise, NCCL as the only -transport. There is no strategy, workspace or autotuner state here, so the -only axis with two execution paths is `sizes`, and both are certified. - -**Who else is on the communicator, inside a serving engine.** Nobody — so the -only calls whose order has to line up are the ones the caller makes. Measured -on this machine by tracing a live 4-rank trtllm serving engine under attention -data parallelism with `NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=INIT,COLL`: one NCCL -communicator per rank, 14,036 collectives on it over 242 forwards (7,018 -all-gathers and 7,018 reduce-scatters, all issued by the model, none by the -runtime), and per-rank sequences bit-identical on (op kind, element count, -stream index) position by position. The runtime's own cross-rank step under -attention DP — agreeing `all_rank_num_tokens` once per forward — is a -**host-side MPI** collective on a communicator it builds itself, congruent to -the session communicator but never identical to it, and never NCCL. That trace -also showed the engine issuing a model's collectives on **two** different -streams (torch's graph-capture stream during capture, the serving stream -otherwise), which is why the stream bullet in *Preconditions* matters. This -paragraph is measured evidence from this repo's own runs, not something this -entry's test carries — the test certifies the pairing rule itself, in isolation. - -**Numerics — this op computes, so the result is not the exact sum.** All of -the following is measured at world size 4, bf16, on payloads with no exactness -property (standard normal cast to bf16): - -- It **is deterministic**. Eight identical eager calls return bitwise - identical tensors; eight replays of a captured call return bitwise identical - tensors and agree bitwise with the eager result; and — in probing, outside - the test — three separate 4-rank processes returned byte-identical results - for the same payloads at 1, 16, 256 and 2048 rows and for a ragged split. A - target's accuracy gate can therefore expect run-to-run reproducibility. -- It is **not** the correctly rounded exact sum. About a third of the elements - differ from an fp32 reduction rounded once to bf16 (measured 0.33 at 1, 8, - 256 and 2048 rows). -- It **is a sequential summation in the input dtype**, in ring order starting - at the destination's successor: for the rank at group position `i` the order - is `x_{i+1} + x_{i+2} + ... + x_{i+G-1} + x_i` (indices mod `G`), each - partial sum rounded to the input dtype. Matched bitwise on every element at - 1, 8 and 256 rows with hidden 2560, and at hidden 64. A consequence worth - stating plainly: **two ranks holding identical rows can get different - answers**, because their orders differ. With rank 0 holding `1.0` and the - other three holding `2^-9`, rank 0 came back with `1.0078125` (its order - accumulates the three small terms first) and rank 2 with `1.0` (its order - starts from the large one and the small terms round away). -- Its distance from the exact sum stays inside the classical bound for a - sequential sum of `G` terms, `(G-1) * u * sum_r |x_r|` elementwise, with - `u = 2^-8` for bf16, `2^-11` for fp16 and `2^-24` for fp32. Largest observed - ratio to that bound: 0.89 for bf16 (over 1, 8, 64, 256 and 2048 rows), 0.82 - for fp16, 0.80 for fp32. -- The ring order is NCCL's algorithm choice for this topology, not a promise - of the op's interface. The entry's test asserts it, so a change shows up as - a test failure rather than as drifting accuracy. - -Because of the exactness of small dyadic values, everything else in this -entry's test is asserted **bitwise** (`rtol=0, atol=0`): its payloads are -multiples of 1/8 bounded so that every partial sum, in any order, is exactly -representable. That is a tightening of the default tolerances, not a -loosening — but it is only available to a test that controls its inputs. A -caller's real activations get the bullets above. - -**CUDA graphs — the surface a decode step actually needs.** Certified at world -size 4, group `[0, 1, 2, 3]`, bf16, under torch's default -`capture_error_mode="global"`, with every rank capturing and replaying the -same graph: - -| form | under capture | -|---|---| -| `sizes = None` (even) | certified at **all 35 per-rank row counts `1..32, 64, 128, 256`** | -| `sizes = [...]` (uneven) | certified — but the split is frozen at capture, see below | - -A trtllm engine instantiates one decode graph per configured batch size — -with `cuda_graph_config.max_batch_size = 256` and no explicit `batch_sizes` -list, 35 of them at `1..32, 64, 128, 256` — keeps all of them alive in one -memory pool, and replays them interleaved as the served batch size moves. -Under attention DP a decode call's per-rank row count *is* that batch size, so -this op's captured input is `G` times it. The whole set is certified: all 35 -captured into one shared pool, none released, then replayed in four orders — -ascending, descending, shuffled, and largest-jump-first (`256, 1, 128, 2, 64, -3, ...`) — each order twice, once one replay at a time with a synchronize and -a full check between graphs, and once with all 35 issued back to back and -nothing synchronizing between them. Every replay's payload is one no earlier -call used, and every result is bitwise equal to the arithmetic reference. The -same is certified with **29 independent reduce-scatters inside each of the 35 -graphs** — 1015 captured collectives in one pool, the shape this checkpoint's -decode graph has (30 layers, layer 0 dense, so 29 expert-parallel MoE calls -each followed by one reduce-scatter). - -A replay re-runs the collective over whatever the input buffer holds at replay -time and writes into the same tensor the capture returned — keep it and read -it after each `replay()`. Replays also stay correct with eager traffic in -between: an eager even reduce-scatter at 1500, 2048 or 8192 rows per rank and -an eager uneven one between every pair of replays, all bitwise correct, which -is the shape a server has when a prefill runs between two decode steps. - -**`sizes` is frozen into a graph; row counts must be padded to capture one.** -`sizes` is a host-side argument, so a capture bakes in the split it was given -and a replay re-runs *that* split. Certified by capturing two graphs with -different sizes vectors into one pool and replaying them in both orders across -two payload rounds: each keeps its own row split and its own output length. -The consequence for an attention-DP target is concrete — a graph-captured -decode step must pad every rank to the same row count and pass `sizes = None`, -because the uneven split cannot vary per replay; the uneven form belongs to -the eager (non-captured) steps. - -**A `group`'s first-ever call cannot be captured**, because it is where the -NCCL communicator is built. Four ranks capturing their first call each raised, -out of the op: - -``` -RuntimeError: Failed, NCCL error ../tensorrt_llm/common/opUtils.cpp:175 -'unhandled cuda error (run with NCCL_DEBUG=INFO for details)' -``` - -which invalidates the capture, so the `with torch.cuda.graph(...)` block then -raises - -``` -AcceleratorError: CUDA error: operation failed due to a previous error -during capture -``` - -(`cudaErrorStreamCaptureInvalidated`). It is survivable, and the fix is one -line: make one eager call per `group` before any capture. After the failure -the same group reduce-scatters correctly eagerly, and a capture taken after -that replays correctly — both certified. One caveat for anyone catching it: -torch's graph context manager ends the capture before it restores the stream, -so a failed capture leaves its own stream current — put the resting one back -with `torch.cuda.set_stream`. - -**Cost: the even split is the cheap one.** Measured on the certified path -(world size 4, bf16, hidden 2560, 200 timed iterations after 20 warm-up calls, -CUDA events, eager), per-rank row count `n`, all times per call: - -| `n` | `sizes = None` | even `sizes` vector | uneven `sizes` vector | -|---|---|---|---| -| 1 | 23.9 us | 24.1 us | 10.2 us (degenerate: `[4, 0, 0, 0]`) | -| 8 | 24.5 us | 24.3 us | 28.1 us | -| 128 | 31.0 us | 30.7 us | 47.2 us | -| 2048 | 163.4 us | 163.6 us | 248.7 us | - -`sizes = None` and an even explicit vector are indistinguishable (ratio -1.00-1.01) and return identical bits, so a caller with even counts loses -nothing by passing them. A genuinely uneven split costs about 1.5x from 128 -rows up — so unlike the sibling all-gather, where the two forms measured -identical, here padding to an even split is a small win on cost as well as -being the only form a graph can replay. Below ~128 rows the call is -launch-bound (24-25 us across a 128x message-size range), which is the regime -a decode step runs in. - -**Splitting along another axis is the caller's problem.** This op only splits -dim 0. trtllm's module-level helper reaches other axes by reshaping around -this call — `chunk`/`split` along the target dim, `reshape`, `cat`, then a -`view` of the result — which is Python-level composition and would have to be -written from catalog entries, not hidden inside a wrapper. For the attention-DP -use (return each rank its own token rows after an expert-parallel MoE call) -dim 0 is already the token axis and no reshape is needed. - -**The forward op exists.** `torch.ops.trtllm.allgather` has the same -`(input, sizes, group)` signature and is the other half of the round trip — -it is a different op and is not this entry. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/reducescatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/reducescatter.py index 101b0735bb5a..24e307836403 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/reducescatter.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/reducescatter.py @@ -1,19 +1,37 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Sum every rank's copy of a tensor and hand each rank back its own rows.""" +"""Sum across ranks and keep this rank's slice, via the trtllm reducescatter op.""" -from typing import List, Optional +from __future__ import annotations + +from typing import List, Optional, Sequence import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from .._op import Arch, Cell, OpWrapper + +#: The dtypes whose NCCL mapping performs the sum the caller means by it. +#: `getDtypeMap` in `cpp/tensorrt_llm/common/opUtils.cpp` is the whole table: +#: these entries land on a native NCCL float or int type, so `ncclSum` is a +#: float or integer sum of the values. Every other entry, and every dtype the +#: table omits, is refused below -- see `is_valid` for the two distinct ways +#: they go wrong. +_SUMMABLE_DTYPES = frozenset( + { + torch.float32, + torch.float16, + torch.bfloat16, + torch.int32, + torch.int64, + torch.int8, + torch.uint8, + } +) + -def reducescatter( - input: torch.Tensor, - sizes: Optional[List[int]], - group: List[int], -) -> torch.Tensor: +class Reducescatter(OpWrapper): """Sum `input` across every rank in `group`, then keep this rank's slice. Every rank passes a tensor of the same full shape; the elementwise sum is @@ -22,36 +40,196 @@ def reducescatter( `None` when the split is even. The result is a freshly allocated tensor on every rank; `input` is left untouched. """ - # Pure-metadata guards, each for a violation measured on this machine to - # end in silent corruption, a wedge, or a process kill rather than an error. - assert input.dim() >= 1, ( - "input must have at least one dimension; a 0-d tensor segfaults inside " - "ReducescatterOp::run_list" - ) - assert input.is_contiguous(), ( - "input must be contiguous; the op reads it as packed memory and a " - "strided view is silently reduced over the wrong elements" - ) - assert input.dtype is not torch.float8_e4m3fn, ( - "float8_e4m3fn is accepted by the op but summed as raw unsigned bytes, " - "not as floats, so the result is meaningless; reduce in bf16/fp16/fp32 " - "and quantize afterwards" + + ARCHS = frozenset({Arch.SM_103}) + + # deepseek-r1 at dep4 always passes `sizes=None`: it pads to a common + # height first. The row counts are the full heights an attention-DP step + # reduces, which is the rank count times a slice. deepseek-v4.1 pads the + # same way; its cells follow. + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode: one row per rank after the split, hidden 7168", + spec=dict(rows_per_rank=1, trailing=(7168,), dtype=torch.bfloat16, sizes=None), + ), + Cell( + why="a small padded batch", + spec=dict(rows_per_rank=32, trailing=(7168,), dtype=torch.bfloat16, sizes=None), + ), + Cell( + why="prefill chunk height", + spec=dict(rows_per_rank=2048, trailing=(7168,), dtype=torch.bfloat16, sizes=None), + ), + # deepseek-v4.1 under attention DP: per-rank rows n (input W n rows) at + # the ends of the swept range 1..16384 (16384 is the largest + # max_num_tokens served) and at 1536, the largest n a decode graph + # captures (batch 256 x next_n 6). The matrix sweeps every n in between + # for both forms. + *( + Cell( + why=f"deepseek-v4.1 {form} n={n}: {role}", + spec=dict(rows_per_rank=n, trailing=trailing, dtype=dtype, sizes=None), + ) + for form, trailing, dtype in ( + ("engram bf16 [Wn, 6144]", (6144,), torch.bfloat16), + ("moe fp32 [Wn, 5120]", (5120,), torch.float32), + ) + for n, role in ( + (1, "decode end of the swept range"), + (1536, "top of the decode capture set"), + (16384, "max_num_tokens, prefill end of the swept range"), + ) + ), ) - if sizes is None: - assert input.shape[0] % len(group) == 0, ( - f"input.shape[0]={input.shape[0]} must be divisible by " - f"len(group)={len(group)} when sizes is None; the op takes " - "shape[0] // len(group) rows per rank and silently drops the remainder" + + note = """ + The ragged form and dtypes other than bf16 (beyond the DeepSeek-V4.1 fp32 + form below) work and are exercised by the op matrix, but no shipped target + reaches them, so they are not cells. CUDA graph behaviour is carried by the + matrix's own named checks rather than as a cell axis. + + DeepSeek-V4.1: under attention DP every rank is padded to + max(all_rank_num_tokens) and reduce-scatters the gathered token set with + `sizes=None`, keeping n rows of a W n-row input (W the group size): on the + engram layers (only with the device-resident engram table) the engram rows + as bf16 [W n, 6144], and on every layer the routed experts' output as fp32 + [W n, 5120]. The MoE runner is chunked at 8192 rows; this call is not, it + takes the whole concatenated output. No metadata. The sum runs in the input + dtype, sequentially in NCCL's ring order, so on real activations the fp32 + form is not the correctly rounded sum (within (W-1) * 2^-24 * sum|x| + elementwise, measured) and the bf16 form carries bf16 rounding per partial + sum; the cells and sweeps use exact payloads, which is what makes the gate + bitwise. The matrix sweeps every per-rank n in 1..16384 for both forms, in + the order an engram layer issues them, ascending and then shuffled, + bitwise; it captures one decode graph per per-rank n of the capture set + (batch G in 1..32, 40..128 step 8, 192, 256; n = G or 6G for next_n 1 or + 6, plus 35 and 160: 84 values up to 1536), each holding the forward's 42 + reduce-scatters, and replays them bitwise against the exact sum and, on + unrestricted payloads, against an eager call on the same buffer. A + captured step must pad every rank to the same n, because `sizes` is frozen + into the graph. The ragged form at V4.1 sizes is certified only at four + dumped splits, eagerly. Not certified: n above 16384; capture at an n + outside that set. The other group layouts the target runs -- world 2 (one + pp2 x dep2 stage), world 4 cut into two groups of 2, world 8 across two + trays, world 8 cut into two groups of 4 -- were measured on the MPI-session + symbol under an mpirun launcher; this matrix runs one group of 4 on the + ProcessGroup symbol and does not certify them. + """ + + def raw_call( + self, + input: torch.Tensor, + sizes: Optional[List[int]], + group: List[int], + process_group: Optional[object] = None, + ) -> torch.Tensor: + """`process_group` selects the communicator flavor, mirroring the + product router in `_torch/distributed/ops.py`: None runs the + MPI-session symbol -- the one a dep4 serving world calls -- and a boxed + c10d ProcessGroup runs the `_pg` symbol the non-MPI backends use. + """ + if process_group is None: + return torch.ops.trtllm.reducescatter(input, sizes, group) + return torch.ops.trtllm.reducescatter_pg(input, sizes, group, process_group) + + def reference( + self, + contributions_to_this_slice: Sequence[torch.Tensor], + ) -> torch.Tensor: + """The sum of what every rank sent to this rank's slice. + + The documented exception to mirroring `raw_call`, for the same reason as + allgather: a collective's output is not a function of what the calling + rank holds. It takes the one block each rank contributed to *this* + position, in ascending rank order -- not their whole inputs, because a + rank can regenerate exactly those WORLD blocks from the shared seed + convention and never has to materialize anybody else's full tensor. + + Accumulated in fp32 (int64 for integer dtypes) and rounded once. The op + sums in the input dtype in whatever order NCCL picks; on a cell's payload + that costs nothing, because every partial sum there is exact in the input + dtype either way -- which is what lets `compare` be bitwise. + """ + first = contributions_to_this_slice[0] + acc_dtype = torch.float32 if first.dtype.is_floating_point else torch.int64 + acc = torch.zeros(first.shape, dtype=acc_dtype, device=first.device) + for block in contributions_to_this_slice: + acc += block.to(acc_dtype) + return acc.to(first.dtype) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """A gate of exactly zero, which the cells earn rather than assume. + + This op computes, so NCCL's summation order and the reference's rank + order could in general differ by rounding. They do not here, because a + cell's payload is drawn from a value set -- multiples of 1/8, bounded -- + on which every partial sum is exactly representable in the input dtype. + That makes the strong gate available, and taking it is the point: any + difference at all is then a wrong reduction or a wrong slice, never + rounding, and a band would have hidden both. + + A cell built from unrestricted random values would need a band instead. + None is; see the matrix's payload builder. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + torch.testing.assert_close(out, ref, rtol=0, atol=0) + + def is_valid( + self, + input: torch.Tensor, + sizes: Optional[List[int]], + group: List[int], + process_group: Optional[object] = None, + ) -> None: + # Each guard is for a violation measured to end in silent corruption, a + # wedge, or a process kill rather than an error. + assert input.dim() >= 1, ( + "input must have at least one dimension; a 0-d tensor segfaults inside " + "ReducescatterOp::run_list" ) - else: - assert len(sizes) == len(group), ( - f"len(sizes)={len(sizes)} must equal len(group)={len(group)}; a short " - "list raises on the highest rank and wedges the others, a long one " - "reduces more rows than the input holds and returns garbage" + assert input.is_contiguous(), ( + "input must be contiguous; the op reads it as packed memory and a " + "strided view is silently reduced over the wrong elements" ) - assert sum(sizes) == input.shape[0], ( - f"sum(sizes)={sum(sizes)} must equal input.shape[0]={input.shape[0]}; " - "the op sizes the reduction from `sizes` alone, so a short split " - "silently drops the trailing rows and a long one reads past the input" + # An allowlist rather than a list of known-bad dtypes, because the two + # ways a dtype goes wrong here are both silent and the second is not + # confined to this call: + # + # mapped to the wrong arithmetic -- float8_e4m3fn and bool both land + # on ncclInt8, so ncclSum adds their raw bytes. fp8 comes back as + # nonsense and bool behaves like OR, neither with an error. + # + # not in the map at all -- float64, quint4x2 and the rest reach + # `TorchUtils::dataType()`, which throws. That call is *inside* the + # group: reducescatterOp.cpp opens with ncclGroupStart() on line 21, + # converts on line 25, and the matching ncclGroupEnd() on line 57 is + # never reached. The communicator is left mid-group, and the + # collectives after this one return invalid data. A guard here turns + # that into a Python error before the group opens. + assert input.dtype in _SUMMABLE_DTYPES, ( + f"{input.dtype} is not summable by this op; it is accepted and then " + f"either summed as raw bytes or thrown on inside an open NCCL group, " + f"which corrupts later collectives. Reduce in one of " + f"{sorted(str(d) for d in _SUMMABLE_DTYPES)} and cast afterwards" ) - return torch.ops.trtllm.reducescatter(input, sizes, group) + if sizes is None: + assert input.shape[0] % len(group) == 0, ( + f"input.shape[0]={input.shape[0]} must be divisible by " + f"len(group)={len(group)} when sizes is None; the op takes " + "shape[0] // len(group) rows per rank and silently drops the remainder" + ) + else: + assert len(sizes) == len(group), ( + f"len(sizes)={len(sizes)} must equal len(group)={len(group)}; a short " + "list raises on the highest rank and wedges the others, a long one " + "reduces more rows than the input holds and returns garbage" + ) + assert sum(sizes) == input.shape[0], ( + f"sum(sizes)={sum(sizes)} must equal input.shape[0]={input.shape[0]}; " + "the op sizes the reduction from `sizes` alone, so a short split " + "silently drops the trailing rows and a long one reads past the input" + ) + + +reducescatter = Reducescatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/bmm_out.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/bmm_out.md deleted file mode 100644 index 9730280f57a7..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/bmm_out.md +++ /dev/null @@ -1,100 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 3} ---- - -# bmm_out - -**Wraps** `torch.ops.trtllm.bmm_out` (one call). - -## Semantics - -Batched matrix multiply written into a caller-provided output buffer: - -``` -out[i, m, n] = sum_k a[i, m, k] * b[i, k, n] for every batch index i -``` - -Accumulation is fp32 for bf16/fp16 inputs: kernel and reference sum the same -exactly-representable products in a different order, so bf16 results agree to -within one rounding step but are **not** bit-identical. The per-element -mismatch rate against `torch.bmm(a.float(), b.float())` tracks K (~7e-6 at -K=16, ~5e-5 at K=64, ~3e-4 at K=576), so whether any one call happens to come -back bit-exact is a matter of how many elements it produces — and some shapes -sit in cuBLAS kernel-selection windows that are bit-identical over millions of -elements. Compare with this test's `atol=1e-3`, which is load-bearing: at -torch's default bf16 `atol` of 1e-5 a prefill-shaped -`(B=8, M=2048, K=512, N=128)` output fails 4 of 10 seeds. There is no -broadcasting: all three tensors are strictly 3D with equal batch sizes. - -Fusion boundary: the single call computes the batched gemm and nothing else — -no bias, no scaling, no activation, no quantization. The caller owns the -allocation of `out`, any transposition of `b` (pass a transpose *view*), and -any packing of head/group dims into the batch dim. Sibling ops exist for -quantized or arch-specialized batched gemms (`fp8_block_scaling_bmm_out`, -`fp4_bmm`, `cute_dsl_bf16_bmm_blackwell`). - -The op's reason to exist over plain `torch.bmm(..., out=)`: it is registered -as an opaque custom op so a torch.compile graph does not break when `out` is -a non-contiguous view. Strided views are first-class for all three arguments -(verified: transposed `out`, transposed `b`, row-strided `a`). - -## Signature - -```python -def bmm_out(a: torch.Tensor, b: torch.Tensor, out: torch.Tensor) -> None -``` - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `a` | `[B, M, K]` (3D only) | bf16 / fp16 / fp32 | any strides (strided views verified) | CUDA | -| `b` | `[B, K, N]` (3D only) | same as `a` | any strides (e.g. `weight.transpose(1, 2)` view) | CUDA | -| `out` | `[B, M, N]` (3D only) | same as `a` | any strides; written in place, never reallocated when the shape is right | CUDA | -| returns | — | — | `None`; the result is the mutation of `out` | — | - -## Metadata consumed - -None. Stateless. - -## Preconditions - -- All three tensors are 3D, on the same CUDA device. 2D inputs raise - (`batch1 must be a 3D tensor`); there is no implicit batch broadcasting. -- Batch sizes and the contraction dim `K` must match between `a` and `b` - (mismatch raises). -- `out.shape == (B, M, N)` exactly. **The op does not validate this**: a - wrong-shaped `out` is silently resized (with only a deprecation - warning). The hazard is real but is **shape/stride corruption, not a - lost write**: an aliased view with room to grow keeps its `data_ptr` and - the result *does* land in the caller's buffer, now under a silently - rewritten shape; a genuine reallocation moves the shared storage, so - sibling views follow it rather than being detached. The wrapper asserts - the shape. -- One dtype across `a`, `b`, `out`. **Mixed input dtypes always raise** — - there is no promotion path and no silent hazard here. The meta check - demands `out` in `b.dtype` while the kernel demands `a.dtype`, so when - `a.dtype != b.dtype` the two can never both be satisfied: all six - combinations over {bf16, fp16, fp32} raise, including the bf16-`a` / - fp32-`b` / fp32-`out` case an earlier revision of this bullet described - as working. The wrapper's single-dtype assert is therefore redundant - rather than load-bearing. -- `out.dtype` must equal the input dtype; a mismatch raises - (`Expected out tensor to have dtype ...`). -- Dtypes verified: bf16, fp16, fp32. float8_e4m3fn raises - (`"baddbmm_cuda" not implemented for 'Float8_e4m3fn'`). fp64 and integer - dtypes are untested (unknown). -- Arbitrary strides are supported for all three tensors, including - zero-copy transpose views of `b` and `out` (verified correct). - -## Notes - -- The op body is exactly one `torch.bmm(a, b, out=out)`; the launched kernel - is torch's cuBLAS strided-batched gemm, so numerical behavior follows the - process-wide torch matmul settings (tf32 flags for fp32, bf16 - reduced-precision-reduction flag). Receipts here were taken under torch - defaults. -- Not arch-gated; TRT-LLM uses it as the bf16 batched-gemm path in MLA - weight-absorption and output projections, with the batch dim carrying - head groups. -- Behavior above was established empirically, not read off the op's - documentation; this entry's test is what holds it. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/bmm_out.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/bmm_out.py index cf0979296def..29ef19dba614 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/bmm_out.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/bmm_out.py @@ -1,21 +1,164 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Batched matmul into a caller-provided output buffer via the trtllm bmm_out op.""" +"""Batched matmul writing into a caller-owned output, via the trtllm bmm_out op.""" + +from __future__ import annotations import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from .._op import Arch, Cell, OpWrapper, assert_within_ulp, true_fp32_matmul -def bmm_out(a: torch.Tensor, b: torch.Tensor, out: torch.Tensor) -> None: - """Compute `out[i] = a[i] @ b[i]` for every batch index, in one bmm_out call.""" - # A wrong-shaped `out` is silently resized and re-allocated by the op, - # detaching it from any buffer the caller aliased it with. - assert a.dim() == 3 and b.dim() == 3 and out.dim() == 3, "a, b, out must be 3D" - assert out.shape == (a.shape[0], a.shape[1], b.shape[2]), ( - "out must be [B, M, N] matching a [B, M, K] and b [B, K, N]" +# DeepSeek-V4.1's grouped low-rank output projection (`wo_a`), in the target's +# layout: a = o.view(T, G, 4096).transpose(0, 1), b = w.view(G, 1024, +# 4096).transpose(1, 2), out = oa.view(T, G, 1024).transpose(0, 1). G is the +# rank's share of the 8 output groups. Rows: the ends of the every-T sweep +# (1..32768) the shared test runs at each G, and at G = 1 the edges of the +# split-K bands (T 1..32 and 65..128). +_DSV41_WO_A_ROWS = { + 1: (1, 32, 33, 64, 65, 128, 129, 32768), + 2: (1, 32768), + 4: (1, 32768), + 8: (1, 32768), +} +_DSV41_CELLS = tuple( + Cell( + why=f"deepseek-v4.1 wo_a G={g} T={t}: [G, T, 4096] @ [G, 4096, 1024], all three " + "strided views; every T 1..32768 swept" + + ("; split-K" if g == 1 and (t <= 32 or 65 <= t <= 128) else ""), + spec=dict( + batch=g, + m=t, + k=4096, + n=1024, + dtype=torch.bfloat16, + transposed_out=True, + transposed_a=True, + transposed_b=True, + ), ) - # Mixed input dtypes are type-promoted instead of rejected, changing the - # required out dtype; the catalog exposes only the single-dtype form. - assert a.dtype == b.dtype == out.dtype, "a, b, out must share one dtype" - torch.ops.trtllm.bmm_out(a, b, out) + for g, rows in _DSV41_WO_A_ROWS.items() + for t in rows +) + + +class BmmOut(OpWrapper): + """`out[i] = a[i] @ b[i]` for every batch index, in place, returns None. + + The point of the op over `torch.bmm` is the caller-owned `out`: both + shipped call sites pass a transposed view of a buffer the rest of the + forward already holds, so the product lands where it is needed with no + copy. That is also the hazard -- see `is_valid`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # deepseek-r1's MLA is the only caller, and it uses the op twice per layer + # in opposite orientations: absorbing the nope half of q into the latent + # space on the way in, and projecting the latent attention output back out + # to v. Batch is the head count (128, replicated under attention DP). + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode, q_nope -> latent: [H, 1, nope] @ [H, nope, kv_lora]", + spec=dict(batch=128, m=1, k=128, n=512, dtype=torch.bfloat16, transposed_out=True), + ), + Cell( + why="decode, latent -> v: [H, 1, kv_lora] @ [H, kv_lora, v_dim]", + spec=dict(batch=128, m=1, k=512, n=128, dtype=torch.bfloat16, transposed_out=True), + ), + Cell( + why="prefill token count, the same absorbing orientation", + spec=dict(batch=128, m=2048, k=128, n=512, dtype=torch.bfloat16, transposed_out=True), + ), + Cell( + why="contiguous out, to separate the op's own correctness from the " + "transposed-view path the targets happen to use", + spec=dict(batch=128, m=64, k=128, n=512, dtype=torch.bfloat16, transposed_out=False), + ), + ) + _DSV41_CELLS + + note = """ + `out` is written, not accumulated into: its prior contents are ignored. + + DeepSeek-V4.1: the grouped low-rank output projection, bf16, K = 4096, + N = 1024, G in {1, 2, 4, 8} (8 / tp under attention TP; 8 under attention + DP and in the DSpark drafter's attention), every operand a strided view of + the target's [T, G * 4096] activation, [G * 1024, 4096] weight and + [T, G * 1024] output buffer -- no copy or transpose kernel runs, and the + result lands in that buffer. No metadata. T is the rank's token count, at + most 16384, or up to 32768 in the decoder bounded-replay pass; the shared + test runs every T in 1..32768 at every G, ascending and shuffled. At each T + nothing is written past the call's rows, and the rows are bitwise the + first T rows of the 32768-row call -- except at G = 1, T 1..32 and + 65..128, where cuBLAS runs a split-K kernel plus a reduction and the + result, inside the gate and repeatable, differs from a larger call's in + up to 1.8e-3 of its elements. Every other call is one kernel. CUDA-graph + replay is bitwise equal to eager at every T of G, 5G and 6G over the + batch set 1..32, 40..128 step 8, 192, 256 (114 counts to 1536), captured + into one shared pool after an eager call at each shape. Not certified for + DeepSeek-V4.1: T above 32768, other G at this K and N, capture at other T + or before an eager call at that shape, fp16 or fp32 in this layout, and + torch matmul flags other than the defaults, which the kernel follows. + """ + + def raw_call(self, a: torch.Tensor, b: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + torch.ops.trtllm.bmm_out(a, b, out) + return out + + def reference(self, a: torch.Tensor, b: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + """The product the op writes into `out`, fp32-accumulated. + + `out` is named to mirror `raw_call`: the op takes a destination, + the reference returns a value, and hiding the argument would hide that + the op mutates one of its inputs. + + TF32 off: torch's default fp32 matmul on this hardware is a TF32 + product, ~1e-3 relative, which would make the reference the inaccurate + side of the comparison. + """ + with true_fp32_matmul(): + return (a.float() @ b.float()).to(a.dtype) + + def is_valid(self, a: torch.Tensor, b: torch.Tensor, out: torch.Tensor) -> None: + # A wrong-shaped `out` is silently resized and re-allocated by the op, + # detaching it from whatever buffer the caller aliased it with -- the + # product then lands somewhere nobody reads, and the forward continues + # with stale data rather than failing. + assert a.dim() == 3 and b.dim() == 3 and out.dim() == 3, "a, b, out must be 3D" + assert out.shape == (a.shape[0], a.shape[1], b.shape[2]), ( + "out must be [B, M, N] matching a [B, M, K] and b [B, K, N]" + ) + # Mixed input dtypes are type-promoted instead of rejected, changing the + # required out dtype; the catalog exposes only the single-dtype form. + assert a.dtype == b.dtype == out.dtype, "a, b, out must share one dtype" + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """One GEMM's worth of accumulation-order difference. + + Kernel and reference consume bit-identical operands and differ only in + the order cuBLAS and torch sum over K. Over random operands a few + outputs land near zero, where that reordering shows up as a relative + difference of 10 on an absolute difference of 1.5e-4 -- which is what + torch's default band rejects, and what no bf16 result could represent + anyway: one ulp at these rows' scale is about 6e-1. + + So: 4 ulp element-wise against the row scale, 2 ulp of relative RMS. + Measured worst case across this entry's cells is 1.5e-4 absolute, or + roughly 3e-4 ulp -- four orders of margin, and a band that a genuinely + wrong product would miss by orders rather than squeak past. + + DeepSeek-V4.1's cells (K = 4096, outputs up to about 360 for + unit-variance operands) take the same gate. Measured on sm_103 over + every T 1..8192 at G 2 and 8, the largest |op - ref| was one bf16 + rounding step, 2.0 at outputs of 256 to 512 -- at most 2 of these ulp + against the row scale -- and at most 2.5e-3 of a call's elements + differed from the reference at all, far inside the 2-ulp relative RMS. + Swapping the weight's two groups lands about 390 away. The shared + test's every-T sweep also holds the absolute band measured there + (rtol 1.6e-2, atol 2e-3). + """ + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +bmm_out = BmmOut() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.md deleted file mode 100644 index 7b32f301dc29..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.md +++ /dev/null @@ -1,129 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 5} ---- - -# cublas_mm - -**Wraps** `torch.ops.trtllm.cublas_mm` (one call). - -## Semantics - -General matrix multiply with an optional fused bias epilogue, executed as a -single cuBLASLt matmul: - -``` -out[m, n] = sum_k mat_a[m, k] * mat_b[k, n] (+ bias[n]) -``` - -Accumulation is fp32 (verified: bf16/fp16 results match an fp32-accumulated -reference exactly at K=256). The result is cast to the output dtype. - -Fusion boundary: the single call computes the gemm and, when `bias` is given, -adds it per output column inside the epilogue. Nothing else happens inside: -no activation, no quantization, no input/output scaling (fp8 inputs are -multiplied with alpha=1; a sibling op `cublas_scaled_mm` exists for the -scaled variant). The caller owns weight transposition metadata (pass a -transpose *view*, see Preconditions) and any flattening of leading dims. - -## Signature - -```python -def cublas_mm( - mat_a: torch.Tensor, - mat_b: torch.Tensor, - bias: torch.Tensor | None = None, - out_dtype: torch.dtype | None = None, - output_buffer_kind: int = 0, - group: list[int] | None = None, -) -> torch.Tensor -``` - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `mat_a` | `[M, K]` (2D only) | bf16 / fp16 / fp32 / float8_e4m3fn | dense row-major (contiguous) | CUDA | -| `mat_b` | `[K, N]` (2D only) | same as `mat_a` | dense column-major: strides `(1, K)`, e.g. `weight.t()` of a contiguous `[N, K]` weight | CUDA | -| `bias` | `[N]` or `None` | must equal the output dtype | contiguous | CUDA | -| `out_dtype` | — | `None` → same as `mat_a`; conversions verified: bf16→fp32, fp8e4m3→bf16 | — | — | -| `output_buffer_kind` | scalar | Python int: 0=DEFAULT, 1=USERBUFFERS, 2=NCCL_WINDOW | — | — | -| `group` | list of ranks or `None` | Python ints | only meaningful with NCCL_WINDOW | — | -| returns | `[M, N]` | `out_dtype` or `mat_a.dtype` | newly allocated, contiguous | CUDA | - -## Metadata consumed - -None. Stateless for the default path. `output_buffer_kind` selects where the -output tensor is allocated: 0 is a plain CUDA allocation; 1 (USERBUFFERS) and -2 (NCCL_WINDOW, together with `group` = the TP rank list) draw from -communication buffer pools so a downstream allreduce can consume the output -in place — those pools must have been initialized by the runtime beforehand. -Single-GPU callers pass `0, None` (2 with `group=None` silently falls back to -a plain allocation). The enum is importable as -`tensorrt_llm.bindings.internal.thop.BufferKind`. - -## Preconditions - -- `mat_a` and `mat_b` are 2D, on the same CUDA device, same dtype. 3D+ - inputs raise; flatten leading dims to `[M, K]` and unflatten after. -- `mat_a` is dense row-major. **The kernel ignores `mat_a` row strides**: a - row-strided view (e.g. a column slice of a wider buffer) is accepted and - produces silently wrong results (verified). The wrapper asserts - contiguity. -- `mat_b` has strides exactly `(1, K)` — a dense column-major buffer, i.e. - the `.t()` view of a contiguous `[N, K]` weight. The op checks - `stride(0) == 1` and raises otherwise, but does **not** check - `stride(1) == K`: the transpose view of a row-strided weight is accepted - and produces silently wrong results (verified). The wrapper asserts both. -- Dtypes: bf16, fp16, fp32 work with `out_dtype=None`; float8_e4m3fn - requires an explicit `out_dtype` (`None` raises a cuBLAS error). - `mat_a.dtype != mat_b.dtype` raises `CUBLAS_STATUS_NOT_SUPPORTED`. -- `bias`, when given: - - shape exactly `[N]`, contiguous, on the same device. **The op does not - validate the bias shape** (a wrong-length bias was accepted); the - wrapper asserts it. - - dtype must equal the *output* dtype (bf16 bias for bf16 out, fp32 bias - for bf16→fp32, bf16 bias for fp8→bf16). A mismatched dtype (e.g. fp32 - bias with bf16 output) is accepted and produces silently wrong results - (verified). - - not supported with fp32 inputs: the bias is accepted and **silently - ignored** (verified: output equals the unbiased product). The wrapper - asserts against this combination. -- `out_dtype` conversions other than the verified ones are not guaranteed: - bf16→fp16 raises a cuBLASLt runtime error on this machine; combinations - not listed above are untested (unknown). - -## Notes - -- The fake (meta) registration of this op accepts N-D `mat_a`, but the real - kernel requires 2D — shape inference under compile can diverge from - eager behavior for N-D inputs. -- Error traces name `cpp/tensorrt_llm/thop/cublasScaledMM.cpp` as the - implementation (shared with the scaled variant); the C++ source is not - shipped in the wheel, so all behavior above was established empirically - on sm_100. -- TRT-LLM routes bf16 linears to this op on sm>=100 because cuBLASLt picks - single-pass cluster-mode kernels for small-M (decode) gemms there; the - op itself is not gated on arch. -- All silently-wrong-result behaviors listed under Preconditions were - observed directly, not inferred from the op's documentation. - - -## The fp32 cell needs its reference pinned, not its tolerance widened - -torch 2.12 defaults `torch.backends.cuda.matmul.fp32_precision` to `tf32` (and -`allow_tf32` to True) on this hardware, so the natural reference -`mat_a.float() @ mat_b.float()` is a **TF32** product. Against a correct -kernel that reads as a ~0.035 absolute error at K=1024 -- about 30x the op's -own -- and it is the *reference* that is wrong. - -Measured on sm_103 against a float64 product on the same inputs: - -| | max abs error vs float64 | -|---|---| -| `cublas_mm` | 1.9e-5 | -| torch, `allow_tf32=False` | 1.9e-5 (**bit-identical to the op**) | -| torch, `allow_tf32=True` | 3.5e-2 | - -So the op is exact to fp32 and this entry's `_ref_mm` now forces TF32 off for -the duration of the reference product. The same trap applies to any entry -whose reference multiplies in fp32; the bf16 and fp16 cells were merely -tolerant enough to hide it. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py index 9a0af503e9c8..ce16fc08c633 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py @@ -2,34 +2,253 @@ # SPDX-License-Identifier: Apache-2.0 """General matmul with optional fused bias via the trtllm cuBLASLt gemm op.""" +from __future__ import annotations + import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from .._op import ULP, Arch, Cell, OpWrapper, assert_within_ulp, true_fp32_matmul + +# The DeepSeek-V4.1 target's calls: (role, K, N, dtype, out_dtype, bias, weight std, rows, what). +# Every call passes a contiguous [M, K] activation and the `.t()` of a contiguous [N, K] weight. +# `rows` are the ends of the every-M sweep the shared test runs at that call, plus the edges of +# the patch embed's refused band. Weight std None is unit-normal (the core calls); the vision +# calls take N(0, 1/K) weights so their outputs are O(1) and a unit-normal bias is not lost. +_DSV41_CALLS = ( + ("compressor-r2", 5120, 1024, torch.float32, None, False, None, (1, 16384), + "fp32 [M, 5120] -> [M, 1024], [kv || score] of a ratio-2 KV source"), + ("router-fp32", 5120, 384, torch.float32, None, False, None, (1, 8192), + "fp32 [M, 5120] -> [M, 384], the router gate on fp32 inputs"), + ("dspark-gate", 5120, 128, torch.float32, None, False, None, (1, 8192), + "fp32 [M, 5120] -> [M, 128], DSpark drafter MoE gate, M = 5G"), + ("markov-head", 256, 129280, torch.float32, None, False, None, (1, 8192), + "fp32 [M, 256] -> [M, 129280], DSpark Markov head, M = G"), + ("compressor-r1", 5120, 512, torch.bfloat16, None, False, None, (1, 16384), + "bf16 [M, 5120] -> [M, 512], compressor of a ratio-1 KV source"), + ("index-k", 512, 128, torch.bfloat16, None, False, None, (1, 16384), + "bf16 [M, 512] -> [M, 128], index-K of the compressed latent"), + ("weights-proj", 5120, 32, torch.bfloat16, None, False, None, (1, 32768), + "bf16 [M, 5120] -> [M, 32], indexer weights_proj"), + ("router", 5120, 384, torch.bfloat16, torch.float32, False, None, (1, 8192), + "bf16 [M, 5120] -> fp32 [M, 384], the router gate's logits"), + ("patch-embed", 588, 1024, torch.bfloat16, None, True, 588**-0.5, (1, 2304, 2433, 17037), + "bf16 [P, 588] -> [P, 1024] + bias; P 2305..2432 is refused, pad to >= 2433"), + ("vit-wqkv", 1024, 3072, torch.bfloat16, None, True, 1024**-0.5, (1, 17037), + "bf16 [P, 1024] -> [P, 3072] + bias, vision block wqkv"), + ("vit-wo", 1024, 1024, torch.bfloat16, None, True, 1024**-0.5, (1, 17037), + "bf16 [P, 1024] -> [P, 1024] + bias, vision block wo"), + ("vit-w1", 1024, 5632, torch.bfloat16, None, False, 1024**-0.5, (1, 17037), + "bf16 [P, 1024] -> [P, 5632], vision block w1 (gate | up)"), + ("vit-w2", 2816, 1024, torch.bfloat16, None, False, 2816**-0.5, (1, 17037), + "bf16 [P, 2816] -> [P, 1024], vision block w2"), + ("aligner-w1", 9216, 5120, torch.bfloat16, None, True, 9216**-0.5, (1, 2048), + "bf16 [C, 9216] -> [C, 5120] + bias, aligner w1"), + ("aligner-w2", 5120, 5120, torch.bfloat16, None, True, 5120**-0.5, (1, 2048), + "bf16 [C, 5120] -> [C, 5120] + bias, aligner w2"), +) # fmt: skip -def cublas_mm( - mat_a: torch.Tensor, - mat_b: torch.Tensor, - bias: torch.Tensor | None = None, - out_dtype: torch.dtype | None = None, - output_buffer_kind: int = 0, - group: list[int] | None = None, -) -> torch.Tensor: - """Return `mat_a @ mat_b (+ bias)`, fp32-accumulated, in one cublas_mm call.""" - # The kernel reads mat_a as dense row-major and mat_b as dense - # column-major; other layouts produce silently wrong results. - assert mat_a.is_contiguous(), "mat_a must be dense row-major [M, K]" - assert mat_b.stride(0) == 1 and mat_b.stride(1) == mat_b.shape[0], ( - "mat_b must be dense column-major [K, N] (e.g. weight.t())" +_DSV41_CELLS = tuple( + Cell( + why=f"deepseek-v4.1 {role} M={m}: {what}; every M {rows[0]}..{rows[-1]} swept", + spec=dict(m=m, k=k, n=n, dtype=dt, bias=bias, out_dtype=out_dt, w_std=w_std), ) - if bias is not None: - out_dt = out_dtype if out_dtype is not None else mat_a.dtype - # A bias in the wrong dtype or shape is accepted by the op and - # produces silently wrong results; with fp32 inputs the bias is - # accepted but silently ignored. - assert mat_a.dtype != torch.float32, "bias is silently ignored for fp32 inputs" - assert bias.dtype == out_dt, "bias dtype must equal the output dtype" - assert bias.shape == (mat_b.shape[1],) and bias.is_contiguous(), ( - "bias must be a contiguous [N] tensor" + for role, k, n, dt, out_dt, bias, w_std, rows, what in _DSV41_CALLS + for m in rows +) + +# The M that `is_valid` refuses at bf16 [M, 588] x [588, 1024] -> bf16 (DeepSeek-V4.1 patch embed). +_PATCH_EMBED_SPLIT_K_ROWS = range(2305, 2433) + + +class CublasMm(OpWrapper): + """`mat_a @ mat_b (+ bias)`, fp32-accumulated, one cublas_mm call. + + `mat_b` is read column-major, which is what a stored weight's `.t()` already + is -- the targets pass `weight.t()` and no transpose kernel runs. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Both targets use this for every dense projection, so the cells are the + # three argument shapes that differ in kind rather than in size: fused bias + # (gpt-oss carries one on qkv, o and the router), no bias (deepseek's + # projections), and an out_dtype that is not the input's (deepseek's router + # asks for fp32 logits). DeepSeek-V4.1's calls follow (`_DSV41_CALLS`): one + # cell per end of each call's swept row range, plus the edges of the patch + # embed's refused band. + CELLS: tuple[Cell, ...] = ( + Cell( + why="gpt-oss qkv with fused bias: hidden 2880 -> (64 + 8 + 8) * 64", + spec=dict(m=16, k=2880, n=5120, dtype=torch.bfloat16, bias=True, out_dtype=None), + ), + Cell( + why="deepseek q_a, no bias: hidden 7168 -> q_lora 1536", + spec=dict(m=16, k=7168, n=1536, dtype=torch.bfloat16, bias=False, out_dtype=None), + ), + Cell( + why="deepseek router: bf16 in, fp32 logits out, 256 experts", + spec=dict( + m=16, k=7168, n=256, dtype=torch.bfloat16, bias=False, out_dtype=torch.float32 + ), + ), + Cell( + why="decode-shaped: one token, the shape every projection sees per step", + spec=dict(m=1, k=7168, n=1536, dtype=torch.bfloat16, bias=False, out_dtype=None), + ), + ) + _DSV41_CELLS + + note = """ + `output_buffer_kind` and `group` are the op's multi-stream and collective + arguments. Both shipped targets leave them at their defaults, so neither is + a cell axis and neither is certified here. + + DeepSeek-V4.1: every call passes bias=None and out_dtype=None except the + router (bf16 in, fp32 logits) and the vision calls (their checkpoint bias, + bf16). No metadata. cuBLAS picks its kernel from M with no algorithm + pinned, and on sm_103 the kernel changes every 16 to 64 rows, so the shared + test sweeps every M from 1 to each call's largest: 16384 at both + compressors and index-K, 32768 at weights_proj, 8192 at the router (the + target splits the routed MoE's rows into chunks of at most 8192) and the + two DSpark calls, 17037 at the vision tower and 2048 at the aligner. The + fp32 calls are IEEE fp32, not TF32. No core call takes a split-K kernel + with bf16 partial sums at any swept M, and at every M at least 0.99 of a + bf16 call's elements are bit-identical to the rounded-once reference. + Hazard: bf16 K = 588, N = 1024 (the patch embed) at M 2305..2432, with or + without bias, takes such a kernel and leaves the bf16 gate on 1.4% or more + of its elements; `is_valid` refuses it, and the first P rows of a call + padded with zero rows to 2433 are in the gate. CUDA-graph replays are + bitwise equal to eager at the row counts the target captures (every batch + G in 1..32, 40..128 step 8, 192, 256, times 1, 3 and 6, and also 4 and 24 + at the router; 5G at the DSpark gate, G at the Markov head) and at the + vision tower's P 1517, 4680, 5325, 8772, 17037 and aligner C 182, 546, + 601, 986, 1902. Rows of one call are not bitwise those of a call with + more rows except at index-K. Not certified for DeepSeek-V4.1: M above the + swept ranges, capture at other row counts or of several calls sharing a + graph pool, fp8 or fp16 inputs, and the patch embed band itself. + """ + + def raw_call( + self, + mat_a: torch.Tensor, + mat_b: torch.Tensor, + bias: torch.Tensor | None = None, + out_dtype: torch.dtype | None = None, + output_buffer_kind: int = 0, + group: list[int] | None = None, + ) -> torch.Tensor: + return torch.ops.trtllm.cublas_mm(mat_a, mat_b, bias, out_dtype, output_buffer_kind, group) + + def reference( + self, + mat_a: torch.Tensor, + mat_b: torch.Tensor, + bias: torch.Tensor | None = None, + out_dtype: torch.dtype | None = None, + output_buffer_kind: int = 0, + group: list[int] | None = None, + ) -> torch.Tensor: + """fp32-accumulated, TF32 off -- see `true_fp32_matmul`. + + fp32 operands are the exception: their reference is the float64 + product, left in float64. The op is IEEE fp32 there, so an fp32 + reference would be the op's own arithmetic in another order, and the + float64 dtype is what tells `compare` to use the fp32 gate. + `output_buffer_kind` and `group` are ignored. + """ + if mat_a.dtype == torch.float32: + ref = mat_a.double() @ mat_b.double() + return ref + bias.double() if bias is not None else ref + with true_fp32_matmul(): + ref = mat_a.float() @ mat_b.float() + if bias is not None: + ref = ref + bias.float() + return ref.to(out_dtype if out_dtype is not None else mat_a.dtype) + + def is_valid( + self, + mat_a: torch.Tensor, + mat_b: torch.Tensor, + bias: torch.Tensor | None = None, + out_dtype: torch.dtype | None = None, + output_buffer_kind: int = 0, + group: list[int] | None = None, + ) -> None: + # The kernel reads mat_a as dense row-major and mat_b as dense + # column-major; other layouts produce silently wrong results. + assert mat_a.is_contiguous(), "mat_a must be dense row-major [M, K]" + assert mat_b.stride(0) == 1 and mat_b.stride(1) == mat_b.shape[0], ( + "mat_b must be dense column-major [K, N] (e.g. weight.t())" ) - return torch.ops.trtllm.cublas_mm(mat_a, mat_b, bias, out_dtype, output_buffer_kind, group) + if bias is not None: + out_dt = out_dtype if out_dtype is not None else mat_a.dtype + # A bias in the wrong dtype or shape is accepted by the op and + # produces silently wrong results; with fp32 inputs the bias is + # accepted but silently ignored. + assert mat_a.dtype != torch.float32, "bias is silently ignored for fp32 inputs" + assert bias.dtype == out_dt, "bias dtype must equal the output dtype" + assert bias.shape == (mat_b.shape[1],) and bias.is_contiguous(), ( + "bias must be a contiguous [N] tensor" + ) + # Measured on sm_103 at DeepSeek-V4.1's vision patch embed: bf16 + # [M, 588] x [588, 1024] -> bf16 at M 2305..2432, with or without bias, + # is the one place cuBLASLt's heuristic picks a split-K kernel whose + # partial sums are stored in bf16. Every partial is rounded to bf16 + # before the sum: at most 0.685 of the elements are bit-identical to + # the rounded-once product (0.9996 or more at every other M) and 1.4% or + # more leave the bf16 gate. The first M rows of a call padded with zero + # rows to 2433 are in the gate. + assert not ( + mat_a.dtype == torch.bfloat16 + and out_dtype in (None, torch.bfloat16) + and mat_a.dim() == 2 + and tuple(mat_b.shape) == (588, 1024) + and mat_a.shape[0] in _PATCH_EMBED_SPLIT_K_ROWS + ), "bf16 K=588 N=1024 at M 2305..2432 takes a bf16-partial split-K kernel; pad M to >= 2433" + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """One GEMM's worth of accumulation-order difference. + + Kernel and reference consume bit-identical operands and differ only in + the order cuBLAS and torch sum over K. Over random operands a few + outputs land near zero, where that reordering shows up as a relative + difference of 10 on an absolute difference of 1.5e-4 -- which is what + torch's default band rejects, and what no bf16 result could represent + anyway: one ulp at these rows' scale is about 6e-1. + + So: 4 ulp element-wise against the row scale, 2 ulp of relative RMS. + Measured worst case across this entry's cells is 1.5e-4 absolute, or + roughly 3e-4 ulp -- four orders of margin, and a band that a genuinely + wrong product would miss by orders rather than squeak past. + DeepSeek-V4.1's bf16 cells, its bf16 -> fp32 router included, take this + gate unchanged; their every-M sweeps in the shared test hold the + tighter checks measured there. + + fp32 operands are scored differently, recognised by the float64 + reference (see `reference`): the largest |out - float64 product| must + be at most 512 fp32 ulp (2^-14) of the output's RMS. Measured on sm_103 + with unit-variance operands, the op sat at most 1.56e-3 from float64 at + K = 5120 over every M to 16384, and 1.05e-4 at K = 256. The output RMS + is sqrt(K) there, so the gate reads 4.4e-3 and 9.8e-4: the absolute + per-K gates those sweeps held (5e-3 and 1e-3), put in the one unit this + method can see and rounded down. A TF32 product of the same operands + lands 0.09 to 0.11 (K = 5120) and 0.022 to 0.027 (K = 256) away, more + than 20x outside. The bf16 scoring cannot see TF32 at all: 0.1 on a row + scale near 250 is a tenth of a bf16 ulp. + """ + if ref.dtype == torch.float64: + assert out.dtype == torch.float32, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + err = (out.double() - ref).abs().max().item() + rms = ref.pow(2).mean().sqrt().item() + gate = 512 * ULP[torch.float32] * rms + assert err <= gate, f"max |out - float64| {err:.3e} exceeds {gate:.3e} (2^-14 of RMS)" + return + # Scored in bf16 ulp even when the output is fp32: every cell's + # operands are bf16, so the fp32 out_dtype the router asks for is a + # wider container for a bf16-limited product, not a more precise one. + # Against fp32's own ulp the same 2.4e-6 difference reads as 5x over. + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0, ulp=ULP[torch.bfloat16]) + + +cublas_mm = CublasMm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/nvfp4_gemm.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/nvfp4_gemm.md deleted file mode 100644 index 98d9f02bec2b..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/nvfp4_gemm.md +++ /dev/null @@ -1,305 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 20} ---- - -# nvfp4_gemm - -**Wraps** `torch.ops.trtllm.nvfp4_gemm` (one call). - -## Semantics - -Dense GEMM in nn.Linear orientation over two **block-scaled NVFP4** operands. -Let `A` be the activation and `B` the weight, both stored as packed e2m1 -nibbles plus one e4m3 scale per 16 contiguous elements along `K`: - -``` -a[m, k] = e2m1(act_fp4 nibble (m, k)) * e4m3(act_sf scale (m, k // 16)) -b[n, k] = e2m1(weight nibble (n, k)) * e4m3(weight_scale scale (n, k // 16)) - -out[m, n] = output_dtype( alpha * sum_k a[m, k] * b[n, k] + bias[n] ) -``` - -i.e. `out = alpha * (a @ b.T) + bias`. The sum is accumulated in **fp32**; -`alpha` (a device fp32 scalar) multiplies the whole accumulator; `bias` (an -optional per-column vector) is added after `alpha`; the result is cast to -`output_dtype` once, at the end. - -Fusion boundary — inside the call: both operands' **block-scale -dequantization**, the fp32 accumulation, the per-tensor `alpha`, the optional -per-column bias, and the output cast. Outside the call, still the caller's: -producing `act_fp4`/`act_sf` (an activation quantizer such as -`torch.ops.trtllm.fp4_quantize`), putting **both** scale buffers into the -swizzled layout below, folding the two per-tensor global scales into `alpha`, -any activation function, any reduction across ranks, and any narrowing of a -padded `N` back to `out_features`. - -### The two global scales live in `alpha` - -An NVFP4 tensor quantized with a per-tensor global scale `g` reconstructs as -`data * sf / g` — this op applies `data * sf` only. With `g_act` for the -activation and `g_w` for the weight, the caller owes - -``` -alpha = 1 / (g_act * g_w) -``` - -For a modelopt/HF checkpoint the two conventions meet like this: the stored -`input_scale` is `amax_act / (448*6)` = `1 / g_act` (the quantizer's -`global_scale` is its reciprocal) and the stored `weight_scale_2` is -`amax_w / (448*6)` = `1 / g_w`, so - -``` -alpha = input_scale * weight_scale_2 # both as stored on disk - = weight_scale_2 / g_act # g_act = the quantizer's global_scale -``` - -`alpha` is a load-time scalar for weight-only-static checkpoints; nothing in -this op derives it. - -### Scale-factor layout (both operands) - -Each scale buffer is **1-D** and holds one e4m3 byte per (row, 16-element -block) in the **128x4 swizzled** order: the order the trtllm NVFP4 activation -quantizer emits with `is_sf_swizzled_layout=True` -(`torch.ops.trtllm.fp4_quantize`), and the order -`torch.ops.trtllm.block_scale_interleave` produces from a row-major -`[rows, cols]` tensor — the latter verified byte-identical to the formula below -on this machine. With `cols = K / 16`, the byte for `(row r, block c)` sits at -flat offset - -``` -(c % 4) -+ (c // 4) * 512 # 4 cols x 128 rows per column group -+ (r % 32) * 16 -+ ((r % 128) // 32) * 4 -+ (r // 128) * 128 * pad_up(cols, 4) -``` - -and the buffer's length is `pad_up(rows, 128) * pad_up(cols, 4)` bytes — -`rows = M` for the activation, `rows = N` for the weight. Bytes at offsets no -real `(r, c)` addresses (row padding up to a multiple of 128, column padding up -to a multiple of 4) are **ignored**: filling them with `0x00` or with `0x7E` -(448) gives bitwise-identical output. - -**A checkpoint that stores `weight_scale` row-major `[N, K/16]` therefore owes a -load-time relayout.** Feeding the linear buffer instead is accepted and silently -wrong — the observed error was of the same order as the output itself. Both -scale buffers are passed as `uint8`; an fp8 `float8_e4m3fn` view of the same -bytes is rejected. - -## Signature - -```python -def nvfp4_gemm( - act_fp4: torch.Tensor, - weight: torch.Tensor, - act_sf: torch.Tensor, - weight_scale: torch.Tensor, - alpha: torch.Tensor, - output_dtype: torch.dtype, - output_buffer_kind: int = 0, - allowed_backends: str = "cutlass,cublaslt,cuda_core", - group: list[int] | None = None, - bias: torch.Tensor | None = None, -) -> torch.Tensor -``` - -| Argument | Shape | Dtype | Layout / value | Device | -|---|---|---|---|---| -| `act_fp4` | `[M, K/2]` (2-D only) | uint8 | contiguous; element `2i` in the low nibble, `2i+1` in the high nibble | CUDA | -| `weight` | `[N, K/2]` (2-D only) | uint8 | contiguous **row-major nn.Linear weight**, not its `.t()` view | CUDA | -| `act_sf` | 1-D, `pad_up(M,128) * pad_up(K/16,4)` | uint8 | contiguous, 128x4 swizzled (above) | CUDA | -| `weight_scale` | 1-D, `pad_up(N,128) * pad_up(K/16,4)` | uint8 | contiguous, 128x4 swizzled (above) | CUDA | -| `alpha` | exactly 1 element (`[1]` or 0-D) | float32 | `1/(g_act * g_w)` | CUDA | -| `output_dtype` | scalar | `torch.dtype` | bfloat16, float16 or float32; the `cutedsl` backend supports bfloat16 only | — | -| `output_buffer_kind` | scalar | Python int | `0` = plain allocation (certified). `1` = userbuffers, `2` = NCCL window — for multi-rank output buffers, not certified here | — | -| `allowed_backends` | scalar | Python str | comma-separated subset of `cutlass`, `cublaslt`, `cuda_core`, `cutedsl`, `marlin`; all compute the same result | — | -| `group` | list of ranks or `None` | Python list | only meaningful with a window output buffer; `None` otherwise | — | -| `bias` | `[N]`, 1-D | **equal to `output_dtype`** | contiguous; added per output column | CUDA | -| returns | `[M, N]` | `output_dtype` | freshly allocated, contiguous | same as `act_fp4` | - -`K` is `2 * act_fp4.shape[1]`. Inputs are never written (verified bitwise), and -the returned buffer aliases nothing. - -`allowed_backends` is a *selection* set, not a composition: exactly one backend -runs per call and every one of them computes the formula above (all verified -equal to the reference on this machine). The op's own schema spells the -parameters `act_fp4, weight, act_sf, weight_scale, alpha, output_dtype, -output_buffer_kind, allowed_backends, group, bias`. - -## Metadata consumed - -The process-global AutoTuner profiling cache -(`tensorrt_llm._torch.autotuner.AutoTuner`). Nothing must be prepared: - -- Outside a tuning context, a cache miss logs a warning and runs the **fallback - tactic**: the `cutlass` backend with its default config if `cutlass` is in - `allowed_backends`, otherwise the first backend listed. -- Inside `tensorrt_llm._torch.autotuner.autotune()` the call profiles every - valid (backend, tactic) pair and caches the winner per M-bucket; one tuning - call at rows `M` fills every power-of-2 bucket `1..last_pow2(M)` for that - `(K, N)` (14 buckets at `M = 8192`, verified at four shapes). Later calls map - their `M` to its bucket. Observed winners on this machine span `cublaslt` and - `cutlass` depending on shape; across the 56 cache entries the four R1 shapes - produce, every winner carried a profiled tactic id rather than the fallback - marker. - -Correctness is certified for both cache states. At the four DeepSeek-R1-0528 -shapes below the tuned tactic's output was **bitwise identical** to the cold -fallback's at `M` in {1, 3, 8, 4097, 8192} — a measured property of this op at -those shapes, not a general guarantee that tactic choice is bit-neutral. No -workspace, runtime object or attention metadata is involved; the op is -otherwise stateless. - -## Preconditions - -Violations marked *silent* were observed on this machine to produce wrong -results without raising; the wrapper rejects each with a metadata assert. -Everything else raises. - -- `K % 32 == 0` and `N % 32 == 0` (16-byte operand lines). Both raise otherwise - — `Expected k to be divisible by 32` / `Expected n to be divisible by 32`. - `M` is unconstrained (1, 3, 7, 9, 127, 128, 129, 1000, 4097, 8192, ... all - verified); `M == 0` raises. This rule is the whole shape domain, and it was - re-measured as such: `K` up to 18432, `N` up to 36864 and `M` up to 8192 - behave exactly like the smaller shapes (see the shape-dependence note below). -- `act_fp4` and `weight` are 2-D, `uint8`, on CUDA, and **contiguous**. A 3-D - activation raises; a non-uint8 view raises; a CPU tensor raises. Contiguity is - *silent*: the `cutlass` and `cuda_core` backends raise, but `cublaslt` and - `cutedsl` ignore the stride and return wrong results, and `cublaslt` is both - in the default backend list and a frequent tuner pick. -- `weight.shape[1] == act_fp4.shape[1]` (same `K`); mismatch raises. -- `act_sf` and `weight_scale` are `uint8`, **contiguous**, and hold the - swizzled layout at the full padded size above. Non-contiguity is *silent* - under `cublaslt`/`cutedsl` (and under `cuda_core` for `weight_scale`); an - unswizzled (row-major linear) buffer of the right size is *silent* on the - default path, and no backend has any way to detect it; a buffer shorter than - the padded size reads out of bounds. Padding bytes may hold anything. -- `alpha` is a float32 CUDA tensor with **exactly one element**. A half or CPU - tensor raises. More than one element is *silent*: element 0 is used and the - rest ignored — this build has no per-token alpha. -- `bias`, when given, is 1-D with exactly `N` elements and dtype equal to - `output_dtype`. Every violation raises. -- `output_dtype` is bfloat16, float16 or float32. The `cutedsl` backend accepts - bfloat16 only and raises `ValueError` otherwise — forcing - `allowed_backends="cutedsl"` with fp16/fp32 raises, and so does an - `autotune()` pass over any list that merely contains `cutedsl` (its runner is - constructed while enumerating tactics); outside tuning such a list is fine, - because the fallback never constructs it. -- `allowed_backends` must be a non-empty comma-separated subset of the five - names; an empty or misspelled string raises `ValueError`. `marlin` is - Hopper-only and raises on sm_100. `cuda_core` requires SM >= 100 and is only - ever *selected* by the tuner for `M <= 8`; forcing it computes correctly up to - `M = 16` and raises `Failed to dispatch cudaCoreGemmLauncher` above that. -- `output_buffer_kind=1` (userbuffers) raises unless a userbuffers workspace of - sufficient size was allocated by the runtime. -- The op is deterministic: repeating a call with the same inputs and the same - tuner state gives bitwise-identical output. - -## Notes - -- Certified on sm_100 (B200) only. The kernels behind every backend are - Blackwell block-scaled MMA paths (the CUTLASS one is instantiated for - `cutlass::arch::Sm100`/`Sm103` and traps on other archs); no receipt is - claimed elsewhere. -- Receipt coverage, shapes: the four DeepSeek-V3-Lite NVFP4 dense linears in - `[out, in]` orientation — `(K, N)` = (2560, 24576), (12288, 2560), - (2560, 6144), (3072, 2560) — at `M` in {1, 2, 8, 64, 1024, 4096} (first - shape) and {1, 8, 64, 1024} (rest); the four DeepSeek-R1-0528 dense linears — - `(K, N)` = (7168, 36864), (18432, 7168), (7168, 4096), (2048, 7168) — each at - `M` in {1, 2, 8, 64, 1024, 4096, 8192}; `M` in {1, 3, 7, 9, 127, 128, 129, - 1000} at (2560, 512); `N = 160` (not a multiple of 128, so the weight scale - buffer pads rows); and a wide block-scale *range* case (scale bytes spanning - 2^-5..2^5, which is what makes the accumulation round) at `K = 12288` and at - `K = 18432`. -- Receipt coverage, everything else: bf16 / fp16 / fp32 outputs — bf16 at every - shape, fp16 at (2560, 512), fp32 at (2560, 512) and at both wide-range shapes - (`K` = 12288 and `K` = 18432); `alpha` in {1, 0.5, 0.03125, 1e-4}; bf16 and - fp32 bias at (2560, 512) only, the R1 shapes having been driven without bias - as the dense path uses none; the default `allowed_backends` - string at every shape above, plus `cutlass`, `cublaslt`, `cuda_core` and - `cutedsl` each forced alone at (2560, 512) and, at every R1 shape, `cutlass` / - `cublaslt` / `cutedsl` forced alone at `M` = 8 and `M` = 8192 with - `cuda_core` alongside them at `M` = 8; the autotuned path at (2560, 6144) and - at each R1 shape (tuned at `M` = 8192, then replayed warm at `M` in - {1, 3, 8, 4097, 8192} and compared bitwise against the same calls on a cold - cache); scale-padding insensitivity; input non-mutation and call-to-call - determinism; the swizzled-vs-linear weight-scale contrast at (2560, 512) and - at both extreme R1 shapes; 20 rejected-domain calls; and the five silent - domains the wrapper guards. -- **Nothing shape-dependent changes past the smaller shapes' maxima.** What is - shape-dependent is the tactic space, not the arithmetic: the number of - cuBLASLt heuristic algorithms varies with `(M, K, N)` (1..8 observed), the - CuTe DSL tile/cluster candidate count varies (120..240 observed), and - `cuda_core` is offered only for `M <= 8`. Measured at the four R1 shapes - against the smaller certified ones, every one of those counts landed inside - the range the smaller shapes already produce, no shape produced an empty - tactic list, the CUTLASS tactic list is 32 configs everywhere (its - enumeration takes no shape at all), and the scale-buffer size formula holds - unchanged at its widest here (`K/16 = 1152` columns, `K = 18432`). At those - shapes `cutlass`, `cublaslt`, `cutedsl` and — at `M <= 8` — `cuda_core` - returned **bitwise identical** outputs, each matching the native-torch - reference, and the warm tuner cache reproduced the cold fallback's bits - exactly. So the domain rule in Preconditions was sufficient on its own at - this scale-up: the shapes' only effect was on which tactic wins. -- Numerics. Against a native-torch reference built from the operand bytes, the - bf16 output matches `torch.testing.assert_close` at **default bf16 - tolerances** (rtol 1.6e-2, atol 1e-5) at every shape above; the observed max - relative deviation was 3.9e-3, i.e. one bf16 ulp. **That figure is against - the unrounded fp32 reference** (measured 3.891e-3 vs a 2^-8 = 3.906e-3 ulp). - The test itself compares against `ref.to(torch.bfloat16)`, so its own - assertions come out **bit-exact** — an instrumented run that records - `assert_close` deviations will therefore report zero here, and that is not a - contradiction of this number. When the block scales sit - in a narrow band the products and partial sums are exactly representable and - the **fp32 output is bit-exact** against an fp64 reference. With block scales - spanning 2^-5..2^5 the fp32 sum does round, and the deviation stays inside - the recursive-summation bound `K * 2^-24 * sum|a_i b_i|` with ~600x margin — - which is what an fp32 accumulator looks like and a bf16/fp16 one would not. - Both regimes were re-measured at `K = 18432`, the largest `K` certified: - narrow band still bit-exact in fp32, wide band inside the same bound with - 715x margin. -- Discrimination at the largest shapes was measured, not assumed: moving a - single e4m3 scale byte, or a single e2m1 nibble, by one code — the smallest - perturbation either operand admits — makes the default-tolerance comparison - fail at `(K, N, M)` = (18432, 7168, 8192) and (7168, 36864, 8192). Feeding - the weight scale row-major (unswizzled) instead of swizzled lands ~0.5x the - output's own magnitude off at those shapes. -- One call is one kernel launch on the certified default path: profiling a - single call showed exactly one device kernel for `cutlass` (bias fused into - its `LinCombPerColBias` epilogue) and for `cublaslt` (bias fused, the - `..._bias_TNT` nvjet variant). Two internal implementation details do not - change the result but are worth knowing: the `cuda_core` backend first - un-swizzles `act_sf` with its own kernel (2 launches), and `cutedsl` adds - `bias` as a post-GEMM elementwise add (2 launches with bias, 1 without). -- `cutedsl` JIT-compiles on first use (seconds) and additionally validates - `alpha.numel() == 1` itself. -- The op's Python body is registered straight through - `torch.library.Library.define/impl` (trtllm's `fast_custom_op`), so - `torch.ops.trtllm.nvfp4_gemm` is an ordinary dispatcher op with a registered - fake kernel returning `[M, N]` in `output_dtype`. There is no autograd - kernel — inference only. -- Upstream (`_torch/modules/linear.py`, `NVFP4LinearMethod`) calls this op with - exactly these operands: `fp4_quantize`'d activations in the **swizzled** - layout, the raw `[N, K/2]` weight, `module.weight_scale` — which its loader - fills with `block_scale_interleave(weight_scale)` — and - `alpha = input_scale * weight_scale_2`. It flattens 3-D activations to - `[M, K]` before the call and reshapes after, and slices `N` back to - `out_features` when the weight was padded. -- Sibling ops with the same operand vocabulary exist: - `torch.ops.trtllm.nvfp4_gemm_cutlass` and - `torch.ops.trtllm.nvfp4_gemm_cublaslt` (single-backend entry points with - their own tuning caches), `torch.ops.trtllm.cuda_core_nvfp4_gemm` and - `torch.ops.trtllm.marlin_nvfp4_gemm` and - `torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell` (the raw per-backend - kernels; the first two take the *unswizzled* activation scales that this op's - `cuda_core`/`marlin` paths derive internally, while the CuTe DSL one takes the - swizzled buffer like this op), `torch.ops.trtllm.fp4_gemm` (the older typed - entry point, also covering - W4A8 MXFP4xMXFP8), `torch.ops.trtllm.fp4_gemm_trtllmgen` and - `torch.ops.trtllm.fp4_bmm`, and `torch.ops.trtllm.nvfp4_gemm_allreduce` - (this GEMM fused with a tensor-parallel allreduce). This entry covers the - unified `nvfp4_gemm` launch only. -- Behaviour under CUDA-graph capture was not exercised (unknown); note that - autotuning and any first-call JIT must complete before capture. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/nvfp4_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/nvfp4_gemm.py index 3e29958228c6..c452d8e7e544 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/nvfp4_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/nvfp4_gemm.py @@ -1,11 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""NVFP4 x NVFP4 dense GEMM in nn.Linear layout via the trtllm unified nvfp4_gemm op.""" +"""Block-scaled NVFP4 matmul via the trtllm nvfp4_gemm op.""" + +from __future__ import annotations import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from .._op import ULP, Arch, Cell, OpWrapper, assert_within_ulp + +_VEC = 16 + +#: e2m1 code -> magnitude. code = (exponent << 1) | mantissa, bit 3 is the sign. +_E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + def _pad_up(value: int, multiple: int) -> int: return -(-value // multiple) * multiple @@ -14,75 +23,200 @@ def _pad_up(value: int, multiple: int) -> int: def _swizzled_scale_numel(rows: int, k: int) -> int: """Bytes a 128x4 swizzled block-scale buffer holds for `rows` x `k`. - The swizzle addresses the padded rectangle, not the real one, so the - buffer is this size whatever `rows` and `k` are -- see nvfp4_gemm.md. + The swizzle addresses the padded rectangle, not the real one, so the buffer + is this size whatever `rows` and `k` are. """ - return _pad_up(rows, 128) * _pad_up(k // 16, 4) - - -def nvfp4_gemm( - act_fp4: torch.Tensor, - weight: torch.Tensor, - act_sf: torch.Tensor, - weight_scale: torch.Tensor, - alpha: torch.Tensor, - output_dtype: torch.dtype, - output_buffer_kind: int = 0, - allowed_backends: str = "cutlass,cublaslt,cuda_core", - group: list[int] | None = None, - bias: torch.Tensor | None = None, -) -> torch.Tensor: - """Return `alpha * (act @ weight.T) (+ bias)` over block-scaled NVFP4 operands, one op call.""" - # Every guard is pure metadata, and every guarded violation was observed on - # this machine to be read wrongly *without raising* by at least one backend - # the default `allowed_backends` string can select: the cutlass path checks - # contiguity and raises, cublaslt does not, and none of the three default - # backends rejects an oversized `alpha`. - assert act_fp4.is_contiguous() and weight.is_contiguous(), ( - "act_fp4 [M, K/2] and weight [N, K/2] must be contiguous; the cublaslt " - "backend ignores strides and returns wrong results" - ) - assert act_sf.is_contiguous() and weight_scale.is_contiguous(), ( - "act_sf and weight_scale must be contiguous; the cublaslt backend " - "ignores strides and returns wrong results" + return _pad_up(rows, 128) * _pad_up(k // _VEC, 4) + + +def _swizzle_index(rows: int, cols: int, device) -> torch.Tensor: + """[rows, cols] scale coordinates -> flat offsets in the swizzled buffer.""" + padded_cols = _pad_up(cols, 4) + r = torch.arange(rows, device=device).view(-1, 1) + c = torch.arange(cols, device=device).view(1, -1) + return ( + (c % 4) + + (c // 4) * (4 * 128) + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) ) - assert alpha.numel() == 1, ( - "alpha must hold exactly one element; extra elements are silently " - "ignored (this build has no per-token alpha)" + + +class Nvfp4Gemm(OpWrapper): + """`alpha * (act @ weight.T) (+ bias)` over block-scaled NVFP4 operands. + + Both operands are packed e2m1, two values to a byte, with one e4m3 scale per + 16 elements of K held in the 128x4 swizzled order. `alpha` folds the two + global scales the quantization divided out back into the product. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # deepseek-r1's dense and shared-expert MLPs are the only callers: hidden + # 7168 into 2 x intermediate, then back down. K and N here are the real + # widths; M is the token count. + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode, gate+up: [1, 7168] @ [4096, 7168].T", + spec=dict(m=1, k=7168, n=4096, bias=False), + ), + Cell( + why="decode, down: [1, 2048] @ [7168, 2048].T", + spec=dict(m=1, k=2048, n=7168, bias=False), + ), + Cell( + why="prefill chunk through the wider of the two", + spec=dict(m=2048, k=7168, n=4096, bias=False), + ), + Cell( + why="an M that is not a multiple of the swizzle's 128-row tile", + spec=dict(m=37, k=7168, n=4096, bias=False), + ), ) - # Contiguity above says nothing about length, and a short scale buffer is - # read past its end rather than rejected: the kernel indexes the *padded* - # rectangle. Both operands are checked because M and N pad independently. - # - # Only for 2-D operands, and only for *under*-length. Anything else about - # the shapes -- wrong rank, zero rows, K disagreement -- the op rejects - # itself and loudly, and preempting that here would swap its RuntimeError - # for an AssertionError that says less. Over-length is not the hazard - # either: the kernel never reads past what the swizzle addresses, and the - # contract already says the padding bytes may hold anything. - if act_fp4.dim() == 2 and weight.dim() == 2: - k = act_fp4.shape[1] * 2 # act_fp4 packs two 4-bit values per byte - floor_act_sf = _swizzled_scale_numel(act_fp4.shape[0], k) - floor_weight_scale = _swizzled_scale_numel(weight.shape[0], k) - assert act_sf.numel() >= floor_act_sf, ( - f"act_sf must hold at least pad_up(M,128) * pad_up(K/16,4) = " - f"{floor_act_sf} bytes for M={act_fp4.shape[0]}, K={k}; got " - f"{act_sf.numel()} -- the kernel would read past its end" + + note = """ + `bias`, `group` and `output_buffer_kind` exist on the op and no shipped + target passes them, so none is a cell axis. + + `allowed_backends` selects among cutlass, cublaslt and cuda_core. The cells + drive the default string, so which backend answers is the op's choice on the + day -- that is deliberate, since a target gets the same choice. + """ + + def raw_call( + self, + act_fp4: torch.Tensor, + weight: torch.Tensor, + act_sf: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + output_buffer_kind: int = 0, + allowed_backends: str = "cutlass,cublaslt,cuda_core", + group: list[int] | None = None, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return torch.ops.trtllm.nvfp4_gemm( + act_fp4, + weight, + act_sf, + weight_scale, + alpha, + output_dtype, + output_buffer_kind, + allowed_backends, + group, + bias, + ) + + def reference( + self, + act_fp4: torch.Tensor, + weight: torch.Tensor, + act_sf: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + output_buffer_kind: int = 0, + allowed_backends: str = "", + group: list[int] | None = None, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + """Dequantize both operands, then one fp32 product. + + Takes the packed form rather than the values it was built from, so the + reference covers the unpacking and the scale layout too -- a swizzle the + kernel and this disagreed about would otherwise never show up. + """ + a = self._dequantize(act_fp4, act_sf) + b = self._dequantize(weight, weight_scale) + out = alpha.float().reshape(()) * (a @ b.T) + if bias is not None: + out = out + bias.float() + return out.to(output_dtype) + + @staticmethod + def _dequantize(packed: torch.Tensor, swizzled_sf: torch.Tensor) -> torch.Tensor: + """[R, K/2] packed e2m1 plus a swizzled scale buffer -> [R, K] fp32.""" + rows, half = packed.shape + k = half * 2 + codes = torch.empty(rows, k, dtype=torch.uint8, device=packed.device) + codes[:, 0::2] = packed & 0xF + codes[:, 1::2] = packed >> 4 + table = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=packed.device) + value = table[(codes & 7).long()] + value = torch.where((codes & 8).bool(), -value, value) + + # Gather the scales back out of the swizzle -- the inverse of the + # scatter that put them there. + cols = k // _VEC + index = _swizzle_index(rows, cols, packed.device).reshape(-1) + scale = swizzled_sf.reshape(-1)[index].reshape(rows, cols) + scale = scale.view(torch.float8_e4m3fn).float() + return value * scale.repeat_interleave(_VEC, dim=-1) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Two fp4 operands accumulated over K, so the band is the operands'. + + The values are exact on both sides -- e2m1 codes and e4m3 scales decode + to the same numbers -- and the only difference is the order K is summed + in. Over a K of 7168 a few outputs cancel toward zero, where torch's + per-element rtol is meaningless. Scored in bf16 ulp because that is what + the output is asked to hold; 8 ulp element-wise and 4 of relative RMS, + the same band a two-GEMM chain gets, because the fp4 operands make each + product coarser than a bf16 one. + """ + assert_within_ulp(out, ref, element_ulp=8.0, rms_ulp=4.0, ulp=ULP[torch.bfloat16]) + + def is_valid( + self, + act_fp4: torch.Tensor, + weight: torch.Tensor, + act_sf: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, + *unused_args: object, + **unused_kwargs: object, + ) -> None: + # Every guard is pure metadata, and every guarded violation was observed + # to be read wrongly *without raising* by at least one backend the + # default `allowed_backends` string can select: cutlass checks + # contiguity and raises, cublaslt does not, and none of the three + # rejects an oversized `alpha`. + assert act_fp4.is_contiguous() and weight.is_contiguous(), ( + "act_fp4 [M, K/2] and weight [N, K/2] must be contiguous; the cublaslt " + "backend ignores strides and returns wrong results" ) - assert weight_scale.numel() >= floor_weight_scale, ( - f"weight_scale must hold at least pad_up(N,128) * pad_up(K/16,4) = " - f"{floor_weight_scale} bytes for N={weight.shape[0]}, K={k}; got " - f"{weight_scale.numel()} -- the kernel would read past its end" + assert act_sf.is_contiguous() and weight_scale.is_contiguous(), ( + "act_sf and weight_scale must be contiguous; the cublaslt backend " + "ignores strides and returns wrong results" ) - return torch.ops.trtllm.nvfp4_gemm( - act_fp4, - weight, - act_sf, - weight_scale, - alpha, - output_dtype, - output_buffer_kind, - allowed_backends, - group, - bias, - ) + assert alpha.numel() == 1, ( + "alpha must hold exactly one element; extra elements are silently " + "ignored (this build has no per-token alpha)" + ) + # Contiguity says nothing about length, and a short scale buffer is read + # past its end rather than rejected: the kernel indexes the *padded* + # rectangle. Both operands are checked because M and N pad + # independently. Only for 2-D operands and only for under-length -- + # wrong rank, zero rows and K disagreement the op rejects itself and + # loudly, and preempting that would swap its error for a worse one. + if act_fp4.dim() == 2 and weight.dim() == 2: + k = act_fp4.shape[1] * 2 # two 4-bit values per byte + floor_act = _swizzled_scale_numel(act_fp4.shape[0], k) + floor_weight = _swizzled_scale_numel(weight.shape[0], k) + assert act_sf.numel() >= floor_act, ( + f"act_sf must hold at least pad_up(M,128) * pad_up(K/16,4) = " + f"{floor_act} bytes for M={act_fp4.shape[0]}, K={k}; got " + f"{act_sf.numel()} -- the kernel would read past its end" + ) + assert weight_scale.numel() >= floor_weight, ( + f"weight_scale must hold at least pad_up(N,128) * pad_up(K/16,4) = " + f"{floor_weight} bytes for N={weight.shape[0]}, K={k}; got " + f"{weight_scale.numel()} -- the kernel would read past its end" + ) + + +nvfp4_gemm = Nvfp4Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml index 7382f31d99a1..03066412aaf0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml @@ -4,9 +4,11 @@ # unique identifier; the directory encodes the category. Synced manually after # catalog changes. # -# The catalog is the complete vocabulary of target modeling code: every call -# that creates or transforms a tensor must be a catalog entry. Only tensor -# metadata reads (.shape/.dtype/.device) and Python control flow live outside. +# The catalog is the vocabulary of target modeling code that is ours: every +# trtllm op a target calls must be an entry. Native torch calls, tensor +# metadata reads (.shape/.dtype/.device) and Python control flow live outside, +# because nothing here would be stating anything about them that torch does +# not already state better. # # AN ENTRY SPANS TWO TREES. The contract (.md) and the wrapper (.py) are here; # the GPU test is at @@ -22,11 +24,15 @@ # importantly they are not duplicates of those -- upstream covers the modules # that wrap these ops, these cover the op itself cell by cell. # -# torch/ entries are hand-added thin mirrors of torch callables — upstream -# owns their correctness, so they carry no contract doc, no test, and no -# arch/precision certification (absence = unconstrained by design; real dtype -# limits, e.g. no float8 arithmetic, are noted in the entry docstring). -# Entries wrapping trtllm ops carry all three. +# Native torch calls are not entries. This catalog used to mirror eleven of +# them -- torch.cat, torch.reshape, Tensor.copy_ and so on -- one file each, +# every file a single forwarding line. They carried no contract doc, no test +# and no certification, because upstream owns their correctness and there was +# nothing left for an entry to say. A target now calls torch directly, which +# is both shorter and the thing a reader already knows how to look up. +# +# An entry is therefore always a trtllm op: something whose behaviour is ours +# to state and ours to prove. # # ── RECEIPT STATUS: all 19 entries certified on sm_103 ──────────────────── # @@ -80,51 +86,6 @@ entries: # ─── torch (native callables, trtllm-free) ──────────────────── - - path: torch/embedding.py - impl: torch.nn.functional.embedding - summary: "Input-embedding lookup: token ids -> rows of the embedding weight" - - - path: torch/concat.py - impl: torch.cat - summary: "Concatenate a list of tensors along a dim" - - - path: torch/split.py - impl: torch.split - summary: "Split a tensor along a dim into equal chunks or explicit sizes" - - - path: torch/add.py - impl: torch.add - summary: "Elementwise add (residual); torch has no float8 arithmetic" - - - path: torch/reshape.py - impl: torch.reshape - summary: "Reshape (covers view: returns a view when layout allows)" - - - path: torch/empty.py - impl: torch.empty - summary: "Uninitialized tensor allocation (scratch/output buffers whose sizing the caller owns)" - - - path: torch/expand.py - impl: torch.Tensor.expand - summary: "Zero-copy broadcast view over singleton dims (stride-0; unsafe to write through)" - - - path: torch/transpose.py - impl: torch.transpose - summary: "Zero-copy swap of two dims (strided view; presents [tokens, heads, dim] to a head-batched gemm)" - - - path: torch/view_dtype.py - impl: torch.Tensor.view - summary: "Bitwise dtype reinterpretation of the same bytes (uint8 scale buffer -> float8_e4m3fn view); no conversion, no copy" - - - path: torch/copy_.py - impl: torch.Tensor.copy_ - summary: "In-place elementwise copy into a tensor or view (canonical slice fill; destroys dst contents)" - - - path: torch/pad.py - impl: torch.nn.functional.pad - summary: "Constant padding of trailing dims (widening an activation to a kernel's operand width)" - - # ─── norm ────────────────────────────────────────────────────── - path: norm/flashinfer_rmsnorm.py impl: torch.ops.trtllm.flashinfer_rmsnorm summary: "RMS normalization over the last dim with elementwise weight scaling" @@ -152,35 +113,11 @@ entries: summary: "NVFP4 x NVFP4 dense GEMM in nn.Linear layout with in-kernel block-scale dequantization: alpha * ([M,K/2] e2m1 + 128x4-swizzled e4m3 scales) @ ([N,K/2] e2m1 + swizzled scales)^T (+ per-column bias) -> fresh [M,N] bf16/fp16/fp32 buffer, fp32 accumulation, backend auto-selected (cutlass/cublaslt/cuda_core/cutedsl)" # ─── attention ───────────────────────────────────────────────── - - path: attention/fused_qk_norm_rope.py - impl: torch.ops.trtllm.fused_qk_norm_rope - summary: "In-place fused per-head QK RMS norm + RoPE (neox/interleaved, YaRN, interleaved mRoPE) on a packed bf16 QKV tensor; v heads untouched" - - - path: attention/mla_rope_generation.py - impl: torch.ops.trtllm.mla_rope_generation - summary: "MLA decode preprocessing: GPT-J RoPE of q_pe + latent [compressed_kv | rope(k_pe)] paged-KV-cache append + decode-FMHA scheduler-buffer fill; bf16 pool writes fused_q's tail, fp8-e4m3 pool writes quant_q_buffer and the two folded bmm scales instead" - - - path: attention/load_paged_kv_cache_for_mla.py - impl: torch.ops.trtllm.load_paged_kv_cache_for_mla - summary: "Gather each context sequence's full [past + new] MLA latent KV from the paged cache into two new contiguous tensors (compressed_kv [T,C], k_pe [T,R]) for KV-reuse context prefill; an fp8-e4m3 pool is dequantized to bf16/fp16/fp32 by kv_scale_quant_orig" - - - path: attention/mla_rope_append_paged_kv_assign_q.py - impl: torch.ops.trtllm.mla_rope_append_paged_kv_assign_q - summary: "MLA context-prefill preprocessing: in-place GPT-J RoPE of q_pe (in q) and k_pe (in latent_cache) at each new token's absolute position + [compressed_kv | rope(k_pe)] paged-latent-cache append; an fp8-e4m3 pool quantizes only the appended row, by kv_scale_orig_quant" - - path: attention/thop_attention.py impl: tensorrt_llm.bindings.internal.thop.attention summary: "Full attention core with fully explicit state (pybind binding, approved exception): paged KV-cache append + causal/padding-masked GQA FMHA over a caller-owned pool (bf16, or fp8-e4m3 with per-tensor kv scales) and explicit length/offset tensors, written into a caller buffer, on either context execution path (use_paged_context_fmha selects packed-QKV context FMHA, or paged-KV context FMHA so a context call may run over a cached prefix — KV-cache reuse and chunked prefill), with optional per-query-head attention sinks (one extra softmax-denominator logit, dropped from the output) and an optional per-call sliding window (attention_window_size keys ending at the query's absolute position; a pure mask — the append stays at absolute positions, so cyclic pool reuse is the caller's page mapping); MLA mode runs context prefill (in-kernel RoPE + latent append), no-append context over explicit K/V (latent_cache=None: cached-KV prefixes and chunked partial passes with softmax-stats output), and generation latent-MQA decode as separate calls over a paged latent pool (bf16 or fp8-e4m3)" # ─── moe ─────────────────────────────────────────────────────── - - path: moe/noaux_tc_op.py - impl: torch.ops.trtllm.noaux_tc_op - summary: "DeepSeek-V3 style MoE routing (noaux_tc): in-kernel sigmoid of the router logits + per-expert bias correction for selection only + optional group-limited top-k (n_group/topk_group; identity at n_group=1) + combine weights gathered from the unbiased sigmoid, renormalized and scaled by routed_scaling_factor, returning weights in the logits' dtype and int32 expert ids" - - - path: moe/fused_moe.py - impl: torch.ops.trtllm.fused_moe - summary: "Full MoE layer over pre-routed tokens: expert permutation + grouped FC1 GEMM over stacked [up; gate] expert weights + gated activation (SwiGLU with optional per-expert alpha/beta/limit, or GeGLU) + grouped FC2 GEMM + routing-weighted fp32 combine, into a fresh or caller-provided [tokens, hidden] buffer (bf16/fp16 unquantized path; expert-parallel slot range selected by ep_size/ep_rank)" - - path: moe/mxe4m3_mxe2m1_block_scale_moe_runner.py impl: torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner summary: "Full MXFP4-weight / MXFP8-activation MoE layer (trtllm-gen W4A8): optional top-k routing over router logits (softmax-then-topk or topk-then-softmax) or caller-supplied top-k ids/weights + grouped FC1 GEMM over pre-shuffled block-scaled [up; gate] expert weights + clamped gated activation + MXFP8 requantization of that activation + grouped FC2 GEMM + routing-weighted fp32 combine, into a fresh or caller-provided [tokens, valid_hidden_size] bf16 buffer (e4m3 hidden states + linear UE8M0 block scales in, expert-parallel window selected by local_expert_offset/local_num_experts)" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.md deleted file mode 100644 index a8e1f3efdb97..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.md +++ /dev/null @@ -1,659 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 15} ---- - -# fp4_block_scale_moe_runner - -**Wraps** `torch.ops.trtllm.fp4_block_scale_moe_runner` (one call). - -## Semantics - -One complete mixture-of-experts layer for **NVFP4 weights over NVFP4 -activations** (E2M1 elements with one E4M3 scale per 16 contiguous K elements -on both operands) — the trtllm-gen "block scale MoE" family's W4A4 member. In -a single call: optional top-k routing over router logits, expert permutation, -the grouped FC1 GEMM, the gated activation, **requantization of that -activation to NVFP4**, the grouped FC2 GEMM, and the routing-weighted combine -back to one row per token. - -Let `T = hidden_states.shape[0]`, `H` = the model hidden size, -`I = intermediate_size`, `K = top_k`, `E = local_num_experts`, -`off = local_expert_offset`. - -### The three scale scalars, and the two global scales behind them - -The kernel's two GEMMs consume raw E2M1 codes and E4M3 block scales; the -**per-tensor global scales** that a modelopt NVFP4 checkpoint carries are not -inside those operands, so they arrive as the three `[E]` fp32 tensors. Write - -- `g1` = the FC1 **activation** global scale (`hidden_states` was quantized - with it), -- `g2` = the FC2 **activation** global scale (the FC1 epilogue quantizes with - it), -- `gw1[e]`, `gw2[e]` = the per-expert weight global scales of the FC1 and FC2 - weights. - -Then, exactly: - -``` -output1_scale_gate_scalar[e] = 1 / (g1 * gw1[e]) # "alpha" of FC1 -output1_scale_scalar[e] = g2 / (g1 * gw1[e]) # that, times g2 -output2_scale_scalar[e] = 1 / (g2 * gw2[e]) # "alpha" of FC2 -``` - -A modelopt checkpoint stores **reciprocals**: its `input_scale` and -`weight_scale_2` scalars are `amax / (448*6)`, so `g1 = 1 / input_scale_fc1`, -`g2 = 1 / input_scale_fc2` (the `down_proj`'s `input_scale`), and - -``` -output1_scale_gate_scalar[e] = input_scale_fc1 * weight_scale_2_fc1[e] -output1_scale_scalar[e] = output1_scale_gate_scalar[e] / input_scale_fc2 -output2_scale_scalar[e] = input_scale_fc2 * weight_scale_2_fc2[e] -``` - -Getting the two `output1_*` scalars the wrong way round is finite and -plausible — measured 177 bf16 ulp off at `H = 256`, `I = 128` and only **48** -(40 ulp RMS) at the R1 shape `H = 7168`, `I = 2048`, in both cases with no -error. It is the *least* loud of the operand mistakes catalogued here, and it -gets quieter at the larger geometry, so a caller cannot expect it to announce -itself. - -### Per token `t` and slot `j < K` - -With `A[t]` = the fp32 value of `hidden_states` **after** dividing out `g1` -(i.e. `code * blockscale / g1`), and `W_up[e]`, `W_gate[e]` (`[I, H]`), -`W_down[e]` (`[H, I]`) the fp32 values of this rank's dequantized NVFP4 -weights (`code * blockscale`, with the weight global scale `gw` divided out — -it lives in the scalars above): - -``` -gid = topk_ids[t, j] # a GLOBAL expert id -skip this slot unless off <= gid < off + E -e = gid - off # index into this rank's weights - -up = A[t] @ W_up[e].T -gate = A[t] @ W_gate[e].T - -# act_type 0 (SwiGlu), the only value certified here: -act = up * gate * sigmoid(gate) # fp32 - -act = nvfp4_requantize(act) # see below — the W4A4 step -y = act @ W_down[e].T -out[t] += topk_weights[t, j] * y # fp32 accumulation -``` - -and `out` is stored as bf16. Note the sigmoid is applied to the **gate** half -and `up` is the plain linear one. - -**The intermediate is NVFP4.** FC2 is an NVFP4 x NVFP4 GEMM, so the FC1 -epilogue quantizes its post-activation output before FC2 reads it. Per token -row and per **16 consecutive intermediate columns** (natural column order, -blocks starting at column 0): - -``` -sf = e4m3_round_to_nearest_even(clamp(g2 * max|act| over the block / 6, max 448)) -act' = e2m1_round_to_nearest_even(act * g2 / sf) * sf / g2 -``` - -— i.e. exactly what `torch.ops.trtllm.fp4_quantize(act, g2, 16, ...)` would -emit, dequantized. `6` is e2m1's largest magnitude, `448` e4m3's largest -finite value; e2m1 rounding is ties-to-even-code (`0.25 -> 0`, `0.75 -> 1`, -`1.25 -> 1`, `1.75 -> 2`, `2.5 -> 2`, `3.5 -> 4`, `5.0 -> 4`), saturating at -`±6`. This entry's test reads the requantized values straight out of the -kernel (a down projection set to the identity, `output2_scale_scalar = 1`) and -matches the formula above **bit-exactly on all 16384 elements**; on the same -data an unrounded (exact) block scale reproduces 23.5% of the elements and the -unquantized activation 0.0%. Modelling this layer without the -intermediate requantization deviates 43 bf16 ulp element-wise / 23 ulp RMS -from the kernel at `H = 256`, `I = 128`, and 28 / 23 at the R1 shape -`H = 7168`, `I = 2048` — it is a real perturbation of the layer output, not a -rounding detail, and the RMS figure barely moves with geometry. An -intermediate element below `1/24` of its block's largest -magnitude rounds to zero (e2m1's smallest nonzero code is `0.5` against a -block top of `6`). - -**FC1 half order is `[up | gate]`.** Before the kernel's row interleave (§ -*Preconditions*) the first `I` rows of the FC1 operand are the up projection -(trtllm's `w3`, HF's `up_proj`) and the last `I` rows the gate projection -(trtllm's `w1`, HF's `gate_proj`). Swapping them produces a plausible-looking, -entirely different result; nothing detects it — this entry's test checks that -a gate/up-swapped reference lands 263 ulp away at `H = 256`, `I = 128` and -278 ulp away at the R1 shape `H = 7168`, `I = 2048`. - -**Fusion boundary.** Inside the call: routing (when driven by logits), -permutation, both GEMMs with on-the-fly NVFP4 dequantization, the gated -activation, the NVFP4 requantization between the GEMMs, and the weighted -combine. Outside, and the caller's job: the router GEMM that produces the -logits and the routing itself when pre-routing (e.g. -`torch.ops.trtllm.noaux_tc_op`), **quantizing the bf16 hidden states to -NVFP4** (`torch.ops.trtllm.fp4_quantize`, § *Preconditions* — there is no -padding step inside this op), all weight preprocessing, any shared/dense -expert branch, the TP all-reduce or EP gather of this call's output, and the -residual add. - -**Nothing is renormalized inside.** `topk_weights` is used exactly as given -(it need not sum to 1). `routed_scaling_factor` is accepted but had no effect -on the certified path. - -**Output.** `do_finalize=True` with `output=None` returns a **one-element -list** whose tensor is a fresh contiguous `[T, H]` **bf16** result. With -`output` given the result is written there (every element overwritten, nothing -outside its rows touched) and the list holds an **empty `[0]` bf16 tensor** — -the caller must read its own buffer. `do_finalize=False` returns three -tensors; see § *Signature*. No input is mutated; two identical calls are -bitwise equal. - -## Signature - -```python -def fp4_block_scale_moe_runner( - routing_logits: Optional[torch.Tensor], - routing_bias: Optional[torch.Tensor], - hidden_states: torch.Tensor, - hidden_states_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, - gemm1_bias: Optional[torch.Tensor], - gemm1_alpha: Optional[torch.Tensor], - gemm1_beta: Optional[torch.Tensor], - gemm1_clamp_limit: Optional[torch.Tensor], - gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, - gemm2_bias: Optional[torch.Tensor], - output1_scale_scalar: torch.Tensor, - output1_scale_gate_scalar: torch.Tensor, - output2_scale_scalar: torch.Tensor, - num_experts: int, - top_k: int, - n_group: Optional[int], - topk_group: Optional[int], - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - routed_scaling_factor: Optional[float], - routing_method_type: int, - do_finalize: bool, - act_type: int = 0, - topk_weights: Optional[torch.Tensor] = None, - topk_ids: Optional[torch.Tensor] = None, - output: Optional[torch.Tensor] = None, - tune_max_num_tokens: int = 8192, - use_dp: bool = False, -) -> list[torch.Tensor] -``` - -Unlike the MXFP4 members of this family the signature carries **no -`valid_hidden_size` / `valid_intermediate_size`**: this op has no padded-vs- -true size distinction at all (§ *Preconditions*), and it gains the three scale -scalars and `do_finalize`. - -**`gemm1_bias`, `gemm1_alpha`, `gemm1_beta`, `gemm1_clamp_limit` and -`gemm2_bias` accept `None`** even though the registered schema spells them -`Tensor`, not `Tensor?`. The schema is generated from the Python custom-op's -annotations; the torch dispatcher turns a `None` argument into an undefined -tensor, the C++ entry point takes `std::optional` for exactly -these five (visible in the exported `FP4BlockScaleMoeRunner::run_moe` symbol), -and it arrives there as an empty optional. Verified at runtime: all five at -`None` are bitwise identical to their neutral values (zero bias, `alpha = 1`, -`beta = 0`, an effectively infinite clamp limit), and each slot is live — a -non-neutral value changes the result. A DeepSeek-V3 style MoE, which has no -expert bias and no clamp, passes `None` for all five. - -### Certified arguments - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `routing_logits` | `None` — the only certified value; see *Notes* | — | — | — | -| `routing_bias` | `None`, or `[num_experts]` and inert | bf16 | contiguous | CUDA | -| `hidden_states` | `[T, H/2]` | **uint8** (2 E2M1 codes per byte, element `2i` in the low nibble) | contiguous | CUDA | -| `hidden_states_scale` | **1-D**, exactly `T * H/16` elements | **float8_e4m3fn** | contiguous, **linear** (row-major `[T, H/16]` flattened) — *not* swizzled | CUDA | -| `gemm1_weights` | `[E, 2*I, H/2]` | **uint8** | contiguous, pre-shuffled | CUDA | -| `gemm1_weights_scale` | `[E, 2*I, H/16]` | **float8_e4m3fn** | contiguous, pre-shuffled + 128x4 swizzled | CUDA | -| `gemm1_bias` / `gemm1_alpha` / `gemm1_beta` / `gemm1_clamp_limit` | `None` | — | — | — | -| `gemm2_weights` | `[E, H, I/2]` | **uint8** | contiguous, pre-shuffled | CUDA | -| `gemm2_weights_scale` | `[E, H, I/16]` | **float8_e4m3fn** | contiguous, pre-shuffled + 128x4 swizzled | CUDA | -| `gemm2_bias` | `None` | — | — | — | -| `output1_scale_scalar` | `[E]` (**local** experts) | **fp32** | contiguous | CUDA | -| `output1_scale_gate_scalar` | `[E]` | **fp32** | contiguous | CUDA | -| `output2_scale_scalar` | `[E]` | **fp32** | contiguous | CUDA | -| `num_experts` | scalar | int, the global routing space; `> top_k` | — | — | -| `top_k` | scalar | int, `0 < top_k < num_experts` | — | — | -| `n_group` / `topk_group` | `None`, `(1,1)`, or (with `routing_method_type = 2`) `(4,2)` / `(8,4)` — all inert | int | — | — | -| `intermediate_size` | scalar | int, **must equal `gemm1_weights.shape[1] // 2`** | — | — | -| `local_expert_offset` / `local_num_experts` | scalars | int, `E = local_num_experts >= 1` | — | — | -| `routed_scaling_factor` | `None`, `1.0`, `2.5` — all inert | float | — | — | -| `routing_method_type` | scalar | int: `0/1/2/4/5/6` all inert on this path | — | — | -| `do_finalize` | scalar | bool | — | — | -| `act_type` | scalar | int: **0** (SwiGlu) only | — | — | -| `topk_weights` | `[T, top_k]` | **bf16** | contiguous | CUDA | -| `topk_ids` | `[T, top_k]` | **int32** | contiguous | CUDA | -| `output` | `[T, H]` or `None` | bf16 | contiguous | CUDA | -| return, `do_finalize=True` | `[[T, H]]`, or `[[0]]` when `output` is given | bf16 | contiguous, newly allocated | same as `hidden_states` | -| return, `do_finalize=False` | see below | | | | - -`torch.ops.trtllm.noaux_tc_op` fed **bf16** router logits produces a directly -consumable `(topk_weights, topk_ids)` pair — bf16 weights and int32 ids, in -that order. Catalog membership of that op is `index.yaml`'s fact alone. - -### `do_finalize=False` - -Returns three tensors: - -0. `[P, H]` bf16 — one row per (token, slot) expert output **before** the - routing-weighted combine, in the kernel's internal permuted/padded order. - `P` depends on `T`, `top_k`, `num_experts` and an internal tile size; it is - always `>= T * top_k`. -1. `[T, top_k]` bf16 — **not written on the pre-routed path.** Observed to - hold uninitialized memory (values of order `1e24`, and NaN once combined). - It is the slot for the routing stage's own weights; do not read it. -2. `[T, top_k]` int32 — the expanded-index -> permuted-row map: token `t`'s - slot `j` expert output is row `out2[t, j]` of output 0. - -Recombining with the caller's own weights, -`out[t] = sum_j topk_weights[t, j] * out0[out2[t, j]]` accumulated in fp32 and -stored as bf16, reproduces the `do_finalize=True` result **bitwise** (verified -here). `output=` is rejected in this mode -(`out_tensor is only supported when do_finalize=true`). **A plain, non -expert-parallel target wants `do_finalize=True`**: the mode exists so an -all-to-all deployment can move the per-slot rows before combining them. - -## Metadata consumed - -None. The op reads no attention metadata, no KV cache, no registered layer and -no module state — every tensor and every scalar it uses is an argument. - -Two process-global caches sit behind it, neither of which changed the result -on any path measured here: - -- an **autotuner profiling cache** that picks the two GEMM tactics, keyed by - `(top_k, intermediate_size, local_num_experts, act_type)` together with the - call's input shapes and a token bucket (the row count rounded down to a - power of two, capped at 8192). Two consequences of that key were measured - here, by counting the profiling sweeps the tuner performed rather than - reading the key: **`local_num_experts` separates entries** — going from 256 - to 64 with everything else fixed forced a fresh sweep, so a cache warmed at - one expert-parallel window size says nothing about another — while - **`local_expert_offset` does not**: after warming the offset-0 64-wide - window, the offset-64 / 128 / 192 windows profiled *nothing*, reusing its - entries. Warming one window of a split therefore warms all four. - - **Cold is the state every call compared against a reference here ran in** — - the tuner then returns the fallback tactic. A serving engine is *not* cold: - the PyTorch runtime profiles this op during warm-up whenever - `enable_autotuner` is set (its default), so a deployed call runs on a tuned - tactic instead. Two direct cold-vs-warm comparisons were made, both by - warming with `tensorrt_llm._torch.autotuner.autotune()` and re-running the - identical operands. At `H = 2560`, `I = 1536`, top-6, `T = 8192`: **bitwise - identical** for `local_num_experts = 72` and all four 18-wide windows. At - the **R1 geometry** (`H = 7168`, `I = 2048`, top-8) the comparison is part - of this entry's test and is broader: warming wrote **28 cache entries** (14 - token buckets x the two `unique_id`s `(8, 2048, 256, 0)` and - `(8, 2048, 64, 0)`), every one of them a **non-fallback** tactic, and all - 85 (expert layout, token count) results — `local_num_experts = 256` plus - the four 64-wide windows, over the full certified `T` column — were - **bitwise unchanged**. *Which* tactics those are is not reproducible: the - tuner selects by measured time, so it is the one figure this test prints - that varies run to run — four runs of the comparison on the same device - recorded between 12 and 15 distinct tactics across the same 28 entries. - The bitwise result therefore held across four different tactic assignments, - which is stronger than one — but it is still a finite sample of tuner picks - on one machine, not a guarantee across every tactic the tuner might select. -- a cache of C++ `FP4BlockScaleMoERunner` objects keyed by `act_type`, each - owning a GPU workspace allocated on first use. - -## Preconditions - -### Sizes - -- **`H` must be a multiple of 256.** With the prepared block-scale layout - below, a hidden size that is only a multiple of 128 makes the kernel read - the FC1 weight scales for the wrong blocks: no error, 300-3400 bf16 ulp - wrong. Measured wrong at `H` = 384, 640 and 896; correct at 256, 512, 768, - 2560 and 7168. The wrapper asserts it. (`H` enters the FC1 GEMM's K axis; - the FC2 K axis, `I`, carries no such rule — every multiple of 64 worked.) -- `I` must be a multiple of **64**, so that `2*I` is a multiple of 128 and - `I/16` a multiple of 4 — the alignment the 128x4 scale swizzle needs on - both operands. Certified `I`: 64, 128, 192, 256, 512, 1536, 2048. -- There is **no padded-vs-true size distinction**. `gemm1_weights.shape[-1] * - 2`, `gemm2_weights.shape[1]` and the width of the returned output are all - the same `H`; a mismatch is rejected (`gemm2_weights_scale has incorrect dim - 1`). `intermediate_size` must equal `gemm1_weights.shape[1] // 2`; anything - else is rejected (`gemm2_weights_scale has incorrect dim 2`, or `No valid - config found for the given problem shape`). -- `0 < top_k < num_experts`. `top_k = 0` is rejected (`only supports - top_k<=32 && top_k>0`), `num_experts == top_k` too (`num_experts must be - greater than top_k`). Certified `top_k`: 1, 2, 4, 6, 8. -- `num_experts` is only the routing space; it need not equal - `local_num_experts`. Certified `num_experts`: 2, 3, 4, 8, 16, 72, 256; - certified `local_num_experts`: 2, 3, 4, 8, 16, 18, 64, 72, 256. Certified - `(local_expert_offset, local_num_experts)` pairs: `(0, E)` for each of - those `E`; `(4, 4)` of `num_experts = 8`; the four-way expert-parallel - split of 72 — `(0, 18)`, `(18, 18)`, `(36, 18)`, `(54, 18)`; and the - four-way split of 256 — `(0, 64)`, `(64, 64)`, `(128, 64)`, `(192, 64)` — - see § *Routing and expert ids*. -- `T >= 1`. Certified `T`: 1, 2, 3, 5, 7, 8, 16, 24, 32, 33, 40, 64, 128, - 256, 1024, 4096, 8192. `T` is the **row count of this call** — - `hidden_states.shape[0]`. It is not `tune_max_num_tokens`, which is an - autotuner bucket hint whose default is also 8192 and which was measured - inert (§ *Notes*); a call passing 8192 rows and a call passing 8 rows with - `tune_max_num_tokens = 8192` are different things, and only the first - certifies this row count. 8192 is trtllm's default `max_num_tokens` - (`llm_args.py`), so it is the widest prefill a stock engine hands this op - in one unchunked call, and it is the top of this list: **a caller whose own - token count exceeds 8192 must chunk into calls this column covers.** The - whole column is certified at the **DeepSeek-R1 routed geometry** - (`H = 7168`, `I = 2048`, `num_experts = 256`, `top_k = 8`, `act_type = 0`, - `do_finalize = True`, pre-routed), for the full 256-expert stack *and* - independently for each of the four 64-wide expert-parallel windows. The - DeepSeek-V3-Lite routed geometry (`H = 2560`, `I = 1536`, top-6) is - certified over a subset of it topping out at the same 8192 rows, for - `local_num_experts = 72` at offset 0 and for the four 18-wide windows — see - § *Routing and expert ids*. Other geometries in the *Certified* lists above - cover only the smaller counts. - -### Activation preparation — NVFP4, linear scales - -`hidden_states` / `hidden_states_scale` are exactly the two outputs of - -```python -# global sf_vec ue8m0 swizzled -data, sf = torch.ops.trtllm.fp4_quantize(x, g1, 16, False, False) -``` - -on the bf16 `[T, H]` hidden states, with `sf` re-viewed as `float8_e4m3fn`. -Concretely, whatever produces them: - -- `hidden_states` is 2-D **uint8** (`hidden_states must be byte` — a - `float8_e4m3fn` view of the same bytes is rejected), contiguous, CUDA, of - width exactly `gemm1_weights.shape[-1]`. -- `hidden_states_scale` is **1-D** (`hidden_states_scale must be 1D`), - **float8_e4m3fn** (`must be fp8` — a uint8 view is rejected), contiguous, - CUDA, and holds exactly `T * H/16` elements (`hidden_states_scale has - incorrect size` otherwise). Element `t * (H/16) + b` is the e4m3 scale of - columns `[16b, 16b+16)` of token `t`. -- The **128x4 swizzled** scale order is a silent wrong answer whenever it - happens to have the same element count (`T % 128 == 0`, `H/16 % 4 == 0` — - true at any graph-friendly batch size) — see *Notes*. Pass - `is_sf_swizzled_layout=False`. -- This op never pads: `H` is used as given, so the quantizer is called on the - full `[T, H]` hidden with no `F.pad` before it. - -### Weight preparation — the caller owns all of it - -Start from the checkpoint's per-expert NVFP4 tensors, laid out **row-major -over the output dim**, K packed two E2M1 codes per byte with the **low nibble -holding the even K index**, and one E4M3 byte per 16 consecutive K elements: - -``` -up (w3 / up_proj) packed [E, I, H/2] uint8 scale [E, I, H/16] e4m3 -gate (w1 / gate_proj) packed [E, I, H/2] uint8 scale [E, I, H/16] e4m3 -down (w2 / down_proj) packed [E, H, I/2] uint8 scale [E, H, I/16] e4m3 -``` - -Then, **per expert**: - -1. **FC1 concat.** Stack the two halves on the row axis as `[up ; gate]`, - giving `[2*I, ...]` for both the packed bytes and the scale bytes. -2. **FC1 row permute** = interleave, then block shuffle: - - interleave: destination row `2i` takes the up half's row `i`, - destination row `2i+1` takes the gate half's row `i`; - - block shuffle (also applied to FC2, which skips the interleave): - within each aligned block of 32 rows, source row `4u + v` - (`0 <= u < 8`, `0 <= v < 4`) moves to destination row `8v + u`. - - Apply the **same** permutation to the weight bytes and the scale bytes, so - row `i` of both still describes the same output channel. Skipping it is a - silent wrong answer (measured 508 ulp at `H = 256`, `I = 128`; 432 at - `H = 7168`, `I = 2048`). -3. **Scale swizzle.** After the row permute, each expert's scale matrix - `[M, C]` (`M % 128 == 0`, `C % 4 == 0`) is rewritten into the trtllm-gen - 128x4 layout: the byte at `(m, c)` moves to flat offset - - ``` - (m // 128) * 512 * (C // 4) + (c // 4) * 512 - + (m % 32) * 16 + ((m % 128) // 32) * 4 + (c % 4) - ``` - - and the flat result is handed to the kernel with its nominal `[E, M, C]` - shape. Weight bytes are **not** swizzled — only scales. Skipping it is a - silent wrong answer (measured 408 ulp at `H = 256`, `I = 128`; 357 at - `H = 7168`, `I = 2048`). - -**No padding is needed anywhere at either certified routed geometry**, and the -swizzle's alignment must be checked per geometry rather than assumed — it is a -property of `H` and `I`, not of the family: - -| geometry | FC1 weights + scales | FC2 weights + scales | `2I % 128`, `H/16 % 4`, `H % 128`, `I/16 % 4` | -|---|---|---|---| -| `H = 2560`, `I = 1536` | `[E, 3072, 1280]` + `[E, 3072, 160]` | `[E, 2560, 768]` + `[E, 2560, 96]` | 0, 0, 0, 0 | -| `H = 7168`, `I = 2048` | `[E, 4096, 3584]` + `[E, 4096, 448]` | `[E, 7168, 1024]` + `[E, 7168, 128]` | 0, 0, 0, 0 | - -At the R1 shapes `torch.ops.trtllm.block_scale_interleave` was checked here to -return exactly `E*M*C` bytes for an `[E, M, C]` scale tensor — measured at -`4096 x 448` and at `7168 x 128` — so the swizzled buffer reshapes straight -back to its nominal `[E, M, C]` shape with nothing appended, and the operand -handed to the kernel is the same size as the one the checkpoint carries. - -The two trtllm helpers `torch.ops.trtllm.shuffle_matrix(x, perm)` (a plain row -gather, `out[i] = x[perm[i]]`) and `torch.ops.trtllm.block_scale_interleave(x)` -(the 128x4 swizzle over a `[E, M, C]` uint8 tensor, returning a flat buffer of -`E * pad_up(M, 128) * pad_up(C, 4)` bytes) produce byte-identical results to -steps 2 and 3; this entry's test asserts that equivalence. - -### Layout and dtypes - -- **Every tensor argument must be contiguous.** A strided view is accepted - silently and read as if dense — see *Notes*. The wrapper asserts this. -- `topk_ids` is `[T, top_k]` **int32** (`topk_ids must be int` — int64 is - rejected) and `topk_weights` is `[T, top_k]` **bf16** (`topk_weights must be - bfloat16` — fp32 and fp16 are both rejected). The two must be given together - (`routing_logits or (topk_ids and topk_weights) must be provided`). -- Both weight-scale tensors must be **float8_e4m3fn** (a uint8 view is - rejected with `must be fp8`). -- The three scale scalars must be **fp32** (`must be float`) with exactly - **`local_num_experts`** elements — not `num_experts` (`has incorrect dim - 0`). -- `output`, when given, must be a contiguous CUDA bf16 tensor of shape exactly - `[T, H]` (`out_tensor must be bfloat16`, `out_tensor dim0 must match - num_tokens`). A leading row-slice of a taller contiguous buffer is valid — - rows past `T` are left bitwise untouched. - -### Routing and expert ids - -- Expert ids outside `[local_expert_offset, local_expert_offset + - local_num_experts)` — including negative ids and ids `>= num_experts` — are - silently dropped; the token's other slots still combine normally. Verified - two ways: a `local_expert_offset = 4`, `local_num_experts = 4` window of - `num_experts = 8` against a reference that masks the same slots, and ids of - `-1`, `num_experts` and `num_experts + 5` all giving the same output as a - dropped slot. -- **The window is a range test plus an index shift, not a validated - partition.** Local weight slot `i` answers for global id - `local_expert_offset + i`; nothing checks the window against - `num_experts`. `local_expert_offset + local_num_experts > num_experts` is - **accepted** — measured at `(60, 18)` of `num_experts = 72`, where the - result matched a reference mapping ids 60..71 onto local slots 0..11 and - the surplus six local slots were simply never addressed — and a window - entirely past the routing space (`local_expert_offset = 72`) returns an - all-zero output. Making the windows tile the routing space is the caller's - arithmetic, not something this op enforces. -- **The four-way expert-parallel split of 72 experts is certified**: - `num_experts = 72` on every call (the routing space is not sharded), - `top_k = 6`, `local_num_experts = 18`, `local_expert_offset` 0 / 18 / 36 / - 54, `H = 2560`, `I = 1536`, `act_type = 0`, `do_finalize = True`, - pre-routed, at T = 1, 8, 256, 1024, 4096 and 8192. Per window — with - `gemm1_weights`, `gemm1_weights_scale`, `gemm2_weights`, - `gemm2_weights_scale` and all three scale scalars sliced to - `[off, off+18)` along the expert axis and made contiguous — the result - matches a torch reference that masks every slot outside the window (worst - 9.97 ulp element-wise, 0.74 ulp RMS; at T = 8192 alone, 9.97 and 0.72), - and every token with **no** slot in the window comes back **bitwise - zero** — at T = 8192 that is ~1300 rows per window (1265-1406 measured - across the four windows: top-6 of 72 lands nothing in a given 18-wide - window for about a sixth of the batch). The four bf16 outputs summed - reproduce the 72-expert result under the same gate (worst 4.86 ulp - element-wise, 0.88 ulp RMS against the fp32 reference; at T = 8192 alone, - 4.86 and 0.84); they are **not** bitwise equal to a single 72-expert - call — measured up to 2.0 ulp element-wise apart at every certified token - count, 8192 included, because each window rounds its own partial to bf16 - before the add. Leaving any one window out of that sum lands >= 124 ulp - RMS away (>= 128 at T = 8192). When every routed slot of a batch happens - to fall inside one window, that window's call **is** bitwise equal to the - 72-expert call and the other three return exactly zero. -- **The four-way expert-parallel split of 256 experts is certified**, the - DeepSeek-R1 routed layout: `num_experts = 256` on every call (the routing - space is not sharded), `top_k = 8`, `local_num_experts = 64`, - `local_expert_offset` 0 / 64 / 128 / 192, `H = 7168`, `I = 2048`, - `act_type = 0`, `do_finalize = True`, pre-routed, over the **whole** - certified `T` column (1 through 8192, § *Sizes*). Sliced the same way — all - four weight/scale operands and all three scale scalars restricted to - `[off, off+64)` along the expert axis and made contiguous — each window - matches a torch reference masking every slot outside it (worst 9.62 ulp - element-wise, 0.74 ulp RMS; at T = 8192 alone, 6.35 and 0.73), and every - token with **no** slot in the window comes back **bitwise zero** — at - T = 8192 that is 764-804 rows per window, i.e. top-8 of 256 lands nothing - in a given 64-wide window for ~10% of the batch (`(3/4)^8`). The four bf16 - outputs summed reproduce the 256-expert result under the same gate (worst - 5.02 ulp element-wise, 0.85 ulp RMS against the fp32 reference; identical - at T = 8192 alone); as with the 18-wide split they are **not** bitwise - equal to a single 256-expert call — up to 2.0 ulp element-wise apart at - every certified token count, because each window rounds its own partial to - bf16 before the add. Leaving any one window out of that sum lands >= 125 - ulp RMS away (>= 128 at T = 8192); a window left at - `local_expert_offset = 0`, a window fed the neighbouring window's weights, - and all four windows issued at offset 0 land 310-363 ulp RMS away. - When every routed slot of a batch falls inside one window, that - window's call **is** bitwise equal to the 256-expert call and the other - three return exactly zero. -- **A repeated expert id inside one token's row is not reliably - deduplicated.** With one geometry it contributed once (the first slot's - weight) at 16 tokens and *twice* at 24 tokens. Supply distinct ids per row; - every routing op in this build does. -- `n_group > 1` requires `routing_method_type = 2` (`Routing kernel with - groups implies DeepSeekV3 routing method`), but on the pre-routed path the - whole grouped-routing configuration is inert: `(n_group, topk_group)` of - `(1,1)`, `(4,2)` and `(8,4)` under `routing_method_type = 2` are all bitwise - identical to `None/None` under `routing_method_type = 1`. -- `routing_method_type = 3` (Llama4) is rejected for `top_k > 1` (`Current - routing kernel (no groups, Llama4) only supports top_k=1`). - -A caller violating none of the above gets the result described under -*Semantics*, to within the bound in *Notes*. - -## Notes - -- **Numerics.** Against a native-torch reference that consumes bit-identical - operands (exact NVFP4 weight and activation dequantization, fp32 GEMM - accumulation, the NVFP4 requantization of the FC1 output modelled exactly, - fp32 combine), the kernel's worst element-wise deviation over every - configuration in this entry's test is **9.97 ulp of the token row's largest - magnitude** (bf16 ulp = 2^-8; worst case: an 18-wide expert-parallel window - of 72 experts, 8192 tokens, `H = 2560`, `I = 1536`), against the test's - 16-ulp element gate; a full-window call sits at 4.75 (3.20 at T = 8192). - Its worst relative RMS deviation over the whole test file is - **0.88 ulp** (those four windows summed), against a 2-ulp aggregate gate. - The **R1 routed geometry** (`H = 7168`, `I = 2048`, 256 experts, top-8) - sits just inside both, over the whole certified token column: 9.62 elt / - 0.74 RMS for a 64-wide window, 4.68 / 0.70 for the 256-expert stack, 5.02 / - 0.85 for the four windows summed. - The element-wise outliers are single intermediate values landing on opposite - sides of an e2m1 rounding boundary — an e2m1 step is coarse (2 mantissa - bits), so one flipped element moves an output row by a visible fraction. - Two things inflate the element figure and not the RMS one: it is an extreme - order statistic over every element compared, and it is normalized by the - reference row's own magnitude, so a window carrying a quarter of the routed - slots divides a similar absolute error by a smaller row. At `H = 2560`, - re-drawn over 144 independent (seed, window) comparisons at T = 1024 / 4096 - it reached **12.0 ulp** while the RMS figure stayed at 0.86; re-drawn again - over 16 fresh seeds at **T = 8192** (16 full-stack + 64 window comparisons) - it reached **16.33 ulp** for a window (p90 9.82, mean 7.15) and 10.41 for - the full stack, while the RMS figure was unchanged from T = 4096 (window - 0.74, full 0.70, four windows summed 0.85). So the element metric's upper - tail *crosses* the 16-ulp gate at the top of the certified token range — it - is a max over `T * H` elements and therefore grows with `T` at constant - accuracy, which is a property of the statistic, not of the kernel. At - `H = 7168` / `I = 2048` the same re-draw (12 fresh weight + activation - seeds, 12 full-stack + 48 window comparisons per token count) has a - **lower** tail: max **12.74** for a window at T = 8192 (p90 8.40, mean - 6.73) and 12.15 at T = 4096, 7.11 / 6.42 for the 256-expert stack, with the - RMS figures flat at 0.74 (window) / 0.70 (full) / 0.86 (sum) at both token - counts — so the element gate keeps ~1.3x headroom on the *distribution* - there, not ~1.0x. The shipped test is seeded and deterministic — two - independent processes on two devices printed byte-identical numbers — and - its T = 8192 draws land at 9.97 / 3.20 (`H = 2560`) and 6.35 / 4.11 - (`H = 7168`). **The RMS figure is the scale-free one and the one a caller - must size its own tolerance from**: it is identical at 4096 and 8192 at both - geometries, and every wrong variant this entry's test constructs sits - 22.5-370 ulp RMS away — 40-370 for every layout, scalar-role or - expert-window mistake (the 40 is an `output1_*` scalar-role swap at the R1 - shape; the same swap is 124 at `H = 256`), and ~23 for the one variant that - is not a mistake in the operands at all (a reference that skips the - FC1-output requantization). -- **Only the pre-routed entry point is certified here.** `routing_logits` was - left at `None` on every certified call and the routing was done outside (see - `torch.ops.trtllm.noaux_tc_op` for the DeepSeek-V3 formula). The in-kernel - routing path — `routing_logits` given, `topk_ids`/`topk_weights` `None`, - `routing_method_type` selecting softmax / renormalize / DeepSeekV3 — is - **not certified by this entry**; nor is `routing_bias`, which was only - observed to be inert on the pre-routed path. -- **A swizzled activation-scale buffer is a silent wrong answer.** - `fp4_quantize(x, g, 16, False, True)` returns the same data bytes and a - scale buffer of `pad_up(T,128) * pad_up(H/16,4)` bytes in 128x4 order. - Whenever that count coincides with the linear one (`T` a multiple of 128 — - exactly the CUDA-graph batch sizes) the size check passes and the kernel - reads scales for the wrong blocks: measured 276 ulp element-wise off at - `T = 128`, `H = 256`. Nothing in the metadata distinguishes the two layouts, - so no guard can catch it — pass `is_sf_swizzled_layout=False`, in a named - constant. -- **Non-contiguous tensors are a silent wrong answer.** The kernel takes raw - data pointers and assumes a dense row-major layout. A strided view of - `hidden_states`, `topk_weights` or any weight/scale operand is accepted - without complaint and reads the wrong elements. The wrapper asserts - contiguity on every tensor argument. -- **`act_type` other than 0 is not certified.** `1` (Relu2) and `2` (Silu) — - the two non-gated activations, which take an `[E, I, H/2]` FC1 operand - rather than `[E, 2*I, H/2]` — run without error when handed a gated-shaped - operand and return something different; so does the undefined value `3`. - Only `0` (SwiGlu) is characterized here. -- **Inert on the certified path** (bitwise identical results): - `routing_method_type` `0/1/2/4/5/6`; `n_group`/`topk_group` as above; - `routed_scaling_factor` `None`/`1.0`/`2.5`; `tune_max_num_tokens` `8192` - and `128` (an autotuner bucket cap); `use_dp` `False` and `True` (an - autotuner token-bucket deflation hint); a zero `routing_bias`. -- **Multi-rank tensor parallelism is not exercised.** `local_expert_offset` / - `local_num_experts` (expert parallelism) **are** certified, including the - four-way splits of 72 and of 256 experts that a 4-rank EP deployment uses — - but every certified call ran in **one process on one device**, the four - windows issued in sequence. This op reads no rank state, holds no - communicator and - launches no collective, so a window is an argument pair rather than a - multi-rank behaviour; summing the per-rank outputs is the caller's - all-reduce and is outside this entry. TP, which instead splits `I` across - ranks and reduces this call's output afterwards, is not certified. -- sm_100 only. The receipt covers sm_100 (B200), the only arch available here; - the installed build carries the assertion `Only SM100f is supported by FP4 - block scale MOE` in `libth_common.so`, so other Blackwell variants are - expected to raise rather than compute something wrong. -- The registered fake (meta) function disagrees with the kernel on **both** - dimensions of the `do_finalize=False` first output: it returns - `[1152, 128]` where the kernel returned `[256, 256]` (`T = 16`, - `top_k = 2`, `num_experts = 8`, `H = 256`). The rows come from a fixed - internal tile of 128; the columns are the **packed** width, because the - helper unpacks only under - `isinstance(hidden_states, Fp4QuantizedTensor)` and the op's schema hands - it a plain `Tensor`, so that branch can never be taken. The - `do_finalize=True` fake path doubles unconditionally and is correct. Eager - callers are unaffected; this is an upstream defect in this build, not - behaviour to rely on. -- Sibling ops in this build address the same job with other operand dtypes: - `torch.ops.trtllm.bf16_mxe2m1_block_scale_moe_runner` (bf16 activations, - MXFP4 weights), `torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner` - (MXFP8 activations, MXFP4 weights), - `torch.ops.trtllm.e4m3_mxe2m1_block_scale_moe_runner`, - `torch.ops.trtllm.fp8_block_scale_moe_runner` and - `torch.ops.trtllm.fp8_fp4_block_scale_moe_runner` (fp8 activations, NVFP4 - weights). The MXFP4 members take an `[E, 2*I, H/32]` **uint8** UE8M0 scale - and carry `valid_hidden_size` / `valid_intermediate_size` parameters this op - does not have, so their prepared weights are **not** interchangeable with - this one's. Catalog membership is `index.yaml`'s fact alone. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py index b2ec6994853f..e44abde79c18 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py @@ -4,126 +4,534 @@ top-k + grouped FC1 GEMM + gated activation + NVFP4 requantization + grouped FC2 GEMM + routing-weighted combine, in one trtllm-gen call.""" -from typing import Optional +from __future__ import annotations + +from typing import Optional, Tuple import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from .._op import Arch, Cell, OpWrapper -def fp4_block_scale_moe_runner( - routing_logits: Optional[torch.Tensor], - routing_bias: Optional[torch.Tensor], - hidden_states: torch.Tensor, - hidden_states_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, - gemm1_bias: Optional[torch.Tensor], - gemm1_alpha: Optional[torch.Tensor], - gemm1_beta: Optional[torch.Tensor], - gemm1_clamp_limit: Optional[torch.Tensor], - gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, - gemm2_bias: Optional[torch.Tensor], - output1_scale_scalar: torch.Tensor, - output1_scale_gate_scalar: torch.Tensor, - output2_scale_scalar: torch.Tensor, - num_experts: int, - top_k: int, - n_group: Optional[int], - topk_group: Optional[int], - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - routed_scaling_factor: Optional[float], - routing_method_type: int, - do_finalize: bool, - act_type: int = 0, - topk_weights: Optional[torch.Tensor] = None, - topk_ids: Optional[torch.Tensor] = None, - output: Optional[torch.Tensor] = None, - tune_max_num_tokens: int = 8192, - use_dp: bool = False, -) -> list[torch.Tensor]: - """Run one NVFP4-weight MoE layer over NVFP4 (e2m1 + e4m3 block scale) activations. - - With `do_finalize=True` returns a one-element list holding the combined - `[num_tokens, hidden]` bf16 result — or an empty `[0]` tensor when `output` - was given, the result having been written there. With `do_finalize=False` - returns three tensors (per-slot expert rows, an unwritten scale buffer, and - the expanded-index -> permuted-row map). - """ - # Pure-metadata guard: the kernel takes raw data pointers and assumes a - # dense row-major layout for every tensor. A strided view is accepted - # without complaint and silently reads (or writes) the wrong elements — - # observed on this machine for hidden_states, topk_weights and the weight - # operands. - for name, tensor in ( - ("routing_logits", routing_logits), - ("routing_bias", routing_bias), - ("hidden_states", hidden_states), - ("hidden_states_scale", hidden_states_scale), - ("gemm1_weights", gemm1_weights), - ("gemm1_weights_scale", gemm1_weights_scale), - ("gemm1_bias", gemm1_bias), - ("gemm1_alpha", gemm1_alpha), - ("gemm1_beta", gemm1_beta), - ("gemm1_clamp_limit", gemm1_clamp_limit), - ("gemm2_weights", gemm2_weights), - ("gemm2_weights_scale", gemm2_weights_scale), - ("gemm2_bias", gemm2_bias), - ("output1_scale_scalar", output1_scale_scalar), - ("output1_scale_gate_scalar", output1_scale_gate_scalar), - ("output2_scale_scalar", output2_scale_scalar), - ("topk_weights", topk_weights), - ("topk_ids", topk_ids), - ("output", output), - ): - if tensor is not None: - assert tensor.is_contiguous(), ( - f"{name} must be contiguous; a strided view is read as if dense " - "and silently produces wrong results" - ) - # Pure-metadata guard: with the block-scale layout this entry contracts - # (row shuffle + 128x4 scale swizzle), a hidden size that is not a multiple - # of 256 makes the kernel read the FC1 weight scales for the wrong blocks. - # Observed on this machine at hidden 384 / 640 / 896: no error, 300-3400 - # bf16 ulp wrong. - assert gemm1_weights.shape[-1] * 2 % 256 == 0, ( - "hidden size (gemm1_weights.shape[-1] * 2) must be a multiple of 256; " - "the prepared block-scale layout is silently misread otherwise" +#: NVFP4 block size: one e4m3 scale per 16 elements along K. +SF_VEC_SIZE = 16 + +#: Largest e2m1 magnitude, and the largest finite e4m3 magnitude. Their product +#: is the global scale a checkpoint's `input_scale` encodes. +E2M1_MAX = 6.0 +E4M3_MAX = 448.0 + +#: `act_type` for SwiGlu, the only activation this entry certifies. +ACT_TYPE_SWIGLU = 0 + +#: The widest unchunked call a stock engine issues: trtllm's default +#: `max_num_tokens`. A caller whose own token count exceeds it must chunk. +MAX_CERTIFIED_TOKENS = 8192 + +#: The 16 e2m1 code points in code order: sign bit 3, exponent bits 2:1, +#: mantissa bit 0. +_E2M1_VALUES = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) + +#: One expert stack's `(codes, block scales)`, unshuffled and unswizzled. +Operand = Tuple[torch.Tensor, torch.Tensor] + +#: DeepSeek-V4.1-Flash's routed geometry. Every rank passes the whole routing +#: space (384 experts, top-6) at H = 5120 and the clamped SwiGLU, whatever its +#: share of the experts or of the intermediate columns. +_V41 = dict(hidden=5120, num_experts=384, top_k=6, swiglu_limit=10.0) + +#: The V4.1 rank forms: (name, local experts, intermediate columns per rank, +#: ranks). An EP window owns `[r * local, (r + 1) * local)` of the experts at the +#: full I; a MoE-TP rank owns columns `[r * I_rank, (r + 1) * I_rank)` of every +#: expert. +_V41_FORMS = ( + ("EP4", 96, 2304, 4), + ("TP4", 384, 576, 4), + ("EP2", 192, 2304, 2), + ("TP2", 384, 1152, 2), + ("EP8", 48, 2304, 8), +) + + +def _v41_cell(form: str, rank: int, tokens: int, g2: str, why: str) -> Cell: + """One V4.1 cell: `form`'s rank `rank` at `tokens` rows. `g2` is "layer" (one + FC2-input global scale in every slot, the target's) or "expert" (`g2[e]`). + The decode cells write a caller `output=` buffer, as the target and the + CUDA-graph tests' eager calls do; the others return a fresh tensor, as the + row-axis sweeps and the per-expert geometry tests ran them.""" + local, inter = {f[0]: f[1:3] for f in _V41_FORMS + (("full", 384, 2304, 1),)}[form] + tp = local == _V41["num_experts"] and form != "full" + return Cell( + why=why, + spec=dict( + model="deepseek-v4.1", + tokens=tokens, + local_expert_offset=0 if tp else rank * local, + local_num_experts=local, + intermediate=inter, + tp_rank=rank if tp else None, + g2=g2, + output=tokens == 1 and g2 == "layer", + **_V41, + ), ) - return torch.ops.trtllm.fp4_block_scale_moe_runner( - routing_logits, - routing_bias, - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - gemm1_bias, - gemm1_alpha, - gemm1_beta, - gemm1_clamp_limit, - gemm2_weights, - gemm2_weights_scale, - gemm2_bias, - output1_scale_scalar, - output1_scale_gate_scalar, - output2_scale_scalar, - num_experts, - top_k, - n_group, - topk_group, - intermediate_size, - local_expert_offset, - local_num_experts, - routed_scaling_factor, - routing_method_type, - do_finalize, - act_type, - topk_weights=topk_weights, - topk_ids=topk_ids, - output=output, - tune_max_num_tokens=tune_max_num_tokens, - use_dp=use_dp, + + +#: The V4.1 cells: every rank form at its decode row count (1) and at the top +#: of the target's eager row axis (8192, one full `_moe_chunks` call), on a +#: different rank at each end, with the layer-level g2 the target passes; then +#: the per-expert g2 form at the full stack's two ends and on an EP4 window. +#: The row axis between the ends is the bucket sweeps' in the entry's test. +_V41_CELLS = tuple( + _v41_cell(form, rank, tokens, "layer", f"deepseek-v4.1 {form} rank {rank}, T={tokens}: {what}") + for form, local, inter, ranks in _V41_FORMS + for rank, tokens, what in ( + (ranks - 1, 1, f"decode, {local} experts x I={inter}, layer g2"), + (0, 8192, "top of the eager row axis, layer g2"), ) +) + ( + _v41_cell("full", 0, 1, "expert", "deepseek-v4.1 full 384-expert stack, T=1, g2[e] per expert"), + _v41_cell( + "full", 0, 8192, "expert", "deepseek-v4.1 full 384-expert stack, T=8192, g2[e] per expert" + ), + _v41_cell("EP4", 1, 1024, "expert", "deepseek-v4.1 EP4 rank 1, T=1024, g2[e] per expert"), +) + + +def _e2m1_rne(a: torch.Tensor) -> torch.Tensor: + """Round to nearest even on the e2m1 grid, saturating at +-6. + + The grid step is 0.5 below 2, 1 below 4 and 2 above -- the binade step of + e2m1's three exponents, the subnormal binade sharing the first normal one's. + `torch.round` is banker's rounding, which is the format's ties-to-even-code + rule: 0.25 -> 0, 0.75 -> 1, 1.25 -> 1, 1.75 -> 2, 2.5 -> 2, 3.5 -> 4, + 5.0 -> 4. + """ + step = torch.where(a.abs() < 2.0, 0.5, torch.where(a.abs() < 4.0, 1.0, 2.0)) + return torch.sign(a) * torch.clamp(torch.round(a.abs() / step) * step, max=E2M1_MAX) + + +def quantize_nvfp4(x: torch.Tensor, global_scale: float) -> torch.Tensor: + """NVFP4 quantize-dequantize of `x` under `global_scale`. + + Per 16 consecutive columns: `sf = e4m3(g * blockmax / 6)` and + `data = e2m1(x * g / sf)`, returning `data * sf` -- the value a downstream + block-scaled MMA multiplies, which is `g` times the reconstruction of `x`. + This models the FC1 epilogue, which the entry's test pins bit-exactly. + """ + rows, cols = x.shape + blocks = x.reshape(rows, cols // SF_VEC_SIZE, SF_VEC_SIZE) + sf = ( + (global_scale * blocks.abs().amax(dim=-1, keepdim=True) / E2M1_MAX) + .clamp(max=E4M3_MAX) + .to(torch.float8_e4m3fn) + .float() + ) + out_scale = torch.where(sf == 0, torch.zeros_like(sf), global_scale / sf) + return (_e2m1_rne(blocks * out_scale) * sf).reshape(rows, cols) + + +def dequantize_nvfp4(codes: torch.Tensor, block_scales: torch.Tensor) -> torch.Tensor: + """One expert's `[N, K]` fp32 weight from its codes and e4m3 block scales.""" + table = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=codes.device) + return table[codes.long()] * block_scales.float().repeat_interleave(SF_VEC_SIZE, dim=1) + + +class Fp4BlockScaleMoeRunner(OpWrapper): + """One NVFP4-weight MoE layer over NVFP4 (e2m1 + e4m3 block-scale) + activations, entered pre-routed. + + The caller owns top-k: `topk_ids` carries global expert ids and + `topk_weights` the combine weights, neither renormalized here. This rank + answers for `[local_expert_offset, local_expert_offset + local_num_experts)` + and contributes zero for every other id, so an expert-parallel split is four + calls whose outputs sum -- which is what the dep4 target does. + + With `do_finalize=True` the result is a one-element list holding the + combined `[T, H]` bf16 rows, or an empty tensor when `output` was given. + With `do_finalize=False` it is three tensors: per-slot expert rows, an + unwritten scale buffer, and the expanded-index -> permuted-row map. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # DeepSeek-R1-0528's routed geometry. The dep4 target gathers four ranks' + # tokens before the expert call and issues one call per 64-wide window, so + # the window cells are the shape it runs and the full stack is the + # single-rank one. + _R1 = dict(hidden=7168, intermediate=2048, num_experts=256, top_k=8) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="dep4's own call: the 64-wide expert window at offset 0, R1 geometry", + spec=dict(tokens=64, local_expert_offset=0, local_num_experts=64, **_R1), + ), + Cell( + why="a window at a non-zero offset -- the ids this rank must answer for move with it", + spec=dict(tokens=64, local_expert_offset=192, local_num_experts=64, **_R1), + ), + Cell( + why="the single-rank shape: all 256 experts local, and one token, the decode row count", + spec=dict(tokens=1, local_expert_offset=0, local_num_experts=256, **_R1), + ), + ) + _V41_CELLS + + note = """ + Certified pre-routed only: `routing_logits` is `None` at every cell, so the + kernel's own routing is not exercised and `routing_method_type`, `n_group`, + `topk_group` and `routed_scaling_factor` are inert -- they are passed + because the signature takes them, not because they do anything here. + + `act_type` 0 (SwiGlu) only. `gemm1_bias`, `gemm1_alpha`, `gemm1_beta`, + `gemm1_clamp_limit` and `gemm2_bias` are `None`. + + `tune_max_num_tokens` is an autotuner bucket hint and was measured inert; + it is not a token count. `T` is `hidden_states.shape[0]`, and 8192 is the + top of the certified column. + + The caller owns the weight layout entirely: FC1 and FC2 weights are + pre-shuffled and their scales additionally 128x4 swizzled. Nothing here + checks that, and getting it wrong is 300-3400 bf16 ulp with no error -- the + entry's test builds the layout twice, once in torch and once through the + trtllm preprocessing ops, and compares them before either reaches the op. + + DeepSeek-V4.1: every rank passes `num_experts` 384, top-6, H = 5120, pre-routed + (`routing_method_type` 1, routing inert), `act_type` 0, `do_finalize=True`, + `output=` a caller buffer, bf16 `topk_weights` summing to the routed scale, and + one of five rank forms: an EP2 / EP4 / EP8 window (192 / 96 / 48 local experts + at offset `r * local`, I = 2304, every weight, scale, scalar and clamp tensor + sliced to the window) or a MoE-TP2 / TP4 rank (all 384 experts at offset 0, + I = 1152 / 576: columns `[r * I, (r + 1) * I)` of every expert, prepared as a + full layer of that width, with the full per-expert scalars and clamp). Two + things differ from the R1 cells. The scalars are per expert: with FC1 / FC2 + weight global scales `gw1[e]`, `gw2[e]` and activation global scales `g1`, + `g2`, `output1_scale_gate_scalar[e] = 1 / (g1 * gw1[e])`, + `output1_scale_scalar[e] = g2 / (g1 * gw1[e])` and `output2_scale_scalar[e] = + 1 / (g2 * gw2[e])`. And `gemm1_clamp_limit` is given: it clamps the two FC1 + accumulators before any scalar (gate from above, up on both sides), so the + model's true-domain limit 10 is passed as `10 / output1_scale_gate_scalar[e]`; + this is exact only because V4.1's up and gate share a weight global scale. + Passing 10 as is, or stacking FC1 `[gate ; up]`, is 256 bf16 ulp RMS wrong with + no error; a window issued at offset 0 is 360, scalars shifted by one expert + 187. The target passes one layer-level `g2` in every slot (a per-expert + `1 / w2.input_scale[e]` saturates the FC1-output block scales on the real + experts); a per-expert `g2` is certified too. The activation scales must be + LINEAR (`fp4_quantize(..., is_sf_swizzled_layout=False)`): the swizzled buffer + has the same byte count whenever T is a multiple of 128 and is silently wrong, + and no metadata tells the two apart. The certified row axis is 1..8192 per call + for every form -- the cells hold its ends; the entry's test sweeps every tuner + bucket's edges, both sides of every token-tile and routing kernel change, the + engine-dumped counts and dep2 / dep4 / dep8's zero-padding skew in between, + cold and autotuned -- and the target cuts larger batches into calls of at most + 8192 rows. Under CUDA-graph capture each form is certified at its captured + counts (`G` and `6G` over the decode batch set, plus a dep member's gathered + counts on an EP window, up to 12288 on EP8 as two chunked calls), replay + bitwise equal to eager and the tactic fixed at capture; the full 384-expert + stack, a per-expert `g2` and counts outside those sets are not certified + captured. + """ + + def raw_call( + self, + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + do_finalize: bool, + act_type: int = 0, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.fp4_block_scale_moe_runner( + routing_logits, + routing_bias, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + gemm1_bias, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, + gemm2_weights, + gemm2_weights_scale, + gemm2_bias, + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + num_experts, + top_k, + n_group, + topk_group, + intermediate_size, + local_expert_offset, + local_num_experts, + routed_scaling_factor, + routing_method_type, + do_finalize, + act_type, + topk_weights=topk_weights, + topk_ids=topk_ids, + output=output, + tune_max_num_tokens=tune_max_num_tokens, + use_dp=use_dp, + ) + + def reference( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + up: Operand, + gate: Operand, + down: Operand, + fc2_global_scale: float | torch.Tensor, + local_expert_offset: int = 0, + local_num_experts: Optional[int] = None, + fc1_weight_global_scale: Optional[torch.Tensor] = None, + fc2_weight_global_scale: Optional[torch.Tensor] = None, + swiglu_limit: Optional[float] = None, + ) -> torch.Tensor: + """The layer, in plain torch over dequantized weights, fp32 throughout. + + The documented exception to mirroring: `raw_call` takes the *prepared* + operands -- packed e2m1 codes, shuffled rows, swizzled scales -- and a + reference given those would be re-deriving the layout rather than the + computation. So this takes the per-expert `(codes, scales)` pairs the + layout was built from, which is where a target's checkpoint starts too, + plus the routing this call was entered with. + + `up`, `gate` and `down` are each `(codes[E, M, K/2], scales[E, M, K/16])` + in their unshuffled, unswizzled form. `topk_ids` carries *global* ids; + this rank contributes only for the ids its window owns, and every other + slot contributes zero -- which is what makes an expert-parallel split + summable. + + The FC1 activation is requantized to NVFP4 under `fc2_global_scale` + before FC2, because the kernel does: the FC1 epilogue emits `g2 * act` + as NVFP4 and FC2's alpha divides it back out. Skipping that step is a + reference that disagrees with a correct kernel by ~23 ulp RMS, so it is + not an optional refinement. + + The DeepSeek-V4.1 form adds three things, each absent by default: + `fc1_weight_global_scale` / `fc2_weight_global_scale` (`[local]` fp32, the + `gw1[e]` / `gw2[e]` each expert's dequantized weight is divided by), + `fc2_global_scale` as a `[local]` tensor (one `g2[e]` per expert; the layer- + level `g2` is the same value in every slot), and `swiglu_limit`, the + true-domain clamp (`gate <= L`, `-L <= up <= L`) applied before the + activation. All per-expert tensors are indexed like the weight operands, + by local expert. A MoE-TP rank is this computation over its own columns of + `up`, `gate` and `down`: its columns start on a multiple of 16, so it + requantizes exactly the FC1-output blocks the full layer does. + """ + num_tokens, hidden = hidden_states.shape + up_c, up_s = up + gt_c, gt_s = gate + dn_c, dn_s = down + local_num_experts = up_c.shape[0] if local_num_experts is None else local_num_experts + out = torch.zeros(num_tokens, hidden, dtype=torch.float32, device=hidden_states.device) + for local_e in range(local_num_experts): + tok, slot = (topk_ids == local_expert_offset + local_e).nonzero(as_tuple=True) + if tok.numel() == 0: + continue + xe = hidden_states[tok] + w_up = dequantize_nvfp4(up_c[local_e], up_s[local_e]) + w_gate = dequantize_nvfp4(gt_c[local_e], gt_s[local_e]) + w_down = dequantize_nvfp4(dn_c[local_e], dn_s[local_e]) + if fc1_weight_global_scale is not None: + w_up = w_up / fc1_weight_global_scale[local_e] + w_gate = w_gate / fc1_weight_global_scale[local_e] + if fc2_weight_global_scale is not None: + w_down = w_down / fc2_weight_global_scale[local_e] + g2 = fc2_global_scale + if isinstance(g2, torch.Tensor): + g2 = float(g2[local_e]) + u = xe @ w_up.t() + g = xe @ w_gate.t() + if swiglu_limit is not None: + g = torch.clamp(g, max=swiglu_limit) + u = torch.clamp(u, min=-swiglu_limit, max=swiglu_limit) + act = u * g * torch.sigmoid(g) + act = quantize_nvfp4(act, g2) / g2 + y = act @ w_down.t() + out.index_add_(0, tok, y * topk_weights[tok, slot].float().unsqueeze(1)) + return out + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """16 ulp per element against the row's own scale, 2 ulp relative RMS. + + Kernel and reference consume bit-identical NVFP4 weights and activations + and model the same FC1-output requantization, so they differ only in + accumulation order and in whether a marginal FC1 value rounds to the + same e2m1 code. An e2m1 code is a coarse step -- two mantissa bits -- so + one flipped intermediate moves the output row by a visible fraction, and + torch's default band cannot express that: its bf16 `atol` of 1e-5 sits + three orders of magnitude below one output ulp of a two-GEMM chain, and + per-element `rtol` is meaningless where cancellation drives `|ref|` to + zero. + + Both numbers are measured, not chosen. Worst over every configuration + the entry's test covers: 9.97 ulp element-wise and 0.88 ulp RMS; the R1 + geometry lands at 9.62 / 0.74 for a 64-wide window and 4.68 / 0.70 for + the 256-expert stack. Only the RMS gate is scale-free -- the element + figure is a max over `T*H` elements, so its tail grows with `T` at + constant accuracy and reaches 16.33 at the top of the certified token + column. Size a caller's tolerance off the RMS number. + + Both were sized against deliberately wrong computations: a gate/up swap + (263-278 ulp), a scalar-role swap (177/48), a missing scale swizzle + (408/357), a missing row shuffle (508/432), and a reference skipping the + FC1 requantization (43/28 element, 23/23 RMS). The smallest RMS distance + any of them reaches is 39.8. + + The DeepSeek-V4.1 cells take the same gate. Measured on sm_103 at the V4.1 + geometry (H 5120, I 2304, 384 experts, top-6, clamp 10, per-expert global + scales), element / RMS: full stack 3.61 / 0.73, an EP4 window 5.75 / 0.86, + an EP2 window 5.82 / 0.76, an EP8 window 8.73 / 0.75, a TP4 rank 7.12 / + 0.82, a TP2 rank 5.59 / 0.70. With the layer-level `g2` over the whole + bucket-by-bucket row axis 1..8192: 9.01 / 1.13 on an EP4 window, 7.57 / + 1.06 on a TP4 rank, 8.00 / 0.76 at most on EP2 / EP8 / TP2; under the dep members' zero-padding skew + 8.87 / 1.11; CUDA-graph replays 13.82 / 1.09 (an EP8 window, the smallest + rows of any layout, over up to 12288 x 5120 elements per graph). The V4.1 + layer's own mistakes sit 128-360 ulp RMS away: the clamp limit passed in + the true domain, FC1 stacked `[gate ; up]`, a window issued at offset 0, + the scalars shifted by one expert, a TP4 rank left out of the sum. + """ + assert out.dtype == ref.dtype == torch.bfloat16, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + ulp = 2.0**-8 + row = ref.float().abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + torch.testing.assert_close(out.float() / row, ref.float() / row, rtol=0.0, atol=16 * ulp) + rel_rms = ( + (out.float() - ref.float()).pow(2).mean().sqrt() + / ref.float().pow(2).mean().sqrt().clamp_min(1e-9) + ).item() / ulp + assert rel_rms <= 2.0, f"relative RMS {rel_rms:.2f} ulp > 2 ulp" + + def is_valid( + self, + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + do_finalize: bool, + act_type: int = 0, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + ) -> None: + """The two ways this op takes an input and answers it wrongly. + + Everything else a caller can get wrong here the op rejects loudly: a + mismatched `intermediate_size`, a `top_k` outside `(0, num_experts)`, a + hidden width that disagrees between the two weight operands. Those are + not restated -- swapping their errors for ours would only make them + worse. + """ + # The kernel takes raw data pointers and assumes a dense row-major + # layout for every tensor. A strided view is accepted without complaint + # and silently reads -- or writes -- the wrong elements. Observed on + # this machine for hidden_states, topk_weights and the weight operands. + for name, tensor in ( + ("routing_logits", routing_logits), + ("routing_bias", routing_bias), + ("hidden_states", hidden_states), + ("hidden_states_scale", hidden_states_scale), + ("gemm1_weights", gemm1_weights), + ("gemm1_weights_scale", gemm1_weights_scale), + ("gemm1_bias", gemm1_bias), + ("gemm1_alpha", gemm1_alpha), + ("gemm1_beta", gemm1_beta), + ("gemm1_clamp_limit", gemm1_clamp_limit), + ("gemm2_weights", gemm2_weights), + ("gemm2_weights_scale", gemm2_weights_scale), + ("gemm2_bias", gemm2_bias), + ("output1_scale_scalar", output1_scale_scalar), + ("output1_scale_gate_scalar", output1_scale_gate_scalar), + ("output2_scale_scalar", output2_scale_scalar), + ("topk_weights", topk_weights), + ("topk_ids", topk_ids), + ("output", output), + ): + if tensor is not None: + assert tensor.is_contiguous(), ( + f"{name} must be contiguous; a strided view is read as if dense " + "and silently produces wrong results" + ) + + # With the block-scale layout this entry contracts -- row shuffle plus + # 128x4 scale swizzle -- a hidden size that is only a multiple of 128 + # makes the kernel read the FC1 weight scales for the wrong blocks. + # Measured wrong at hidden 384 / 640 / 896: no error, 300-3400 bf16 ulp. + assert gemm1_weights.shape[-1] * 2 % 256 == 0, ( + "hidden size (gemm1_weights.shape[-1] * 2) must be a multiple of 256; " + "the prepared block-scale layout is silently misread otherwise" + ) + + +fp4_block_scale_moe_runner = Fp4BlockScaleMoeRunner() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fused_moe.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fused_moe.md deleted file mode 100644 index 2fa0f34c4814..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fused_moe.md +++ /dev/null @@ -1,426 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 17} ---- - -# fused_moe - -**Wraps** `torch.ops.trtllm.fused_moe` (one call). - -## Semantics - -One complete mixture-of-experts layer over already-routed tokens: expert -permutation, the grouped FC1 GEMM, the gated activation, the grouped FC2 -GEMM, and the routing-weighted combine back to one row per token — all -inside a single call. - -Let `T = input.shape[0]`, `H = hidden size`, `I = intermediate size per -expert`, `K = top-k`, `E = fc1_expert_weights.shape[0]` (the experts this -rank holds) and `first = ep_rank * E` (this rank's first *global* expert -id). For every token `t` and slot `j < K`: - -``` -g = token_selected_experts[t, j] # a GLOBAL expert id -skip this slot unless first <= g < first + E -e = g - first # index into the weight tensors - -up = input[t] @ fc1_expert_weights[e, :I, :].T (+ fc1_expert_biases[e, :I]) -gate = input[t] @ fc1_expert_weights[e, I:, :].T (+ fc1_expert_biases[e, I:]) - -# activation_type 5 (Swiglu) and 7 (SwigluBias) — identical on this path: -gate = min(gate, swiglu_limit[e]) # only when swiglu_limit given -up = clamp(up, -swiglu_limit[e], swiglu_limit[e]) -act = gate * sigmoid(swiglu_alpha[e] * gate) * (up + swiglu_beta[e]) - # swiglu_alpha defaults to 1.0, swiglu_beta to 0.0, no clamp by default, - # so the default is plain SwiGLU: silu(gate) * up -# activation_type 6 (Geglu): -act = gelu(gate) * up - -act = round_to(input.dtype)(act) # materialized between the GEMMs -y = act @ fc2_expert_weights[e].T (+ fc2_expert_biases[e]) -out[t] += token_final_scales[t, j] * y # fp32 accumulation -``` - -and finally `out` is cast to `output_dtype`. Both GEMMs accumulate in -fp32; only the FC1 activation result and the final store are in the -low-precision dtype. `token_final_scales=None` means every selected slot -combines with weight `1.0`. - -**FC1 row order is `[up | gate]`.** The first `I` rows of -`fc1_expert_weights[e]` are the up projection (trtllm's `w3`, HF's -`up_proj`) and the last `I` rows are the gate projection (trtllm's `w1`, -HF's `gate_proj`) — the activation's sigmoid is applied to the *second* -half. `fc1_expert_biases[e]` is split the same way. Swapping the halves -produces a plausible-looking, entirely different result; nothing detects -it. - -**Expert parallelism.** `ep_size`/`ep_rank` do not slice anything: they -only shift which global ids this rank answers for. This rank owns -`[ep_rank * E, (ep_rank + 1) * E)`; slots routed anywhere else contribute -nothing to its output, so summing the outputs of all `ep_size` ranks -reproduces the single-rank result. Certified at `ep_size = 2` with -`E = 8`, and at `ep_size = 4` with `E = 64` over a 256-expert global stack -(all four windows driven, each against its own reference, and their sum -against a 256-expert native-torch reference — see *Notes* for the -numbers). A rank whose window catches no slot at all returns exactly -zero. With `ep_size = 1` (the default) the ids are plain indices into the -weight tensors. `tp_size`/`tp_rank` are a separate axis and are **not** -certified here (held at `1`/`0`): under tensor parallelism a rank passes -its own `I/tp_size` slice of both weight tensors and the reduction across -ranks happens outside this call. - -**No routing, no reduction.** The router GEMM and the top-k selection that -produce `token_selected_experts` / `token_final_scales` are the caller's -(a sibling op, `torch.ops.trtllm.renorm_moe_routing_op`, produces exactly -that pair), as are any shared/dense expert branch, the TP all-reduce or EP -gather of this call's output, and the residual add. Nothing is -renormalized inside: the combine weights are used exactly as given (they -need not sum to 1, and may be negative). - -**Output.** The call returns a Python list. With `out_tensor=None` it is -`[out]`, one freshly allocated contiguous `[T, H]` tensor in -`output_dtype`. With `out_tensor` given the result is written into that -buffer (every element overwritten, nothing outside its rows touched) and -the returned list is **empty** — the caller uses its own buffer. Inputs -are never mutated. - -## Signature - -```python -def fused_moe( - input: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: Optional[torch.Tensor], - fc1_expert_weights: torch.Tensor, - fc1_expert_biases: Optional[torch.Tensor], - fc2_expert_weights: torch.Tensor, - fc2_expert_biases: Optional[torch.Tensor], - output_dtype: torch.dtype, - quant_scales: List[torch.Tensor], - input_sf: Optional[torch.Tensor] = None, - swizzled_input_sf: bool = True, - swiglu_alpha: Optional[torch.Tensor] = None, - swiglu_beta: Optional[torch.Tensor] = None, - swiglu_limit: Optional[torch.Tensor] = None, - tp_size: int = 1, tp_rank: int = 0, - ep_size: int = 1, ep_rank: int = 0, - cluster_size: int = 1, cluster_rank: int = 0, - enable_alltoall: bool = False, - use_deepseek_fp8_block_scale: bool = False, - use_w4_group_scaling: bool = False, - use_int8_woq_per_channel: bool = False, - use_mxfp8_act_scaling: bool = False, - min_latency_mode: bool = False, - use_fused_finalize: bool = True, - tune_max_num_tokens: int = 8192, - tuner_num_tokens: Optional[int] = None, - tuner_top_k: Optional[int] = None, - activation_type: int = 5, - unpadded_hidden_size: Optional[int] = None, - out_tensor: Optional[torch.Tensor] = None, - use_dynamic_fc2_scale: bool = False, - use_mxfp8_weight_scaling: bool = False, - # routed-expert LoRA, per-request and slot-indexed families - fc1_lora_ranks=None, fc1_lora_weight_ptrs=None, - fc2_lora_ranks=None, fc2_lora_weight_ptrs=None, - gated_lora_ranks=None, gated_lora_weight_ptrs=None, - host_request_types=None, host_context_lengths=None, - lora_max_low_rank: int = 0, - fc1_slot_lora_ranks=None, fc1_slot_lora_weight_ptrs=None, - fc2_slot_lora_ranks=None, fc2_slot_lora_weight_ptrs=None, - gated_slot_lora_ranks=None, gated_slot_lora_weight_ptrs=None, - token_to_slot=None, -) -> List[torch.Tensor] -``` - -### Certified arguments - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `input` | `[T, H]` | bf16 or fp16 | contiguous | CUDA | -| `token_selected_experts` | `[T, K]` | **int32** | contiguous | CUDA | -| `token_final_scales` | `[T, K]` or `None` | **fp32** | contiguous | CUDA | -| `fc1_expert_weights` | `[E, 2I, H]`, rows `[up; gate]` | = `input.dtype` | contiguous | CUDA | -| `fc1_expert_biases` | `[E, 2I]` or `None` | = `input.dtype` | contiguous | CUDA | -| `fc2_expert_weights` | `[E, H, I]` | = `input.dtype` | contiguous | CUDA | -| `fc2_expert_biases` | `[E, H]` or `None` | = `input.dtype` | contiguous | CUDA | -| `output_dtype` | scalar | must equal `input.dtype` here | — | — | -| `quant_scales` | `[]` (empty list) | — | — | — | -| `swiglu_alpha` / `swiglu_beta` / `swiglu_limit` | `[E]` or `None` | **fp32** | contiguous | CUDA | -| `ep_size` / `ep_rank` | scalars | Python int, `0 <= ep_rank < ep_size` | — | — | -| `use_fused_finalize` | scalar | bool — both values accepted, giving identical bits; not switchable per call (see *Metadata consumed*) | — | — | -| `activation_type` | scalar | int: 5 (Swiglu), 6 (Geglu), 7 (SwigluBias) | — | — | -| `out_tensor` | `[T, H]` or `None` | = `output_dtype` | contiguous | CUDA | -| returns `[0]` (absent when `out_tensor` given) | `[T, H]` | `output_dtype` | contiguous, newly allocated | same as `input` | - -Biases are all-or-nothing: pass both or neither (see *Preconditions*). - -### Arguments held inert (not certified) - -`input_sf`, `swizzled_input_sf`, `use_deepseek_fp8_block_scale`, -`use_w4_group_scaling`, `use_int8_woq_per_channel`, -`use_mxfp8_act_scaling`, `use_mxfp8_weight_scaling`, -`use_dynamic_fc2_scale` — the quantized-activation/weight paths (fp8 -per-tensor and block scale, NVFP4, MXFP4/MXFP8, INT8 weight-only), which -also populate `quant_scales`. `min_latency_mode` — NVFP4-only (it is -rejected here) and returns a different tensor list. -`enable_alltoall`, `tuner_num_tokens`, -`tuner_top_k`, `cluster_size`, `cluster_rank` — multi-rank -dispatch/smart-router paths. `tune_max_num_tokens` — an autotuner bucket -cap; left at its default. `unpadded_hidden_size` — for callers that padded -`H`; left at `None`. The whole LoRA family (`fc1_lora_*`, `fc2_lora_*`, -`gated_lora_*`, `host_request_types`, `host_context_lengths`, -`lora_max_low_rank`, `*_slot_lora_*`, `token_to_slot`). `tp_size` / -`tp_rank` are forwarded but were only exercised at `1` / `0`. - -## Metadata consumed - -None. The op reads no attention metadata, no KV cache, no registered -layer, and no module state — every tensor and every scalar it uses is an -argument. - -Two process-global caches sit behind it. The runner cache is result-neutral; -the tactic cache is not: - -- an **autotuner profiling cache** that picks the two GEMM tactics. Cold — - the state most calls certified here run in — both tactics are the `-1` - fallback, and nothing is printed (the cache-miss message is a - `warning_once`, below trtllm's default `error` severity). - **Tactic choice is not bitwise neutral.** After one `autotune()` - pass the same inputs at `T=256, H=512, I=256, E=32, K=4` came back - different in 83447 of 131072 elements (2.96 ulp of the token row's largest - magnitude), and clearing the cache restored the cold bits exactly. It is a - different accumulation order, not a different computation — cold and hot - both pass this entry's gate against the torch reference (2.4 and 2.0 ulp). - Almost all of the spread is the FC2 GEMM: over the tactic space at that - shape, `gemm1`'s 209 tactics yield 2 distinct outputs, `gemm2`'s 309 yield - 101. **Serving runs on the hot side**, so a target exercises tactics a cold - receipt never saw: `PyExecutor` warmup calls `_run_autotuner_warmup` - (`_torch/pyexecutor/model_engine.py:1103,1278`), which wraps a real forward - in `autotune()` (`:1293`) unless `enable_autotuner=False` (default `True`, - `llmapi/llm_args.py:4669`). - - **The hot side is certified at `E=64, H=7168, I=2048, K=8`**, two ways: - - - The *whole tactic space*, driven directly. The tactic population is a - property of the build, not of the geometry — the count comes from - `get_tactic_num(gemm_idx)`, which takes no shape argument — and that was - measured rather than assumed: the tactic space captured at - `E=64, H=7168, I=2048` has exactly the same size as one captured at - `E=32, H=512, I=256`, 209 x 309 = 64581 configurations. Every one of - `gemm2`'s 309 tactics and `gemm1`'s 209 passes this entry's gate at - `T=256`, and all 309 `gemm2` tactics pass at `T=8192`; none failed to - run. The distinct-output counts match the small shape's exactly — 101 - from `gemm2`, 2 from `gemm1` — and the worst deviation over the whole - space is **3.71 ulp** element-wise at `T=256` and **3.98 ulp** at - `T=8192` (relative RMS 1.42 both), so the 8 / 4 ulp bounds below hold at - this cell with 2.0x and 2.8x margin. - - One real `autotune()` pass at that geometry: it fills exactly **28** - cache entries (2 tunable GEMMs x 14 power-of-2 token buckets, 1 … 8192), - every one carrying a profiled tactic rather than the `-1` fallback. The - warm results move the bits (36172254 of 58720256 elements at `T=8192`) - and still pass at 3.88 ulp; clearing the cache restores the cold bits - exactly at every token count. The pass costs ~70 s at this geometry. -- a cache of C++ `FusedMoeRunner` objects keyed by - `(activation dtype, weight dtype, output dtype, quant flags)`, each - holding a GPU **workspace** buffer that is allocated on first use and - grown as needed. It is never freed by the op; - `tensorrt_llm._torch.custom_ops.torch_custom_ops.MoERunner.clear_all_workspaces()` - releases the workspaces of every cached runner. **`use_fused_finalize` is - baked into that object** and is not part of the key, so the first call in - a process for a given dtype set fixes it for every later call — a caller - cannot flip it per call, and the entry's own test therefore runs on a - runner built with `True`. On this bf16 path the value makes no difference - either way: two fresh processes driving `True` then `False`, and `False` - then `True`, produced four bitwise identical outputs at - `T=2048, H=512, I=256, E=32, K=8` (a probe, not part of the test), and - `True` and `False` agree bitwise inside the test at - `T=2048, E=64, H=7168, I=2048, K=8`. - -Two identical calls return bitwise identical results **within one tuner -state** (verified at `T=256, H=512, I=256, E=32, K=4`, cold and hot -separately; and at `E=64, H=7168, I=2048, K=8` for `T` in {1, 256, 8192}, -where a cold call repeated after an intervening tuning pass and a cache -clear returned the same bits); across a cold/hot change they differ by the -ulp above. - -## Preconditions - -Shapes and layout: - -- `input` is a **2D contiguous CUDA** tensor with `T >= 1`. 3D input raises - `input must be 2D`; a strided view raises `input must be contiguous`; a - CPU tensor raises `input must be a CUDA tensor`; `T == 0` fails a C++ - assertion (`Assertion failed: input_activations`). -- `H % 8 == 0` and `I % 8 == 0`, else `hidden_size must be divisible - by 8 for weights` / `inter_size must be divisible by 8 for weights`. - `H = 8` and `I = 8` work. -- `fc1_expert_weights` is `[E, 2I, H]` and `fc2_expert_weights` is - `[E, H, I]`, both 3D, contiguous, CUDA, same dtype as `input`, same `E`. - Violations raise (`must be 3D`, `must be contiguous`, - `must be a CUDA tensor`, `fc1_expert_weights and fc2_expert_weights must - have the same number of experts.`). Certified: `E` in - {1, 2, 7, 8, 16, 32, **64**, 128, 256} and `(H, I)` in {(8,8), (16,8), - (128,64), (192,96), (256,128), (512,256), (1024,512), (2048,768), - **(7168,2048)**}. The two enumerations are not a product: `(7168, 2048)` - was driven at `E = 64` only, and `E = 64` only at `(7168, 2048)`. Nothing - in the permutation or gather path keys on the expert count, or on which - index a given expert sits at: permuting the 64-expert stack and - relabelling the ids to match reproduces the output **bitwise**. -- `token_selected_experts` is `[T, K]` **int32** contiguous CUDA; - `token_final_scales`, when given, is `[T, K]` **fp32** contiguous CUDA - with the same `K`. Other dtypes raise (`token_selected_experts dtype is - Long, while Int is expected`, `token_final_scales.value() dtype is - BFloat16, while Float is expected`), as do mismatched `T` or `K`. - Certified `K`: 1, 2, 3, 4, 8, 16. **A token's `K` ids must be distinct - when `T > 256`** — a repeat there is accepted and then computed wrong, or - faults; see *Notes*. At `T <= 256` a repeat is honoured once per slot, so - that expert's output is added twice with its two combine weights. -- Biases are optional but **paired**: passing exactly one of - `fc1_expert_biases` / `fc2_expert_biases` raises - `RuntimeError: bad optional access`. Both must be `[E, 2I]` and `[E, H]` - contiguous CUDA tensors in `input.dtype` (fp32/fp16 bias against a bf16 - input raises). -- `out_tensor`, when given, must be a contiguous CUDA tensor of shape - exactly `[T, H]` in `output_dtype`; wrong dtype, wrong shape, or a - strided view each raise. A leading row-slice of a larger contiguous - buffer is a valid `out_tensor` — rows past `T` are left bitwise - untouched. - -Dtypes and modes: - -- Activations and weights must be the **same** dtype: bf16 or fp16. - fp32 activations, or weights whose dtype differs from `input`, raise - `Could not construct fused moe op with the requested input combination - ...`. -- On this unquantized path `output_dtype` **must equal** `input.dtype`. - A mismatch does **not** raise — see *Notes*. The wrapper asserts it. -- `quant_scales` must carry the scale set of whatever quantization is - requested; on the unquantized path pass `[]` (a non-empty list is - accepted and ignored). fp8-e4m3 inputs with `[]` raise `Expecting 4 - quant scales for fp8 quantization`. -- `swiglu_alpha` / `swiglu_beta` / `swiglu_limit`, when given, are fp32 - CUDA tensors with exactly `E` elements (`swiglu_alpha must have - num_experts_on_rank elements.`, `... must be a CUDA tensor`, `... dtype - is BFloat16, while Float is expected`). -- `activation_type` must be a gated type (5 Swiglu, 6 Geglu, 7 - SwigluBias) for the `[E, 2I, H]` FC1 layout. The non-gated values of the - same enum (1 Identity, 2 Gelu, 3 Relu, 4 Silu, 8 Relu2) — and any - unrecognised value — expect `[E, I, H]` instead, so with a gated layout - they raise `fc1_expert_weights inter size must be equal to - fc2_expert_weights inter size.` (verified for 1, 2, 4, 8 and 99). -- `0 <= ep_rank < ep_size` (`Assertion failed: ep_rank < ep_size`). -- `min_latency_mode=True` is NVFP4-only and fails here with `Assertion - failed: use_fp4 == true`. `cluster_size > 1` raises `smart_router is - supported in min_latency mode`. -- `tuner_num_tokens` and `tuner_top_k` must both be `None` unless - `enable_alltoall=True`, and both must be set when it is; the op's own - Python `assert` fires otherwise (bare `AssertionError`). -- Any LoRA tensor implies `lora_max_low_rank > 0` (`MoE LoRA requires - lora_max_low_rank > 0; got 0`), and the per-request and slot-indexed - families are mutually exclusive. - -A caller violating none of the above gets the result described under -*Semantics*, to within a few low-precision ulps (see *Notes*). - -## Notes - -- **`output_dtype != input.dtype` is a silent wrong answer.** On the - unquantized path the output buffer is allocated with `output_dtype` - while the epilogue stores elements in the activation dtype, so the - caller reads reinterpreted bits: bf16 input with `output_dtype=fp32` - returns finite, plausibly-scaled values that are pairs of bf16 results - glued into fp32 words (element `j` carries roughly the true element - `2j+1`). No error, no NaN. `use_fused_finalize=False` behaves the same. - The wrapper carries the one guard this entry needs: when - `input.dtype in (bf16, fp16)` and `fc1_expert_weights.dtype == - input.dtype`, `output_dtype` must equal `input.dtype`. The guard is - scoped so it never fires on the quantized paths, where a differing - `output_dtype` is the norm. -- **Expert ids outside this rank's range are silently dropped**, not - clamped and not rejected: with `ep_size=1, E=8`, ids `8`, `100` and `-1` - each contribute nothing, and the token's other slots still combine - normally. A routing bug therefore shows up as a quietly weaker token, - never as an error. Checking ids costs a device-to-host sync, so the - wrapper does not. -- **Three quantization flags are accepted and ignored** on the bf16 path: - `use_w4_group_scaling`, `use_mxfp8_act_scaling` and - `use_dynamic_fc2_scale` return bitwise the unquantized result. The other - three (`use_deepseek_fp8_block_scale`, `use_int8_woq_per_channel`, - `use_mxfp8_weight_scaling`) raise. -- **Numerics.** Against a native-torch reference that consumes the same - low-precision operands (fp32 GEMM accumulation, the FC1 activation - rounded to the input dtype, fp32 combine), the kernel's worst - element-wise deviation over the covered configurations was **2.8 ulp of - the token row's largest magnitude** (bf16 ulp = 2^-8, fp16 ulp = 2^-11) - and its relative RMS deviation **1.2 ulp**. Both were calibrated on the - fallback tactic; walking the *whole* tactic space reaches **3.98 ulp** - element-wise and **1.43 ulp** RMS (`gemm2`, `T = 8192`, at - `E=64, H=7168, I=2048`) — still inside the test's 8 ulp / 4 ulp gates, - with 2.0x and 2.8x margin. That is rounding noise, not - a different computation: the same comparison against a gate/up-swapped - reference lands at 176+ ulp element-wise and ~190 ulp RMS at the small - shapes and 284 / 197 at `E=64, H=7168, I=2048`, and against a - one-expert-dropped reference at 997 / 212 there (measured while - calibrating the bound; the entry's test asserts only that such a - reference is rejected). Default `torch.testing.assert_close` tolerances - do not fit this op — their bf16 `atol` of 1e-5 is three orders of - magnitude below one output ulp of a two-GEMM chain. -- **Scale is inert.** `(H, I) = (7168, 2048)` at `E = 64` is 3.5x the - hidden size and 2.7x the intermediate size of anything else certified - here, and behaves exactly like the small shapes: same tactic population, - same distinct-output counts, same ulp band, no tactic failing to run. - The deviation grows only with token count, and only slightly — 1.12 ulp - at `T = 1` to 2.54 at `T = 8192`, the same trend the smaller shapes show. -- **Four-way expert parallelism**, measured at `E = 64` per rank over a - 256-expert global stack with global ids in `[0, 256)`, at - `T` in {1, 1024, 8192}: each rank's output matches its own reference to - at most 2.69 ulp, and the four outputs summed in fp32 match a 256-expert - native-torch reference to 3.53 ulp element-wise / 1.34 ulp RMS. The sum - sits slightly wider than any single rank because each rank rounds its own - partial to bf16 before the add. A rank whose window catches no slot - returns exactly zero, and the same weights driven under the wrong - `ep_rank` land 90x outside the RMS gate — so the tiling check is a real - check and not a formality. -- **Repeated expert ids in one token's row are a defect past `T = 256`, - not a supported input.** At `T <= 256` a repeat is honoured once per - slot, as *Preconditions* says. At `T >= 257` the op accepts the input and - then either returns values ~200-460 ulp wrong on essentially every row, - or dies with `CUDA error: an illegal memory access was encountered` from - inside the call. Which of the two happens depends on whether the - out-of-bounds address is mapped; both were seen at the same shape. - `compute-sanitizer` names the fault: an invalid 4-byte global read in - `tensorrt_llm::_v1::kernels::cutlass_kernels::finalizeMoeRoutingKernel`, - launched from `finalizeMoeRoutingKernelLauncher` inside - `CutlassMoeFCRunner<...>::gemm2`. The boundary is exactly 256 tokens - (`T = 256` clean, `T = 257` faults) and matches the largest block size the - single-CTA expert-map builder is instantiated for — - `fusedBuildExpertMapsSortFirstTokenDispatch<256, …>` in - `libth_common.so`, whose fallback is - `threeStepBuildExpertMapsSortFirstToken`. `use_fused_finalize` does not - change it. This is **not specific to the new geometry**: reproduced at - `(E,H,I,K) = (8,256,128,4)`, `(32,512,256,8)` and `(128,2048,768,8)` — - cells this entry already certified — as well as at `(64,7168,2048,8)`. - Top-k selection returns distinct ids by construction, so a caller taking - its ids from a selection op cannot reach this; a caller synthesizing - `token_selected_experts` itself can. -- `activation_type=7` (SwigluBias) computed bitwise the same result as - `5` (Swiglu) for every input tried, with and without - `swiglu_alpha/beta/limit`. -- Geglu (6) matches a `gelu(gate) * up` reference; the tanh-approximate - and exact gelu variants are indistinguishable at bf16 output resolution, - so which one the kernel uses is **not** pinned by this entry. -- `T = 8192` (one token past the default `tune_max_num_tokens` bucket cap) - is certified, at `(2048, 768, E=128)` and at `(7168, 2048, E=64)`; - `T = 8193` was also observed correct during an earlier run's probing but - is not in the entry's test. -- **What `(7168, 2048)` at `E = 64` does *not* cover.** bf16 only (no - fp16), no biases, no `swiglu_alpha`/`beta`/`limit`, `activation_type = 5` - only (no Geglu, no SwigluBias), `tp_size = 1`, and `K = 8` only. Those - axes are certified at the smaller shapes and were not re-driven here. -- Sibling ops in this build address the same job differently: - `torch.ops.trtllm.moe_custom_op` takes a registered layer by string id - instead of explicit weights, and `torch.ops.trtllm.fp8_block_scale_moe_runner` - / `torch.ops.trtllm.fp4_block_scale_moe_runner` are the trtllm-gen - block-scaled MoE entry points. Catalog membership is `index.yaml`'s fact - alone. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fused_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fused_moe.py deleted file mode 100644 index ca528377a573..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fused_moe.py +++ /dev/null @@ -1,133 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Fused mixture-of-experts layer: expert permutation + grouped FC1/FC2 GEMMs -with a gated activation between them + routing-weighted combine, in one call.""" - -from typing import List, Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* - - -def fused_moe( - input: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: Optional[torch.Tensor], - fc1_expert_weights: torch.Tensor, - fc1_expert_biases: Optional[torch.Tensor], - fc2_expert_weights: torch.Tensor, - fc2_expert_biases: Optional[torch.Tensor], - output_dtype: torch.dtype, - quant_scales: List[torch.Tensor], - input_sf: Optional[torch.Tensor] = None, - swizzled_input_sf: bool = True, - swiglu_alpha: Optional[torch.Tensor] = None, - swiglu_beta: Optional[torch.Tensor] = None, - swiglu_limit: Optional[torch.Tensor] = None, - tp_size: int = 1, - tp_rank: int = 0, - ep_size: int = 1, - ep_rank: int = 0, - cluster_size: int = 1, - cluster_rank: int = 0, - enable_alltoall: bool = False, - use_deepseek_fp8_block_scale: bool = False, - use_w4_group_scaling: bool = False, - use_int8_woq_per_channel: bool = False, - use_mxfp8_act_scaling: bool = False, - min_latency_mode: bool = False, - use_fused_finalize: bool = True, - tune_max_num_tokens: int = 8192, - tuner_num_tokens: Optional[int] = None, - tuner_top_k: Optional[int] = None, - activation_type: int = 5, # ActivationType.Swiglu - unpadded_hidden_size: Optional[int] = None, - out_tensor: Optional[torch.Tensor] = None, - use_dynamic_fc2_scale: bool = False, - use_mxfp8_weight_scaling: bool = False, - fc1_lora_ranks: Optional[torch.Tensor] = None, - fc1_lora_weight_ptrs: Optional[torch.Tensor] = None, - fc2_lora_ranks: Optional[torch.Tensor] = None, - fc2_lora_weight_ptrs: Optional[torch.Tensor] = None, - gated_lora_ranks: Optional[torch.Tensor] = None, - gated_lora_weight_ptrs: Optional[torch.Tensor] = None, - host_request_types: Optional[torch.Tensor] = None, - host_context_lengths: Optional[torch.Tensor] = None, - lora_max_low_rank: int = 0, - fc1_slot_lora_ranks: Optional[torch.Tensor] = None, - fc1_slot_lora_weight_ptrs: Optional[torch.Tensor] = None, - fc2_slot_lora_ranks: Optional[torch.Tensor] = None, - fc2_slot_lora_weight_ptrs: Optional[torch.Tensor] = None, - gated_slot_lora_ranks: Optional[torch.Tensor] = None, - gated_slot_lora_weight_ptrs: Optional[torch.Tensor] = None, - token_to_slot: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """Run one MoE layer over pre-routed tokens. - - Returns `[out]` with `out` a fresh `[num_tokens, hidden_size]` tensor in - `output_dtype`, or `[]` when `out_tensor` is given (written in place). - """ - # Pure-metadata guard: on the unquantized high-precision path the store - # is done in the activation dtype while the output tensor is allocated - # with output_dtype, so a mismatch reinterprets the written bits — observed - # on this machine to return plausible-looking wrong values, never to raise. - if input.dtype in (torch.bfloat16, torch.float16) and (fc1_expert_weights.dtype == input.dtype): - assert output_dtype == input.dtype, ( - f"output_dtype ({output_dtype}) must equal input.dtype ({input.dtype}) " - "on the unquantized path; a mismatch is written as raw activation-dtype " - "bits into an output_dtype buffer and silently gives wrong values" - ) - return torch.ops.trtllm.fused_moe( - input, - token_selected_experts, - token_final_scales, - fc1_expert_weights, - fc1_expert_biases, - fc2_expert_weights, - fc2_expert_biases, - output_dtype, - quant_scales, - input_sf=input_sf, - swizzled_input_sf=swizzled_input_sf, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - swiglu_limit=swiglu_limit, - tp_size=tp_size, - tp_rank=tp_rank, - ep_size=ep_size, - ep_rank=ep_rank, - cluster_size=cluster_size, - cluster_rank=cluster_rank, - enable_alltoall=enable_alltoall, - use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale, - use_w4_group_scaling=use_w4_group_scaling, - use_int8_woq_per_channel=use_int8_woq_per_channel, - use_mxfp8_act_scaling=use_mxfp8_act_scaling, - min_latency_mode=min_latency_mode, - use_fused_finalize=use_fused_finalize, - tune_max_num_tokens=tune_max_num_tokens, - tuner_num_tokens=tuner_num_tokens, - tuner_top_k=tuner_top_k, - activation_type=activation_type, - unpadded_hidden_size=unpadded_hidden_size, - out_tensor=out_tensor, - use_dynamic_fc2_scale=use_dynamic_fc2_scale, - use_mxfp8_weight_scaling=use_mxfp8_weight_scaling, - fc1_lora_ranks=fc1_lora_ranks, - fc1_lora_weight_ptrs=fc1_lora_weight_ptrs, - fc2_lora_ranks=fc2_lora_ranks, - fc2_lora_weight_ptrs=fc2_lora_weight_ptrs, - gated_lora_ranks=gated_lora_ranks, - gated_lora_weight_ptrs=gated_lora_weight_ptrs, - host_request_types=host_request_types, - host_context_lengths=host_context_lengths, - lora_max_low_rank=lora_max_low_rank, - fc1_slot_lora_ranks=fc1_slot_lora_ranks, - fc1_slot_lora_weight_ptrs=fc1_slot_lora_weight_ptrs, - fc2_slot_lora_ranks=fc2_slot_lora_ranks, - fc2_slot_lora_weight_ptrs=fc2_slot_lora_weight_ptrs, - gated_slot_lora_ranks=gated_slot_lora_ranks, - gated_slot_lora_weight_ptrs=gated_slot_lora_weight_ptrs, - token_to_slot=token_to_slot, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/mxe4m3_mxe2m1_block_scale_moe_runner.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/mxe4m3_mxe2m1_block_scale_moe_runner.md deleted file mode 100644 index 168e1448fb77..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/mxe4m3_mxe2m1_block_scale_moe_runner.md +++ /dev/null @@ -1,512 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 18} ---- - -# mxe4m3_mxe2m1_block_scale_moe_runner - -**Wraps** `torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner` (one call). - -## Semantics - -One complete mixture-of-experts layer for **MXFP4 (E2M1 + per-32 E8M0 block -scale) weights over MXFP8 (E4M3 + per-32 E8M0 block scale) activations** — the -trtllm-gen "block scale MoE" family's W4A8 member. In a single call: optional -top-k routing over router logits, expert permutation, the grouped FC1 GEMM, the -clamped gated activation, **requantization of that activation to MXFP8**, the -grouped FC2 GEMM, and the routing-weighted combine back to one row per token. - -Let `T = hidden_states.shape[0]`, `H = valid_hidden_size` (the model's true -hidden size), `I = valid_intermediate_size` (the true per-rank intermediate -size), `K = top_k`, `E = local_num_experts`, `off = local_expert_offset`, and -let `W_up[e]`, `W_gate[e]` (`[I, H]`) and `W_down[e]` (`[H, I]`) be the fp32 -values of this rank's dequantized MXFP4 weights (§ *Preconditions* fixes how -they are packed). - -The activations arrive **already quantized**: `hidden_states` holds e4m3 -elements and `hidden_states_scale` one UE8M0 byte per 32 consecutive columns, -so the value the kernel computes with is - -``` -hidden[t, k] = hidden_states[t, k].float() * 2^(sf[t, k // 32] - 127) -``` - -where `sf` is `hidden_states_scale` read as `[T, hidden_width / 32]` row-major -(§ *Preconditions*). - -**Routing** — two mutually exclusive entry points, both certified: - -- *Routed here*: pass `routing_logits` `[T, num_experts]`, leave - `topk_weights`/`topk_ids` at `None`. `routing_method_type` selects - - `0` (Default): `p = softmax_fp32(logits)` over **all** experts, then - `(w, id) = top_k(p)` — the combine weights are those softmax - probabilities and do **not** sum to 1; - - `1` (Renormalize): `(v, id) = top_k(logits)` first, then - `w = softmax_fp32(v)` over the `K` selected logits — the weights **do** - sum to 1. This is the gpt-oss / Qwen3 style. -- *Pre-routed*: pass `topk_ids` `[T, K]` int32 and `topk_weights` `[T, K]` - **bf16**, leave `routing_logits` at `None`. The kernel's routing stage is - bypassed entirely; `routing_method_type` is still validated but never - used (values `0/1/2/4/5/6` all give bitwise identical output). - -Passing both is not an error: the op announces it uses the pre-routed pair. - -**Per token `t` and slot `j < K`:** - -``` -g = topk_ids[t, j] # a GLOBAL expert id -skip this slot unless off <= g < off + E -e = g - off # index into this rank's weights - -up = hidden[t] @ W_up[e].T (+ gemm1_bias_up[e]) -gate = hidden[t] @ W_gate[e].T (+ gemm1_bias_gate[e]) - -# gemm1_clamp_limit (per expert), applied only when the tensor is given: -gate = min(gate, limit[e]) -up = clamp(up, -limit[e], limit[e]) - -# act_type 0 (SwiGlu) — the only kernel that exists on this path: -act = (up + beta[e]) * gate * sigmoid(alpha[e] * gate) - # alpha defaults to 1.0 when gemm1_alpha is None, - # beta defaults to 0.0 when gemm1_beta is None, - # so the default is plain SwiGLU: up * silu(gate) - -act = mx_requantize(act) # see below — the W4A8 step -y = act @ W_down[e].T (+ gemm2_bias[e]) -out[t] += topk_weights[t, j] * y # fp32 accumulation -``` - -and `out` is stored as bf16. The gpt-oss clamped GLU is exactly -`alpha = 1.702`, `beta = 1.0`, `limit = 7.0`. - -**The intermediate is MXFP8, on the OCP scale.** FC2 is an MXFP8 x MXFP4 GEMM, -so the FC1 epilogue quantizes its post-activation output before FC2 reads it. -Per token row and per **32 consecutive intermediate columns** (natural column -order, blocks starting at column 0 of the padded intermediate): - -``` -amax = max |act| over the 32-element block -e = floor(log2(amax)) - 8 # E8M0 byte = e + 127; e = -127 if amax == 0 -act' = round_to_nearest_even(clamp(act / 2^e, -448, +448)) * 2^e -``` - -This is the OCP MX scale — `floor` of the block max's exponent, with the block -max itself **saturating** to `448 * 2^e` whenever its mantissa exceeds 1.75 -(18–21% of blocks measured, clipping that one element by up to 12.5%). It is -**not** the round-up scale `torch.ops.trtllm.mxfp8_quantize` applies to -activations, and modelling it as such is wrong for exactly those blocks. An -intermediate element below roughly `2^-18` of its block's largest magnitude -rounds to zero. This entry's test reads the requantized values straight out of -the kernel (a down projection set to the identity) and matches the formula -above bit-exactly on all 16384 + 46080 elements of two geometries -(`H = I = 512` and `H = I = 2880`); on the same data the round-up scale -matches 99.2% of elements and the *un*quantized activation 6%. - -**FC1 half order is `[up | gate]`.** Before the kernel's row interleave (§ -*Preconditions*) the first `I_pad` rows of the FC1 operand are the up -projection (trtllm's `w3`, HF's `up_proj`) and the last `I_pad` rows the gate -projection (trtllm's `w1`, HF's `gate_proj`). The sigmoid is applied to the -**gate** half. Swapping them produces a plausible-looking, entirely different -result; nothing detects it — this entry's test checks that a gate/up-swapped -reference lands far outside the numerical gate (measured 246 ulp per element). - -**Fusion boundary.** Inside the call: routing (when driven by logits), -permutation, both GEMMs with on-the-fly MXFP4/MXFP8 dequantization, the clamped -gated activation, the MXFP8 requantization between the GEMMs, and the weighted -combine. Outside, and the caller's job: the router GEMM that produces -`routing_logits`, **adding the router bias to those logits** (see *Notes* — -`routing_bias` is a no-op here), **quantizing the bf16 hidden states to MXFP8** -(`torch.ops.trtllm.mxfp8_quantize`, § *Preconditions* — the hidden-width -padding happens inside that call, not here), all weight preprocessing, any -shared/dense expert branch, the TP all-reduce or EP gather of this call's -output, and the residual add. - -**Nothing is renormalized inside.** `topk_weights` is used exactly as given -(need not sum to 1, may be negative). `routed_scaling_factor` is accepted but -had no effect on any certified path. - -**Output.** With `output=None` the call returns a fresh contiguous -`[T, valid_hidden_size]` **bf16** tensor (the output is bf16 regardless of the -fp8 input). With `output` given, the result is written into that buffer (every -element overwritten, nothing outside its rows touched) and the call returns an -**empty `[0]` bf16 tensor** — the caller must read its own buffer. No input is -mutated; two identical calls are bitwise equal. - -## Signature - -```python -def mxe4m3_mxe2m1_block_scale_moe_runner( - routing_logits: Optional[torch.Tensor], - routing_bias: Optional[torch.Tensor], - hidden_states: torch.Tensor, - hidden_states_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, - gemm1_bias: Optional[torch.Tensor], - gemm1_alpha: Optional[torch.Tensor], - gemm1_beta: Optional[torch.Tensor], - gemm1_clamp_limit: Optional[torch.Tensor], - gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, - gemm2_bias: Optional[torch.Tensor], - num_experts: int, - top_k: int, - n_group: Optional[int], - topk_group: Optional[int], - intermediate_size: int, - valid_hidden_size: Optional[int], - valid_intermediate_size: Optional[int], - local_expert_offset: int, - local_num_experts: int, - routed_scaling_factor: Optional[float], - routing_method_type: int, - act_type: int, - topk_weights: Optional[torch.Tensor] = None, - topk_ids: Optional[torch.Tensor] = None, - output: Optional[torch.Tensor] = None, - tune_max_num_tokens: int = 8192, - use_dp: bool = False, -) -> torch.Tensor -``` - -The signature is the bf16-activation sibling op's plus `hidden_states_scale` -in position 4; every other argument means the same thing. - -Derived sizes used below (`pad_up(x, a) = ceil(x / a) * a`): - -``` -I_pad = pad_up(I, 128) # FC1 rows per half, FC2 K axis -H1_pad = pad_up(H, 512) # FC1 K axis == hidden_states width -H2_pad = pad_up(H, 128) # FC2 rows -``` - -For gpt-oss-120b (`H = I = 2880`): `I_pad = 2944`, `H1_pad = 3072`, -`H2_pad = 2944`. - -### Certified arguments - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `routing_logits` | `[T, num_experts]` or `None` | bf16 or fp32 (bitwise identical results) | contiguous | CUDA | -| `routing_bias` | must be `None` — see *Notes* | — | — | — | -| `hidden_states` | `[T, H1_pad]` | **float8_e4m3fn** | contiguous | CUDA | -| `hidden_states_scale` | **1-D**, exactly `T * H1_pad/32` elements | **uint8** (E8M0) | contiguous, **linear** (row-major `[T, H1_pad/32]` flattened) — *not* swizzled | CUDA | -| `gemm1_weights` | `[E, 2*I_pad, H1_pad/2]` | **uint8** (2 E2M1 codes per byte) | contiguous, pre-shuffled | CUDA | -| `gemm1_weights_scale` | `[E, 2*I_pad, H1_pad/32]` | **uint8** (E8M0) | contiguous, pre-shuffled + swizzled | CUDA | -| `gemm1_bias` | `[E, 2*I_pad]` or `None` | **fp32** | contiguous, pre-shuffled | CUDA | -| `gemm1_alpha` / `gemm1_beta` / `gemm1_clamp_limit` | `[E]` or `None` | **fp32** | contiguous | CUDA | -| `gemm2_weights` | `[E, H2_pad, I_pad/2]` | **uint8** | contiguous, pre-shuffled | CUDA | -| `gemm2_weights_scale` | `[E, H2_pad, I_pad/32]` | **uint8** (E8M0) | contiguous, pre-shuffled + swizzled | CUDA | -| `gemm2_bias` | `[E, H2_pad]` or `None` | **fp32** | contiguous, pre-shuffled | CUDA | -| `num_experts` | scalar | int, the global routing space; `> top_k` | — | — | -| `top_k` | scalar | int, `0 < top_k < num_experts` | — | — | -| `n_group` / `topk_group` | `None` | — | — | — | -| `intermediate_size` | scalar | int, **must equal `I_pad`** | — | — | -| `valid_hidden_size` | scalar (not `None`) | int, `H`; multiple of 32 with `pad_up(H, 128) == H2_pad` | — | — | -| `valid_intermediate_size` | scalar or `None` | int, multiple of 32, `<= I_pad` | — | — | -| `local_expert_offset` / `local_num_experts` | scalars | int, `E = local_num_experts >= 1` | — | — | -| `routed_scaling_factor` | `None` | — | — | — | -| `routing_method_type` | scalar | int: **0** (Default) or **1** (Renormalize) | — | — | -| `act_type` | scalar | int: **0** (SwiGlu) only | — | — | -| `topk_weights` | `[T, top_k]` or `None` | **bf16** | contiguous | CUDA | -| `topk_ids` | `[T, top_k]` or `None` | **int32** | contiguous | CUDA | -| `output` | `[T, valid_hidden_size]` or `None` | bf16 | contiguous | CUDA | -| return | `[T, valid_hidden_size]`, or `[0]` when `output` is given | bf16 | contiguous, newly allocated | same as `hidden_states` | - -Biases are independently optional: `gemm1_bias` and `gemm2_bias` may each be -`None`, in any combination. - -### Arguments held inert (not certified) - -- `routing_bias` — see *Notes*; the wrapper rejects it on every path where it - was observed to be a no-op. -- `n_group` / `topk_group` — the grouped (DeepSeek-V3 style) routing path, - reachable only with `routing_method_type = 2`. -- `routing_method_type` values other than `0` and `1`. `4` (RenormalizeNaive) - and `6` (SigmoidRenorm) do run on the logits entry point here, but their - weight formulas are not pinned; `2` (DeepSeekV3), `3` (Llama4) and - `5` (MiniMax2) were not exercised on that entry point at all. On the - pre-routed entry point every one of `0/1/2/4/5/6` is inert. -- `routed_scaling_factor` — changed nothing on either certified entry point - (`None`, `1.0` and `2.5` give bitwise identical results); it belongs to the - grouped-routing path. -- `tune_max_num_tokens` — an autotuner bucket cap; `8192` and `128` give - bitwise identical results. Left at its default. -- `use_dp` — an autotuner token-bucket deflation hint for data-parallel - deployments; `True`/`False` give bitwise identical results at `ep_size = 1`. - Left at `False`. -- Multi-rank tensor parallelism. `local_expert_offset` / `local_num_experts` - (expert parallelism) **are** certified; TP, which instead splits `I` across - ranks and reduces this call's output afterwards, is not exercised here. - -## Metadata consumed - -None. The op reads no attention metadata, no KV cache, no registered layer -and no module state — every tensor and every scalar it uses is an argument. - -Two process-global caches sit behind it, neither of which changes the result: - -- an **autotuner profiling cache** that picks the two GEMM tactics. Cold — - the state every call certified here ran in — it takes a fallback tactic; - tactic choice is a performance knob only. -- a cache of C++ `MxE4m3MxE2m1BlockScaleMoERunner` objects keyed by - `act_type`, each owning a GPU workspace allocated on first use. - -## Preconditions - -### Activation preparation — MXFP8, linear scales, padded hidden - -`hidden_states` / `hidden_states_scale` are exactly the two outputs of - -```python -# swizzled_layout alignment -data, sf = torch.ops.trtllm.mxfp8_quantize(x, False, 512) -``` - -on the bf16 `[T, H]` hidden states, and that pairing is certified: for -`H = 2880` it yields `data [T, 3072] float8_e4m3fn` and a 1-D -`sf [T * 96] uint8`, which is what this op wants. Concretely, whatever -produces them: - -- `hidden_states` is 2-D, **float8_e4m3fn** (a uint8 view of the same bytes is - rejected: `hidden_states must be Float8_e4m3fn`), contiguous, CUDA, with - width exactly `gemm1_weights.shape[-1] * 2 == H1_pad`. `alignment` must - therefore be the FC1 K alignment (512 for this weight family); quantizing - with `alignment=32` leaves the tensor at width `H` and is rejected (`the - third dimension of weights must be equal to hidden_size`). The op never - pads the hidden itself — the widening `2880 -> 3072` happens inside - `mxfp8_quantize`, and the caller does no `F.pad` at all. -- `hidden_states_scale` is **1-D** (`hidden_states_scale must be 1D`), - **uint8** (`must be UInt8`), contiguous, CUDA, and holds exactly - `T * H1_pad/32` bytes (`hidden_states_scale has incorrect size` otherwise — - one byte too few or too many both raise). Byte `t * (H1_pad/32) + b` is the - E8M0 scale (`value = 2^(byte - 127)`) of columns `[32b, 32b+32)` of token - `t`. The **128x4 swizzled** scale order is a silent wrong answer, not an - error, whenever it happens to have the same byte count (`T % 128 == 0`) — - see *Notes*. -- Columns `[H, H1_pad)` multiply zero-valued padded weights; `mxfp8_quantize` - writes zero data bytes and zero scale bytes there, but any **finite** - content is inert (verified with e4m3 `400.0` under a live scale byte). A - NaN there propagates into the output. -- `T >= 1`; `T = 0` is rejected. Certified `T`: 1, 2, 4, 6, 8, 10, 12, 16, 17, - 24, 32, 128, 256, 1024, 8192. - -### Weight preparation — the caller owns all of it - -The activation dtype changes none of it: this is the same prepared weight -layout `torch.ops.trtllm.bf16_mxe2m1_block_scale_moe_runner` consumes, and in -this build both ops are fed by the same trtllm weight-loading class -(`MXFP4WeightTRTLLMGenFusedMoEMethod`, which the W4A16 and W4A8-MXFP8 methods -inherit without overriding a single padding, shuffle or swizzle step). One -prepared expert stack therefore serves both. - -Start from the checkpoint's per-expert MXFP4 tensors, laid out **row-major -over the output dim**, K packed two E2M1 codes per byte with the **low nibble -holding the even K index**, and one E8M0 byte (`value = 2^(byte - 127)`) per -32 consecutive K elements: - -``` -up (w3 / up_proj) packed [E, I, H/2] uint8 scale [E, I, H/32] uint8 -gate (w1 / gate_proj) packed [E, I, H/2] uint8 scale [E, I, H/32] uint8 -down (w2 / down_proj) packed [E, H, I/2] uint8 scale [E, H, I/32] uint8 -up/gate bias [E, I] , down bias [E, H] -> both must reach the kernel as fp32 -``` - -(Checkpoints that store gate and up interleaved on one `2I` axis must -de-interleave first.) Then, **per expert**: - -1. **Zero-pad.** `up`/`gate`: rows `I -> I_pad`, byte columns - `H/2 -> H1_pad/2`, scale columns `H/32 -> H1_pad/32`. `down`: rows - `H -> H2_pad`, byte columns `I/2 -> I_pad/2`, scale columns - `I/32 -> I_pad/32`. Biases: `I -> I_pad` and `H -> H2_pad`. -2. **FC1 concat.** Stack the padded halves on the row axis as - `[up ; gate]`, giving `[2*I_pad, ...]`. Same for the bias. -3. **FC1 row permute** = interleave, then block shuffle: - - interleave: destination row `2i` takes the up half's row `i`, - destination row `2i+1` takes the gate half's row `i`; - - block shuffle (also applied to FC2, which skips step 2/3's interleave): - within each aligned block of 32 rows, source row `4u + v` - (`0 <= u < 8`, `0 <= v < 4`) moves to destination row `8v + u`. - - Apply the **same** permutation to the weight bytes, the scale bytes and - the fp32 bias, so row `i` of all three still describe the same output - channel. A bias not permuted with its weights is silently wrong. -4. **Scale swizzle.** After the row permute, each expert's scale matrix - `[M, C]` (`M % 128 == 0`, `C % 4 == 0`) is rewritten into the trtllm-gen - 128x4 layout: the byte at `(m, c)` moves to flat offset - - ``` - (m // 128) * 512 * (C // 4) + (c // 4) * 512 - + (m % 32) * 16 + ((m % 128) // 32) * 4 + (c % 4) - ``` - - and the flat result is handed to the kernel with its nominal `[E, M, C]` - shape. Weight bytes and biases are **not** swizzled — only scales. - -The two trtllm helpers `torch.ops.trtllm.shuffle_matrix(x, perm)` (a plain -row gather, `out[i] = x[perm[i]]`) and -`torch.ops.trtllm.block_scale_interleave(x)` (the 128x4 swizzle over a -`[E, M, C]` uint8 tensor, returning a flat buffer of -`E * pad_up(M, 128) * pad_up(C, 4)` bytes) produce byte-identical results to -steps 3 and 4; this entry's test asserts that equivalence. - -### Shapes and layout - -- `intermediate_size` must equal `gemm1_weights.shape[1] // 2`; anything else - is rejected (`No valid config found for the given problem shape`). -- `valid_hidden_size` is the **output width**, unrelated to the widened - activation: it stays at the model's true hidden `H` (2880 for gpt-oss) even - though `hidden_states` is `H1_pad` (3072) wide. It must be a multiple of 32 - satisfying `pad_up(valid_hidden_size, 128) == gemm2_weights.shape[1]`. - `None` means "use the `hidden_states` width", which now fails whenever - `H1_pad != H2_pad` (`gemm2_weights_scale has incorrect dim 1`) — so on this - op it must always be passed explicitly. Setting it to `H2_pad` yields the - wider output whose tail columns are all zero; other values are rejected - (`gemm2_weights_scale has incorrect dim 1` or `No valid config found`). -- `valid_intermediate_size` must be a multiple of 32 and at most - `intermediate_size`; `None` means `intermediate_size`. It is a bandwidth - hint: the kernel only reads the first `valid_intermediate_size` intermediate - columns. Since padded columns carry zero weight, every value `>= I` gives - the same result — but a **smaller** value silently truncates the layer. -- `topk_ids` is `[T, top_k]` **int32** contiguous CUDA and `topk_weights` is - `[T, top_k]` **bf16** contiguous CUDA; the two must be given together - (`routing_logits or (topk_ids and topk_weights) must be provided`), their - row count must match `hidden_states` and their column count must match - `top_k`. int64 ids and fp32 weights are both rejected. -- `routing_logits` is `[T, num_experts]`, bf16 or fp32, contiguous CUDA. A - column count other than `num_experts` is rejected. -- `gemm1_alpha` / `gemm1_beta` / `gemm1_clamp_limit`, when given, are fp32 - CUDA tensors with exactly **`local_num_experts`** elements (not - `num_experts`), indexed by local expert. -- Both bias tensors, when given, must be **fp32** (bf16 is rejected); both - weight-scale tensors must be **uint8** (an int8 view is rejected). -- `output`, when given, must be a contiguous CUDA bf16 tensor of shape - exactly `[T, valid_hidden_size]`; wrong dtype and wrong shape are both - rejected. A leading row-slice of a taller contiguous buffer is valid — rows - past `T` are left bitwise untouched. -- **Every tensor argument must be contiguous.** A strided view is accepted - silently and read as if dense — see *Notes*. The wrapper asserts this. - -### Sizes and counts - -- `0 < top_k < num_experts` (`num_experts must be greater than top_k`; - `top_k = 0` is rejected). Certified `top_k`: 1, 2, 3, 4. -- `num_experts` is only the routing space; it need not equal - `local_num_experts` (certified at `num_experts = 8`, - `local_num_experts = 4`, `local_expert_offset = 4`). Certified - `num_experts`: 2, 3, 4, 5, 8, 16, 128; certified `local_num_experts`: - 2, 3, 4, 5, 8, 16, 128. -- Certified geometries `(H, I)`: (512, 128), (512, 256), (512, 512), - (640, 128), (1024, 512), (2048, 512), (2880, 1024), (2880, 2880). `H` and - `I` need not be multiples of the kernel's alignments — that is what the - padding is for — but both must be multiples of 32, since they are passed as - `valid_hidden_size` / `valid_intermediate_size` (a `valid_hidden_size` of - 500 was rejected with `No valid config found`). -- `local_expert_offset + local_num_experts > num_experts` is **not** checked. -- Expert ids outside `[local_expert_offset, local_expert_offset + - local_num_experts)` — including negative ids and ids `>= num_experts` — are - silently dropped; the token's other slots still combine normally. -- A repeated expert id inside one token's row contributes **once**, with the - weight of its **first** occurrence; the duplicate slot's weight is - discarded. - -### Dtypes and modes - -- `hidden_states` must be float8_e4m3fn; bf16 is rejected - (`hidden_states must be Float8_e4m3fn`). bf16 hidden states belong to the - sibling op `torch.ops.trtllm.bf16_mxe2m1_block_scale_moe_runner`, whose - signature has no `hidden_states_scale` parameter at all. -- `act_type` must be `0`. `1` (Relu2) and `2` (Silu) fail with - `No kernel found for the given options: mDtypeA: MxE4m3, mDtypeB: MxE2m1 - ...` — the non-gated activations have no cubin in this MXFP8 x MXFP4 family. - `SwigluBias` is not a separate value: the per-expert `alpha`/`beta`/ - `clamp_limit` tensors turn `act_type = 0` into it. -- `hidden_states` must be 2-D; a 3-D `[1, T, H1_pad]` view is rejected. -- sm_100 only. The receipt covers sm_100 (B200), the only arch available - here; the installed build carries the assertion `Only SM100f is supported - by MXFP4 block scale MOE` in this op's C++ entry point, so other Blackwell - variants are expected to raise rather than compute something wrong. - -A caller violating none of the above gets the result described under -*Semantics*, to within the bound in *Notes*. - -## Notes - -- **Numerics.** Against a native-torch reference that consumes bit-identical - operands (exact MXFP4 weight and MXFP8 activation dequantization, fp32 GEMM - accumulation, the MXFP8 requantization of the FC1 output modelled exactly, - fp32 combine), the kernel's worst element-wise deviation over every - configuration in this entry's test was **2.0 ulp of the token row's largest - magnitude** (bf16 ulp = 2^-8, worst case: 128 experts, 8192 tokens, - `H = I = 2880`) and its worst relative RMS deviation **0.87 ulp**. Dropping - the intermediate requantization from that reference — i.e. modelling this - layer the way a bf16-intermediate MoE would be modelled — costs a factor of - seven: 13.5 ulp element-wise and 8.5 ulp RMS (~3% relative) at the gpt-oss - geometry. That gap is a property of the kernel, not of the test: it is what - the extra e4m3 rounding between the two GEMMs does to the layer's output. -- **`routing_bias` is silently ignored** for `routing_method_type` 0, 1, 4 - and 6, and on the pre-routed entry point for every routing method: a bias - of `+1e3 / -1e3` on two experts leaves the output bitwise unchanged. - A checkpoint whose router carries a bias (gpt-oss does) must therefore have - it **added into `routing_logits`** before the call — that is, use the - router `nn.Linear`'s own bias and pass `routing_bias=None`. The wrapper - asserts `routing_bias is None` on exactly those paths, because the failure - is otherwise invisible: expert selection quietly ignores the bias and the - model degrades without any error. The argument is presumably live for the - grouped (`routing_method_type = 2`) path, which is not certified here. -- **A swizzled activation-scale buffer is a silent wrong answer.** - `mxfp8_quantize(x, True, 512)` returns the same data bytes and a scale - buffer of `pad_up(T,128) * pad_up(cols,4)` bytes in 128x4 order. Whenever - that count coincides with the linear one (`T` a multiple of 128, `cols` a - multiple of 4 — true for gpt-oss at any graph-friendly batch size) the size - check passes and the kernel reads scales for the wrong blocks: measured - 260 ulp element-wise off at `T = 128`, `H = I = 2880`. Nothing in the - metadata distinguishes the two layouts, so no guard can catch it — pass - `swizzled_layout=False`. -- **Non-contiguous tensors are a silent wrong answer.** The kernel takes raw - data pointers and assumes a dense row-major layout. A strided view of - `hidden_states` (measured 479 ulp off), of `hidden_states_scale` (311 ulp, - using a same-length stride-2 1-D view), of any of the six - weight/scale/bias tensors, of `routing_logits`, `topk_weights`, `topk_ids` - or `output` is accepted without complaint and reads (or writes) the wrong - elements. The wrapper asserts contiguity on every tensor argument. -- **The FC1 accumulation precision is not pinned by this entry**, and cannot - be: the epilogue's e4m3 requantization (3 mantissa bits) erases any - difference between an fp32 and a bf16 intermediate accumulator. What *is* - pinned is the value FC2 consumes — bit-exactly, per the formula in - *Semantics*. -- Sibling ops in this build address the same job with other operand dtypes: - `torch.ops.trtllm.bf16_mxe2m1_block_scale_moe_runner` (bf16 activations, - same prepared weight layout, no `hidden_states_scale` parameter), - `torch.ops.trtllm.fp4_block_scale_moe_runner` (NVFP4) and - `torch.ops.trtllm.fp8_block_scale_moe_runner`. Catalog membership is - `index.yaml`'s fact alone. - - -## The FC1 epilogue's block-scale recipe is architecture-specific - -trtllm-gen ships one cubin per architecture, and the two differ **bit-exactly** -in how the FC1 epilogue picks the e8m0 scale when it requantizes its activation -output to MXFP8 for FC2. Measured with an identity down-projection, which reads -the intermediate out element by element rather than inferring it from output -noise (`test_intermediate_is_mxfp8_quantized`): - -| Arch | e8m0 exponent | Name | -|---|---|---| -| sm_100 | `floor(log2(amax)) - 8` | OCP scale | -| sm_103 | `ceil(log2(amax / 448))` | round-up scale | - -Both were verified bit-exact on their own architecture (0 mismatched elements) -and each *refutes* the other's, so the cases genuinely separate the recipes. -The round-up form is what `torch.ops.trtllm.mxfp8_quantize` has always used, so -sm_103 brings the MoE epilogue into agreement with the standalone quantizer. - -This is the whole of the sm_100 -> sm_103 numerical difference for this entry. -Before the reference was made architecture-aware, 10 of 19 cells failed -- -relative RMS ~5.5 ulp against a 4 ulp gate, max abs ~0.07 against 0.031. With -the correct recipe in force all 19 pass **at the original tolerances**, which -is what identifies the recipe as the sole cause rather than one contributor. - -A future cubin that changes recipe again will fail the bit-exact test rather -than drift quietly; record the new recipe in `_SCALE_RECIPE_BY_SM`, and do not -widen a tolerance instead. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/mxe4m3_mxe2m1_block_scale_moe_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/mxe4m3_mxe2m1_block_scale_moe_runner.py index 52e3056cdb35..9f9ef4d402ea 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/mxe4m3_mxe2m1_block_scale_moe_runner.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/mxe4m3_mxe2m1_block_scale_moe_runner.py @@ -4,126 +4,531 @@ top-k) + grouped FC1 GEMM + clamped gated activation + MXFP8 requantization + grouped FC2 GEMM + routing-weighted combine, in one trtllm-gen call.""" -from typing import Optional +from __future__ import annotations + +from typing import Optional, Tuple import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -# routing_method_type values for which the kernel never reads routing_bias: -# 0 Default, 1 Renormalize, 4 RenormalizeNaive, 6 SigmoidRenorm. Observed on -# this machine: a +-1e3 bias leaves the result bitwise unchanged. +from .._op import Arch, Cell, OpWrapper + +#: routing_method_type values for which the kernel never reads routing_bias: +#: 0 Default, 1 Renormalize, 4 RenormalizeNaive, 6 SigmoidRenorm. Observed on +#: this machine: a +-1e3 bias leaves the result bitwise unchanged. _ROUTING_BIAS_IGNORED = (0, 1, 4, 6) +#: routing_method_type the gpt-oss target passes, and `act_type` for SwiGlu. +ROUTING_METHOD_RENORMALIZE = 1 +ACT_TYPE_SWIGLU = 0 + +#: The clamped gated activation's fixed constants: `(up + beta) * gate * +#: sigmoid(alpha * gate)` after `gate.clamp(max=limit)`, `up.clamp(+-limit)`. +#: gpt-oss carries these two in its HF reference rather than its checkpoint -- +#: every expert on every layer uses the same alpha and beta -- which is why +#: `bind_glu` below builds them from a count rather than taking them as +#: arguments: there is nothing a caller could supply that varies. +GLU_ALPHA = 1.702 +GLU_BETA = 1.0 + +#: MX block size: one e8m0 scale per 32 elements along K. +SF_VEC_SIZE = 32 + +#: Largest finite e4m3 magnitude. +E4M3_MAX = 448.0 + +#: The 16 e2m1 code points in code order. +_E2M1_VALUES = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) + +#: One expert stack's `(codes, e8m0 exponents, bias)`. +Operand = Tuple[torch.Tensor, torch.Tensor, torch.Tensor] -def mxe4m3_mxe2m1_block_scale_moe_runner( - routing_logits: Optional[torch.Tensor], - routing_bias: Optional[torch.Tensor], - hidden_states: torch.Tensor, - hidden_states_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, - gemm1_bias: Optional[torch.Tensor], - gemm1_alpha: Optional[torch.Tensor], - gemm1_beta: Optional[torch.Tensor], - gemm1_clamp_limit: Optional[torch.Tensor], - gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, - gemm2_bias: Optional[torch.Tensor], - num_experts: int, - top_k: int, - n_group: Optional[int], - topk_group: Optional[int], - intermediate_size: int, - valid_hidden_size: Optional[int], - valid_intermediate_size: Optional[int], - local_expert_offset: int, - local_num_experts: int, - routed_scaling_factor: Optional[float], - routing_method_type: int, - act_type: int, - topk_weights: Optional[torch.Tensor] = None, - topk_ids: Optional[torch.Tensor] = None, - output: Optional[torch.Tensor] = None, - tune_max_num_tokens: int = 8192, - use_dp: bool = False, -) -> torch.Tensor: - """Run one MXFP4-weight MoE layer over MXFP8 (e4m3 + UE8M0) activations. - - Returns a fresh `[num_tokens, valid_hidden_size]` bf16 tensor, or — when - `output` is given — an empty `[0]` tensor, the result having been written - into `output`. +#: The DeepSeek-V4.1 DSpark drafter's routed experts: 128, all local, top-3, +#: H = 5120 and I = 2304 (neither padded), the clamp at 10 per expert, no +#: alpha / beta, no biases; pre-routed, into a caller `output`, inside the +#: autotuner's `fallback_tactic()` scope. +_DSPARK = dict( + hidden=5120, + intermediate=2304, + num_experts=128, + top_k=3, + swiglu_limit=10.0, + model="deepseek-v4.1", + routed_by_kernel=False, + local_expert_offset=0, + local_num_experts=128, + pinned=True, + output=True, +) + +#: The drafter's row counts: the ends of the eager sweep 1..1280 (T = 5G for up +#: to 256 drafting requests), the first draft block, both sides of the two +#: token-tile changes of the fallback tactic, and the top of the eager column -- +#: which was run as a plain call (no pin, fresh `output`), so its cell is one. +_DSPARK_CELLS = tuple( + Cell(why=f"deepseek-v4.1 DSpark drafter T={tokens}: {what}", spec=dict(_DSPARK, tokens=tokens)) + for tokens, what in ( + (1, "bottom of the sweep; FC1 static scheduler, FC2 split-K 2"), + (5, "one request's draft block, the smallest captured T"), + (383, "top of the 8-token tile"), + (384, "bottom of the 16-token tile"), + (725, "top of the 16-token tile"), + (726, "bottom of the 32-token tile"), + (1280, "top of the sweep: 256 requests x 5"), + ) +) + ( + Cell( + why="deepseek-v4.1 DSpark drafter T=8192: top of the eager column, unpinned, fresh output", + spec=dict(_DSPARK, tokens=8192, pinned=False, output=False), + ), +) + +#: The FC1 epilogue's block-scale recipe is architecture-specific, and the +#: difference is bit-exact rather than a tolerance: trtllm-gen ships one cubin +#: per architecture. Measured on each with an identity down-projection reading +#: the intermediate out element by element. +#: +#: sm_100 e8m0 = floor(log2(amax)) - 8 "OCP scale" +#: sm_103 e8m0 = ceil(log2(amax / 448)) "round-up scale" +#: +#: The round-up form is the one ``torch.ops.trtllm.mxfp8_quantize`` has always +#: used, so sm_103 makes the MoE epilogue and the standalone quantizer agree. +#: Both are named so that a future cubin switching back cannot pass silently. +OCP_SCALE, ROUND_UP_SCALE = "ocp", "round_up" +SCALE_RECIPE_BY_SM = {(10, 0): OCP_SCALE, (10, 3): ROUND_UP_SCALE} + + +def scale_recipe() -> str: + """The block-scale recipe this device's FC1 epilogue uses. + + Refuses an architecture nobody measured rather than guessing one of the + two: the difference is bit-exact, so a wrong guess is a reference that + disagrees with a correct kernel by ~10 ulp RMS and reads as a kernel bug. """ - # Pure-metadata guard: the kernel takes raw data pointers and assumes a - # dense row-major layout for every tensor. A strided view is accepted - # without complaint and silently reads the wrong elements (observed on - # this machine for every tensor argument listed here, hidden_states_scale - # included). - for name, tensor in ( - ("routing_logits", routing_logits), - ("routing_bias", routing_bias), - ("hidden_states", hidden_states), - ("hidden_states_scale", hidden_states_scale), - ("gemm1_weights", gemm1_weights), - ("gemm1_weights_scale", gemm1_weights_scale), - ("gemm1_bias", gemm1_bias), - ("gemm1_alpha", gemm1_alpha), - ("gemm1_beta", gemm1_beta), - ("gemm1_clamp_limit", gemm1_clamp_limit), - ("gemm2_weights", gemm2_weights), - ("gemm2_weights_scale", gemm2_weights_scale), - ("gemm2_bias", gemm2_bias), - ("topk_weights", topk_weights), - ("topk_ids", topk_ids), - ("output", output), - ): - if tensor is not None: - assert tensor.is_contiguous(), ( - f"{name} must be contiguous; a strided view is read as if dense " - "and silently produces wrong results" - ) - # Pure-metadata guard: routing_bias is a no-op for every non-grouped - # routing method, and for any routing method once topk_ids/topk_weights - # carry the routing. A caller expecting the bias to shift expert selection - # gets a silently different model. - if routing_bias is not None: - assert topk_ids is None, ( - "routing_bias is ignored when topk_ids/topk_weights are given " - "(routing has already happened); fold it into the router logits" + sm = torch.cuda.get_device_capability() + recipe = SCALE_RECIPE_BY_SM.get(sm) + assert recipe is not None, ( + f"the FC1 epilogue's requantization recipe is not certified on sm_{sm[0]}{sm[1]}; " + "run the identity-down-projection probe and record it before trusting this entry" + ) + return recipe + + +def _block_scale(amax: torch.Tensor, recipe: str) -> torch.Tensor: + """The per-32-column e8m0 scale, under the named recipe.""" + if recipe == OCP_SCALE: + exp = torch.floor(torch.log2(amax)) - 8.0 + else: + exp = torch.ceil(torch.log2(amax / E4M3_MAX)) + exp = torch.where(amax == 0, torch.full_like(amax, -127.0), exp) + return torch.exp2(exp.clamp(-127.0, 127.0)) + + +def quantize_mxfp8(act: torch.Tensor, recipe: Optional[str] = None) -> torch.Tensor: + """The MXFP8 requantization the FC1 epilogue applies to its own output. + + Per 32 consecutive intermediate columns: this architecture's block scale, + then round to nearest even into e4m3 saturating at +-448. Returns the + dequantized fp32 value FC2 actually consumes, which is what a reference has + to model -- skipping it lands ~10 ulp RMS away from a correct kernel. + """ + rows, cols = act.shape + blocks = act.reshape(rows, cols // SF_VEC_SIZE, SF_VEC_SIZE) + amax = blocks.abs().amax(dim=-1, keepdim=True) + scale = _block_scale(amax, recipe or scale_recipe()) + q = (blocks / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn).float() + return (q * scale).reshape(rows, cols) + + +def dequantize_mxfp4(codes: torch.Tensor, exponents: torch.Tensor) -> torch.Tensor: + """One expert's `[N, K]` fp32 weight from its e2m1 codes and e8m0 exponents.""" + table = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=codes.device) + scale = torch.exp2(exponents.float() - 127.0).repeat_interleave(SF_VEC_SIZE, dim=1) + return table[codes.long()] * scale + + +class Mxe4m3Mxe2m1BlockScaleMoeRunner(OpWrapper): + """One MXFP4-weight MoE layer over MXFP8 (e4m3 + UE8M0) activations, with + gpt-oss's clamped gated activation. + + Unlike its NVFP4 sibling this op will route for you: given `routing_logits` + it computes top-k itself, and given `topk_ids` / `topk_weights` it takes the + caller's. The gpt-oss target passes logits and lets the kernel route. + + `valid_hidden_size` and `valid_intermediate_size` are the *unpadded* widths. + Both GEMMs want their rows aligned -- 128 on each, and 512 more on FC1's K + axis -- so the operands are padded and these two say how much of the result + is real. The returned rows are `valid_hidden_size` wide. + + Returns a fresh `[num_tokens, valid_hidden_size]` bf16 tensor, or an empty + `[0]` tensor when `output` was given. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # gpt-oss-120b's expert geometry: 128 experts, H = I = 2880, top-4, and the + # clamped SwiGlu constants the checkpoint carries (alpha 1.702, beta 1.0, + # clamp limit 7.0). The target runs tp1, so every expert is local. + _GPT_OSS = dict(hidden=2880, intermediate=2880, num_experts=128, top_k=4) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="the gpt-oss target's own call: router entry point, all 128 experts local", + spec=dict( + tokens=64, + routed_by_kernel=True, + local_expert_offset=0, + local_num_experts=128, + **_GPT_OSS, + ), + ), + Cell( + why="one token -- the decode row count, where the routing has a single row to place", + spec=dict( + tokens=1, + routed_by_kernel=True, + local_expert_offset=0, + local_num_experts=128, + **_GPT_OSS, + ), + ), + Cell( + why="pre-routed with an expert window: the caller owns top-k and this rank owns 32 ids", + spec=dict( + tokens=64, + routed_by_kernel=False, + local_expert_offset=32, + local_num_experts=32, + **_GPT_OSS, + ), + ), + ) + _DSPARK_CELLS + + note = """ + `act_type` 0 (SwiGlu) only, with the clamp: `gemm1_alpha`, `gemm1_beta` and + `gemm1_clamp_limit` are per-local-expert fp32 vectors, and gpt-oss sets them + to 1.702 / 1.0 / 7.0. Passing `None` for all three is the unclamped form, + which this entry does not certify against the target. + + The FC1 epilogue's requantization recipe is per architecture and the + difference is bit-exact, not a tolerance -- see `scale_recipe`. That is the + reason `ARCHS` is load-bearing on this entry rather than bookkeeping. + + `tune_max_num_tokens` is an autotuner bucket hint and was measured inert; it + is not a token count. + + The caller owns the weight layout: both GEMMs' weights are pre-shuffled and + their scales 128x4 swizzled, padded up to the alignments above. Nothing here + checks it, and the entry's test builds the layout twice -- once in torch, + once through the trtllm preprocessing ops -- and compares them first. + + DeepSeek-V4.1: the DSpark drafter calls this once per stage on its draft-block + rows, `T = 5G` for the `G` drafting requests on the rank (up to 1280 at + `max_batch_size` 256; never chunked, never gathered): 128 experts all local, + top-3, H = 5120, I = 2304 (`valid_*` equal to the padded widths -- nothing is + padded), `gemm1_clamp_limit` 10 per expert with `gemm1_alpha`, `gemm1_beta` + and both biases `None` (so the activation is `silu(min(gate, 10)) * + clamp(up, -10, 10)`), pre-routed bf16 weights, `routing_method_type` 1, + `hidden_states` from `mxfp8_quantize(x, False, 512)` (linear scales), and + `output=` a `[T, 5120]` bf16 buffer. The call is made inside the autotuner's + `fallback_tactic()` scope: tactic choice changes output bits (in the gates, + but not bitwise), and inside that scope every rank takes the fallback whatever + its cache holds -- `autotune(tune_mode=False)` alone still takes a cached + tactic and is not certified here. The fallback's token tile is 8 up to 383, 16 + at 384..725 and 32 from 726; the cells hold those edges and the sweep's ends, + and the entry's test runs every T = 1..1280 eagerly (kernel map, skewed + routing at every run's edges, bitwise repeatability, rows independent of T + from T = 2) and captures and replays the pinned call at every `T = 5G` over + the 46 captured batch sizes. The routing weight multiplies FC2's output, so a + reference in the checkpoint's own order (weight into the activation before + FC2's MXFP8 quantization) is a different computation, about 10 ulp RMS away. + Not certified for V4.1: `T` above 1280 except 8192 (no capture there), an + unpinned call, an expert window, alpha / beta or biases at this geometry. + """ + + def bind_glu(self, num_experts: int, clamp_limit: float, device: torch.device) -> None: + """Bind gpt-oss's clamped-GLU epilogue vectors, via `bind_const`. + + `gemm1_alpha` and `gemm1_beta` are this op's own fixture: a fixed pair + the checkpoint never carries, so the op manufactures the per-expert + vectors itself from `GLU_ALPHA`/`GLU_BETA` rather than taking them as + tensors a caller built. `clamp_limit` does come from the checkpoint -- + gpt-oss's `swiglu_limit` -- but is still a scalar here: every local + expert shares the one value, and broadcasting it into the per-expert + `gemm1_clamp_limit` vector `raw_call` wants is this op's business, not + a caller's. + """ + self.bind_const( + gemm1_alpha=torch.full((num_experts,), GLU_ALPHA, dtype=torch.float32, device=device), + gemm1_beta=torch.full((num_experts,), GLU_BETA, dtype=torch.float32, device=device), + gemm1_clamp_limit=torch.full( + (num_experts,), float(clamp_limit), dtype=torch.float32, device=device + ), ) - assert routing_method_type not in _ROUTING_BIAS_IGNORED, ( - f"routing_bias is silently ignored for routing_method_type=" - f"{routing_method_type}; add it to routing_logits before the call" + + def raw_call( + self, + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + valid_hidden_size: Optional[int], + valid_intermediate_size: Optional[int], + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + act_type: int, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + ) -> torch.Tensor: + """Run the layer. Returns a fresh `[num_tokens, valid_hidden_size]` + bf16 tensor, or an empty `[0]` tensor when `output` was given.""" + return torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner( + routing_logits, + routing_bias, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + gemm1_bias, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, + gemm2_weights, + gemm2_weights_scale, + gemm2_bias, + num_experts, + top_k, + n_group, + topk_group, + intermediate_size, + valid_hidden_size, + valid_intermediate_size, + local_expert_offset, + local_num_experts, + routed_scaling_factor, + routing_method_type, + act_type, + topk_weights=topk_weights, + topk_ids=topk_ids, + output=output, + tune_max_num_tokens=tune_max_num_tokens, + use_dp=use_dp, ) - return torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner( - routing_logits, - routing_bias, - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - gemm1_bias, - gemm1_alpha, - gemm1_beta, - gemm1_clamp_limit, - gemm2_weights, - gemm2_weights_scale, - gemm2_bias, - num_experts, - top_k, - n_group, - topk_group, - intermediate_size, - valid_hidden_size, - valid_intermediate_size, - local_expert_offset, - local_num_experts, - routed_scaling_factor, - routing_method_type, - act_type, - topk_weights=topk_weights, - topk_ids=topk_ids, - output=output, - tune_max_num_tokens=tune_max_num_tokens, - use_dp=use_dp, - ) + + def reference( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + up: Operand, + gate: Operand, + down: Operand, + alpha: Optional[torch.Tensor] = None, + beta: Optional[torch.Tensor] = None, + clamp_limit: Optional[torch.Tensor] = None, + local_expert_offset: int = 0, + local_num_experts: Optional[int] = None, + ) -> torch.Tensor: + """The layer, in plain torch over dequantized weights, fp32 throughout. + + The documented exception to mirroring: `raw_call` takes the *prepared* + operands -- packed e2m1 codes, shuffled and padded rows, swizzled + scales -- and a reference given those would be re-deriving the layout + rather than the computation. So this takes the per-expert + `(codes, e8m0 exponents, bias)` triples the layout was built from, at + their valid (unpadded) widths, plus the routing the call resolved to. + + `topk_ids` carries *global* ids; this rank contributes only for the ids + its window owns and zero elsewhere, which is what makes an + expert-parallel split summable. + + The activation is gpt-oss's clamped GLU: + `(up + beta) * gate * sigmoid(alpha * gate)`, with `gate` clamped above + at `clamp_limit` and `up` clamped to both sides -- an asymmetry, not a + typo. Its output is requantized to MXFP8 before FC2 because the kernel + does; skipping that is ~10 ulp RMS away from a correct kernel. + """ + num_tokens, hidden = hidden_states.shape + up_c, up_s, up_b = up + gt_c, gt_s, gt_b = gate + dn_c, dn_s, dn_b = down + if local_num_experts is None: + local_num_experts = up_c.shape[0] + xf = hidden_states.float() + out = torch.zeros(num_tokens, hidden, dtype=torch.float32, device=hidden_states.device) + for local_e in range(local_num_experts): + tok, slot = (topk_ids == local_expert_offset + local_e).nonzero(as_tuple=True) + if tok.numel() == 0: + continue + xe = xf[tok] + u = xe @ dequantize_mxfp4(up_c[local_e], up_s[local_e]).t() + up_b[local_e] + g = xe @ dequantize_mxfp4(gt_c[local_e], gt_s[local_e]).t() + gt_b[local_e] + if clamp_limit is not None: + lim = float(clamp_limit[local_e]) + g = g.clamp(max=lim) + u = u.clamp(-lim, lim) + a = 1.0 if alpha is None else float(alpha[local_e]) + b = 0.0 if beta is None else float(beta[local_e]) + act = quantize_mxfp8((u + b) * g * torch.sigmoid(a * g)) + y = act @ dequantize_mxfp4(dn_c[local_e], dn_s[local_e]).t() + dn_b[local_e] + out.index_add_(0, tok, y * topk_weights[tok, slot].float().unsqueeze(1)) + return out + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """8 ulp per element against the row's own scale, 4 ulp relative RMS. + + Kernel and reference consume bit-identical mxfp4 weights and MXFP8 + activations and model the same FC1-output requantization, so they differ + only in accumulation order and in whether a marginal FC1 value rounds to + the same e4m3 code. Torch's default band cannot express that: its bf16 + `atol` of 1e-5 sits three orders of magnitude below one output ulp of a + two-GEMM chain, and per-element `rtol` is meaningless where cancellation + drives `|ref|` to zero. + + Both numbers are measured. Worst over every configuration the entry's + test covers: 2.0 ulp element-wise (128 experts, 8192 tokens, H=I=2880) + and 0.87 ulp RMS -- so this gate sits at 4x and 4.6x the observed + maxima. Both were sized against deliberately wrong computations: a + reference skipping the FC1 requantization lands at 19.6 / 10.6 ulp, and + a gate/up swap or an unshuffled operand past 100. + + The DeepSeek-V4.1 drafter cells take the same gate. Measured on sm_103 at + that geometry (H 5120, I 2304, 128 experts, top-3, clamp 10, no biases): + 1.98 ulp element-wise / 0.71 ulp RMS at worst over every eager T = 1..1280 + and the skewed-routing reruns, 2.20 / 0.71 over the CUDA-graph replays. + Its wrong references -- no clamp, the clamp at 7, gate/up swapped, alpha / + beta 1.702 / 1.0, no FC1 requantization -- all land outside the gate, and + the checkpoint's weight-before-FC2 order sits above 6 ulp RMS. + """ + assert out.dtype == ref.dtype == torch.bfloat16, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + ulp = 2.0**-8 + row = ref.float().abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + torch.testing.assert_close(out.float() / row, ref.float() / row, rtol=0.0, atol=8 * ulp) + rel_rms = ( + (out.float() - ref.float()).pow(2).mean().sqrt() + / ref.float().pow(2).mean().sqrt().clamp_min(1e-9) + ).item() / ulp + assert rel_rms <= 4.0, f"relative RMS {rel_rms:.2f} ulp > 4 ulp" + + def is_valid( + self, + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + valid_hidden_size: Optional[int], + valid_intermediate_size: Optional[int], + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + act_type: int, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + ) -> None: + """The two ways this op takes an input and answers it wrongly. + + Everything else a caller can get wrong here the op rejects loudly -- a + mismatched `intermediate_size`, a `top_k` outside its range, a weight + shape that disagrees with the declared valid widths. Restating those + would only swap their errors for worse ones. + """ + # Pure-metadata guard: the kernel takes raw data pointers and assumes a + # dense row-major layout for every tensor. A strided view is accepted + # without complaint and silently reads the wrong elements (observed on + # this machine for every tensor argument listed here, hidden_states_scale + # included). + for name, tensor in ( + ("routing_logits", routing_logits), + ("routing_bias", routing_bias), + ("hidden_states", hidden_states), + ("hidden_states_scale", hidden_states_scale), + ("gemm1_weights", gemm1_weights), + ("gemm1_weights_scale", gemm1_weights_scale), + ("gemm1_bias", gemm1_bias), + ("gemm1_alpha", gemm1_alpha), + ("gemm1_beta", gemm1_beta), + ("gemm1_clamp_limit", gemm1_clamp_limit), + ("gemm2_weights", gemm2_weights), + ("gemm2_weights_scale", gemm2_weights_scale), + ("gemm2_bias", gemm2_bias), + ("topk_weights", topk_weights), + ("topk_ids", topk_ids), + ("output", output), + ): + if tensor is not None: + assert tensor.is_contiguous(), ( + f"{name} must be contiguous; a strided view is read as if dense " + "and silently produces wrong results" + ) + # Pure-metadata guard: routing_bias is a no-op for every non-grouped + # routing method, and for any routing method once topk_ids/topk_weights + # carry the routing. A caller expecting the bias to shift expert selection + # gets a silently different model. + if routing_bias is not None: + assert topk_ids is None, ( + "routing_bias is ignored when topk_ids/topk_weights are given " + "(routing has already happened); fold it into the router logits" + ) + assert routing_method_type not in _ROUTING_BIAS_IGNORED, ( + f"routing_bias is silently ignored for routing_method_type=" + f"{routing_method_type}; add it to routing_logits before the call" + ) + + +mxe4m3_mxe2m1_block_scale_moe_runner = Mxe4m3Mxe2m1BlockScaleMoeRunner() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/noaux_tc_op.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/noaux_tc_op.md deleted file mode 100644 index a70d6556364c..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/noaux_tc_op.md +++ /dev/null @@ -1,243 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 16} ---- - -# noaux_tc_op - -**Wraps** `torch.ops.trtllm.noaux_tc_op` (one call). - -## Semantics - -DeepSeek-V3 style MoE routing (`noaux_tc` is DeepSeek's `topk_method` name for -auxiliary-loss-free routing with a bias correction term): per-token expert -selection by a *bias-corrected* sigmoid score, with combine weights taken from -the *uncorrected* sigmoid. - -Let `L` be one row of `router_logits` (`num_experts` values) and `b` the -per-expert correction bias. For every token, with the precision of each step -annotated: - -``` -s = sigmoid(L) # fp32; see "The kernel's sigmoid" -c = s + b # fp32; SELECTION SCORE ONLY - -# group stage — active only when n_group > 1 (identity at n_group == 1): -# split c into n_group contiguous groups of num_experts/n_group experts -# gscore[g] = sum of the two largest c inside group g -# keep the topk_group groups with the largest gscore, set c to -inf elsewhere - -ids = indices of the topk largest c, ordered by decreasing c # int32 -w = s[ids] # from BARE s — the bias is NOT in here -tot = sum(w) # fp32 accumulation -w = fp64(w) / (fp64(tot) + eps) * routed_scaling_factor # fp64 -``` - -The call returns `(w, ids)` — **weights first, ids second**. - -The renormalization and the scaling are evaluated in **fp64** — pinned by -observation: an all-fp32 reference disagrees with the kernel on a third of the -elements at `routed_scaling_factor = 2.5`, an fp64 one is bit-exact. Only the -sum `tot` is accumulated in fp32. `eps` is a guard term that keeps a row whose -selected scores are all exactly zero at exactly zero instead of `NaN`; HF uses -`1e-20` and that value reproduces the kernel everywhere, but the constant -itself is not observable from outside (this kernel's sigmoid returns either -exactly `0` or at least `2^-25`, so any `eps <= 1e-20` behaves identically). - -Three consequences that a reader of the argument name `scores` would get -wrong, each verified here: - -- The first argument is **raw router logits**, not scores. The sigmoid is - applied *inside* the call. Passing pre-sigmoided values silently routes on - `sigmoid(sigmoid(x))`. -- The bias enters **selection only**. The returned weights are gathered from - the bias-free sigmoid, so they are not the values that were ranked. -- The renormalization (`norm_topk_prob: true` in HF configs) and the - `routed_scaling_factor` multiply are both **done inside**. There is no flag - to skip either; a model with `norm_topk_prob: false` cannot use this op. - Each returned row sums to `routed_scaling_factor` — except a row whose - selected scores all saturated to zero, which sums to 0 (see below). - -**The kernel's sigmoid.** The kernel evaluates `0.5 * tanh(0.5 * x) + 0.5` in -fp32, not `1 / (1 + exp(-x))`. The two are algebraically identical, and over -the logit range a router actually produces (`|x| <= 8`) they agree well inside -dtype tolerance — a `torch.sigmoid` reference reproduces this op there, which -is what makes it a drop-in for the HF DeepSeek-V3 gate. But the tanh form -saturates early: it returns **exactly 1.0** for `x >= ~17` and **exactly 0.0** -for `x <= ~-18.5`, and its relative error against the exponential form grows -through the negative tail (measured 4.3e-5 at `x = -8`, 2.6e-2 at `x = -15`). -A row whose *selected* experts all sit below `~-18.5` therefore comes back as -all-zero weights (the `eps` turns `0/0` into `0`, not `NaN`) where the -exponential form would return `routed_scaling_factor / topk` each. A torch -reference built on `0.5 * torch.tanh(0.5 * x) + 0.5` reproduced this op at -every shape, dtype and logit range in the test — out to `|logits| ~ 40`. - -**Ordering and ties.** Within a row the slots are sorted by decreasing -selection score `c`, and equal `c` values are ordered by *increasing* expert -index — i.e. the result equals a stable descending sort of `c` truncated to -`topk`. This differs from `torch.topk`, which on this machine emits the -*larger* index first among equals; the difference is visible whenever the -logits are bf16/fp16 (exact ties are common) or the bias is coarse. Note the -returned *weights* are therefore **not** monotone: they are `s` read out in -`c` order. - -**Fusion boundary.** The call is routing only. The router GEMM -(`hidden_states @ gate.weight.T`) that produces `router_logits` is the -caller's, and so is everything downstream: no expert permutation, no expert -GEMMs, no combine, no shared-expert branch. The two returned tensors are -shaped and typed as the `topk_weights` / `topk_ids` pair that trtllm's -pre-routed MoE runners take (see *Notes* for the dtype those runners demand). - -## Signature - -```python -def noaux_tc_op( - router_logits: torch.Tensor, - bias: torch.Tensor, - n_group: int, - topk_group: int, - topk: int, - routed_scaling_factor: float, -) -> tuple[torch.Tensor, torch.Tensor] -``` - -The wrapper renames the op's first two schema arguments (`scores`, `bias`) to -`router_logits`, `bias`; they are passed positionally and unchanged. - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `router_logits` | `[num_tokens, num_experts]` | bf16 / fp16 / fp32 | **contiguous** (see Preconditions) | CUDA | -| `bias` | `[num_experts]` | bf16 / fp16 / fp32, independent of `router_logits` except one combination (see Preconditions) | **contiguous** | CUDA, same device as `router_logits` | -| `n_group` | scalar | Python int, `>= 1` | — | — | -| `topk_group` | scalar | Python int | — | — | -| `topk` | scalar | Python int, `0 <= topk <= min(32, num_experts)` | — | — | -| `routed_scaling_factor` | scalar | Python float (any finite value, including 0 and negative) | — | — | -| returns `[0]` (`topk_weights`) | `[num_tokens, topk]` | **`router_logits.dtype`** | contiguous, newly allocated | CUDA (same device as `router_logits`) | -| returns `[1]` (`topk_ids`) | `[num_tokens, topk]` | **int32** | contiguous, newly allocated | CUDA (same device as `router_logits`) | - -The weight dtype follows `router_logits`, **not** `bias` — a bf16 logits / -fp32 bias call returns bf16 weights. Neither input is mutated (verified -bitwise) and neither is retained; both outputs are fresh allocations. - -## Metadata consumed - -None. Stateless — no attention metadata, no cache, no workspace, no -process-global tuner state. Two identical calls return bitwise identical -tensors, and the call runs on the ambient CUDA stream (verified: an -alternate-stream call returns the same result). - -## Preconditions - -- `router_logits` is a **2D** CUDA tensor. 1D and 3D inputs raise - `RuntimeError: scores must be a 2D Tensor`; a CPU tensor raises - `NotImplementedError` from the dispatcher. -- `bias` is a **1D** CUDA tensor with exactly `num_experts` elements, on the - same device as `router_logits`. Anything else raises `RuntimeError: bias - must be 1D with length == number of experts`; a CPU bias with CUDA logits - raises `RuntimeError: scores and bias must be CUDA tensors`. The bias is - mandatory — there is no "no bias" mode; pass a zero tensor. -- **Both tensors must be contiguous. The kernel ignores strides**: it - addresses `router_logits` as a dense `[num_tokens, num_experts]` row-major - buffer and `bias` as a dense `[num_experts]` buffer, each from `data_ptr()`. - A column slice of a wider buffer, a transposed view, or a strided bias is - silently routed on the wrong elements and returns a plausible-looking wrong - answer — it never raises. The wrapper asserts both; call `.contiguous()` - first (e.g. when the router logits are a slice of a fused projection). -- Dtypes: `router_logits` and `bias` are each bf16, fp16 or fp32. Anything - else raises `ValueError: Invalid dtype, only supports float16, float32, and - bfloat16` (or the `Invalid bias dtype` variant). The two are independent - **except** that `router_logits` fp16 + `bias` bf16 is rejected with `Invalid - bias dtype`, even though the corresponding kernel is compiled into this - build — an op-level dispatch gap, not a kernel limit. The other eight - combinations all work. -- `0 <= topk <= 32`. `topk > 32` raises `RuntimeError: topk should be smaller - than or equal to 32 for now`; a negative `topk` raises from the output - allocation. `topk == 0` is accepted and returns two `[num_tokens, 0]` - tensors. -- **`topk <= num_experts` is not enforced by the op.** With `topk > - num_experts` the kernel reads past the end of each row and emits expert ids - `>= num_experts` at non-zero weight (an out-of-bounds read; the values are - meaningless but **finite and deterministic** — at `num_experts = 3, - topk = 16` eight independent processes, three of them churning the caching - allocator first, each returned 104 of 128 ids out of range, identical - weights, and no NaN, so the read stays inside the logits tensor's own - buffer rather than picking up allocator residue). It never raises. The - wrapper asserts this. -- `num_experts % n_group == 0`, else `RuntimeError: num_experts should be - divisible by n_group`. `n_group <= 32`, else `RuntimeError: n_group should - be smaller than or equal to 32 for now`. -- **`n_group >= 1` is not checked.** `n_group == 0` reaches a host-side - `num_experts % n_group` and raises **SIGFPE, killing the process** — not a - Python exception, so it cannot be caught. The wrapper does not guard it - (guarding would require a non-metadata policy call); callers must never - pass 0. -- Configuration support, enforced inside the kernel with `RuntimeError: - [TensorRT-LLM][ERROR] Assertion failed: invokeNoAuxTc: unsupported - configuration (n_group=..., num_experts=..., topk_group=..., topk=...)`: - - `n_group == 1` (the ungrouped case): requires only `num_experts <= 1024`. - `topk_group` is then **completely ignored** — values `0, 1, 2, 5, 32, 100` - all give bitwise identical results. `num_experts = 1025` is rejected. - - `n_group > 1`: requires `1 <= topk_group <= n_group`, `topk <= 8`, - `num_experts <= 256`, and `num_experts / n_group <= 32` (experts per - group). `topk_group == 0` is *accepted* without error at `n_group > 1` but - is not characterized by this entry — do not use it. - - `topk_group == n_group` keeps every group and is bitwise identical to - `n_group = 1`. -- `num_tokens == 0` is accepted and returns empty `[0, topk]` tensors. - -A caller violating none of the above gets the result described under -*Semantics*. - -## Notes - -- **Certified surface.** Passing: `num_tokens` - in `{0, 1, 2, 4, 7, 8, 16, 64, 128, 256, 512, 1024, 2048, 4096, 8192}`; - `num_experts` in `{1, 2, 7, 8, 16, 32, 64, 72, 100, 128, 256, 257, 512, - 1024}`; `topk` in `{0, 1, 2, 3, 4, 6, 8, 16, 31, 32}`; all eight accepted - (logits, bias) dtype pairs; `routed_scaling_factor` in `{-1, 0, 0.5, 1, 2, - 2.5, 3, 1000}`; ungrouped and the grouped configurations `(num_experts, - n_group, topk_group, topk)` = `(256,8,4,8)`, `(128,4,2,6)`, `(72,8,2,6)`, - `(72,4,2,6)`, `(64,8,4,8)`. -- **Numerics: the reference of *Semantics* is not an approximation, it is the - kernel.** Expert ids matched exactly everywhere. Weights were bit-identical - except for a last-bit rounding step that appears when the kernel's fp32 - accumulation order for `tot` differs from `torch.sum`'s — measured at 1 - element in 49152 (8192 tokens x 72 experts, bf16) and 4 in ~10^7 over a - soak. Every element stayed within **one ulp** of the reference, so this - entry's test gates at `rtol = finfo(dtype).eps, atol = 0` (an order of - magnitude tighter than `assert_close`'s defaults for fp32) plus a - bit-exactness budget. Getting the fp64 detail of *Semantics* wrong is not a - rounding difference: an all-fp32 normalization misses a third of the - elements at `routed_scaling_factor = 2.5`. -- **The `register_fake` meta function disagrees with the kernel.** Under - `FakeTensorMode` (torch.compile / export / AD tracing) the first output is - given `bias.dtype`; eager execution gives `router_logits.dtype`. They differ - whenever the two input dtypes differ — e.g. fp32 logits with a bf16 bias - traces as bf16 and runs as fp32. Eager callers are unaffected. This is an - upstream defect in this build, not a behaviour to rely on. -- **Downstream dtype match.** `torch.ops.trtllm.fp4_block_scale_moe_runner`, - the NVFP4 pre-routed MoE runner, was observed here to reject a fp32 - `topk_weights` (`RuntimeError: topk_weights must be bfloat16.`) and an int64 - `topk_ids` (`RuntimeError: topk_ids must be int`), and to accept the - bf16/int32 pair. Since this op's weight dtype follows `router_logits`, - feeding it **bf16** router logits produces a directly consumable pair with - no cast. (Catalog membership of that runner is `index.yaml`'s fact alone.) -- Sigmoid, bias correction, group scoring and selection are computed in fp32 - regardless of the input dtype (the normalization then steps up to fp64); - only the final weights are cast back to `router_logits.dtype`. -- `-inf` logits map to score `0` exactly, like any logit below `~-18.5`, so a - `-inf`-masked expert loses to every expert with a positive score (verified: - none was selected while `topk` fit in the unmasked experts). It is *not* a - hard mask — with more slots than positively-scored experts the masked ones - come back at weight 0. `+inf` saturates to 1 like any logit above `~17`. - `NaN` logits *are* selected and yield `NaN` weights. -- A hand-written PyTorch path with the same semantics exists in this build - (`tensorrt_llm._torch.modules.fused_moe.routing.Deepseekv3RoutingImpl`) and - is what trtllm falls back to outside the supported configurations above; it - uses `torch.sigmoid`, so it differs from this op in the saturating tails and - in tie-breaking. -- Related routing ops in this build: `torch.ops.trtllm.renorm_moe_routing_op` - and `torch.ops.trtllm.default_moe_routing_op` (softmax routing, no bias), - and the block-scale MoE runners' built-in `routing_method_type = 2` - (DeepSeekV3) path, which is a different code path from this op. Catalog - membership is `index.yaml`'s fact alone. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/noaux_tc_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/noaux_tc_op.py deleted file mode 100644 index 904495aee2c2..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/noaux_tc_op.py +++ /dev/null @@ -1,49 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""DeepSeek-V3 style MoE routing: sigmoid + bias-corrected group-limited top-k.""" - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* - - -def noaux_tc_op( - router_logits: torch.Tensor, - bias: torch.Tensor, - n_group: int, - topk_group: int, - topk: int, - routed_scaling_factor: float, -) -> tuple[torch.Tensor, torch.Tensor]: - """Route each token to `topk` experts by biased sigmoid score. - - Selection uses `sigmoid(router_logits) + bias`; the combine weights are - gathered from the *unbiased* `sigmoid(router_logits)`, renormalized to sum - to 1 and scaled by `routed_scaling_factor`. - - Returns `(topk_weights [T, topk] in router_logits.dtype, topk_ids - [T, topk] int32)`, both newly allocated and contiguous. - """ - # Pure-metadata guards for three domains the op does not police. All three - # were observed on this machine to return a plausible-looking wrong answer - # instead of raising: - # - the kernel addresses `router_logits` as a dense [T, num_experts] - # row-major buffer and `bias` as a dense [num_experts] buffer, both from - # data_ptr(), ignoring strides; - # - topk > num_experts makes the kernel read past the end of a row and - # emit expert ids >= num_experts at non-zero weight. - assert router_logits.is_contiguous(), ( - "router_logits must be contiguous; a strided view is read as a dense " - "[num_tokens, num_experts] buffer and silently gives wrong routing" - ) - assert bias.is_contiguous(), ( - "bias must be contiguous; a strided view is read as a dense " - "[num_experts] buffer and silently gives wrong routing" - ) - assert topk <= router_logits.shape[-1], ( - f"topk ({topk}) exceeds num_experts ({router_logits.shape[-1]}); the " - "kernel reads out of bounds and emits out-of-range expert ids" - ) - return torch.ops.trtllm.noaux_tc_op( - router_logits, bias, n_group, topk_group, topk, routed_scaling_factor - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_fused_add_rmsnorm.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_fused_add_rmsnorm.md deleted file mode 100644 index b8e61d3bc7e1..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_fused_add_rmsnorm.md +++ /dev/null @@ -1,98 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 3} ---- - -# flashinfer_fused_add_rmsnorm - -**Wraps** `torch.ops.trtllm.flashinfer_fused_add_rmsnorm` (one call). - -## Semantics - -Fused residual add + RMS normalization, in-place on both tensor arguments. -One call computes, per row: - -``` -h = fp32(x) + fp32(residual) # fp32 accumulation -residual = cast(h, dtype) # overwritten in place -x = cast(h / sqrt(mean(h^2) + eps) * weight, dtype) # overwritten -``` - -The normalization reads the fp32 `h`, not the rounded `residual` output: -the add, the squared-mean reduction, and the weight scaling all happen in -fp32 inside the kernel before the final cast back to the input dtype. - -Fusion boundary: residual add + rmsnorm + weight scaling, nothing else. -No `(1 + weight)` gemma-style scaling (see -`flashinfer_gemma_fused_add_rmsnorm`) and no quantization (see -`flashinfer_fused_add_rmsnorm_quant`). The caller owns everything before -(producing `x`, e.g. an attention/MLP output) and after (consuming the -normed `x` and the updated `residual`). - -## Signature - -```python -def flashinfer_fused_add_rmsnorm( - x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float -) -> None -``` - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `x` | `[num_tokens, hidden]` | fp16 / bf16 / fp32 | last-dim stride must be 1; row stride may exceed `hidden` | CUDA | -| `residual` | `[num_tokens, hidden]` | same as `x` | same constraint as `x` | CUDA (same device) | -| `weight` | `[hidden]` | same as `x` | contiguous | CUDA (same device) | -| `eps` | scalar | Python float | — | — | -| returns | — | — | — | — | - -Returns `None`. Both `x` and `residual` are mutated in place: on return, -`residual` holds the pre-norm sum `x + residual` (the residual stream for -the next layer) and `x` holds the normalized, weight-scaled output. - -## Metadata consumed - -None. Stateless. - -## Preconditions - -- `x` and `residual` are 2D on the same CUDA device, same shape, same - dtype. The kernel is compiled for `[M, H]`; do not pass 3D tensors (it - would read `M = shape[0]`), but the guard is **loud, not silent**: the - CuTe compiled-kernel argument check raises `ValueError: Mismatched - Tensor on argument #0 ... expected ndim=2` before launch. Nothing is - silently skipped; the precondition stands, its - earlier justification did not. -- `weight.shape == (hidden,)`, `weight.dtype == x.dtype`, contiguous. -- Dtype is one of fp16, bf16, fp32. `float64`, `int8`, `uint8` and - `float8_e5m2` raise `KeyError`, but **`float8_e4m3fn` does not** — the - CuTe DSL path's dtype table maps it and returns a plausible fp8 - rmsnorm, so a caller who forgets to dequantize gets a result instead of - an error (measured on the certified path, 2026-07-28). The uncertified - CUDA-JIT path (`FLASHINFER_USE_CUDA_NORM=1`) does reject it — this - bullet described that path. -- `x.stride(-1) == 1` and `residual.stride(-1) == 1`. If either tensor is - non-contiguous (row stride > `hidden`), a strided kernel variant is - selected; its symbolic row stride is declared divisible by the kernel - vector size (up to 8 elements for bf16/fp16, 4 for fp32) and data - pointers are assumed 16-byte aligned, so arbitrary odd row strides or - misaligned slices are outside the contract. -- `hidden` has no divisibility requirement: sizes not divisible by the - 128-bit vector width (111, 1152) were verified correct on this machine. -- The caller must not need the pre-call contents of `x` or `residual` - afterward — both are destroyed. - -## Notes - -- The op is registered only when flashinfer is importable - (`IS_FLASHINFER_AVAILABLE`), and routes to the CuTe DSL kernel - (`fused_add_rmsnorm_cute`); a CUDA JIT fallback exists behind - `FLASHINFER_USE_CUDA_NORM=1` but is not what these receipts certify. -- Programmatic dependent launch (PDL) is controlled by the env var - `TRTLLM_ENABLE_PDL` (default enabled) inside the trtllm custom op; - it affects scheduling only, not results. -- The kernel uses fast-math `rsqrt`; results matched the fp32 torch - reference at default `assert_close` tolerances for all three dtypes. -- TRT-LLM's own `RMSNorm` module routes to this op only for fp16/bf16 - inputs, but fp32 works and passed the test on sm_100. -- Contiguous inputs with `num_tokens * hidden > 2^31 - 1` are - transparently routed to the strided (int64-offset) kernel variant. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_fused_add_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_fused_add_rmsnorm.py index 92fa4acffc5f..e84c58b437e9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_fused_add_rmsnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_fused_add_rmsnorm.py @@ -1,14 +1,106 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""In-place fused residual add + RMS normalization via the flashinfer kernel.""" +"""In-place fused residual add + RMS normalization via the flashinfer kernel. + +Everything a target author needs in order to decide whether this op covers +their case is below, and all of it is checked: `reference` is what CI compares +the kernel against, and `CELLS` is what CI drives. +""" + +from __future__ import annotations import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from .._op import Arch, Cell, OpWrapper + + +class FlashinferFusedAddRmsnorm(OpWrapper): + """`residual += x; x = rmsnorm(residual) * weight`, in place, returns + `(x, residual)` -- the same two tensors, mutated and handed back so a call + site can show what it touched without allocating. + + Both tensors are written. The residual output is the fp32 sum rounded once + to its own dtype; the normalization reads that sum in fp32, not the rounded + value, which is what `reference` mirrors. + + Fusion boundary: add, normalization and weight scaling. No gemma `(1 + w)` + form, no quantization. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The two hidden sizes the shipped targets normalize over, at a decode + # shape and a prefill shape. Both targets call this op only on the + # residual stream, so there is no third width to cover and no 3-D case: + # unlike plain rmsnorm this one never runs on a per-head view. + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode-shaped, gpt-oss-120b hidden", + spec=dict(shape=(1, 2880), dtype=torch.bfloat16, eps=1e-5), + ), + Cell( + why="decode-shaped, deepseek-r1 hidden", + spec=dict(shape=(1, 7168), dtype=torch.bfloat16, eps=1e-6), + ), + Cell( + why="prefill-shaped: enough rows to cross the kernel's grid tiling", + spec=dict(shape=(2048, 7168), dtype=torch.bfloat16, eps=1e-6), + ), + ) + + note = """ + Registered only when flashinfer is importable (`IS_FLASHINFER_AVAILABLE`). + + PDL is controlled by `TRTLLM_ENABLE_PDL` (default on) inside the trtllm + custom op; it changes scheduling, not results, and is not a cell axis. + + Both targets pass contiguous tensors here -- the residual stream is a + freshly allocated buffer -- so the strided-row path this kernel also has + is not certified by this entry. + """ + + def raw_call( + self, x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float + ) -> tuple[torch.Tensor, torch.Tensor]: + """Writes `x` and `residual` in place and returns both, in that order + -- the op destroys each argument's prior contents, so a caller that + wants them has to have kept its own copy beforehand.""" + torch.ops.trtllm.flashinfer_fused_add_rmsnorm(x, residual, weight, eps) + return x, residual + + def reference( + self, x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float + ) -> tuple[torch.Tensor, torch.Tensor]: + """The two values the op writes, as (x_out, residual_out). + + Returns rather than mutates so a cell can drive both sides from the + same inputs; the kernel's in-place write is what `raw_call` does. + """ + h = x.float() + residual.float() + normed = h * torch.rsqrt(h.pow(2).mean(dim=-1, keepdim=True) + eps) + return (normed * weight.float()).to(x.dtype), h.to(residual.dtype) + + def is_valid( + self, x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float + ) -> None: + # Only the cases the op takes and answers wrongly. Rank and width + # disagreements it rejects itself, loudly. + assert x.dtype is not torch.float8_e4m3fn, ( + "float8_e4m3fn is accepted by the kernel's dtype table and returns " + "a plausible fp8 result; dequantize before normalizing" + ) + assert x.stride(-1) == 1 and residual.stride(-1) == 1, ( + f"the normalized dim must be contiguous; got strides " + f"{x.stride(-1)} and {residual.stride(-1)}" + ) + # Aliasing is the trap specific to the in-place form: the kernel writes + # both tensors, so one buffer passed twice is overwritten mid-kernel + # and the result is neither value. + assert x.data_ptr() != residual.data_ptr(), ( + "x and residual must be distinct buffers; the op writes both" + ) + -def flashinfer_fused_add_rmsnorm( - x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float -) -> None: - """In-place: residual += x; x = rmsnorm(residual) * weight. Returns None.""" - torch.ops.trtllm.flashinfer_fused_add_rmsnorm(x, residual, weight, eps) +flashinfer_fused_add_rmsnorm = FlashinferFusedAddRmsnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.md deleted file mode 100644 index 490f1d9f71ac..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.md +++ /dev/null @@ -1,72 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 4} ---- - -# flashinfer_rmsnorm - -**Wraps** `torch.ops.trtllm.flashinfer_rmsnorm` (one call). - -## Semantics - -Root-mean-square normalization over the last dimension, with elementwise -weight scaling: - -``` -out[..., i] = x[..., i] / sqrt(mean(x[..., :]^2) + eps) * weight[i] -``` - -The squared-mean reduction and normalization are accumulated in fp32 -inside the kernel; the result is cast back to the input dtype. - -Fusion boundary: the single call computes normalization and weight scaling -only. There is no residual add (see `flashinfer_fused_add_rmsnorm` for -that), no `(1 + weight)` gemma-style scaling (see `flashinfer_gemma_rmsnorm`), -and no quantization. The caller owns everything else. - -## Signature - -```python -def flashinfer_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor -``` - -| Argument | Shape | Dtype | Layout | Device | -|---|---|---|---|---| -| `x` | `[num_tokens, hidden]` or `[batch, num_heads, head_dim]` | fp16 / bf16 / fp32 | last-dim stride must be 1; non-contiguous row strides are handled | CUDA | -| `weight` | `[hidden]` (or `[head_dim]` for 3D input) | same as `x` | contiguous | CUDA (same device as `x`) | -| `eps` | scalar | Python float | — | — | -| returns | same shape as `x` | same as `x` | newly allocated | CUDA | - -The output is a new tensor (`x` is not mutated). - -## Metadata consumed - -None. Stateless. - -## Preconditions - -- `x` is 2D or 3D on a CUDA device; the normalized dim is the last one. -- `weight.shape == (x.shape[-1],)` and `weight.dtype == x.dtype`. -- `x.stride(-1) == 1` (rows may be strided, e.g. a column slice of a wider - buffer; the kernel selects a strided code path in that case). -- Dtype is one of fp16, bf16, fp32. `float64`, `int8`, `uint8` and - `float8_e5m2` raise `KeyError`, but **`float8_e4m3fn` does not** — the - CuTe DSL path's dtype table maps it and returns a plausible fp8 - rmsnorm, so a caller who forgets to dequantize gets a result instead of - an error (measured on the certified path, 2026-07-28). The uncertified - CUDA-JIT path (`FLASHINFER_USE_CUDA_NORM=1`) does reject it — this - bullet described that path. -- `hidden` (last dim) has no divisibility requirement: sizes not divisible - by the 128-bit vector width (e.g. 111, 1152) fall back to a smaller - vector size and were verified correct on this machine. - -## Notes - -- The op is registered only when flashinfer is importable - (`IS_FLASHINFER_AVAILABLE`). -- Programmatic dependent launch (PDL) is controlled by the env var - `TRTLLM_ENABLE_PDL` (default enabled) inside the trtllm custom op; - it affects scheduling only, not results. -- TRT-LLM's own `RMSNorm` module routes to this op only for fp16/bf16 - inputs, but fp32 works and passed the test on sm_100. -- Weight scaling is plain `w * normed(x)`, not the gemma `(1 + w)` form. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.py index 6082c216094d..b15e2d342337 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.py @@ -1,12 +1,170 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""RMS normalization over the last dim via the flashinfer rmsnorm kernel.""" +"""RMS normalization over the last dim via the flashinfer rmsnorm kernel. + +Everything a target author needs in order to decide whether this op covers +their case is below, and all of it is checked: `reference` is what CI +compares the kernel against, and `CELLS` is what CI drives. Read the cell +list for the certified range. +""" + +from __future__ import annotations import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from .._op import Arch, Cell, OpWrapper + +# DeepSeek-V4.1's layouts, all bf16 at eps 1e-20: +# (id tag, width, parent row stride or None for contiguous, column offset, call site). +# The q / kv views slice one [T, 1792] qkv buffer; the DSpark drafter's ring +# write slices one [M, 1536] buffer per stage. Each layout is certified at +# every row count 1..16384 by the row sweep in the test; the cells hold the +# sweep's two ends, and at one row a column view takes the compact kernel. +_DSV41_LAYOUTS = ( + ("hidden 5120", 5120, None, 0, "attn / ffn / final norm"), + ("q 1280@0/1792", 1280, 1792, 0, "q_norm on qkv[:, :1280]"), + ("kv 512@1280/1792", 512, 1792, 1280, "kv_norm on qkv[:, 1280:]"), + ("comp 512", 512, None, 0, "compressed-latent norm"), + ("idx 128", 128, None, 0, "index-K norm"), + ("ring0 512@0/1536", 512, 1536, 0, "DSpark ring kv_norm, stage 0"), + ("ring1 512@512/1536", 512, 1536, 512, "DSpark ring kv_norm, stage 1"), + ("ring2 512@1024/1536", 512, 1536, 1024, "DSpark ring kv_norm, stage 2"), +) +_DSV41_ROWS = (1, 16384) +_DSV41_VIT_ROWS = (1, 51111) + +_DSV41_CELLS = tuple( + Cell( + why=f"deepseek-v4.1 {tag} rows={rows}: {site}, eps 1e-20, " + f"{'one row' if rows == 1 else 'top'} of the row sweep 1..16384", + spec=dict( + shape=(rows, width), + dtype=torch.bfloat16, + eps=1e-20, + row_stride=stride, + col_offset=offset, + data="eps_classes", + ), + ) + for tag, width, stride, offset, site in _DSV41_LAYOUTS + for rows in _DSV41_ROWS +) + tuple( + Cell( + why=f"deepseek-v4.1 vision 1024 rows={rows}: vision tower norm1 / norm2 / " + f"final norm, eps 1e-6, {'one row' if rows == 1 else 'top'} of the row sweep 1..51111", + spec=dict(shape=(rows, 1024), dtype=torch.bfloat16, eps=1e-6, data="vit"), + ) + for rows in _DSV41_VIT_ROWS +) + + +class FlashinferRmsnorm(OpWrapper): + """`x / sqrt(mean(x^2, -1) + eps) * weight`, one op call, new tensor out. + + Fusion boundary: normalization and weight scaling only. No residual add + (`flashinfer_fused_add_rmsnorm`), no gemma `(1 + w)` form + (`flashinfer_gemma_rmsnorm`), no quantization. `x` is not mutated. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The widths the shipped targets normalize over. gpt-oss reaches this op + # only on its residual stream; deepseek reaches it there and on the two + # lora ranks its MLA projects through, which is why the narrow cells are + # here at all -- a 512-wide row is one 128-bit vector per lane and takes a + # different path through the kernel than a 7168-wide one. + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode-shaped, gpt-oss-120b hidden 2880", + spec=dict(shape=(1, 2880), dtype=torch.bfloat16, eps=1e-5), + ), + Cell( + why="decode-shaped, deepseek-r1 hidden 7168", + spec=dict(shape=(1, 7168), dtype=torch.bfloat16, eps=1e-6), + ), + Cell( + why="prefill-shaped: enough rows to cross the kernel's grid tiling", + spec=dict(shape=(2048, 7168), dtype=torch.bfloat16, eps=1e-6), + ), + Cell( + why="deepseek q_lora rank 1536, the q_norm this op is called on", + spec=dict(shape=(16, 1536), dtype=torch.bfloat16, eps=1e-6), + ), + Cell( + why="deepseek kv_lora rank 512, the narrowest row any target asks for", + spec=dict(shape=(16, 512), dtype=torch.bfloat16, eps=1e-6), + ), + # DeepSeek-V4.1: see `_DSV41_LAYOUTS`. Optional spec keys: `row_stride` + # and `col_offset` make `x` a column view of a wider parent, `data` + # names the row-magnitude pattern (the test's `_build` documents both). + *_DSV41_CELLS, + ) + + note = """ + Registered only when flashinfer is importable (`IS_FLASHINFER_AVAILABLE`). + + PDL is controlled by `TRTLLM_ENABLE_PDL` (default on) inside the trtllm + custom op; it changes scheduling, not results, and is not a cell axis. + + fp16 and fp32 inputs are accepted by the op and were measured to work, + but no shipped target uses them, so they are not cells. TRT-LLM's own + `RMSNorm` module routes here for fp16/bf16 only. + + `float8_e4m3fn` is the trap this entry cannot guard: the CuTe DSL path's + dtype table maps it and returns a plausible fp8 rmsnorm rather than + raising, so a caller who forgets to dequantize gets a result. `is_valid` + rejects it below. The CUDA-JIT path + (`FLASHINFER_USE_CUDA_NORM=1`) does raise, but this entry does not + certify that path. + + DeepSeek-V4.1: bf16 at eps 1e-20 on the layouts in `_DSV41_LAYOUTS`, + metadata none. Every row count 1..16384 is certified at each layout: the + 16384-row call meets the reference, every prefix of it is bitwise the op + on that many rows, nothing is written past the last row, and no kernel + variant compiles per row count. Rows mix unit, eps-scale and zero + magnitudes, so a kernel that clamped or ignored the tiny eps fails the + gate; zero rows come back exactly 0. Column views need a 16-byte-aligned + row stride and start; a one-row view reports contiguous and takes the + compact kernel, which needs a 128-byte-aligned start (column offset a + multiple of 64 elements, as every certified site is). A misaligned view + raises `ValueError`; it is never computed wrongly. A column view's output + is a new contiguous `[rows, width]` tensor (at one row it keeps the view's + row stride). Any width-512 view with a 16-byte-aligned row stride up to + 2048 and column offset was measured at 13 row counts up to 160. CUDA graphs + are certified, after an eager call at the same shape, at the 126 per-rank + row counts {G, 3G, 5G, 6G} over the 46 captured batch sizes up to 256 for + the five core layouts, and at 473 counts up to 1280 for the ring views. + The vision tower calls it at width 1024, eps 1e-6, contiguous `[P, 1024]` + with +-2000 massive channels, every P 1..51111, captured at P = 1517, 4680, + 5325, 8772 and 17037. Not certified for V4.1: fp16 / fp32 at eps 1e-20; + rows whose squares underflow fp32 (|x| below about 1e-19, flushed to zero); + more than 16384 rows (51111 at width 1024); width 1024 at another eps or as + a view; capture at other row counts, or before the shape's kernel variant + has compiled; 3-D input at eps 1e-20. + """ + + def raw_call(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + return torch.ops.trtllm.flashinfer_rmsnorm(x, weight, eps) + + def reference(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """fp32-accumulated; the kernel reduces in fp32 and rounds once at the end.""" + xf = x.float() + normed = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (normed * weight.float()).to(x.dtype) + + def is_valid(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> None: + # Only the case the op takes and answers wrongly. Shape and dtype + # mismatches it rejects itself, loudly, and duplicating that here + # would only swap its error for a worse one. + assert x.dtype is not torch.float8_e4m3fn, ( + "float8_e4m3fn is accepted by the kernel's dtype table and returns " + "a plausible fp8 rmsnorm; dequantize before normalizing" + ) + assert x.stride(-1) == 1, ( + f"the normalized dim must be contiguous; got stride {x.stride(-1)}" + ) + -def flashinfer_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: - """Return `x / sqrt(mean(x^2, dim=-1) + eps) * weight` as a new tensor.""" - return torch.ops.trtllm.flashinfer_rmsnorm(x, weight, eps) +flashinfer_rmsnorm = FlashinferRmsnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.md deleted file mode 100644 index f8e81a788774..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.md +++ /dev/null @@ -1,340 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 19} ---- - -# fp4_quantize - -**Wraps** `torch.ops.trtllm.fp4_quantize` (one call). - -## Semantics - -Block-scaled FP4 quantization of an activation tensor: one call converts -values to 4-bit **e2m1** elements (two packed per output byte) plus one -scale byte per `sf_vec_size` **contiguous elements along the last dim**. -Nothing else is fused in — no normalization, no activation, no transpose, -and (unlike `mxfp8_quantize`) **no width padding**: the last dim is used -exactly as given, so any padding a consumer needs is the caller's own -`pad` before the call. The caller owns producing the input and owns both -outputs afterwards. - -The op has two modes, selected by `sf_use_ue8m0`: - -| Mode | `sf_vec_size` | `sf_use_ue8m0` | `global_scale` | Scale byte | -|---|---|---|---|---| -| **NVFP4** | 16 | `False` | required | e4m3 (UE4M3 in practice) | -| MXFP4 | 32 | `True` | optional | UE8M0 power of two | - -Any other combination raises. **Only the NVFP4 mode is certified here**; -what follows describes it. - -Let `M` = product of all leading dims, `K` = last dim, `cols = K / 16`. -For row `m` and block `b` of 16 elements, with `g` = the `global_scale` -value: - -``` -vecmax = max |x[m, 16b : 16b+16]| # exact in the input dtype -sf[m, b] = e4m3_round_to_nearest_even( g * vecmax / 6 ) # saturates at 448 -out_scale = g / float(sf[m, b]) # 0 when vecmax == 0 -data[m, i] = e2m1_round_to_nearest_even( x[m, i] * out_scale ) # saturates at +-6 -``` - -`6` is e2m1's largest magnitude and `448` is e4m3's largest finite value, -so the canonical `g = 448 * 6 / amax(x)` keeps every block's `sf` inside -e4m3's finite range. - -**Mind the direction.** A modelopt NVFP4 checkpoint stores its per-tensor -`input_scale` as `amax / (448*6)` — the **reciprocal** of this `g` — so a -caller passes `1 / input_scale` here, never `input_scale`. Passing it -straight is accepted and silently destructive: measured on `64 x 7168` -bf16 with `amax = 5.1875`, it drove **96.4 % of the scale bytes to zero** -(those blocks read back as all-zero, the `g`-underflow path in -*Preconditions*) while the surviving blocks saturated — max reconstructed -magnitude 6.07 against a true 5.1875. The result is neither an error nor -an obviously dead tensor. - -**Dequantization is exactly `data * sf / g`.** A consumer GEMM folds -`1/g` into its own `alpha` (typically `alpha = (amax_input / (448*6)) * -weight_scale_2`); this op emits `data` and `sf` only, and never `g` or -`1/g`. - -e2m1 encodes 8 magnitudes — `0, 0.5, 1, 1.5, 2, 3, 4, 6` — as -`code = (exponent << 1) | mantissa`, with bit 3 the sign; `-0.0` keeps -its sign bit (code 8). Rounding is **round-to-nearest, ties to even -code**, saturating: `0.25 -> 0`, `0.75 -> 1.0`, `1.25 -> 1.0`, -`1.75 -> 2.0`, `2.5 -> 2.0`, `3.5 -> 4.0`, `5.0 -> 4.0`, and anything -above `6` clamps to `6`. - -Because `sf` has only a 3-bit mantissa, `sf` may round *below* `g*vecmax/6`, -which makes `out_scale` slightly larger than `6/vecmax` and lets the -block's largest element clip to `6`. That is the format, not an error: it -happened in about half the blocks of random bf16 input, and the -reconstruction still obeys the bound below. - -**Tie non-determinism (the one place this op is not bit-reproducible from -a plain fp32 model).** The kernel forms `out_scale` through two -`rcp.approx.ftz.f32` reciprocals (~2^-23 relative error each) rather than -an exact division. An element whose exact `x * out_scale` sits within a -few fp32 ulps of an e2m1 midpoint (`0.25, 0.75, 1.25, 1.75, 2.5, 3.5, -5.0`) can therefore land on either neighbouring code. Every other element -matches the formula above bit for bit, and the scale bytes are always -exact. The op itself is deterministic: repeating the same call gives the -same bytes. - -**What sets how many such elements there are: the global scale, not the -shape.** A near-tie needs the block's `sf` to come out *exactly* -`g * vecmax / 6` — then `out_scale` is exactly `6 / vecmax` and -`6 * x / vecmax` is an exact small rational, which lands on a midpoint -for a sizeable minority of the block's lanes. Measured on this machine -(bf16, `g = 448*6/amax`, the canonical convention): - -- **the width is inert.** One 16,515,072-element sample reshaped to - `K` = 7168, 18432 and 2048 gave the *identical* counts — 34,405 - near-tie elements and 12,091 kernel/reference disagreements in all - three. Scale blocks are 16 *contiguous* elements, so a reshape never - moves a block boundary; `M` and `K` reach the arithmetic only through - `amax`, hence through `g`. -- **`g` is the whole story.** With that `g`, 0.68 % of blocks had an - exactly representable `sf`. Multiplying `g` by 1.001 left **no** block - with one, and the same call then matched the reference **bit for bit**, - zero near-ties — while multiplying by 2 (a power of two, so `sf` scales - exactly and `out_scale` is unchanged) kept the population at 33,113. -- across the 24 bf16 shapes of the R1 sweep, near-ties were **0.03 % to - 0.60 %** of all elements and the kernel disagreed with an exact-fp32 - reference on **0.00 % to 0.14 %** — always by one adjacent e2m1 code - with the sign unchanged. - -So the population is a property of the *exactly calibrated* case. A -serving call passes a static per-tensor `input_scale` rather than -`448*6/amax` of the tensor in hand, which is the regime measured above as -bit-exact — but that was one `g` offset (0.1 %), not a proof for every -static scale, and no run here used a real checkpoint's scale. - -### Output layouts - -`data` is **identical byte for byte under both `is_sf_swizzled_layout` -values** — only the scale buffer's size and element order change. - -- `is_sf_swizzled_layout=False` (**linear**): `M * cols` bytes, row-major. - Byte `b` of row `m` is at flat offset `m * cols + b`; `sf.view(M, cols)` - is exactly the per-block scale matrix. -- `is_sf_swizzled_layout=True` (**128x4 swizzled**): - `pad_up(M, 128) * pad_up(cols, 4)` bytes. The scale of `(m, b)` sits at - flat offset - - ``` - (b % 4) - + (b // 4) * 512 # 4 cols x 128 rows per column group - + (m % 32) * 16 - + ((m % 128) // 32) * 4 - + (m // 128) * 128 * pad_up(cols, 4) - ``` - - Every offset not addressed by a real `(m, b)` pair — row padding up to - `pad_up(M,128)`, column padding up to `pad_up(cols,4)` — is `0x00`. - -## Signature - -```python -def fp4_quantize( - input: torch.Tensor, - global_scale: Optional[torch.Tensor], - sf_vec_size: int, - sf_use_ue8m0: bool = False, - is_sf_swizzled_layout: bool = True, -) -> tuple[torch.Tensor, torch.Tensor] -``` - -| Argument | Shape | Dtype | Layout / value | Device | -|---|---|---|---|---| -| `input` | `[..., K]`, rank >= 2 | bfloat16 or float16 (e4m3 accepted, uncertified) | contiguous | CUDA | -| `global_scale` | exactly 1 element (`[1]` or 0-D) | float32 | positive, finite; `448*6/amax` by convention | CUDA | -| `sf_vec_size` | scalar | Python int | `16` (NVFP4). `32` only with `sf_use_ue8m0=True` | — | -| `sf_use_ue8m0` | scalar | Python bool | `False` for NVFP4 e4m3 scales | — | -| `is_sf_swizzled_layout` | scalar | Python bool | `True` = 128x4 swizzled scale order, `False` = linear | — | -| returns `data` | `[..., K/2]` (leading dims preserved) | uint8 | contiguous, freshly allocated; element `2i` in the low nibble, `2i+1` in the high nibble | same as `input` | -| returns `sf` | 1-D: `[M*cols]` (linear) or `[pad_up(M,128)*pad_up(cols,4)]` (swizzled) | uint8 | contiguous, freshly allocated; reinterpret as `float8_e4m3fn` to read the values | same as `input` | - -`M` = product of `input`'s leading dims, `cols = K / sf_vec_size`. -`input` is not modified. Leading dims are collapsed for the scale buffer -in both layouts (a `[B, S, K]` input yields a scale buffer laid out for -`B*S` rows), while `data` keeps `input`'s leading dims. - -The op's own schema spells the parameters `globalScale`, `sfVecSize`, -`sfUseUE8M0`, `isSfSwizzledLayout`; all five are positional here. - -## Metadata consumed - -None. Stateless — no runtime, attention metadata, or workspace. - -## Preconditions - -- `input` is on CUDA, contiguous, rank >= 2, dtype bfloat16 / float16 / - float8_e4m3fn. Violations raise: a CPU tensor, a non-contiguous view - (column slice or transpose), a 1-D tensor and an fp32 tensor were each - observed to raise on this machine rather than compute something wrong. -- `K % sf_vec_size == 0`. Raises otherwise. This is the whole width - domain, and it was re-measured as such: `K` = 2048, 7168 and 18432 — - the last 1.5x wider than any previously certified width — behave - exactly like 2560 / 3072 / 12288, at every row count from 1 to 8192. - Nothing about the op is a function of `K` beyond `cols = K / 16` (see - *Notes*: one kernel, no width-keyed dispatch). -- `sf_vec_size == 16` with `sf_use_ue8m0=False`, or `sf_vec_size == 32` - with `sf_use_ue8m0=True`. Every other pairing (including - `sf_vec_size=8` or `0`) raises. -- `global_scale` is required when `sf_use_ue8m0=False` (passing `None` - raises), must be a **float32 CUDA** tensor (a half tensor and a CPU - tensor each raise), and must hold **exactly one element**. A tensor - with more elements is *not* rejected by the op: it silently reads - element 0 and applies it to every row — this build has no per-token - global scale. The wrapper asserts the element count for that reason. -- `global_scale` must be **positive and finite**: - - `g = 0` yields an all-zero scale buffer and every data lane at `+6`; - a consumer dividing by `g` gets NaN. - - `g < 0` yields **negative** e4m3 scale bytes (e.g. `0xB0`), which - every consumer reads as unsigned UE4M3 — silently wrong. -- Every block must satisfy `2^-10 < g * vecmax / 6 <= 448` (or - `vecmax == 0`). Both ends fail silently: - - **Above 448** (a stale/mis-calibrated static `input_scale`, or a - runtime `amax` above the calibration one): `sf` clamps to `0x7E` - (448) and `out_scale` collapses to `g/448`. Only lanes above - `5 * 448 / g` then clip to `+-6`; the rest quantize normally against - the collapsed scale and still dequantize close to their true value — - at 1.17x overshoot a hand-built block came back - `[6, 4, 2, 1, 0.5, 0, 0, 0]`, not all-`+-6`. - **Do not look for an all-`+-6`, signs-only block as the signature**: - a mis-calibrated block looks ordinary, and only ~100x overshoot - produces the saturated form. - - **At or below `2^-10`** (`g` far too small): `sf` rounds to `0x00`, - the reciprocal of that zero scale is `+inf`, so every lane saturates - to `+6` (exact zeros included, via `0 * inf = NaN` which the e2m1 - cast saturates). Because `sf` is zero, a consumer reads the whole - block back as zeros. One binade higher (`g*vecmax/6 = 2^-9`, the - e4m3 minimum subnormal) is already the normal, exact path. - - An all-zero block is fine: scale byte `0x00`, all-zero data bytes. -- `input` must be **finite**. Non-finite lanes are accepted silently and - are destructive: - - a `+inf`/`-inf` lane drives the block max to infinity, so the scale - saturates at 448 and every other lane in the block is divided by 448 - — small lanes round to zero; the inf lane itself becomes `+6`; - - a NaN lane is excluded from the block max, so the scale stays correct - and the NaN lane saturates to `+6` without disturbing its neighbours. -- `M == 0` (zero rows) is accepted and returns empty `data` and `sf`. -- The caller owns both returned buffers; nothing else writes them. - -## Notes - -- Certified on sm_100 (B200) only, and only for the **NVFP4** mode - (`sf_vec_size=16`, `sf_use_ue8m0=False`) with bfloat16 and float16 - inputs. The kernel behind this op is Blackwell-targeted, and the - installed binary shows it directly: the shipped library carries cubins - of this kernel for sm_80, sm_86, sm_89, sm_90a, sm_100f and sm_120f, - but only the sm_100f and sm_120f ones have a body (32 and 36 registers, - 1024 B shared). The sm_80/86/89 cubins are **4 registers and no shared - memory** and the sm_90a one 4 registers — i.e. the conversion is - compiled away below Blackwell, leaving a kernel that launches, writes - nothing and returns success. A pre-Blackwell arch is therefore a - silent-garbage risk rather than a raise. No receipt is claimed there, - and sm_120 was not run. -- **Not certified** (accepted by the op, no observed-behaviour claims - here): the MXFP4 mode (`sf_vec_size=32`, `sf_use_ue8m0=True`, UE8M0 - scale bytes) beyond its argument validation, and float8_e4m3fn input - (which routes to a different in-kernel conversion). Also uncertified, - though the shape rule covers them: `M > 8192`; float16 and rank-3 - inputs at `K` in {2048, 7168, 18432} (only bf16, rank-2 ran there); - and CUDA-graph capture, which nothing here exercises. -- Against a native-torch reference built from the formula above, the - kernel was **bit-exact on every scale byte** and on every data nibble - outside the near-tie window described in Semantics; every element - inside that window differed by at most one adjacent e2m1 code with the - sign unchanged. Shapes run: - - **linear** layout, bf16: `T x K` for every `T` in - {1, 2, 7, 64, 1023, 1024, 2048} crossed with `K` in - {2560, 3072, 12288} except `2048 x 12288`; plus `T` in - {1, 2, 7, 64, 1023, 1024, 2048, **8192**} crossed with `K` in - {**2048**, **7168**, **18432**}; plus one 16,515,072-element sample - reshaped to `8064 x 2048`, `2304 x 7168` and `896 x 18432`. - - **swizzled** layout, bf16: `1 x 2560`, `7 x 2560`, `129 x 2560`, - `1024 x 2560`, `200 x 3072`, `3 x 112` (`cols = 7`, the only case - exercising column padding), and the full `T` x `K` grid above at - `K` in {2048, 7168, 18432}. Every call listed here was checked - offset-by-offset against the same input's linear call, and its data - bytes asserted equal to the linear call's. - - fp16: `T` in {1, 512, 1024} at `K = 2560`, linear only. A 3-D - `2 x 5 x 2560` bf16 input: linear offset-checked, swizzled checked for - buffer size and data equality only. - The bf16 `K` in {2048, 7168, 18432} grid is the DeepSeek-R1 activation - set — model hidden, dense-MLP intermediate, shared-expert intermediate - — with `T` reaching `max_num_tokens = 8192`. At those widths - `cols = K/16` is 128 / 448 / 1152, all already multiples of 4, so the - swizzled buffer there carries row padding but no column padding. -- Reconstruction bound, verified elementwise on `1024 x 2560` bf16 input - in fp32 and on `1024 x 18432` bf16 input in **float64**: - `|data*sf/g - x| <= max(0.25 * sf/g, |x|/4)`. The two terms are the - e2m1 subnormal half-step carried back through the block scale, and the - half-step of the coarsest e2m1 binade (which also covers the clipping - of a block max when `sf` rounds down). Every operand is exact in - float64, so that form of the check is the mathematical bound; a - consumer reconstructing in **fp32** can land up to half an fp32 ulp - outside it (measured 3.0e-8 at `1024 x 18432`) purely from its own - rounding. -- **One kernel — nothing about rows or width selects a different one.** - Every call of this op in this build launches exactly one CUDA kernel, - `tensorrt_llm::_v1::kernels::quantize_with_block_size`, and it is the same instantiation at 7 rows and at 8192, at - widths 2048 / 7168 / 18432 (multiples of 512) and at 112 (not), in both - scale layouts and for bf16 and fp16 alike — measured with the CUDA - profiler in this entry's test. Consistent with the library's symbol - table, inspected separately, which carries 14 instantiations of that one - kernel template and no TMA variant at all. - A TMA high-throughput variant selected on `M >= 1024 && width % - 512 == 0` **does** exist in FlashInfer's copy of this kernel family, - which ships in this install as readable source — but it is reachable - only through `torch.ops.trtllm.tunable_fp4_quantize`, never through - this op. Treat that vendored source as a *newer* upstream revision, not - as this build: its `invokeFP4Quantization` takes `enable_pdl`, - `use_row_wise_scale` and `inverse_scale` parameters that the installed - symbol does not have, and it accepts a per-token `globalScale` that - this build silently ignores (see *Preconditions*). -- Consumer pairing, read from the TensorRT-LLM call paths in this install - (each consumer's own contract is authoritative; this entry certifies - only what each layout *contains*): - - the NVFP4 dense GEMM family (`torch.ops.trtllm.nvfp4_gemm` and the - backends it dispatches to) is fed the **swizzled** buffer — - `is_sf_swizzled_layout=True`, the op's default; its CUDA-core backend - explicitly un-swizzles before use; - - the trtllm-gen NVFP4 block-scale MoE runner - (`torch.ops.trtllm.fp4_block_scale_moe_runner`) is fed the **linear** - buffer — `is_sf_swizzled_layout=False`, then viewed as - `[num_tokens, cols]` — matching how `mxfp8_quantize` feeds the - trtllm-gen W4A8 runner. - So a target running an NVFP4 dense GEMM and an NVFP4 trtllm-gen MoE - needs two different calls to this op, not one shared result. Both - layouts are certified at `K = 7168`, which is where a DeepSeek-R1-shaped - target needs exactly that pair of calls. -- **How discriminating the certification above is**, measured rather than - assumed — the check tolerates one adjacent e2m1 code inside the - near-tie window and nothing else, so the window is its only blind spot, - and it is 0.03–0.60 % of elements. Driven through the same comparison at - `1024 x 7168`: a single data nibble moved one code *outside* the window - fails it; two codes *inside* the window fails it; a sign flip inside the - window fails it; a single scale byte moved one code fails it (the scale - bytes have no carve-out at all); one code inside the window passes, and - is visible only as one extra disagreement in the count. A swizzled scale - buffer handed over as if it were linear is a different byte string even - when the two have the same length (rows a multiple of 128), so the - layout check is not vacuous either. -- Sibling quantizers with the same scale-layout vocabulary exist - (`torch.ops.trtllm.mxfp8_quantize` for fixed-32-block MXFP8, - `torch.ops.trtllm.fp8_quantize_1x128` for 1x128 block-scaled fp8), as - do fp4 variants that fuse extra work into the same launch - (`torch.ops.trtllm.fp4_quantize_with_residual`, - `torch.ops.trtllm.fp4_batched_quantize`) and a helper that converts an - existing linear scale buffer to the swizzled order - (`torch.ops.trtllm.block_scale_interleave`). - `torch.ops.trtllm.tunable_fp4_quantize` is a different thing again: an - autotuned chooser between this op's kernel and FlashInfer's, so it can - reach the TMA kernel described above. This entry covers the plain - `fp4_quantize` launch only. -- The op allocates both outputs itself: it is functional, with no `out=` - parameter and no in-place mode. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py index 058aa6997c70..bf011dcc29f5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py @@ -1,7 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Block-scaled FP4 activation quantization: packed e2m1 data + one scale byte per -sf_vec_size contiguous elements (NVFP4: 16-element blocks, e4m3 scales).""" +"""NVFP4 quantization: packed e2m1 data plus one e4m3 scale per 16 elements.""" + +from __future__ import annotations from typing import Optional @@ -9,29 +10,347 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from .._op import Arch, Cell, OpWrapper + +_E4M3_MAX = 448.0 +_E2M1_MAX = 6.0 + +#: Midpoints between consecutive e2m1 magnitudes. A value landing exactly on one +#: is a rounding tie; `_TIE_UP[i]` is True when rounding *up* yields the even +#: code, which is the direction round-to-nearest-even takes. +_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_TIE_UP = (False, True, False, True, False, True, False) + +#: The kernel builds its per-block output scale through two `rcp.approx.ftz.f32` +#: reciprocals, each ~2^-23 relative. A value within a few fp32 ulps of a +#: midpoint can therefore round to either neighbour, and no reference can say +#: which. Everything further out than this must match bit for bit. Largest +#: deviation observed on this machine: 7.9e-8 relative, ~12x inside the window. +TIE_WINDOW = 2.0**-20 + +#: MXFP4 mode (`sf_vec_size=32`, `sf_use_ue8m0=True`): the reference floors each +#: 32-block's max at `6 * 2^-126`, so its smallest scale is 2^-126 (UE8M0 byte +#: 0x01). The op has no floor; see `reference`. +_MX_AMAX_FLOOR = 6 * 2.0**-126 +#: The reference multiplies the block max by this fp32 constant and takes +#: ceil(log2(.)) from the product's IEEE bits. +_INV6_F32 = 1.0 / 6.0 + + +#: DeepSeek-V4.1's routed-expert input: bf16 NVFP4 with a static calibrated `[1]` scale. +_V41_MOE = dict(dtype=torch.bfloat16, vec=16, gs="static") +#: DeepSeek-V4.1's index K / index Q quantizer: bf16 MXFP4, linear, no global scale, and +#: a power-of-two magnitude per 32-block so the UE8M0 scales spread. +_V41_INDEX = dict(dtype=torch.bfloat16, vec=32, swizzled=False, ue8m0=True, gs=None, block_mag=True) + + +def _pad_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple -def fp4_quantize( - input: torch.Tensor, - global_scale: Optional[torch.Tensor], - sf_vec_size: int, - sf_use_ue8m0: bool = False, - is_sf_swizzled_layout: bool = True, -) -> tuple[torch.Tensor, torch.Tensor]: - """Quantize to FP4: returns (packed e2m1 data [..., K/2] uint8, block scales uint8, 1-D). - One scale per `sf_vec_size` contiguous elements along the last dim. - NVFP4 is `sf_vec_size=16, sf_use_ue8m0=False` (e4m3 scale bytes, scaled by - `global_scale`); `is_sf_swizzled_layout` selects the 128x4 swizzled scale - order (True) or the row-major linear order (False). +class Fp4Quantize(OpWrapper): + """Quantize to FP4: `(packed e2m1 data [..., K/2] uint8, scale bytes, 1-D)`. + + One scale per `sf_vec_size` contiguous elements of the last dim. NVFP4 is + `sf_vec_size=16` with `sf_use_ue8m0=False`, which makes the scale an e4m3 + byte derived from the block maximum and the caller's `global_scale`. MXFP4 + is `sf_vec_size=32` with `sf_use_ue8m0=True`: the scale is a UE8M0 power of + two of the block maximum alone, and `global_scale` is not read. Any other + pairing raises. + + `is_sf_swizzled_layout` selects the 128x4 swizzled scale order (True) or the + row-major linear order (False); the buffer is 1-D either way, so the layout + is something a consumer has to be told, not something its shape shows. """ - # Pure-metadata guard: the kernel loads one scalar from global_scale and - # ignores every element past the first, so a per-token [num_tokens] tensor - # is silently applied as global_scale[0] to every row -- observed on this - # machine to return a valid-looking result, never to raise. - assert global_scale is None or global_scale.numel() == 1, ( - "global_scale must hold exactly one element; extra elements are " - "silently ignored and global_scale[0] is applied to every row" - ) - return torch.ops.trtllm.fp4_quantize( - input, global_scale, sf_vec_size, sf_use_ue8m0, is_sf_swizzled_layout + + ARCHS = frozenset({Arch.SM_103}) + + # The first five cells are deepseek-r1's. It reaches this op from both layouts: the + # swizzled one feeding nvfp4_gemm in the dense and shared-expert MLPs, the + # linear one feeding the routed-expert path. Widths are its hidden size and + # its intermediate. + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode, hidden 7168, swizzled scales -- the dense MLP's input", + spec=dict(shape=(1, 7168), dtype=torch.bfloat16, vec=16, swizzled=True), + ), + Cell( + why="prefill chunk, hidden 7168, swizzled", + spec=dict(shape=(2048, 7168), dtype=torch.bfloat16, vec=16, swizzled=True), + ), + Cell( + why="linear scales -- the routed-expert path asks for this order", + spec=dict(shape=(2048, 7168), dtype=torch.bfloat16, vec=16, swizzled=False), + ), + Cell( + why="the activation width, 2 x moe_intermediate, between the two GEMMs", + spec=dict(shape=(512, 4096), dtype=torch.bfloat16, vec=16, swizzled=True), + ), + Cell( + why="a row count that is not a multiple of the swizzle's 128-row tile", + spec=dict(shape=(37, 7168), dtype=torch.bfloat16, vec=16, swizzled=True), + ), + # deepseek-v4.1 routed-expert input: NVFP4 at its hidden 5120, with the + # layer's static `[1]` g1 = 1 / input_scale calibrated on another batch + # (`gs="static"`). The target asks for linear scales; both layouts ran. + # Rows 1..8192 are swept by test_v41_moe_input_k5120_every_row_count and + # the prefix test; these cells are the sweep's ends and the grid edge. + Cell( + why="deepseek-v4.1 MoE in K5120 linear: 1 row, static g1 [1]", + spec=dict(shape=(1, 5120), swizzled=False, **_V41_MOE), + ), + Cell( + why="deepseek-v4.1 MoE in K5120 linear: 609 rows, past the launcher's 608-row grid cap " + "(152 SMs x 4 blocks), so rows go through the grid-stride loop", + spec=dict(shape=(609, 5120), swizzled=False, **_V41_MOE), + ), + Cell( + why="deepseek-v4.1 MoE in K5120 linear: 8192 rows, the target's MoE chunk size", + spec=dict(shape=(8192, 5120), swizzled=False, **_V41_MOE), + ), + Cell( + why="deepseek-v4.1 MoE in K5120 swizzled: 1 row, 127 rows of padding", + spec=dict(shape=(1, 5120), swizzled=True, **_V41_MOE), + ), + Cell( + why="deepseek-v4.1 MoE in K5120 swizzled: 8192 rows", + spec=dict(shape=(8192, 5120), swizzled=True, **_V41_MOE), + ), + # deepseek-v4.1 index K [T, 128] and index Q [T, 32, 128]: MXFP4, linear, + # global_scale None -- the checkpoint's `fp4_act_quant(x, 32)`. Each + # 32-block gets its own power-of-two magnitude in 2^-20..2^20 + # (`block_mag=True`) so the UE8M0 scales spread. Rows 1..32768 are swept by + # the prefix test; these cells are the sweep's ends and the grid edges. + Cell( + why="deepseek-v4.1 index K MXFP4 [1, 128]", + spec=dict(shape=(1, 128), **_V41_INDEX), + ), + Cell( + why="deepseek-v4.1 index K MXFP4 [19457, 128]: past the launcher's 19456-row grid cap " + "(152 SMs x 128 blocks) at K 128", + spec=dict(shape=(19457, 128), **_V41_INDEX), + ), + Cell( + why="deepseek-v4.1 index K MXFP4 [32768, 128]: the certified row maximum", + spec=dict(shape=(32768, 128), **_V41_INDEX), + ), + Cell( + why="deepseek-v4.1 index Q MXFP4 [1, 32, 128]: rank 3, leading dims collapse in sf", + spec=dict(shape=(1, 32, 128), **_V41_INDEX), + ), + Cell( + why="deepseek-v4.1 index Q MXFP4 [609, 32, 128]: 19488 rows, past the 19456-row grid cap", + spec=dict(shape=(609, 32, 128), **_V41_INDEX), + ), + Cell( + why="deepseek-v4.1 index Q MXFP4 [32768, 32, 128]: the certified token maximum", + spec=dict(shape=(32768, 32, 128), **_V41_INDEX), + ), ) + + note = """ + `sf_use_ue8m0=True` selects power-of-two scales instead of e4m3. It is + certified only in the DeepSeek-V4.1 form below: `sf_vec_size=32`, bf16, + `K = 128`, linear layout. + + DeepSeek-V4.1: three call forms, no metadata. (1) The routed-expert input + `[r, 5120]` bf16, NVFP4, linear, with the layer's static `[1]` fp32 scale + `g1 = 1 / max_e max(w1.input_scale[e], w3.input_scale[e])` -- the reciprocal + of the checkpoint's `input_scale`; passing `input_scale` itself is accepted + and zeroes most scale bytes. The target splits the MoE rows into 8192-row + chunks and hands each over as a contiguous row view at an offset of a + multiple of 8192 rows. (2) Index K `[r, 128]` and (3) index Q `[r, 32, 128]`, + `fp4_quantize(x, None, 32, True, False)`, bit for bit the checkpoint's + `fp4_act_quant(x, 32)` with no tie window; `data` feeds + `indexer_k_cache_scatter_op` as `k_data` and `sf.view(torch.int32).view(T, 1)` + as `k_scale`. Certified eagerly at every row count 1..8192 for (1) (both + layouts) and 1..32768 for (2) and (3), each prefix bitwise equal to the + longest call, one kernel instantiation per mode at every probed row count; + the 8192-row chunk views of a 65536-row input; and the three calls captured + together in one CUDA graph at n in {1..32, 64, 96, 128, 1024, 8192}, and as a + whole decode step at every graph batch G of {1..32, 40..128 step 8, 192, 256} + with next_n 1 and 6 (rows G, 3G, 6G, 4G, 24G) -- a replay reads the inputs and + `global_scale` from device memory and rewrites the same two output buffers. + Hazards: V4.1's own g1 lands in the NVFP4 tie regime (0.19% near-ties, + 0.033% resolved the other way, one adjacent code each), and it crosses the + `g * vecmax / 6 <= 448` bound on real inputs (column block 153 of the BOS and + a few other rows in layers 20 and 28-37, up to 6.1x); the op clamps those + blocks to 448 silently. In MXFP4 mode a block whose max is at or below + `3 * 2^-126` (in practice an all-zero padding row) gets scale byte 0x00 where + `fp4_act_quant` writes 0x01; both decode to zero. Not certified for V4.1: + MXFP4 with float16 input, the swizzled layout, `K != 128`, non-finite input + or zero rows; either mode above its row maximum in one call; CUDA-graph + capture at row counts outside the sets above or of any other call form. + """ + + def raw_call( + self, + input: torch.Tensor, + global_scale: Optional[torch.Tensor], + sf_vec_size: int, + sf_use_ue8m0: bool = False, + is_sf_swizzled_layout: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.fp4_quantize( + input, global_scale, sf_vec_size, sf_use_ue8m0, is_sf_swizzled_layout + ) + + def reference( + self, + input: torch.Tensor, + global_scale: Optional[torch.Tensor], + sf_vec_size: int, + sf_use_ue8m0: bool = False, + is_sf_swizzled_layout: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Native-torch NVFP4, with the scales in the requested layout. + + Per block: `vecmax = max|x|`, the scale is `e4m3(global_scale * vecmax / + 6)`, and the data is the input times `global_scale / scale` rounded to + e2m1. torch's fp32 -> e4m3 cast emits NaN above 464 where the kernel's + saturates, so the scale is clamped first -- between 448 and 464 both + round to 448 anyway. + + With `sf_use_ue8m0=True` (MXFP4) the reference is DeepSeek-V4.1's + `fp4_act_quant(x, 32)` instead, and `global_scale` is ignored: per + 32-block, `s = 2^ceil(log2(max(vecmax, 6 * 2^-126) * fp32(1/6)))`, the + scale byte is `log2(s) + 127`, and the data is `x / s` (exact, `s` is a + power of two) rounded to e2m1, ties to even. `data` keeps the input's + leading dims. The op has no `6 * 2^-126` floor: a block whose max is at + or below `3 * 2^-126` gets scale byte 0x00 and signed-zero data, where + this writes 0x01 (and, for a max above 2^-128, nonzero codes). Those + are the only blocks where the two differ. + """ + if sf_use_ue8m0: + packed, sf = self._mx_quantize(input, sf_vec_size) + else: + packed, sf, _ = self._quantize(input, global_scale, sf_vec_size) + if not is_sf_swizzled_layout: + return packed, sf.reshape(-1) + return packed, self._swizzle(sf) + + def scaled( + self, input: torch.Tensor, global_scale: Optional[torch.Tensor], sf_vec_size: int + ) -> torch.Tensor: + """The pre-rounding product, which is what `near_tie` reads. + + Exposed because a caller cannot reconstruct it from the op's outputs: + rounding is the step that loses it. + """ + return self._quantize(input, global_scale, sf_vec_size)[2] + + @staticmethod + def near_tie(scaled: torch.Tensor) -> torch.Tensor: + """Elements whose rounding this entry refuses to predict. + + Inside `TIE_WINDOW` of an e2m1 midpoint the kernel's approximate + reciprocals decide the direction, so a code that differs there is not a + defect. Outside it, a difference is. + """ + mid = torch.tensor(_E2M1_MIDPOINTS, device=scaled.device, dtype=torch.float32) + a = scaled.abs().unsqueeze(-1) + return ((a - mid).abs() <= TIE_WINDOW * a).any(-1) + + def is_valid( + self, + input: torch.Tensor, + global_scale: Optional[torch.Tensor], + sf_vec_size: int, + sf_use_ue8m0: bool = False, + is_sf_swizzled_layout: bool = True, + ) -> None: + # The kernel loads one scalar from global_scale and ignores every + # element past the first, so a per-token [num_tokens] tensor is silently + # applied as global_scale[0] to every row -- observed to return a + # valid-looking result, never to raise. + assert global_scale is None or global_scale.numel() == 1, ( + "global_scale must hold exactly one element; extra elements are " + "silently ignored and global_scale[0] is applied to every row" + ) + assert input.shape[-1] % sf_vec_size == 0, ( + f"the quantized dim must be a whole number of {sf_vec_size}-element " + f"blocks; got {input.shape[-1]}" + ) + assert input.is_contiguous(), "input must be contiguous; the kernel reads it flat" + + # ── the arithmetic, shared by reference() and scaled() ────────────────── + + def _quantize( + self, input: torch.Tensor, global_scale: Optional[torch.Tensor], vec: int + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + gs = 1.0 if global_scale is None else global_scale.float().reshape(()) + k = input.shape[-1] + rows = input.numel() // k + flat = input.reshape(rows, k) + vecmax = flat.abs().reshape(rows, k // vec, vec).amax(-1).float() + sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) + out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) + scaled = (flat.float().reshape(rows, k // vec, vec) * out_scale.unsqueeze(-1)).reshape( + rows, k + ) + codes = self._e2m1_codes(scaled) + packed = codes[:, 0::2] | (codes[:, 1::2] << 4) + return packed, sf.view(torch.uint8), scaled + + @staticmethod + def _mx_quantize(input: torch.Tensor, vec: int) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP4: (packed data [..., K/2], UE8M0 scale bytes [rows, K/vec]).""" + k = input.shape[-1] + rows = input.numel() // k + xf = input.reshape(rows, k // vec, vec).float() + amax = xf.abs().amax(-1).clamp_min(_MX_AMAX_FLOOR) + # ceil(log2(amax / 6)) from the fp32 product's bits: the exponent, plus one + # unless the mantissa is zero (an exact power of two). + bits = (amax * torch.tensor(_INV6_F32, dtype=torch.float32, device=amax.device)).view( + torch.int32 + ) + e = ((bits >> 23) & 0xFF) - 127 + ((bits & 0x7FFFFF) != 0).to(torch.int32) + s = ((e + 127) << 23).view(torch.float32) + q = (xf / s.unsqueeze(-1)).clamp(-_E2M1_MAX, _E2M1_MAX) + # e2m1 steps are 0.5 below 2, 1 on [2, 4) and 2 on [4, 6]; torch.round is + # half-to-even on each binade's grid. + a = q.abs() + r = torch.where( + a < 2, + torch.round(a * 2) / 2, + torch.where(a < 4, torch.round(a), torch.round(a / 2) * 2), + ) + # value -> code: {0, 0.5, 1, 1.5} -> 2r; {2, 3} -> r + 2; {4, 6} -> r/2 + 4 + code = torch.where(r < 2, r * 2, torch.where(r < 4, r + 2, r / 2 + 4)).to(torch.uint8) + # The sign comes from the input, so -0.0 and a negative lane rounding to 0 are code 8. + codes = (code | (torch.signbit(xf).to(torch.uint8) << 3)).reshape(rows, k) + packed = codes[:, 0::2] | (codes[:, 1::2] << 4) + return packed.reshape(*input.shape[:-1], k // 2), (e + 127).to(torch.uint8) + + @staticmethod + def _e2m1_codes(v: torch.Tensor) -> torch.Tensor: + """fp32 -> 4-bit e2m1 codes (bit 3 = sign), round-to-nearest-even.""" + mid = torch.tensor(_E2M1_MIDPOINTS, device=v.device, dtype=torch.float32) + tie_up = torch.tensor(_TIE_UP, device=v.device, dtype=torch.bool) + a = v.abs().unsqueeze(-1) + greater = (a > mid).sum(-1) # strict: ties provisionally round down + up = ((a == mid) & tie_up).any(-1) # ties whose even side is up + code = (greater + up).to(torch.uint8) + negative = (v < 0) | ((v == 0) & torch.signbit(v)) + return code | (negative.to(torch.uint8) << 3) + + @staticmethod + def _swizzle(sf: torch.Tensor) -> torch.Tensor: + """Scatter [rows, cols] scale bytes into the 128x4-swizzled buffer.""" + rows, cols = sf.shape + padded_cols, padded_rows = _pad_up(cols, 4), _pad_up(rows, 128) + r = torch.arange(rows, device=sf.device).view(-1, 1) + c = torch.arange(cols, device=sf.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * (4 * 128) + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=sf.device) + out[index.reshape(-1)] = sf.reshape(-1) + return out + + +fp4_quantize = Fp4Quantize() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/mxfp8_quantize.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/mxfp8_quantize.md deleted file mode 100644 index 83275c688373..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/mxfp8_quantize.md +++ /dev/null @@ -1,169 +0,0 @@ ---- -receipts: - sm_103: {status: passed, tests: 12} ---- - -# mxfp8_quantize - -**Wraps** `torch.ops.trtllm.mxfp8_quantize` (one call). - -## Semantics - -Dynamic MXFP8 (OCP microscaling FP8) quantization of an activation tensor: -one call converts bf16/fp16 values to e4m3 elements plus one UE8M0 -(power-of-two) scale per **32 contiguous elements along the last dim**. -Nothing else is fused in — no normalization, no activation, no transpose. -The caller owns producing the input (e.g. bf16 hidden states) and owns -both outputs afterwards. - -Let `M` = product of all leading dims, `K` = last dim, and -`padded_k = ceil(K / alignment) * alignment`. For row `m` and block -`b` of 32 elements: - -``` -amax = max |x[m, 32b : 32b+32]| # fp32; exact for bf16/fp16 inputs -scale = smallest power of two >= amax / 448 # 448 = e4m3 max finite -sf[m, b] = log2(scale) + 127 # UE8M0 byte, so scale == 2^(sf-127) -data[m,i] = e4m3_round_to_nearest_even( x[m, i] / scale ) for i in the block -``` - -The scale is rounded **up** (toward +inf) to the next power of two, so -the block max always lands at or below e4m3's 448 and no element -saturates. Dequantization is exactly `data.float() * 2^(sf - 127)`; -per element the reconstruction error is at most `2^-4` relative (e4m3's -3-bit mantissa) plus `2^-10 * scale` absolute (e4m3 subnormals). - -Padding is part of the op, not the caller's job: - -- **Column padding** (`padded_k > K`, i.e. `alignment` wider than `K`): - columns `K .. padded_k` of `data` are written as the zero byte `0x00` - and their scale bytes are `0x00`. The valid columns are bit-identical - to what an unpadded call produces. -- **Row padding** (swizzled layout only): the scale buffer is sized for - `pad_up(M, 128)` rows; the bytes belonging to rows `M .. pad_up(M,128)` - are `0x00`. The `data` tensor is never row-padded. - -`data` is identical byte-for-byte under both `swizzled_layout` values — -only the scale buffer's size and element order change: - -- `swizzled_layout=False` (**linear**): `M * cols` bytes, row-major, where - `cols = padded_k / 32`. Byte `b` of row `m` is at flat offset - `m * cols + b`; `sf.view(M, cols)` is exactly the per-block scale matrix. -- `swizzled_layout=True` (**128x4 swizzled**): `pad_up(M,128) * pad_up(cols,4)` - bytes. The scale of `(m, b)` sits at flat offset - - ``` - (b % 4) - + (b // 4) * 512 # 4 cols x 128 rows per column group - + (m % 32) * 16 - + ((m % 128) // 32) * 4 - + (m // 128) * 128 * pad_up(cols, 4) - ``` - - Every offset not addressed by a real `(m, b)` pair is `0x00` padding. - -## Signature - -```python -def mxfp8_quantize( - input: torch.Tensor, - swizzled_layout: bool = True, - alignment: int = 32, -) -> tuple[torch.Tensor, torch.Tensor] -``` - -| Argument | Shape | Dtype | Layout / value | Device | -|---|---|---|---|---| -| `input` | `[..., K]`, rank >= 2 | bfloat16 or float16 | contiguous | CUDA | -| `swizzled_layout` | scalar | Python bool | `False` = linear scale order, `True` = 128x4 swizzled | — | -| `alignment` | scalar | Python int | positive multiple of 32; pads `K` up to `ceil(K/alignment)*alignment` | — | -| returns `data` | `[..., padded_k]` (leading dims preserved) | float8_e4m3fn | contiguous, freshly allocated | same as `input` | -| returns `sf` | 1-D: `[M*cols]` (linear) or `[pad_up(M,128)*pad_up(cols,4)]` (swizzled) | uint8 | contiguous, freshly allocated | same as `input` | - -`M` = product of `input`'s leading dims, `cols = padded_k / 32`. `input` -is not modified. Leading dims are collapsed for the scale buffer in both -layouts (a `[B, S, K]` input yields a scale buffer laid out for `B*S` -rows), while `data` keeps `input`'s leading dims. - -The op's own schema names the second parameter `swizzedLayout` (upstream -spelling); it is positional here. - -## Metadata consumed - -None. Stateless — no runtime, attention metadata, or workspace. - -## Preconditions - -- `input` is on CUDA, contiguous, rank >= 2, dtype bfloat16 or float16. - Violations raise: a CPU tensor, a non-contiguous view (column slice or - transpose), a 1-D tensor and an fp32 tensor were each observed to raise - on this machine rather than compute something wrong. -- `K % 32 == 0` (the block size). Raises otherwise. -- `alignment % 32 == 0` and `alignment > 0`. A non-multiple raises; - `alignment = 0` is **not** an exception — it kills the process with - SIGFPE (integer division by zero in the padding computation). There is - no upper bound: `alignment` may exceed `K` (`K=2880, alignment=512` - gives `padded_k=3072`) or equal the block size (`alignment=32`, no pad). - A **negative** multiple of 32 is the quiet one: it passes the - `% 32 == 0` guard, raises nothing, and **silently truncates** — at - `K = 2880` the values `-32 / -64 / -512` return `padded_k` - `2816 / 2752 / 2048`, matching `((K + a - 1) / a) * a` exactly, i.e. the - round-up above turning into a round-*down* under truncating division. - The columns - that survive are bit-exact against the reference and the scale tensor is - self-consistent, so the result looks entirely well-formed; the tail of - every row is simply gone. Check the sign before the modulus. -- `input` must be **finite**. Non-finite lanes are accepted silently and - are destructive: - - a `+inf`/`-inf` lane drives the block scale to the E8M0 finite max - (byte 254 = 2^127); the inf lane becomes NaN and **every other lane in - that 32-element block is zeroed**; - - a NaN lane is excluded from the block max, so the scale stays correct - and the NaN stays confined to its own lane. -- Every 32-element block must have `amax == 0` or `amax > 448 * 2^-127` - (~2.63e-36). An all-zero block is fine (scale byte `0x00`, all-zero - data). A block whose max is nonzero but at or below that threshold - silently produces garbage: scale byte `0x00`, every nonzero lane - saturated to +-448, and **every exact-zero lane turned into NaN**. Real - bf16 hidden states never reach this range; a block of scaled-down - denormals does. -- `M == 0` (zero rows) is accepted and returns empty `data` and `sf`. -- The caller owns both returned buffers; nothing else writes them. - -## Notes - -- Certified on sm_100 (B200) only. The kernel behind this op is - Blackwell-targeted: the same TensorRT-LLM quantization kernel is - vendored as readable source elsewhere in this install (flashinfer's - bundled `nv_internal` copy) and compiles its body only under - `__CUDA_ARCH__ >= 1000`, the empty branch returning without writing the - output buffers. A pre-Blackwell arch is therefore a silent-garbage risk - rather than a raise. No receipt is claimed there. -- Against a native-torch reference built from the formula above, the - kernel was **bit-exact** on every element and every scale byte across - `T x 2880` and `T x 3072` bf16 inputs for `T` in {1, 2, 7, 64, 1024, - 8192} and `alignment` in {32, 128, 512}, plus fp16 inputs and a 3-D - input. The scale math runs through approximate reciprocals (the - flush-to-zero surface above is their footprint), yet the result did not - shift even when `amax` is exactly `448 * 2^k` — the only value where a - 1-ulp difference could flip the round-up (verified for `k = -2, 0, 1, 3`). -- Scale-byte dynamic range from bf16/fp16 inputs: the largest finite bf16 - magnitude (3.39e38) yields byte 247, so the E8M0 finite max (254) is - reachable only from a non-finite input. -- Consumer pairing. The trtllm-gen W4A8 MXFP4xMXFP8 MoE runner - `torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner` takes this op's - two outputs as its `hidden_states` / `hidden_states_scale` and expects - the **linear** scale buffer (`swizzled_layout=False`), sized exactly - `M * padded_k/32` bytes, with `alignment` equal to that kernel's - input-hidden alignment — 512 for the trtllm-gen MXFP4 weight path, which - is what pads gpt-oss's 2880 hidden to 3072 inside this call (so the - caller does not pre-pad). The swizzled layout is what the CUTLASS-side - mxfp8 GEMM/MoE consumers take instead. Those pairings are fixed by the - consuming ops' own contracts; this entry certifies only what each layout - contains. -- Sibling quantizers with the same scale-layout vocabulary exist - (`torch.ops.trtllm.fp4_quantize` for NVFP4/MXFP4 with a 16- or - 32-element vector, `torch.ops.trtllm.fp8_quantize_1x128` for 1x128 - block-scaled fp8); this entry covers the fixed-32-block mxfp8 path only. -- The op allocates both outputs itself: it is functional, with no `out=` - parameter and no in-place mode. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/mxfp8_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/mxfp8_quantize.py index 511d3051fef8..5956f33f60c5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/mxfp8_quantize.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/mxfp8_quantize.py @@ -1,21 +1,239 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Dynamic MXFP8 quantization: bf16/fp16 -> e4m3 data + per-32-element UE8M0 block scales.""" +"""MXFP8 quantization: e4m3 data plus one UE8M0 scale per 32 elements.""" + +from __future__ import annotations import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from .._op import Arch, Cell, OpWrapper + +_E4M3_MAX = 448.0 +_BLOCK = 32 + + +def _pad_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +# DeepSeek-V4.1's dense FP8 linear inputs, (K, which linear): each quantized +# with the default call (swizzled, alignment 32) and handed unchanged to the +# MXFP8 GEMM. The wo_b input is 8 / tp * 1024 per rank. Every row count +# 1..32768 is certified at each K by the row sweep in the test; the cells hold +# the sweep's two ends. +_DSV41_K = ( + (5120, "hidden (wqkv_a, shared w1 || w3)"), + (1280, "q_lora (wq_b, indexer wq_b)"), + (2048, "wo_b under attention TP4"), + (8192, "wo_b under attention DP"), + (2304, "shared w2"), + (6144, "Engram wkv"), + (4096, "wo_b under attention TP2"), + (1024, "wo_b under attention TP8"), +) +_DSV41_ROWS = (1, 32768) +# The DSpark drafter's two calls the core target does not make, (tag, K, +# swizzled, alignment, what): every row count 1..1280 is certified, the cells +# hold the two ends. +_DSV41_DSPARK = ( + ("MoE in", 5120, False, 512, "the MXFP4 MoE runner's linear-layout input"), + ("main_proj", 15360, True, 32, "main_proj's input"), +) +_DSV41_DSPARK_ROWS = (1, 1280) + +_DSV41_CELLS = tuple( + Cell( + why=f"deepseek-v4.1 K={k} M={m}: {site} input, swizzled, alignment 32, " + f"{'one row' if m == 1 else 'top'} of the row sweep 1..32768", + spec=dict( + shape=(m, k), dtype=torch.bfloat16, swizzled=True, alignment=32, data="block_scaled" + ), + ) + for k, site in _DSV41_K + for m in _DSV41_ROWS +) + tuple( + Cell( + why=f"deepseek-v4.1 DSpark {tag} M={m}: {what}, K {k}, " + f"{'swizzled' if swizzled else 'linear'}, alignment {alignment}, " + f"{'one row' if m == 1 else 'top'} of the row sweep 1..1280", + spec=dict( + shape=(m, k), + dtype=torch.bfloat16, + swizzled=swizzled, + alignment=alignment, + data="block_scaled", + ), + ) + for tag, k, swizzled, alignment, what in _DSV41_DSPARK + for m in _DSV41_DSPARK_ROWS +) -def mxfp8_quantize( - input: torch.Tensor, - swizzled_layout: bool = True, - alignment: int = 32, -) -> tuple[torch.Tensor, torch.Tensor]: - """Quantize to MXFP8: returns (e4m3 data [..., pad_up(K, alignment)], uint8 UE8M0 scales, 1D). - One scale per 32 contiguous elements along the last dim; `swizzled_layout` - selects the 128x4 swizzled scale order (True) or the row-major linear - order (False). +class Mxfp8Quantize(OpWrapper): + """Quantize to MXFP8: `(e4m3 data [..., pad_up(K, alignment)], uint8 scales)`. + + One scale per 32 contiguous elements of the last dim. The scale is the + smallest power of two at or above `amax / 448`, carried as its E8M0 byte + (`exponent + 127`) rather than a float. + + `swizzled_layout` selects the 128x4 swizzled scale order (True) or the + row-major linear order (False). The scale tensor is 1-D either way; the + layout is what a consumer kernel expects, not something the shape shows. """ - return torch.ops.trtllm.mxfp8_quantize(input, swizzled_layout, alignment) + + ARCHS = frozenset({Arch.SM_103}) + + # gpt-oss asks for the *linear* layout with the 512-element alignment its + # MoE FC1 needs -- hidden 2880 pads to 3072. The swizzled cell at its width + # is here because the two layouts are separate code paths in the kernel; + # DeepSeek-V4.1, whose cells follow, is the target that takes the swizzled + # layout. + CELLS: tuple[Cell, ...] = ( + Cell( + why="gpt-oss hidden 2880, linear scales, aligned to 512 -> padded 3072", + spec=dict(shape=(16, 2880), dtype=torch.bfloat16, swizzled=False, alignment=512), + ), + Cell( + why="decode-shaped: one token through the same padding", + spec=dict(shape=(1, 2880), dtype=torch.bfloat16, swizzled=False, alignment=512), + ), + Cell( + why="prefill-shaped, and enough rows to cross the swizzle's 128-row tile", + spec=dict(shape=(2048, 2880), dtype=torch.bfloat16, swizzled=False, alignment=512), + ), + Cell( + why="the swizzled scale order, the other of the kernel's two paths", + spec=dict(shape=(160, 2880), dtype=torch.bfloat16, swizzled=True, alignment=512), + ), + # DeepSeek-V4.1: see `_DSV41_K` and `_DSV41_DSPARK`. `data` names the + # input's value pattern (the test's `_build` documents it). + *_DSV41_CELLS, + ) + + note = """ + Blocks whose amax is below `448 * 2^-127` take a flush-to-zero path in the + kernel that `reference` does not model; no shipped target produces one, + since the residual stream it quantizes never gets that small. + + Leading dims are collapsed: a [B, T, K] input quantizes as [B*T, K], and + under the swizzled layout the collapsed rows are padded as one matrix. + + DeepSeek-V4.1: every dense FP8 linear of the core target calls it as + `mxfp8_quantize(x)` (swizzled, alignment 32) on contiguous bf16 `[M, K]` + at the eight K of `_DSV41_K` and feeds both outputs unchanged to the MXFP8 + GEMM. Metadata none. The row count is the rank's token count, and the + bounded decoder replay of layers 21-39 can reach 32768 rows uncut; every + M 1..32768 is certified at each K, bit-exact in every data byte, every + scale byte at its 128x4 offset and every row-padding byte. All eight K are + multiples of 128, so no scale column is padded; rows pad to 128. The + DSpark drafter adds its MoE input (linear, alignment 512, K 5120, the + layout the MXFP4 MoE runner consumes) and main_proj's input (swizzled, + alignment 32, K 15360), each certified at every M 1..1280. CUDA graphs, + after an eager call at each shape: the eight K at the 114 per-rank row + counts {G, 5G, 6G} over the 46 captured batch sizes up to 256, plus 1935; + the two drafter calls at 176 counts (every M 1..160 plus G and 5G). + Data in the certification has block maxima spread over 2^-24..2^24 so a + misplaced scale byte is visible. Not certified for V4.1: fp16 input; more + than 32768 rows; another K in the swizzled layout; an alignment above 32 + at the eight K; the linear layout beyond the MoE-input call; capture at + other row counts, at a batch size outside that list, or with no eager call + first. + """ + + def raw_call( + self, + input: torch.Tensor, + swizzled_layout: bool = True, + alignment: int = 32, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.mxfp8_quantize(input, swizzled_layout, alignment) + + def reference( + self, + input: torch.Tensor, + swizzled_layout: bool = True, + alignment: int = 32, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Native-torch MXFP8, returning the scales in the requested layout. + + Columns beyond the input's K, up to the aligned width, are zero -- an + all-zero block quantizes to scale byte 0 and zero data. + """ + k = input.shape[-1] + rows = input.numel() // k + padded_k = _pad_up(k, alignment) + + padded = torch.zeros(rows, padded_k, dtype=torch.float32, device=input.device) + padded[:, :k] = input.reshape(rows, k).float() + blocks = padded.view(rows, padded_k // _BLOCK, _BLOCK) + amax = blocks.abs().amax(dim=-1) + + # E8M0 of amax/448 rounded toward +inf: take the fp32 biased exponent + # and bump it whenever any mantissa bit is set, i.e. whenever the value + # is not already a power of two. amax is exact in fp32 for bf16 and + # fp16 inputs, so this is the exact round-up rather than an estimate. + bits = (amax / _E4M3_MAX).view(torch.int32) + byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) + scale = torch.exp2(byte.to(torch.float32) - 127.0).unsqueeze(-1) + data = (blocks / scale).view(rows, padded_k).to(torch.float8_e4m3fn) + + byte = byte.to(torch.uint8) + if not swizzled_layout: + return data, byte.reshape(-1) + return data, self._swizzle(byte) + + @staticmethod + def _swizzle(byte: torch.Tensor) -> torch.Tensor: + """Scatter [rows, cols] scale bytes into the 128x4-swizzled buffer.""" + rows, cols = byte.shape + padded_cols = _pad_up(cols, 4) + padded_rows = _pad_up(rows, 128) + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * (4 * 128) + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Bit-exact. + + Both sides compute the same power-of-two scale and the same e4m3 + rounding, so any difference is a defect rather than accumulated error. + fp8 and uint8 both compare as bytes. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + assert torch.equal(out.view(torch.uint8), ref.view(torch.uint8)), ( + "quantization is deterministic; a mismatch is a defect, not drift" + ) + + def is_valid( + self, + input: torch.Tensor, + swizzled_layout: bool = True, + alignment: int = 32, + ) -> None: + # A last dim that is not a whole number of blocks leaves a partial + # block the kernel reads past the end of, and an alignment that is not + # a multiple of the block size pads to a width the scale count cannot + # describe. Neither is rejected by the op. + assert alignment % _BLOCK == 0, ( + f"alignment must be a multiple of the {_BLOCK}-element scale block; got {alignment}" + ) + assert input.shape[-1] % _BLOCK == 0, ( + f"the quantized dim must be a whole number of {_BLOCK}-element blocks; " + f"got {input.shape[-1]}" + ) + assert input.is_contiguous(), "input must be contiguous; the kernel reads it flat" + + +mxfp8_quantize = Mxfp8Quantize() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/__init__.py deleted file mode 100644 index 9c12fb0b373a..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Thin mirrors of torch callables. - -Upstream owns their correctness, so these carry no contract, no test and -no receipt by design. They exist so a target's forward can satisfy the -closed-vocabulary rule without leaving the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/add.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/add.py deleted file mode 100644 index 9eb889284740..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/add.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Elementwise addition (thin mirror of torch.add). - -torch does not implement float8 arithmetic; inputs are expected in the -high-precision residual dtypes (fp32/bf16/fp16). -""" - -import torch - - -def add(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - """Return `a + b`; identical semantics to torch.add.""" - return torch.add(a, b) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/concat.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/concat.py deleted file mode 100644 index 8b60d8219c04..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/concat.py +++ /dev/null @@ -1,10 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Tensor concatenation (thin mirror of torch.cat).""" - -import torch - - -def concat(tensors: list[torch.Tensor], dim: int = 0) -> torch.Tensor: - """Concatenate `tensors` along `dim`; identical semantics to torch.cat.""" - return torch.cat(tensors, dim=dim) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/copy_.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/copy_.py deleted file mode 100644 index 0c32b729b4a9..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/copy_.py +++ /dev/null @@ -1,16 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""In-place copy into a tensor or view (thin mirror of torch.Tensor.copy_). - -Writes `src` into `dst` element-wise (broadcasting and dtype conversion follow -torch semantics) and returns `dst`. The canonical way to fill a slice of a -larger buffer (e.g. `copy_(k[..., :nope], k_nope)`); the destination's -pre-call contents are destroyed. -""" - -import torch - - -def copy_(dst: torch.Tensor, src: torch.Tensor) -> torch.Tensor: - """Copy `src` into `dst` in place; identical semantics to Tensor.copy_.""" - return dst.copy_(src) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/embedding.py deleted file mode 100644 index d9057a48f7e2..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/embedding.py +++ /dev/null @@ -1,11 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Input-embedding lookup (thin mirror of torch.nn.functional.embedding).""" - -import torch -import torch.nn.functional as F - - -def embedding(input_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - """Look up rows of `weight` by token id: [*] ids -> [*, H] embeddings.""" - return F.embedding(input_ids, weight) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/empty.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/empty.py deleted file mode 100644 index 4f0c5c805193..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/empty.py +++ /dev/null @@ -1,15 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Uninitialized tensor allocation (thin mirror of torch.empty). - -The general-purpose allocation primitive for scratch and output buffers -whose sizing the caller owns; layer-aware output allocation has dedicated -entries (e.g. create_attn_outputs). Contents are garbage until written. -""" - -import torch - - -def empty(shape: list[int], dtype: torch.dtype, device: torch.device | str) -> torch.Tensor: - """Return an uninitialized tensor; identical semantics to torch.empty.""" - return torch.empty(shape, dtype=dtype, device=device) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/expand.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/expand.py deleted file mode 100644 index 7265fe27216e..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/expand.py +++ /dev/null @@ -1,15 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Broadcast view over singleton dims (thin mirror of torch.Tensor.expand). - -Returns a zero-copy view with stride 0 on the expanded dims; only singleton -dims can be expanded, and -1 keeps a dim unchanged. Writing through the view -is unsafe (aliased elements) — copy first if mutation is needed. -""" - -import torch - - -def expand(x: torch.Tensor, sizes: list[int]) -> torch.Tensor: - """Return `x` broadcast to `sizes`; identical semantics to Tensor.expand.""" - return x.expand(sizes) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/pad.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/pad.py deleted file mode 100644 index e02e16737c0c..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/pad.py +++ /dev/null @@ -1,16 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Constant padding of a tensor's trailing dims (thin mirror of torch.nn.functional.pad). - -Layout glue: widens a tensor to a kernel's required operand width by -appending constant-valued columns (the padded region multiplies zero-valued -padded weights in block-scaled GEMMs, so the fill value must be finite). -""" - -import torch -import torch.nn.functional as F - - -def pad(x: torch.Tensor, padding: list[int], value: float = 0.0) -> torch.Tensor: - """Return `x` padded with `value` per the last-dim-first `padding` pairs.""" - return F.pad(x, padding, mode="constant", value=value) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/reshape.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/reshape.py deleted file mode 100644 index f167787c3b52..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/reshape.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Tensor reshape (thin mirror of torch.reshape). - -Covers view use cases too: returns a view when the layout allows, copies -otherwise — unlike Tensor.view, it never fails on non-contiguous input. -""" - -import torch - - -def reshape(x: torch.Tensor, shape: list[int]) -> torch.Tensor: - """Reshape `x` to `shape`; identical semantics to torch.reshape.""" - return torch.reshape(x, shape) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/split.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/split.py deleted file mode 100644 index 6653e4e6b1f0..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/split.py +++ /dev/null @@ -1,12 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Tensor split (thin mirror of torch.split).""" - -import torch - - -def split( - tensor: torch.Tensor, split_size_or_sections: int | list[int], dim: int = 0 -) -> tuple[torch.Tensor, ...]: - """Split `tensor` along `dim`; identical semantics to torch.split.""" - return torch.split(tensor, split_size_or_sections, dim=dim) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/transpose.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/transpose.py deleted file mode 100644 index 1a6a2235a94a..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/transpose.py +++ /dev/null @@ -1,15 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Zero-copy swap of two dimensions (thin mirror of torch.transpose). - -Returns a strided view sharing storage with the input; the canonical way to -present a `[tokens, heads, dim]` activation to a batched-gemm entry whose -batch axis is the head axis. -""" - -import torch - - -def transpose(x: torch.Tensor, dim0: int, dim1: int) -> torch.Tensor: - """Return `x` with `dim0` and `dim1` swapped; identical semantics to torch.transpose.""" - return torch.transpose(x, dim0, dim1) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/view_dtype.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/view_dtype.py deleted file mode 100644 index 59d33774d4d7..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/view_dtype.py +++ /dev/null @@ -1,16 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Bitwise dtype reinterpretation of a tensor (thin mirror of torch.Tensor.view(dtype)). - -Reinterprets the same bytes under another dtype of equal itemsize — no -conversion, no copy. The canonical bridge between a producer that emits a -byte buffer (e.g. packed block scales as uint8) and a consumer that demands -the typed view of those bytes (float8_e4m3fn). -""" - -import torch - - -def view_dtype(x: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: - """Reinterpret `x`'s bytes as `dtype`; identical semantics to Tensor.view(dtype).""" - return x.view(dtype) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/explain.py b/tensorrt_llm/_torch/_experimental/modeling_v2/explain.py index 683a11400d38..22ee3804757d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/explain.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/explain.py @@ -58,6 +58,9 @@ def build_parser() -> argparse.ArgumentParser: def main(argv: Optional[list] = None) -> int: args = build_parser().parse_args(argv) + # Registers the in-tree custom configs (e.g. deepseek_v41) with + # transformers, as the engine does before it reads a checkpoint. + import tensorrt_llm._torch.configs # noqa: F401 from tensorrt_llm._torch.model_config import ModelConfig world_size = args.tp * args.pp diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py index 65e903151a7a..d26b4c4c51e4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py @@ -2,50 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """ModelingV2 target: deepseek-r1-0528-nvfp4 / sm_103 / dep4 — self-contained modeling code. -Flat single-entry forward assembled from catalog entries only; every call -that creates or transforms a tensor is a catalog entry, everything else is -tensor-metadata reads and Python control flow. - -DeepSeek-R1-0528, modelopt NVFP4 export: 61 layers, hidden 7168, 128 query -heads, `q_lora_rank` 1536, layers 0-2 dense (intermediate 18432), layers 3-60 -MoE with 256 routed experts at top-8 plus one shared expert (intermediate -2048). The quantization is `nvfp4_moe_only`-shaped — every `self_attn*` and -`lm_head` is excluded from NVFP4, so attention, router, embedding and lm_head -are bf16 and only the MLP path is NVFP4 — **and the KV cache is fp8-e4m3** -(`hf_quant_config.json`: `kv_cache_quant_algo: FP8`). - -The query path is a LoRA pair (`q_lora_rank: 1536`), routing is group-limited -(`n_group 8`, `topk_group 4`), and the rope table is YaRN-scaled with the -attention temperature carried in `q_scaling` rather than in the table. - -**Layer 61 — the checkpoint's bf16 MTP module — is a second forward path this -file also carries, and it exists only when `speculative_config` turns it on.** -Under the target's identity config nothing below MTP is declared, layer 61's -790 keys stay a predicted non-load in the weight manifest, and the shell's -forward is a plain call to the inherited base. `MTPLayer` below is the -module's forward; the checkpoint ships no reference implementation of it and -neither does transformers, so what it computes is stated where it is built. - -`dep4` is `tensor_parallel_size: 4` plus `moe_expert_parallel_size: 4` plus -`enable_attention_dp: true`: the requests are split, not the heads. Everything -but the MoE is replicated and runs over a rank's own tokens, so there is no -attention-side collective; the MoE stays expert-parallel, rank `r` holding -experts `[64r, 64r+64)` whole, which costs one `comm/allgather` and one -`comm/reducescatter` per MoE layer. - -Weights are target-owned: a flat ParameterDict declared here (HF [out, in] -storage so checkpoint rows copy in unchanged, plus the kernel-ready expert -stacks weights.py builds during the load), with the column-major GEMM views, -the MLA absorption operands, the rope table and every NVFP4 call scalar -derived once after load. The registration shell inherits -DecoderModelForCausalLM for lm_head, packed-batch logits gathering, and the -meta-init/load/post-load hooks. - -Everything this forward relies on that is not visible in the call itself is -stated at the call: which scales are read and which are inert, why the two -decode producers are ordered rather than concurrent, why the gather precedes -the router GEMM, why the reduce-scatter is crossed in bf16, and which bounds -are certification boundaries rather than tuning knobs. +DeepSeek-R1-0528, modelopt NVFP4 export: 61 layers, hidden 7168, 128 query heads, layers 0-2 +dense and 3-60 MoE (256 routed experts at top-8 plus a shared expert). Only the MLP path is +NVFP4; attention, router, embedding and lm_head are bf16, and the KV cache is fp8-e4m3. Layer +61 is the checkpoint's bf16 MTP module (`MTPLayer`), live only under a `speculative_config`. +`dep4` is attention DP with the 256 experts split four ways, so only the MoE needs a collective. """ import math @@ -55,51 +16,32 @@ from torch import nn from transformers import PretrainedConfig +from tensorrt_llm._torch._experimental.modeling_v2._target import Phase, Target, phase_of +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import advance_step_generation from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( # noqa: E501 + FlashinferSiluAndMul, flashinfer_silu_and_mul, ) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.load_paged_kv_cache_for_mla import ( # noqa: E501 - load_paged_kv_cache_for_mla, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.mla_rope_append_paged_kv_assign_q import ( # noqa: E501 - mla_rope_append_paged_kv_assign_q, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.mla_rope_generation import ( - mla_rope_generation, -) from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( - thop_attention, + ThopAttention, ) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import allgather -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.reducescatter import reducescatter -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.bmm_out import bmm_out -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.nvfp4_gemm import nvfp4_gemm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import Allgather +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.reducescatter import Reducescatter +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.bmm_out import BmmOut +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import CublasMm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.nvfp4_gemm import Nvfp4Gemm from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( # noqa: E501 - fp4_block_scale_moe_runner, + Fp4BlockScaleMoeRunner, ) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fused_moe import fused_moe -from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.noaux_tc_op import noaux_tc_op from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 - flashinfer_fused_add_rmsnorm, + FlashinferFusedAddRmsnorm, ) from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( - flashinfer_rmsnorm, + FlashinferRmsnorm, ) from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( - fp4_quantize, + Fp4Quantize, ) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.add import add -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.concat import concat -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.copy_ import copy_ -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.embedding import embedding -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.empty import empty -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.expand import expand -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.pad import pad -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.reshape import reshape -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.split import split -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.transpose import transpose -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.view_dtype import view_dtype from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.model_config import ModelConfig @@ -111,86 +53,14 @@ from tensorrt_llm._torch.speculative import get_spec_worker from . import weights as _weights - -# The GPU architecture this target IS. Routing will not send another one here, -# but a direct instantiation could, and the certification is per arch: this -# assert is what the version pin used to be. In-tree the version moves with -# the code, so pinning it is meaningless; the architecture does not. -_SM = (10, 3) - - -#: Every trtllm op this target reaches for, in its forward and in the weight -#: load. Declared here, asserted in tests/unittest/_torch/modeling_v2: a symbol -#: that does not exist is a fact of the build, and the place to find that out -#: is a machine with the extension built rather than every import of this -#: module. -REQUIRED_TRTLLM_OPS = ( - "cublas_mm", - "bmm_out", - "nvfp4_gemm", - "flashinfer_rmsnorm", - "flashinfer_fused_add_rmsnorm", - "flashinfer_silu_and_mul", - "fp4_quantize", - "noaux_tc_op", - "fp4_block_scale_moe_runner", - "fused_moe", - "mla_rope_generation", - "mla_rope_append_paged_kv_assign_q", - "load_paged_kv_cache_for_mla", - "allgather", - "reducescatter", - # Not a forward call: the load-time expert/scale relayout in weights.py - # depends on it, so it belongs to the same contract. - "block_scale_interleave", +from .ops.attention.load_paged_kv_cache_for_mla import LoadPagedKvCacheForMla # noqa: E501 +from .ops.attention.mla_rope_append_paged_kv_assign_q import ( # noqa: E501 + MlaRopeAppendPagedKvAssignQ, ) +from .ops.attention.mla_rope_generation import MlaRopeGeneration +from .ops.moe.fused_moe import FusedMoe +from .ops.moe.noaux_tc_op import NoauxTcOp -# Metadata fields consumed each step (sourcing mirrors the in-tree -# FallbackFmha for this trtllm version; existence checked at first forward). -# The tail group is read only by the first-forward contract check: this -# target holds those features at the MLA columns' inert values, and the -# asserts make that honest instead of silently dropping an enabled feature. -_STEP_FIELDS = ( - "all_rank_num_tokens", - "kv_lens_cuda_runtime", - "kv_lens_runtime", - "host_total_kv_lens", - "prompt_lens_cuda_runtime", - "prompt_lens_cpu_runtime", - "host_request_types_runtime", - "kv_cache_block_offsets", - "host_kv_cache_pool_pointers", - "host_kv_cache_pool_mapping", - "effective_workspace", - "tokens_per_block", - "max_num_requests", - "max_context_length", - "max_seq_len", - "num_contexts", - "num_ctx_tokens", - "num_seqs", - "trtllm_gen_jit_warmup", - "effective_beam_width", - "cache_indirection", - "block_ids_per_seq", - "is_cross", - "is_spec_decoding_enabled", - "use_spec_decoding", - "flash_mla_tile_scheduler_metadata", - "flash_mla_num_splits", - # Both are engine-prepared per-instance constants (max_num_sequences - # defaults to max_num_requests; the tree-mask flag is set from - # is_spec_dec_dynamic_tree, and this target's MTP is a linear tree), so - # they project like the rest. - "max_num_sequences", - "force_prepare_spec_dec_tree_mask", -) - -# The cached-prefix context group. The engine only creates these attributes -# when it prepares the metadata for MLA context over reused blocks — under -# trtllm's default kv_cache_config that is on, and with block reuse disabled -# they are absent entirely — so their existence selects the context flavor -# rather than being a hard requirement. _CACHED_CTX_FIELDS = ( "enable_context_mla_with_cached_kv", "ctx_cached_token_indptr", @@ -201,14 +71,8 @@ def _build_step_args(md: TrtllmAttentionMetadata) -> dict: - """Project the prepared metadata onto the batch state both MLA attention - calls of one forward share; built once per forward. CUDA-graph classes: - every tensor here is an engine-owned persistent buffer refreshed in place - (reference class), and every Python int is an engine-construction or - per-capture constant (host-derived class — a decode-only graph always - sees num_contexts == 0). Phase-specific arguments (q/k/v, output, head - geometry, the scheduler buffers, the fp8 decode buffers) are passed at the - call sites.""" + """Project the prepared metadata onto the batch state both MLA attention calls of one + forward share.""" return dict( sequence_length=md.kv_lens_cuda_runtime, host_past_key_value_lengths=md.kv_lens_runtime, @@ -235,33 +99,13 @@ def _build_step_args(md: TrtllmAttentionMetadata) -> dict: ) -# Feature groups this target holds inert, shared by both MLA calls: the -# contract's MLA columns plus every not-certified group at its listed inert -# value. Causal mask, beam width 1, no sinks, no spec-dec mask / sparse / -# cross / relative-bias / mRoPE / helix / FlashMLA. `quant_mode` and -# `q_scaling` are *not* here — both are derived from this checkpoint's config -# in __init__ and joined on in `self._call`. Neither is -# `predicted_tokens_per_seq`: it is 1 on the context call and the generation -# call's own query-tokens-per-sequence, which is 1 for an ordinary decode step -# and `runtime_draft_len + 1` under MTP, so it is passed per call site. -# -# Both kv scale tensors stay None over the fp8 latent pool: the generation -# call reads neither, and both context flavors are correct only at s = 1.0, -# which None is read as exactly. This checkpoint's k_scale/v_scale are 1.0, -# asserted after load. -# -# position_embedding_type=8 selects the in-kernel GPT-J rope of the MLA path; -# rope_dim / rope_base / the two tables are per-instance. The seven scalars -# below (rope_scale_type, rope_scale, the two m-scales, and the two position -# windows in the step args) are measured inert on the MLA path — the table's -# content is the only rope input the op reads, and the whole YaRN blend lives -# there. They are held at neutral values rather than at the config's `yarn` -# names, whose trtllm enum coding is not derivable from this target's sources. -# The MLA KV pool stores no residual tail: the op accepts 0 or rope_size -# and rejects non-zero unless the pool is FP4, and this checkpoint's is -# fp8-e4m3. The in-tree caller passes a literal 0 on the same path. _KV_RESIDUAL_DIM = 0 +# Both kv scale tensors stay None: None is read as a scale of exactly 1.0, which is what this +# checkpoint's k_scale/v_scale hold and the only value the context flavors are correct at. +# The MLA path reads the rope table and nothing else — rope_scale_type, rope_scale, the two +# m-scales and the two position windows are silently ignored, so the YaRN blend must live in +# the table. _CALL_INERT = dict( output_sf=None, out_scale=None, @@ -311,234 +155,77 @@ def _build_step_args(md: TrtllmAttentionMetadata) -> dict: num_sparse_topk=None, flash_mla_tile_scheduler_metadata=None, flash_mla_num_splits=None, - # Held at the op's defaults, which are also what the in-tree caller passes - # on this path. kv_norm_weight is not merely a default: non-None would fold - # the kv_a_layernorm into the KV kernel and make it read latent_cache - # RAW, and this target normalizes that itself -- passing the weight would - # normalize twice. skip_correction is a lossy trtllm-gen MLA option - # (SM100/SM103, off by default upstream); enabling it is a caller's - # business, not the identity assembly's. + # kv_norm_weight must stay None: a weight here folds kv_a_layernorm into the KV kernel, + # which this target already applies itself — a silent double normalization. kv_norm_weight=None, kv_norm_eps=1e-6, skip_correction_threshold=0.0, ) -# QuantMode's fp8-KV-cache bit — the value every MLA entry certifies for an -# fp8-e4m3 latent pool. Bits outside the KV-cache group ride along unread -# (`1152` = | FP8_1x128_128x128 and `384` = | FP8_QDQ were measured -# bit-identical to a bare 128 on every MLA flavor and on the gather), so the -# derivation below maps the checkpoint's declared KV algo onto this bit rather -# than reconstructing the engine's full bitmask. _QUANT_MODE_FP8_KV = 128 -# NVFP4 block size: one e4m3 scale per 16 contiguous elements along K. _SF_VEC = 16 -# The MoE runner reads the activation scales as a *linear* (row-major) -# buffer; the dense GEMM reads the 128x4-swizzled one. The two have the same -# byte count whenever num_tokens is a multiple of 128 — every decode -# CUDA-graph batch — and the wrong one is then accepted silently, so both -# layouts are spelled out rather than left to the quantizer's default. +# The MoE runner reads the activation scales as a linear buffer and the dense GEMM reads the +# 128x4-swizzled one; the two have the same byte count whenever num_tokens is a multiple of 128 +# — every decode CUDA-graph batch — and the wrong one is then accepted silently. _SF_LINEAR = False _SF_SWIZZLED = True -# Routing already happened outside (noaux_tc_op), so the runner's routing -# method and group configuration are inert on this pre-routed path — including -# this checkpoint's (n_group 8, topk_group 4), measured bitwise identical to -# None/None at routing_method_type 1. _ROUTING_METHOD_INERT = 1 -# Largest token count the expert call is allowed to see in one invocation. -# `fp4_block_scale_moe_runner` is certified over its whole token column — up -# to 8192 — at this checkpoint's routed geometry (H 7168, I 2048, 256 experts -# top-8) for the full stack and for each of the four 64-wide expert-parallel -# windows; `noaux_tc_op` is certified to num_tokens 8192. Under dep4 the -# gathered token set reaches `4 * max_num_tokens` = 32768, so the call is -# chunked rather than run outside those columns. Routing and the activation -# quantization ride along per chunk — both are per-token, so chunking is a -# no-op for them mathematically. _MOE_MAX_T = 8192 _ACT_TYPE_SWIGLU = 0 -# `fused_moe`'s activation enum is a *different* enum from the trtllm-gen -# runner's above: 5 is Swiglu there, 0 is not a gated type at all. +# A different enum from the trtllm-gen runner's above: 5 is Swiglu here, 0 is not a gated type. _FUSED_MOE_ACT_SWIGLU = 5 -# The MTP layer's `eh_proj` operand order, at one place so flipping it is a -# one-line experiment. True = `concat(enorm(e), hnorm(h))`, the embedding block -# first — measured on this checkpoint's own weights by two independent -# statistics, which contradicts the -# DeepSeek-V3 report's `M_k[RMSNorm(h); RMSNorm(e)]` notation and agrees with -# the parameter's name. Nothing in the checkpoint's metadata pins it and -# nothing downstream detects a flip: the drafts are simply rejected, so the -# acceptance rate is what confirms it end to end. +# True = `concat(enorm(e), hnorm(h))`, which contradicts the DeepSeek-V3 report's +# `M_k[RMSNorm(h); RMSNorm(e)]`; a flip is silent, only the draft acceptance rate shows it. _MTP_EMBED_BLOCK_FIRST = True -# thop_attention's certified q_scaling values over an fp8 latent pool are 1.0 -# and DeepSeek-R1's YaRN temperature; the config-derived value is checked -# against the latter so a config change lands outside the certified column -# loudly rather than silently. -_CERTIFIED_YARN_Q_SCALING = 0.5336594 def _moe_chunk_sizes(total: int) -> list[int]: - """Split `total` gathered rows into expert-call chunks of at most - `_MOE_MAX_T` rows. Host arithmetic only — no tensor is created.""" + """Split `total` gathered rows into expert-call chunks of at most `_MOE_MAX_T` rows.""" full, rest = divmod(total, _MOE_MAX_T) return [_MOE_MAX_T] * full + ([rest] if rest else []) def _mtp_dp_rows(all_rank_num_tokens, rank: int, dp_size: int, rows: int) -> int: - """The uniform row count the MTP layer pads its token block to before its - MoE round trip: `max` over the group's per-rank counts. - - **The list is a parameter and there is no metadata argument, on purpose.** - Inside the draft loop `attn_metadata.all_rank_num_tokens` is the *trunk's* - list from draft step 1 onwards — the worker leaves it in place for the - whole loop and passes the correct basis in as the `all_rank_num_tokens` - keyword instead (`spec_metadata.all_rank_num_tokens` at step 0, then - `spec_metadata.subseq_all_rank_num_tokens`, which is the per-rank - *sequence* count). Both collective contracts certify that calls pair by - position and that at equal byte counts a divergence is silent — every rank - wrong in 98-99% of elements, bitwise reproducibly, with no hang — so a - helper that *could* reach the metadata is the whole hazard. This one - cannot. `_dp_rows` on the trunk keeps its own shape; the two are not - interchangeable.""" + """The uniform row count the MTP layer pads its token block to before its MoE round trip. + + Callers must pass the worker's own basis, never `attn_metadata.all_rank_num_tokens`, which + from draft step 1 onwards is still the trunk's list: at equal byte counts the collectives + diverge silently on every rank, with no hang.""" counts = [int(n) for n in all_rank_num_tokens] - assert len(counts) == dp_size, (counts, dp_size) - assert counts[rank] == rows, ( - f"rank {rank} holds {rows} MTP rows but the worker's " - f"all_rank_num_tokens says {counts[rank]} ({counts}); the collectives " - "would gather the wrong split" - ) return max(counts) def _yarn_mscale(factor: float, mscale: float) -> float: - """DeepSeek's YaRN magnitude scaling `m(x)`: 1.0 at or below factor 1, - `0.1 * x * ln(factor) + 1` above it.""" + """DeepSeek's YaRN magnitude scaling `m(x)`.""" if factor <= 1.0: return 1.0 return 0.1 * mscale * math.log(factor) + 1.0 -class ModelingV2Core(DecoderModel): +class DeepseekV3Dep4ModelingV2(DecoderModel): def __init__(self, model_config: ModelConfig): super().__init__(model_config) cfg = model_config.pretrained_config - assert cfg is not None - - # This target IS this topology and this geometry — assert, never - # adapt. Read the topology off the mapping the engine built, never - # off the path segment: the segment names the intent. - assert torch.cuda.get_device_capability() == _SM, ( - f"target certified on sm_{_SM[0]}{_SM[1]}, running on " - f"sm_{''.join(map(str, torch.cuda.get_device_capability()))}" - ) + # The first-step derivation runs once and always; the contract probe that + # shared this flag now lives on each target, behind its own. + self._derive_pending = True + mapping = model_config.mapping - assert mapping.world_size == 4, f"dep4 target, world size {mapping.world_size}" - assert mapping.tp_size == 4, f"dep4 target, tp_size {mapping.tp_size}" - assert mapping.pp_size == 1, "pipeline parallelism is not implemented" - assert mapping.moe_ep_size == 4, f"dep4 target, ep {mapping.moe_ep_size}" - assert mapping.moe_tp_size == 1, ( - "the routed experts are split by expert parallelism only; a " - f"moe_tp_size of {mapping.moe_tp_size} splits them a second way" - ) - assert mapping.enable_attention_dp, ( - "attention data parallelism is what the dep4 segment declares; " - "without it attention would be head-split (that is tep4)" - ) self.rank = mapping.rank self.ep_rank = mapping.moe_ep_rank self.ep_size = mapping.moe_ep_size - # The collective group: at pp_size 1 with one expert-parallel split, - # every rank participates in both halves of the MoE round trip, so - # the group is the whole world. self.dp_group = list(range(mapping.world_size)) self.dp_size = mapping.world_size - dt = model_config.torch_dtype - assert dt == torch.bfloat16, f"bf16 target, engine resolved {dt}" - assert cfg.torch_dtype == torch.bfloat16, cfg.torch_dtype - assert not cfg.tie_word_embeddings, "untied lm_head" - assert not cfg.attention_bias, "no q/kv/o bias anywhere" - assert cfg.hidden_act == "silu", "SwiGLU over a silu gate" - # The NVFP4 recipe covers the MLP only; the latent KV pool is fp8-e4m3 - # per the checkpoint's own quant config, which is what selects - # quant_mode on every MLA call. Derived, not hard-coded: a checkpoint - # declaring no KV quantization is a bf16-pool target and a different - # assembly. - kv_algo = model_config.quant_config.kv_cache_quant_algo - assert kv_algo is not None and str(kv_algo).upper().endswith("FP8"), ( - f"this target is the fp8-e4m3 latent-pool assembly; the checkpoint " - f"declares kv_cache_quant_algo {kv_algo!r}" - ) self.quant_mode = _QUANT_MODE_FP8_KV - self.num_layers = cfg.num_hidden_layers - self.hidden = cfg.hidden_size - self.eps = cfg.rms_norm_eps - self.vocab = cfg.vocab_size - # The checkpoint ships `num_nextn_predict_layers` extra decoder layers - # past `num_hidden_layers` for multi-token prediction — on this one a - # single bf16 layer 61 with its own 256 experts, embedding, eh_proj, - # two extra norms and an output head. self.mtp_layers = int(getattr(cfg, "num_nextn_predict_layers", 0) or 0) - # Whether this engine drafts. The checkpoint decides the *mode* - # (`num_nextn_predict_layers: 1` -> MTP-Eagle one-model, one layer - # replayed `max_draft_len` times); the caller's - # `speculative_config` decides whether it runs at all, and the engine - # has already resolved that onto `model_config.spec_config` by the time - # the model is built. With it absent — the target's identity config — - # layer 61 is not part of this model at all: its keys stay a predicted - # non-load in the weight manifest and nothing below is declared. spec_config = getattr(model_config, "spec_config", None) self.mtp_enabled = spec_config is not None - if self.mtp_enabled: - assert self.mtp_layers == 1, ( - "this target implements the one-layer MTP-Eagle module the " - f"checkpoint ships; num_nextn_predict_layers is {self.mtp_layers}" - ) - # MLA geometry. num_key_value_heads is 128 in this config but MLA has - # no separate KV heads: the context call runs Hq == Hkv == heads over - # head_size nope+rope, the generation call runs Hq == heads against a - # single latent KV head of width kv_lora+rope. Attention is - # replicated under DP, so every rank runs the whole head set. - self.heads = cfg.num_attention_heads - self.nope = cfg.qk_nope_head_dim - self.rope = cfg.qk_rope_head_dim - self.v_dim = cfg.v_head_dim - self.kv_lora = cfg.kv_lora_rank - self.q_lora = cfg.q_lora_rank - self.qk_dim = self.nope + self.rope - self.lat_dim = self.kv_lora + self.rope - # The query is a LoRA pair here: q_a_proj -> q_a_layernorm -> q_b_proj. - # A checkpoint with q_lora_rank null projects directly and needs the - # single-q_proj path instead (that is the deepseek-v3-lite sibling). - assert isinstance(self.q_lora, int) and self.q_lora > 0, ( - "this target implements the q-LoRA query path; a checkpoint with " - f"q_lora_rank {self.q_lora!r} needs the direct q_proj path" - ) - assert cfg.num_key_value_heads == cfg.num_attention_heads, "MLA: Hkv == Hq" - # thop_attention aborts the process on a head_size its FMHA kernels - # do not carry, so both call shapes are pinned here. - assert self.qk_dim == 192 and self.lat_dim == 576 and self.v_dim == 128, ( - f"certified MLA head dims are 192 (context) / 576 (generation) / " - f"128 (v), got {self.qk_dim} / {self.lat_dim} / {self.v_dim}" - ) - # The MLA generation phase compiles a decode kernel per head count, so - # the count is a certification axis rather than a free shape — and over - # an fp8-e4m3 latent pool 128 is the *only* certified count. - assert self.heads == 128, ( - f"thop_attention certifies MLA over an fp8 latent pool at 128 query " - f"heads only; this topology yields {self.heads}" - ) - - # RoPE: GPT-J interleaved pairs over the rope slice, YaRN-scaled. The - # engine hands either the flat pre-migration fields or the - # transformers-5.x rope_parameters dict; both shapes are resolved and - # every scalar the table depends on is asserted rather than defaulted. rope_cfg = getattr(cfg, "rope_scaling", None) or getattr(cfg, "rope_parameters", None) - assert isinstance(rope_cfg, dict), f"YaRN rope config expected, got {rope_cfg!r}" - rope_kind = rope_cfg.get("type", rope_cfg.get("rope_type")) - assert rope_kind == "yarn", f"this target builds a YaRN table, got {rope_kind!r}" theta = rope_cfg.get("rope_theta", getattr(cfg, "rope_theta", None)) - assert theta is not None and float(theta) > 0.0, "rope theta" self.theta = float(theta) self.rope_factor = float(rope_cfg["factor"]) self.rope_orig_max = int(rope_cfg["original_max_position_embeddings"]) @@ -546,259 +233,51 @@ def __init__(self, model_config: ModelConfig): self.beta_slow = float(rope_cfg["beta_slow"]) mscale = float(rope_cfg["mscale"]) mscale_all_dim = float(rope_cfg["mscale_all_dim"]) - assert self.rope_factor > 1.0 and self.rope_orig_max > 0, rope_cfg - assert getattr(cfg, "rope_interleave", True), ( - "the MLA ops apply GPT-J (interleaved-pair) rope in kernel" - ) - self.max_pos = cfg.max_position_embeddings - # The table's amplitude and the softmax temperature are the two halves - # of YaRN's magnitude correction, and they go to different places: the - # amplitude multiplies cos/sin (exactly 1.0 whenever the config's two - # m-scales agree), while the temperature — built from mscale_all_dim, - # as the reference model does — is folded into the op's q_scaling as - # 1/m^2. Putting either in the other's place is silently wrong. + # YaRN's magnitude correction splits in two: the amplitude multiplies cos/sin, the + # temperature folds into q_scaling as 1/m^2. Either in the other's place is silently wrong. self.rope_amplitude = _yarn_mscale(self.rope_factor, mscale) / _yarn_mscale( self.rope_factor, mscale_all_dim ) temperature = _yarn_mscale(self.rope_factor, mscale_all_dim) self.q_scaling = 1.0 / (temperature * temperature) - assert abs(self.q_scaling - _CERTIFIED_YARN_Q_SCALING) < 1e-6, ( - f"q_scaling {self.q_scaling} is outside thop_attention's certified " - f"fp8-pool column (1.0 and {_CERTIFIED_YARN_Q_SCALING})" - ) - # Per-call constants: the inert groups above plus the two values this - # checkpoint derives. self._call = dict(_CALL_INERT, quant_mode=self.quant_mode, q_scaling=self.q_scaling) - # MLP structure: the first `first_k_dense_replace` layers are dense, - # the rest are routed MoE plus a shared-expert pair fused into one - # dense linear pair of width n_shared * moe_intermediate_size. Both - # dense shapes are replicated under attention DP and run over this - # rank's own tokens. - self.dense_layers = cfg.first_k_dense_replace - assert 0 < self.dense_layers < self.num_layers, "mixed dense/MoE stack" - assert cfg.moe_layer_freq == 1, "every layer past the dense prefix is MoE" - self.num_experts = cfg.n_routed_experts - self.topk = cfg.num_experts_per_tok - self.moe_inter = cfg.moe_intermediate_size - self.shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts - self.dense_inter = cfg.intermediate_size - assert 0 < self.topk < self.num_experts, "MoE top-k bound" - # noaux_tc_op does the whole gate in-kernel and is not configurable: - # sigmoid, bias correction for selection only, renormalization, the - # routed_scaling_factor multiply. A checkpoint that wants any of those - # differently cannot use it. - assert cfg.topk_method == "noaux_tc", cfg.topk_method - assert cfg.scoring_func == "sigmoid", cfg.scoring_func - assert cfg.norm_topk_prob, "noaux_tc_op always renormalizes" - self.n_group = cfg.n_group - self.topk_group = cfg.topk_group - # Four hard limits of the op's grouped path, each checked here because - # the op reports any violation as one opaque "unsupported - # configuration". - assert self.n_group > 1 and 1 <= self.topk_group <= self.n_group, ( - f"grouped routing needs 1 <= topk_group <= n_group, got " - f"{self.topk_group} / {self.n_group}" - ) - assert self.num_experts % self.n_group == 0, "experts per routing group" - assert self.num_experts <= 256 and self.num_experts // self.n_group <= 32, ( - "noaux_tc_op's grouped path takes at most 256 experts and 32 per " - f"group; this config has {self.num_experts} in {self.n_group} groups" - ) - assert self.topk <= 8, "noaux_tc_op's grouped path takes at most top-8" - self.routed_scale = float(cfg.routed_scaling_factor) - # Expert parallelism: the routing space stays global and every rank - # runs the full top-k over the *gathered* token set, but a rank holds - # only its own window of experts and the kernel drops every slot - # outside it. The four windows' outputs sum to the whole layer — that - # is what the reduce-scatter completes. - assert self.num_experts % mapping.moe_ep_size == 0, "experts per rank" - self.local_experts = self.num_experts // mapping.moe_ep_size + num_experts = cfg.n_routed_experts + self.local_experts = num_experts // mapping.moe_ep_size self.expert_offset = self.local_experts * self.ep_rank - assert self.local_experts == 64, ( - f"fp4_block_scale_moe_runner certifies the four-way split of 256 " - f"experts (windows of 64); this topology yields {self.local_experts}" - ) - # fp4_block_scale_moe_runner's block-scale layout rules. - assert self.hidden % 256 == 0, "MoE hidden must be a multiple of 256" - assert self.moe_inter % 64 == 0, "MoE intermediate must be a multiple of 64" - # nvfp4_gemm needs K and N multiples of 32 on both dense linears - # (gate_up: K=hidden, N=2*inter; down: K=inter, N=hidden), and - # flashinfer_silu_and_mul vectorizes the up half from element - # `inter`, so that offset must be 16-element aligned too. The 128x4 - # scale swizzle adds two more: the gate_up scale matrix has 2*inter - # rows (a multiple of 128) and the down one inter/16 columns (a - # multiple of 4) — together, inter must be a multiple of 64. - assert self.hidden % 32 == 0, "nvfp4_gemm operand width" - for inter in (self.dense_inter, self.shared_inter): - assert inter % 64 == 0, ( - "MLP intermediate must be a multiple of 64: nvfp4_gemm " - "operand width, the silu_and_mul half offset, and the 128x4 " - "scale swizzle's row/column alignment" - ) - # Weight declaration. HF [out, in] row-major so checkpoint rows copy - # in unchanged, except kv_b_proj (row-regrouped at load, see - # weights.py) and the expert stacks (declared in the MoE runner's - # kernel-ready shuffled/swizzled layout, at the rank's window). - # Meta-init intercepts torch.empty here — real CUDA storage arrives - # when the engine materializes the registry. - def P(*shape, dtype=dt): - return nn.Parameter(torch.empty(*shape, dtype=dtype), requires_grad=False) - - u8, f32 = torch.uint8, torch.float32 - w = nn.ParameterDict() - for i in range(self.num_layers): - w[f"l{i}_norm1"] = P(self.hidden) - w[f"l{i}_qa"] = P(self.q_lora, self.hidden) - w[f"l{i}_q_norm"] = P(self.q_lora) - w[f"l{i}_qb"] = P(self.heads * self.qk_dim, self.q_lora) - w[f"l{i}_kva"] = P(self.lat_dim, self.hidden) - w[f"l{i}_kv_norm"] = P(self.kv_lora) - w[f"l{i}_kvb"] = P(self.heads * (self.nope + self.v_dim), self.kv_lora) - w[f"l{i}_o"] = P(self.hidden, self.heads * self.v_dim) - # The checkpoint's calibrated fp8 KV-cache scales. They are loaded - # rather than skipped so `derive_after_load` can assert the value - # the whole fp8 MLA path depends on. - w[f"l{i}_k_scale"] = P(1, dtype=f32) - w[f"l{i}_v_scale"] = P(1, dtype=f32) - w[f"l{i}_norm2"] = P(self.hidden) - inter = self.dense_inter if i < self.dense_layers else self.shared_inter - w[f"l{i}_mlp_gu_w"] = P(2 * inter, self.hidden // 2, dtype=u8) - w[f"l{i}_mlp_gu_s"] = P(2 * inter * (self.hidden // _SF_VEC), dtype=u8) - w[f"l{i}_mlp_dn_w"] = P(self.hidden, inter // 2, dtype=u8) - w[f"l{i}_mlp_dn_s"] = P(self.hidden * (inter // _SF_VEC), dtype=u8) - for name in ( - "isc1", - "isc1_up", - "ws2_1", - "ws2_1_up", - "isc2", - "ws2_2", - ): - w[f"l{i}_mlp_{name}"] = P(1, dtype=f32) - if i < self.dense_layers: - continue - e, mi = self.local_experts, self.moe_inter - w[f"l{i}_router"] = P(self.num_experts, self.hidden) - # fp32 on this checkpoint (the reference model keeps the - # correction bias in fp32 whatever the rest of the weights are); - # noaux_tc_op takes bf16 logits against an fp32 bias and returns - # weights in the *logits* dtype, which is what the MoE runner - # demands. - w[f"l{i}_router_bias"] = P(self.num_experts, dtype=f32) - w[f"l{i}_fc1_w"] = P(e, 2 * mi, self.hidden // 2, dtype=u8) - w[f"l{i}_fc1_s"] = P(e, 2 * mi, self.hidden // _SF_VEC, dtype=u8) - w[f"l{i}_fc2_w"] = P(e, self.hidden, mi // 2, dtype=u8) - w[f"l{i}_fc2_s"] = P(e, self.hidden, mi // _SF_VEC, dtype=u8) - # The per-expert NVFP4 scalars stay whole on every rank: they - # cost 6 floats per expert, and the shared-expert activation - # scale is asserted against the max over *all* routed experts, - # which a window could not see. The window is sliced out in - # derive_after_load, where the kernel's [local_num_experts] - # operands are built. - for name in ( - "isc1", - "isc1_up", - "ws2_1", - "ws2_1_up", - "isc2", - "ws2_2", - ): - w[f"l{i}_e_{name}"] = P(self.num_experts, dtype=f32) - w["final_norm"] = P(self.hidden) - w["embed"] = P(self.vocab, self.hidden) - # The MTP module at layer index `num_hidden_layers`, declared only when - # `speculative_config` turned drafting on. Its attention block is - # byte-identical in geometry to a trunk layer's; its MLP path is the - # same structure at a different **dtype** — `hf_quant_config.json` - # carries `model.layers.61*` as one wildcard entry in its - # `exclude_modules` list, so every weight here is bf16 while the - # trunk's MLP is NVFP4. That is the export's choice about where - # accuracy is worth the bytes, so the stacks are declared bf16 and fed - # to the unquantized `fused_moe` rather than re-quantized at load to - # reuse the trunk's expert vocabulary. `embed_tokens` and - # `shared_head.head` are *not* declared: both are bitwise copies of the - # trunk's `model.embed_tokens` / `lm_head`, and the draft-model - # container points at those instead — 1.85 GB per rank saved. - if self.mtp_enabled: - e, mi = self.local_experts, self.moe_inter - w["mtp_enorm"] = P(self.hidden) - w["mtp_hnorm"] = P(self.hidden) - w["mtp_eh"] = P(self.hidden, 2 * self.hidden) - w["mtp_norm1"] = P(self.hidden) - w["mtp_qa"] = P(self.q_lora, self.hidden) - w["mtp_q_norm"] = P(self.q_lora) - w["mtp_qb"] = P(self.heads * self.qk_dim, self.q_lora) - w["mtp_kva"] = P(self.lat_dim, self.hidden) - w["mtp_kv_norm"] = P(self.kv_lora) - w["mtp_kvb"] = P(self.heads * (self.nope + self.v_dim), self.kv_lora) - w["mtp_o"] = P(self.hidden, self.heads * self.v_dim) - w["mtp_k_scale"] = P(1, dtype=f32) - w["mtp_v_scale"] = P(1, dtype=f32) - w["mtp_norm2"] = P(self.hidden) - w["mtp_router"] = P(self.num_experts, self.hidden) - w["mtp_router_bias"] = P(self.num_experts, dtype=f32) - # `fused_moe`'s stacked layout: `[E, 2I, H]` with the **up** rows - # first and the gate rows last (the opposite half order from the - # dense gate_up linear below, which flashinfer_silu_and_mul reads - # gate-first), and `[E, H, I]` for FC2. No interleave, no 32-row - # block shuffle, no swizzle — those belong to the trtllm-gen - # block-scale runner the trunk uses, not to this one. - w["mtp_fc1"] = P(e, 2 * mi, self.hidden) - w["mtp_fc2"] = P(e, self.hidden, mi) - w["mtp_sh_gu"] = P(2 * self.shared_inter, self.hidden) - w["mtp_sh_dn"] = P(self.hidden, self.shared_inter) - w["mtp_head_norm"] = P(self.hidden) - self.w = w - - self._attn: list | None = None - self._mlp: list | None = None - self._moe: list | None = None + self.w = _weights.MODEL_WEIGHTS.declare(_weights.MODEL_WEIGHTS.dims(self)) + self._mtp: dict | None = None - self._next_norm: list | None = None self._rope: dict | None = None self._rope_positions = 0 self._side_stream: torch.cuda.Stream | None = None self._cached_ctx = False - self._step_contract_checked = False def _rope_tables(self, device, positions: int) -> dict: - """The duplicated-layout GPT-J rope table the MLA ops read: per - position, `rope` (cos, sin) pairs whose second half duplicates the - first, flattened to `[1, positions * rope * 2]` fp32, plus the - `[rope/2]` inverse-frequency vector from the same construction. - - `positions` is a row count, not a model property: every row depends - only on its own index, so a longer table is the same table with more - rows and rebuilding one at a larger size changes no existing row. - - The inverse frequencies carry this checkpoint's **YaRN** blend: the - interpolated (factor-divided) frequency for the low-frequency half of - the spectrum, the original one for the high-frequency half, ramped - between the two correction dimensions. Everything else about the rope - configuration is inert on the MLA path — the table's content is the - only rope input the ops read — so this construction is the whole of - it, and a table built from the unscaled theta is a silently wrong - model that diverges with position. Built in fp64 on the host, rounded - once.""" - half = self.rope // 2 + """The duplicated-layout GPT-J rope table the MLA ops read, plus its inverse-frequency + vector. + + `positions` is a row count, not a model property: every row depends only on its own + index, so rebuilding at a larger size changes no existing row.""" + rope_dim = self.model_config.pretrained_config.qk_rope_head_dim + half = rope_dim // 2 d = torch.arange(half, dtype=torch.float64) - freq = self.theta ** (2.0 * d / self.rope) + freq = self.theta ** (2.0 * d / rope_dim) two_pi = 2.0 * math.pi log_theta = math.log(self.theta) low = max( 0.0, math.floor( - self.rope + rope_dim * math.log(self.rope_orig_max / (self.beta_fast * two_pi)) / (2.0 * log_theta) ), ) high = min( - self.rope - 1.0, + rope_dim - 1.0, math.ceil( - self.rope + rope_dim * math.log(self.rope_orig_max / (self.beta_slow * two_pi)) / (2.0 * log_theta) ), @@ -807,327 +286,437 @@ def _rope_tables(self, device, positions: int) -> dict: inv = ramp / (self.rope_factor * freq) + (1.0 - ramp) / freq ang = torch.arange(positions, dtype=torch.float64)[:, None] * inv[None, :] cos, sin = ang.cos() * self.rope_amplitude, ang.sin() * self.rope_amplitude - table = torch.empty(positions, self.rope, 2, dtype=torch.float64) + table = torch.empty(positions, rope_dim, 2, dtype=torch.float64) table[:, :half, 0] = cos table[:, half:, 0] = cos table[:, :half, 1] = sin table[:, half:, 1] = sin return { - "rotary_cos_sin": table.reshape(1, positions * self.rope * 2) + "rotary_cos_sin": table.reshape(1, positions * rope_dim * 2) .float() .to(device) .contiguous(), "rotary_inv_freq": inv.float().to(device).contiguous(), - "rope_dim": self.rope, + "rope_dim": rope_dim, "rope_base": self.theta, } - def derive_after_load(self) -> None: - """Post-load derivation: column-major GEMM views (`.t()` is - zero-copy), the two MLA absorption operands split out of the - row-regrouped kv_b_proj, the rope table, and every NVFP4 call scalar - folded from the checkpoint's per-tensor `input_scale` / - `weight_scale_2` pairs — the routed ones sliced to this rank's expert - window, which is where the kernel's `[local_num_experts]` operands - come from. Meta is over here, so real tensors may be created.""" - w = self.w - device = w["final_norm"].device - self._rope_positions = self.max_pos + def post_load_weights(self) -> None: + """Build the two targets, now that the weights are real. + + Named for the shell hook that drives it. Not an override: that hook lives on + DecoderModelForCausalLM, and a core derives from DecoderModel, which has none. + + The rope tables, the MTP operands and the side stream stay on the core rather + than a target: `MTPLayer` reads the first two, and a CUDA stream created inside + a graph capture is not capturable, so it is made once, here. + """ + cfg = self.model_config.pretrained_config + device = self.w["final_norm"].device + self._rope_positions = cfg.max_position_embeddings self._rope = self._rope_tables(device, self._rope_positions) - window = slice(self.expert_offset, self.expert_offset + self.local_experts) - - attn, mlp, moe, nxt = [], [], [], [] - hn = self.heads * self.nope - for i in range(self.num_layers): - # The fp8 latent pool is written at 1/s and read at s, and the two - # roles live in different ops with no relation checked anywhere in - # the chain. This target passes None for both, which every op reads - # as exactly 1.0 — correct only for a checkpoint calibrated at 1.0, - # and additionally the only value the fp8 MLA *context* path is - # self-consistent at (it quantizes q/k/v at 1.0 while applying - # s^2/s regardless). So the checkpoint's own scales are checked - # rather than assumed. - for role in ("k_scale", "v_scale"): - s = w[f"l{i}_{role}"] - assert torch.equal(s, torch.ones_like(s)), ( - f"layer {i}: {role} is {s.item()}, not 1.0; the fp8 MLA " - "context path is only correct at a KV scaling factor of " - "1.0, and this assembly passes no scale tensors" - ) - kvb = w[f"l{i}_kvb"] - attn.append( - ( - w[f"l{i}_qa"].t(), - w[f"l{i}_q_norm"], - w[f"l{i}_qb"].t(), - w[f"l{i}_kva"].t(), - w[f"l{i}_kv_norm"], - # k_b [H, nope, C] absorbs into q_nope; v_b_t [H, C, v] - # expands the latent attention output. Both are views of - # the row-regrouped kv_b_proj. - kvb[:hn].reshape(self.heads, self.nope, self.kv_lora), - transpose(kvb[hn:].reshape(self.heads, self.v_dim, self.kv_lora), 1, 2), - kvb.t(), - w[f"l{i}_o"].t(), - w[f"l{i}_norm2"], - ) - ) - nxt.append(w[f"l{i + 1}_norm1"] if i + 1 < self.num_layers else w["final_norm"]) - # The fused gate_up GEMM assumes one activation scale and one - # weight global scale for both halves; the checkpoint stores them - # per projection, so the equality the fusion rests on is asserted. - assert torch.equal(w[f"l{i}_mlp_isc1"], w[f"l{i}_mlp_isc1_up"]), ( - f"layer {i}: gate/up input_scale differ; the fused gate_up " - "GEMM needs one activation scale" - ) - assert torch.equal(w[f"l{i}_mlp_ws2_1"], w[f"l{i}_mlp_ws2_1_up"]), ( - f"layer {i}: gate/up weight_scale_2 differ; the fused gate_up GEMM needs one alpha" - ) - # The checkpoint stores reciprocals: `input_scale = amax/(448*6) - # = 1/g_act` and `weight_scale_2 = 1/g_w`, so the quantizer's - # global scale is `1/input_scale` and the GEMM's alpha is their - # product — both straight off disk, no further reciprocal. - isc1 = w[f"l{i}_mlp_isc1"] - isc2 = w[f"l{i}_mlp_isc2"] - mlp.append( - ( - w[f"l{i}_mlp_gu_w"], - w[f"l{i}_mlp_gu_s"], - (isc1 * w[f"l{i}_mlp_ws2_1"]).contiguous(), - (1.0 / isc1).contiguous(), - w[f"l{i}_mlp_dn_w"], - w[f"l{i}_mlp_dn_s"], - (isc2 * w[f"l{i}_mlp_ws2_2"]).contiguous(), - (1.0 / isc2).contiguous(), - ) - ) - if i < self.dense_layers: - moe.append(None) - continue - assert torch.equal(w[f"l{i}_e_isc1"], w[f"l{i}_e_isc1_up"]), ( - f"layer {i}: expert gate/up input_scale differ" - ) - assert torch.equal(w[f"l{i}_e_ws2_1"], w[f"l{i}_e_ws2_1_up"]), ( - f"layer {i}: expert gate/up weight_scale_2 differ" - ) - # One quantization of the gathered hidden states feeds every - # expert on every rank, so the routed FC1 activation scale must - # be a single value — and the same value everywhere, or the four - # windows would not sum to the whole layer. The shared expert - # sees every token where each routed expert sees only its own - # subset, and the checkpoint's shared-expert input_scale is - # exactly the max over *all* routed ones — the conservative - # choice that cannot saturate an activation block. Asserted over - # the full 256, which is why the scalars are loaded whole. - e_isc1 = w[f"l{i}_e_isc1"] - assert torch.equal(e_isc1.max(), w[f"l{i}_mlp_isc1"][0]), ( - f"layer {i}: shared-expert input_scale is not the max over " - "the routed experts; the shared activation quantization " - "would saturate" - ) - e_isc2 = w[f"l{i}_e_isc2"][window] - gate1 = (w[f"l{i}_mlp_isc1"][0] * w[f"l{i}_e_ws2_1"][window]).contiguous() - moe.append( - ( - w[f"l{i}_router"].t(), - w[f"l{i}_router_bias"], - w[f"l{i}_fc1_w"], - view_dtype(w[f"l{i}_fc1_s"], torch.float8_e4m3fn), - w[f"l{i}_fc2_w"], - view_dtype(w[f"l{i}_fc2_s"], torch.float8_e4m3fn), - # output1_scale_scalar, output1_scale_gate_scalar, - # output2_scale_scalar — the FC1 alpha, that alpha times - # the per-expert FC2 activation global scale, and the FC2 - # alpha, each [local_num_experts]. Swapping the first two - # is finite and silent. - (gate1 / e_isc2).contiguous(), - gate1, - (e_isc2 * w[f"l{i}_e_ws2_2"][window]).contiguous(), - # The routed FC1 activation global scale: the same - # 1/input_scale the shared-expert GEMM quantizes with, - # in the linear scale layout the MoE runner reads. - (1.0 / w[f"l{i}_mlp_isc1"]).contiguous(), - ) - ) - self._attn, self._mlp, self._moe, self._next_norm = attn, mlp, moe, nxt if self.mtp_enabled: self._mtp = self._derive_mtp() - # The side stream the shared-expert / dense-MLP branch runs on. Built - # here because a stream must exist before the first forward: creating - # one inside a CUDA-graph capture is not a capturable operation, and - # the runtime's first eager forwards are already past this point. self._side_stream = torch.cuda.Stream(device=device) + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + + def _select_target(self, attn_metadata) -> Target: + """Which target runs this step.""" + return self._decode if phase_of(attn_metadata) is Phase.DECODE else self._prefill + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None = None, + position_ids: torch.IntTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + lora_params: dict | None = None, + **kwargs, + ) -> torch.Tensor: + """Derive what the first step fixes, then route.""" + advance_step_generation() + self._derive_first_step(attn_metadata) + return self._select_target(attn_metadata).forward( + attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + lora_params=lora_params, + **kwargs, + ) def _derive_mtp(self) -> dict: - """The MTP layer's operand set, derived exactly as a trunk layer's is: - column-major GEMM views (`.t()` is zero-copy), the two MLA absorption - operands split out of the row-regrouped kv_b_proj, and the same fp8 KV - scale check. No NVFP4 scalars — this module is bf16 throughout, so its - expert stacks go to `fused_moe` as they are stored.""" + """The MTP layer's operand set, derived exactly as a trunk layer's is.""" + nl = self.model_config.pretrained_config.num_hidden_layers + cfg = self.model_config.pretrained_config + heads = cfg.num_attention_heads + nope = cfg.qk_nope_head_dim + v_dim = cfg.v_head_dim + kv_lora = cfg.kv_lora_rank w = self.w - hn = self.heads * self.nope - for role in ("k_scale", "v_scale"): - s = w[f"mtp_{role}"] - assert torch.equal(s, torch.ones_like(s)), ( - f"MTP layer: {role} is {s.item()}, not 1.0; the fp8 MLA " - "context path is only correct at a KV scaling factor of 1.0, " - "and this assembly passes no scale tensors" - ) - kvb = w["mtp_kvb"] + hn = heads * nope + kvb = w[f"l{nl}_kvb"] return { - "enorm": w["mtp_enorm"], - "hnorm": w["mtp_hnorm"], - "eh": w["mtp_eh"].t(), - "norm1": w["mtp_norm1"], - "qa": w["mtp_qa"].t(), - "q_norm": w["mtp_q_norm"], - "qb": w["mtp_qb"].t(), - "kva": w["mtp_kva"].t(), - "kv_norm": w["mtp_kv_norm"], - "k_b": kvb[:hn].reshape(self.heads, self.nope, self.kv_lora), - "v_b_t": transpose(kvb[hn:].reshape(self.heads, self.v_dim, self.kv_lora), 1, 2), + "enorm": w[f"l{nl}_enorm"], + "hnorm": w[f"l{nl}_hnorm"], + "eh": w[f"l{nl}_eh"].t(), + "norm1": w[f"l{nl}_norm1"], + "qa": w[f"l{nl}_qa"].t(), + "q_norm": w[f"l{nl}_q_norm"], + "qb": w[f"l{nl}_qb"].t(), + "kva": w[f"l{nl}_kva"].t(), + "kv_norm": w[f"l{nl}_kv_norm"], + "k_b": kvb[:hn].reshape(heads, nope, kv_lora), + "v_b_t": torch.transpose(kvb[hn:].reshape(heads, v_dim, kv_lora), 1, 2), "kvb": kvb.t(), - "o": w["mtp_o"].t(), - "norm2": w["mtp_norm2"], - "router": w["mtp_router"].t(), - "router_bias": w["mtp_router_bias"], - "fc1": w["mtp_fc1"], - "fc2": w["mtp_fc2"], - "sh_gu": w["mtp_sh_gu"].t(), - "sh_dn": w["mtp_sh_dn"].t(), - "head_norm": w["mtp_head_norm"], + "o": w[f"l{nl}_o"].t(), + "norm2": w[f"l{nl}_norm2"], + "router": w[f"l{nl}_router"].t(), + "router_bias": w[f"l{nl}_router_bias"], + "fc1": w[f"l{nl}_fc1"], + "fc2": w[f"l{nl}_fc2"], + "sh_gu": w[f"l{nl}_sh_gu"].t(), + "sh_dn": w[f"l{nl}_sh_dn"].t(), + "head_norm": w[f"l{nl}_head_norm"], } - def _check_step_contract(self, md, position_ids) -> None: - """First-forward fail-fast: the metadata fields this target consumes - must exist (private trtllm surface, pinned by version), the paged - latent pool must be the single pool the MLA entries are certified - over at the page size their fp8 column covers, the rope table must - cover every position the engine admits, and every feature this target - holds inert must actually be off. Everything checked is fixed at - engine construction — once per model instance is sound.""" - missing = [name for name in _STEP_FIELDS if not hasattr(md, name)] - assert not missing, f"metadata fields missing: {missing}" - # position_ids is not consumed: the MLA ops derive each context - # token's position from its row index within the sequence and each - # generation token's from sequence_length - 1. Checked anyway so a - # layout change upstream is loud rather than silent. - assert position_ids.dtype == torch.int32, position_ids.dtype - # Narrower than the gpt-oss target's bound on purpose. That one reaches - # its KV cache through thop_attention, which is certified at one pool - # and two; this one goes through the MLA family -- - # load_paged_kv_cache_for_mla, mla_rope_generation, - # mla_rope_append_paged_kv_assign_q -- and none of those entries has a - # multi-pool cell. Every layer here is the same attention-window class, - # so a manager gives them one group and this holds; if that ever - # changes, certify the MLA entries before widening it. - pools = {row[0] for row in md.host_kv_cache_pool_mapping.tolist()} - assert pools == {0}, ( - f"multi-pool KV addressing is not certified for the MLA entries; " - f"layer->pool ids {sorted(pools)}" - ) - # Every MLA entry's fp8-e4m3 column is page 32 only (the bf16 columns - # also carry 64). 32 is what a default KvCacheConfig produces. - assert md.tokens_per_block == 32, ( - f"the fp8 latent-pool column of every MLA entry is certified at " - f"tokens_per_block 32; this engine built {md.tokens_per_block}" - ) - # The rope table must cover every position the engine admits: a short - # table is read out of bounds with no check, and `rope_max_positions`, - # the argument that looks like it bounds this, is one of the inert - # seven. `max_position_embeddings` is the right size under the identity - # config, but **a speculative_config inflates the engine's max_seq_len - # past it** — measured 163840 -> 163848 at `max_draft_len: 3`, which is - # more than the `max_draft_len - 1` extra KV tokens per sequence the - # runtime reference documents, so the engine's own number is taken - # rather than a formula. Growing here is exact: every row of the table - # depends only on its own position, so the rows the identity config - # uses are bit-identical either way. This runs on the first forward, - # before any CUDA-graph capture. + def _derive_first_step(self, md) -> None: + """Grow the rope table to cover every position the engine admits, and fix the context + flavor this engine construction runs. + + Unconditional, unlike the targets' contract probe: both of these decide what the + model computes, not whether it is well-formed. They used to hang off that probe + and so ran only under TRTLLM_MODELING_V2_VALIDATE, which the accuracy tests set + and a served engine does not -- leaving production with a rope table sized to + `max_position_embeddings` and `_cached_ctx` stuck at False. + """ + if not self._derive_pending: + return + # A table shorter than the engine's max_seq_len is read out of bounds with no check, and + # `rope_max_positions` is one of the arguments the MLA path ignores. `speculative_config` + # inflates max_seq_len past `max_position_embeddings`, so take the engine's own number. if md.max_seq_len > self._rope_positions: self._rope_positions = md.max_seq_len self._rope = self._rope_tables(self.w["final_norm"].device, self._rope_positions) - # Context flavor, fixed at engine construction: with the cached-KV - # surface present the target runs append -> gather -> up-project -> - # explicit-K/V FMHA, which serves reused and fresh sequences alike; - # without it (block reuse off) no context sequence can carry a - # prefix, and the fresh-prefill flavor with the in-kernel rope and - # append is the whole context path. The two cache ops reject any - # index dtype but int64. self._cached_ctx = all(hasattr(md, n) for n in _CACHED_CTX_FIELDS) and bool( md.enable_context_mla_with_cached_kv ) - if self._cached_ctx: - for name in ("ctx_cached_token_indptr", "ctx_kv_indptr"): - t = getattr(md, name) - assert t.dtype == torch.int64 and t.is_cuda, (name, t.dtype, t.device) - assert md.effective_beam_width == 1, "beam search is not implemented" - assert md.cache_indirection is None, "beam search is not implemented" - assert md.block_ids_per_seq is None, "the FlashMLA layout is not implemented" - assert md.flash_mla_tile_scheduler_metadata is None, "FlashMLA is not implemented" - assert md.flash_mla_num_splits is None, "FlashMLA is not implemented" - assert not md.is_cross, "cross attention is not implemented" - # The spec-dec **mask** machinery, which is a different thing from - # speculative decoding being on. Under a linear-tree MTP on a trtllm-gen - # arch the - # backend computes `is_spec_decoding_enabled and (not trtllm_gen_arch - # or is_spec_dec_dynamic_tree)` and gets False, so every - # `spec_decoding_*` tensor stays None — which is exactly the inert - # group `_CALL_INERT` holds and every MLA column certifies. Drafting - # itself reaches the attention ops through `predicted_tokens_per_seq` - # alone. This assert is the precondition those inert values rest on: it - # fires on a tree draft, or on a pre-Blackwell arch where the gating - # does not force the mask off, and either would need the mask surface - # certified first. - assert not md.is_spec_decoding_enabled and not md.use_spec_decoding, ( - "the spec-decoding mask surface is live; this target holds the " - "whole spec_decoding_* group at its inert values, which is only " - "valid while the Blackwell linear-tree gating keeps it off" - ) - self._step_contract_checked = True + self._derive_pending = False def _dp_rows(self, md, num_tokens: int) -> int: - """The uniform row count every rank pads its token block to before the - MoE round trip: `max` over the group's per-rank token counts. - - Attention DP gives each rank a different batch by construction, so the - split is engine state, not something a rank can derive locally: - `attn_metadata.all_rank_num_tokens` is where the engine publishes it, - as a plain list of host ints identical on every rank. Every rank - therefore computes the same maximum, and both collectives run in their - uniform form (`sizes=None`). - - **Both collectives are run uniform on purpose, and the ragged form is - not used at all.** The even split is cheaper for the reduce-scatter at - these row counts, and `sizes` is a host argument baked into a - CUDA-graph capture, so only the uniform form is replayable. - - CUDA-graph classification: this returns a host int, but under capture - the counts are uniform (the engine pads the decode batch to a captured - size on every rank), so the padding is zero rows and the graph - contains no padding at all — and no host value reaches either - collective, whose row counts come from tensor shapes the graph - fixes.""" + """The uniform row count every rank pads its token block to before the MoE round trip.""" counts = [int(n) for n in md.all_rank_num_tokens] - assert len(counts) == self.dp_size, (counts, self.dp_size) - assert counts[self.rank] == num_tokens, ( - f"rank {self.rank} holds {num_tokens} rows but the engine's " - f"all_rank_num_tokens says {counts[self.rank]} ({counts}); the " - "collectives would gather the wrong split" - ) return max(counts) - def _dense_mlp(self, x, params, dt): - """One NVFP4 SwiGLU MLP over this rank's own tokens: fused gate_up - GEMM, silu_and_mul, down GEMM. Replicated weights, so the result is - complete — nothing to reduce. Both quantizations emit the - 128x4-swizzled scale buffer nvfp4_gemm consumes.""" - gu_w, gu_s, gu_alpha, gu_g, dn_w, dn_s, dn_alpha, dn_g = params - xq, xsf = fp4_quantize(x, gu_g, _SF_VEC, False, _SF_SWIZZLED) - gu = nvfp4_gemm(xq, gu_w, xsf, gu_s, gu_alpha, dt) + +class PrefillTarget(Target): + """The general case: context rows, and possibly generation rows beside them.""" + + def __init__(self, core: "DeepseekV3Dep4ModelingV2") -> None: + """Bind every op this target calls, once, against the real weights.""" + super().__init__(core) + cfg = core.model_config.pretrained_config + + num_layers = cfg.num_hidden_layers + + heads = cfg.num_attention_heads + + nope = cfg.qk_nope_head_dim + + rope_dim = cfg.qk_rope_head_dim + + v_dim = cfg.v_head_dim + + kv_lora = cfg.kv_lora_rank + + q_lora = cfg.q_lora_rank + + qk_dim = nope + rope_dim + lat_dim = kv_lora + rope_dim + + eps = cfg.rms_norm_eps + + dense_layers = cfg.first_k_dense_replace + + num_experts = cfg.n_routed_experts + + topk = cfg.num_experts_per_tok + + moe_inter = cfg.moe_intermediate_size + + w = core.w + + window = slice(core.expert_offset, core.expert_offset + core.local_experts) + + self._rms0 = FlashinferRmsnorm() + + self._rms0.bind_const(weight=w["l0_norm1"], eps=eps) + + self._qa = CublasMm() + + self._q_lora_norm = FlashinferRmsnorm() + + self._q_lora_norm.bind_const(eps=eps) + + self._qb = CublasMm() + + self._kva = CublasMm() + + self._kv_lora_norm = FlashinferRmsnorm() + + self._kv_lora_norm.bind_const(eps=eps) + + self._kvb = CublasMm() + + self._o = CublasMm() + + self._norm2 = FlashinferFusedAddRmsnorm() + + self._norm2.bind_const(eps=eps) + + self._next_norm = FlashinferFusedAddRmsnorm() + + self._next_norm.bind_const(eps=eps) + + self._absorb = BmmOut() + + self._expand = BmmOut() + + self._attn_ctx = ThopAttention() + + self._attn_ctx.bind_const( + num_heads=heads, + num_kv_heads=heads, + head_size=qk_dim, + q_lora_rank=q_lora, + kv_lora_rank=kv_lora, + qk_nope_head_dim=nope, + qk_rope_head_dim=rope_dim, + v_head_dim=v_dim, + is_fused_qkv=False, + attention_input_type=1, + q_pe=None, + cu_q_seqlens=None, + cu_kv_seqlens=None, + fmha_scheduler_counter=None, + predicted_tokens_per_seq=1, + **_CALL_INERT, + quant_mode=core.quant_mode, + q_scaling=core.q_scaling, + ) + + self._attn_gen = ThopAttention() + + self._attn_gen.bind_const( + num_heads=heads, + num_kv_heads=1, + head_size=lat_dim, + q_lora_rank=q_lora, + kv_lora_rank=kv_lora, + qk_nope_head_dim=nope, + qk_rope_head_dim=rope_dim, + v_head_dim=kv_lora, + is_fused_qkv=True, + attention_input_type=2, + k=None, + v=None, + **_CALL_INERT, + quant_mode=core.quant_mode, + q_scaling=core.q_scaling, + ) + + self._mla_append = MlaRopeAppendPagedKvAssignQ() + + self._mla_append.bind_const( + head_num=heads, + nope_size=nope, + rope_size=rope_dim, + lora_size=kv_lora, + kv_scale_orig_quant=None, + residual_dim=_KV_RESIDUAL_DIM, + beam_width=1, + quant_mode=core.quant_mode, + ) + + self._mla_load = LoadPagedKvCacheForMla() + + self._mla_load.bind_const( + kv_scale_quant_orig=None, + kv_lora_rank=kv_lora, + qk_rope_head_dim=rope_dim, + beam_width=1, + quant_mode=core.quant_mode, + ) + + self._mla_gen = MlaRopeGeneration() + + self._mla_gen.bind_const( + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + kv_cache_scale_orig_quant=None, + out_scale=None, + block_ids_per_seq=None, + helix_tensor_params=[None, None], + num_heads=heads, + num_kv_heads=1, + head_size=lat_dim, + residual_dim=_KV_RESIDUAL_DIM, + beam_width=1, + quant_mode=core.quant_mode, + q_scaling=core.q_scaling, + q_lora_rank=q_lora, + kv_lora_rank=kv_lora, + qk_nope_head_dim=nope, + qk_rope_head_dim=rope_dim, + v_head_dim=kv_lora, + rope_append=True, + ) + + self._allgather = Allgather() + + self._allgather.bind_const(sizes=None, group=core.dp_group) + + self._reducescatter = Reducescatter() + + self._reducescatter.bind_const(sizes=None, group=core.dp_group) + + self._noaux = NoauxTcOp() + + self._noaux.bind_const( + n_group=cfg.n_group, + topk_group=cfg.topk_group, + topk=topk, + routed_scaling_factor=float(cfg.routed_scaling_factor), + ) + + self._router = CublasMm() + + self._mlp_gu_quant = Fp4Quantize() + + self._mlp_gu_quant.bind_const( + sf_vec_size=_SF_VEC, sf_use_ue8m0=False, is_sf_swizzled_layout=_SF_SWIZZLED + ) + + self._mlp_gu_gemm = Nvfp4Gemm() + + self._mlp_dn_quant = Fp4Quantize() + + self._mlp_dn_quant.bind_const( + sf_vec_size=_SF_VEC, sf_use_ue8m0=False, is_sf_swizzled_layout=_SF_SWIZZLED + ) + + self._mlp_dn_gemm = Nvfp4Gemm() + + self._moe_quant = Fp4Quantize() + + self._moe_quant.bind_const( + sf_vec_size=_SF_VEC, sf_use_ue8m0=False, is_sf_swizzled_layout=_SF_LINEAR + ) + + self._moe_runner = Fp4BlockScaleMoeRunner() + + self._moe_runner.bind_const( + routing_logits=None, + routing_bias=None, + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_bias=None, + num_experts=num_experts, + top_k=topk, + n_group=None, + topk_group=None, + intermediate_size=moe_inter, + local_expert_offset=core.expert_offset, + local_num_experts=core.local_experts, + routed_scaling_factor=None, + routing_method_type=_ROUTING_METHOD_INERT, + do_finalize=True, + act_type=_ACT_TYPE_SWIGLU, + ) + + hn = heads * nope + + for i in range(num_layers): + kvb = w[f"l{i}_kvb"] + self._qa.bind_layered(i, mat_b=w[f"l{i}_qa"].t()) + self._q_lora_norm.bind_layered(i, weight=w[f"l{i}_q_norm"]) + self._qb.bind_layered(i, mat_b=w[f"l{i}_qb"].t()) + self._kva.bind_layered(i, mat_b=w[f"l{i}_kva"].t()) + self._kv_lora_norm.bind_layered(i, weight=w[f"l{i}_kv_norm"]) + self._kvb.bind_layered(i, mat_b=kvb.t()) + self._o.bind_layered(i, mat_b=w[f"l{i}_o"].t()) + self._norm2.bind_layered(i, weight=w[f"l{i}_norm2"]) + next_w = w[f"l{i + 1}_norm1"] if i + 1 < num_layers else w["final_norm"] + self._next_norm.bind_layered(i, weight=next_w) + self._absorb.bind_layered(i, b=kvb[:hn].reshape(heads, nope, kv_lora)) + self._expand.bind_layered( + i, b=torch.transpose(kvb[hn:].reshape(heads, v_dim, kv_lora), 1, 2) + ) + self._attn_ctx.bind_layered(i, local_layer_idx=i) + self._attn_gen.bind_layered(i, local_layer_idx=i) + self._mla_append.bind_layered(i, layer_idx=i) + self._mla_load.bind_layered(i, layer_idx=i) + self._mla_gen.bind_layered(i, layer_idx=i) + + # The checkpoint stores reciprocals (`input_scale` = 1/g_act, `weight_scale_2` = + # 1/g_w), so the quantizer's global scale is 1/input_scale and the GEMM alpha is + # their product — no further reciprocal anywhere. + isc1 = w[f"l{i}_mlp_isc1"] + isc2 = w[f"l{i}_mlp_isc2"] + self._mlp_gu_quant.bind_layered(i, global_scale=(1.0 / isc1).contiguous()) + self._mlp_gu_gemm.bind_layered( + i, + weight=w[f"l{i}_mlp_gu_w"], + weight_scale=w[f"l{i}_mlp_gu_s"], + alpha=(isc1 * w[f"l{i}_mlp_ws2_1"]).contiguous(), + ) + self._mlp_dn_quant.bind_layered(i, global_scale=(1.0 / isc2).contiguous()) + self._mlp_dn_gemm.bind_layered( + i, + weight=w[f"l{i}_mlp_dn_w"], + weight_scale=w[f"l{i}_mlp_dn_s"], + alpha=(isc2 * w[f"l{i}_mlp_ws2_2"]).contiguous(), + ) + + if i < dense_layers: + continue + self._router.bind_layered(i, mat_b=w[f"l{i}_router"].t()) + self._noaux.bind_layered(i, bias=w[f"l{i}_router_bias"]) + # One quantization of the gathered hidden states feeds every expert on every rank, + # so the routed FC1 activation scale has to be a single value and the same value on + # all four ranks, or the windows no longer sum to the whole layer; the shared + # expert's `input_scale` is the max over all routed ones, so it is the safe one. + e_isc2 = w[f"l{i}_e_isc2"][window] + gate1 = (w[f"l{i}_mlp_isc1"][0] * w[f"l{i}_e_ws2_1"][window]).contiguous() + self._moe_quant.bind_layered(i, global_scale=(1.0 / w[f"l{i}_mlp_isc1"]).contiguous()) + self._moe_runner.bind_layered( + i, + gemm1_weights=w[f"l{i}_fc1_w"], + gemm1_weights_scale=w[f"l{i}_fc1_s"].view(torch.float8_e4m3fn), + gemm2_weights=w[f"l{i}_fc2_w"], + gemm2_weights_scale=w[f"l{i}_fc2_s"].view(torch.float8_e4m3fn), + output1_scale_scalar=(gate1 / e_isc2).contiguous(), + output1_scale_gate_scalar=gate1, + output2_scale_scalar=(e_isc2 * w[f"l{i}_e_ws2_2"][window]).contiguous(), + ) + if core.mtp_enabled: + self._bind_mtp() + + def step_args(self, md: TrtllmAttentionMetadata) -> dict: + """The batch state both MLA calls of one step share.""" + return _build_step_args(md) + + def _dense_mlp(self, x, layer: int, dt): + """One NVFP4 SwiGLU MLP over this rank's own tokens, used for both the dense layers' MLP + and every MoE layer's shared expert.""" + xq, xsf = self._mlp_gu_quant(x, layer=layer) + gu = self._mlp_gu_gemm(act_fp4=xq, act_sf=xsf, output_dtype=dt, layer=layer) act = flashinfer_silu_and_mul(gu) - aq, asf = fp4_quantize(act, dn_g, _SF_VEC, False, _SF_SWIZZLED) - return nvfp4_gemm(aq, dn_w, asf, dn_s, dn_alpha, dt) + aq, asf = self._mlp_dn_quant(act, layer=layer) + return self._mlp_dn_gemm(act_fp4=aq, act_sf=asf, output_dtype=dt, layer=layer) def forward( self, @@ -1138,580 +727,427 @@ def forward( lora_params: dict | None = None, **kwargs, ) -> torch.Tensor: - attn_w, mlp_w, moe_w = self._attn, self._mlp, self._moe - next_norm = self._next_norm - assert attn_w is not None and mlp_w is not None and moe_w is not None, ( - "load_weights must run before forward" + core = self.core + self._check_step_contract(attn_metadata) + + cfg = core.model_config.pretrained_config + heads = cfg.num_attention_heads + nope = cfg.qk_nope_head_dim + rope_dim = cfg.qk_rope_head_dim + v_dim = cfg.v_head_dim + kv_lora = cfg.kv_lora_rank + qk_dim = nope + rope_dim + lat_dim = kv_lora + rope_dim + num_layers = cfg.num_hidden_layers + dense_layers = cfg.first_k_dense_replace + + md = attn_metadata + rope = core._rope + + step = self.step_args(md) + self._attn_ctx.bind_const(**step, **rope) + self._attn_gen.bind_const(**step, **rope) + self._mla_append.bind_const( + cos_sin_cache=rope["rotary_cos_sin"], + tokens_per_block=md.tokens_per_block, + attention_window_size=md.max_seq_len, ) - assert next_norm is not None, "load_weights must run before forward" - assert position_ids is not None - assert isinstance(attn_metadata, TrtllmAttentionMetadata) - # Inputs this target does not implement must fail loudly, not be - # silently dropped (unlike runtime-owned features, which pass through). - assert lora_params is None, "LoRA is not implemented by this target" - # The trunk does not read `spec_metadata` even when the engine drafts: - # on MTP_EAGLE_ONE_MODEL the runtime sets `layers_to_capture = ()`, so - # no hidden-state capture hook is owed and the shell keeps the argument - # to itself. Keeping this assert is a stronger guarantee than deleting - # it. - assert kwargs.get("spec_metadata") is None, ( - "spec_metadata reached the trunk; the shell owns the draft loop and must not forward it" + self._mla_load.bind_const( + tokens_per_block=md.tokens_per_block, attention_window_size=md.max_seq_len + ) + # Explicit subsets rather than **step: the three MLA helpers' raw_calls + # do not take the projection's engine-bookkeeping keys. + self._mla_gen.bind_const( + rotary_cos_sin=rope["rotary_cos_sin"], + tokens_per_block=md.tokens_per_block, + attention_window_size=md.max_seq_len, + sequence_length=md.kv_lens_cuda_runtime, + host_past_key_value_lengths=md.kv_lens_runtime, + host_context_lengths=md.prompt_lens_cpu_runtime, + num_contexts=md.num_contexts, + kv_cache_block_offsets=md.kv_cache_block_offsets, + host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, ) - md = attn_metadata - if not self._step_contract_checked: - self._check_step_contract(md, position_ids) - # Read after the contract check: that is where a table too short for - # the engine's admitted max_seq_len is regrown. - rope = self._rope - assert rope is not None, "load_weights must run before forward" - step = _build_step_args(md) num_ctx = md.num_contexts tc = md.num_ctx_tokens - # The paged-cache address book, shared by the two MLA preprocessing - # ops and both attention calls. Under attention DP each rank owns a - # pool holding only its own requests' latent rows. - tokens_per_block = md.tokens_per_block - block_offsets = md.kv_cache_block_offsets - pool_ptrs = md.host_kv_cache_pool_pointers - pool_map = md.host_kv_cache_pool_mapping - assert tokens_per_block is not None, "paged KV cache is required" - assert block_offsets is not None and pool_ptrs is not None, "paged KV cache is required" - assert pool_map is not None, "paged KV cache is required" if inputs_embeds is None: - assert input_ids is not None - h = embedding(input_ids, self.w["embed"]) + h = nn.functional.embedding(input_ids, core.w["embed"]) else: h = inputs_embeds num_tokens = h.shape[0] gen = num_tokens - tc dt = h.dtype dev = h.device - # Query tokens per generation sequence, which is the MLA generation - # call's `predicted_tokens_per_seq`. 1 for an ordinary decode step; - # under MTP a generation request arrives carrying its whole draft - # chain, `runtime_draft_len + 1` rows, and that count is the *only* - # thing expressing the taller query block to the attention ops. Derived - # from the metadata rather than from a spec object, so it costs nothing - # when MTP is off and is a per-capture host constant when it is on. + # Query rows per generation sequence (the whole draft chain under MTP): the only thing + # that tells the attention ops the query block is taller, and what aligns the + # bottom-right within-block causal mask. gen_seqs = md.num_seqs - num_ctx gen_p = gen // gen_seqs if gen_seqs else 1 - assert gen_p * gen_seqs == gen, ( - f"{gen} generation rows do not divide over {gen_seqs} generation " - "sequences; a ragged per-sequence draft length cannot be expressed " - "to the MLA generation call" - ) - # The context sequences' cached+new latent-KV row count: what the - # cache gather returns and what the context FMHA attends over. - ctx_kv_tokens = int(md.host_total_kv_lens[0]) if tc else 0 - # Every rank pads its token block to the group-wide maximum, so both - # MoE collectives run in their uniform form; the padded rows are - # sliced off after the reduce-scatter and never reach the residual - # stream. Zero rows quantize to an all-zero NVFP4 block (scale byte - # 0x00, zero data) and route like any other token, so they perturb - # nothing but wasted expert work. - dp_rows = self._dp_rows(md, num_tokens) + ctx_kv_tokens = int(md.host_total_kv_lens[0]) + dp_rows = core._dp_rows(md, num_tokens) pad_rows = dp_rows - num_tokens + if core._cached_ctx: + # The ctx_* metadata fields exist only in this flavor. + self._mla_append.bind_const( + num_contexts=num_ctx, + cu_ctx_cached_kv_lens=md.ctx_cached_token_indptr, + cu_seq_lens=md.ctx_kv_indptr, + max_input_uncached_seq_len=int(md.max_ctx_seq_len), + kv_cache_block_offsets=md.kv_cache_block_offsets, + host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, + ) + self._mla_load.bind_const( + out_dtype=dt, + num_contexts=num_ctx, + num_ctx_kv_tokens=ctx_kv_tokens, + max_ctx_kv_len=int(md.max_ctx_kv_len), + cu_ctx_kv_lens=md.ctx_kv_indptr, + kv_cache_block_offsets=md.kv_cache_block_offsets, + host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, + ) - attn_out = empty([num_tokens, self.heads * self.v_dim], dt, dev) - attn_ctx, attn_gen = split(attn_out, [tc, gen], 0) - x = flashinfer_rmsnorm(h, self.w["l0_norm1"], self.eps) + attn_out = torch.empty([num_tokens, heads * v_dim], dtype=dt, device=dev) + attn_ctx, attn_gen = torch.split(attn_out, [tc, gen], 0) + # One set of attention scratch for all layers: every tensor here is fully + # rewritten by layer i's producer before layer i's consumer reads it, and + # the whole attention path runs on the main stream, so layer i+1 cannot + # overwrite what layer i still reads. 9 allocations per step, not 9 x 61. + # `k`'s rope half is written only in the cached-ctx flavor; the fresh + # flavor never reads it, in either allocation scheme. + tkv = ctx_kv_tokens + k = torch.empty([tkv, heads, qk_dim], dtype=dt, device=dev) + k_nope_dst, k_pe_dst = torch.split(k, [nope, rope_dim], -1) + fused_q = torch.empty([gen, heads, lat_dim], dtype=dt, device=dev) + fq_nope, _ = torch.split(fused_q, [kv_lora, rope_dim], -1) + fq_nope_t = torch.transpose(fq_nope, 0, 1) + lat_out = torch.empty([gen, heads * kv_lora], dtype=dt, device=dev) + lat_out_t = torch.transpose(torch.reshape(lat_out, [gen, heads, kv_lora]), 0, 1) + attn_gen_t = torch.transpose(torch.reshape(attn_gen, [gen, heads, v_dim]), 0, 1) + cu_q = torch.empty([gen + 1], dtype=torch.int32, device=dev) + cu_kv = torch.empty([gen + 1], dtype=torch.int32, device=dev) + counter = torch.empty([1], dtype=torch.uint32, device=dev) + quant_q = torch.empty([gen, heads, lat_dim], dtype=torch.float8_e4m3fn, device=dev) + bmm1_scale = torch.empty([2], dtype=torch.float32, device=dev) + bmm2_scale = torch.empty([1], dtype=torch.float32, device=dev) + scratch = dict( + cu_q_seqlens=cu_q, + cu_kv_seqlens=cu_kv, + fmha_scheduler_counter=counter, + mla_bmm1_scale=bmm1_scale, + mla_bmm2_scale=bmm2_scale, + quant_q_buffer=quant_q, + predicted_tokens_per_seq=gen_p, + ) + self._mla_gen.bind_const(**scratch) + self._attn_gen.bind_const(**scratch) + x = self._rms0(h) residual = h - for i in range(self.num_layers): - ( - w_qa, - w_q_norm, - w_qb, - w_kva, - w_kv_norm, - k_b, - v_b_t, - w_kvb, - w_o, - w_n2, - ) = attn_w[i] - # The q-LoRA pair: down-project to q_lora_rank, RMS-norm the - # latent, up-project to the per-head [nope | rope] rows. - q = cublas_mm(flashinfer_rmsnorm(cublas_mm(x, w_qa), w_q_norm, self.eps), w_qb) - kva = cublas_mm(x, w_kva) - ckv_raw, k_pe = split(kva, [self.kv_lora, self.rope], -1) - ckv = flashinfer_rmsnorm(ckv_raw, w_kv_norm, self.eps) - latent = concat([ckv, k_pe], -1) - q_ctx, q_gen = split(q, [tc, gen], 0) - latent_ctx, latent_gen = split(latent, [tc, gen], 0) - - if tc: - if self._cached_ctx: - # Reuse-capable context: rotate q_pe/k_pe in place and - # append this step's latent rows (quantized to e4m3 on the - # way into the pool, at the write-side scale — None = 1.0), - # read each sequence's whole [cached + new] latent range - # back out (dequantized at the read-side scale — None = - # 1.0), up-project it, and attend over explicit K/V. One - # call serves a batch mixing reused and fresh sequences — - # a fresh one is the cached_s == 0 case. A cached prefix - # therefore pays fp8 twice: the gather widens the pool's - # e4m3 rows to bf16 and the FMHA below quantizes the - # up-projected result straight back. - mla_rope_append_paged_kv_assign_q( - q_ctx, - latent_ctx, - num_ctx, - md.ctx_cached_token_indptr, - md.ctx_kv_indptr, - int(md.max_ctx_seq_len), - rope["rotary_cos_sin"], - self.heads, - self.nope, - self.rope, - self.kv_lora, - block_offsets, - pool_ptrs, - pool_map, - None, - _KV_RESIDUAL_DIM, - i, - tokens_per_block, - md.max_seq_len, - 1, - self.quant_mode, - ) - ckv_full, k_pe_full = load_paged_kv_cache_for_mla( - dt, - num_ctx, - ctx_kv_tokens, - int(md.max_ctx_kv_len), - md.ctx_kv_indptr, - block_offsets, - pool_ptrs, - pool_map, - None, - i, - self.kv_lora, - self.rope, - tokens_per_block, - md.max_seq_len, - 1, - self.quant_mode, - ) - latent_arg = None - else: - # Block reuse off: no context sequence can carry a - # prefix, so the fresh-prefill flavor does the rope and - # the append inside the attention call and the latent - # rows never leave registers. - assert ctx_kv_tokens == tc, ( - "a context sequence arrived with a cached KV prefix " - "while the cached-KV metadata surface is off" - ) - ckv_full, _ = split(ckv, [tc, gen], 0) - k_pe_full = None - latent_arg = latent_ctx - tkv = ctx_kv_tokens - # kv is packed [all heads' k_nope | all heads' v]: the context - # FMHA hard-codes V's row stride as the full packed width and - # reads the column block, so V must stay this split view. - kv = cublas_mm(ckv_full, w_kvb) - k_nope, v_view = split(kv, [self.heads * self.nope, self.heads * self.v_dim], -1) - k = empty([tkv, self.heads, self.qk_dim], dt, dev) - k_nope_dst, k_pe_dst = split(k, [self.nope, self.rope], -1) - copy_(k_nope_dst, reshape(k_nope, [tkv, self.heads, self.nope])) - if k_pe_full is not None: - # k_pe came back from the pool already rotated; every - # query head shares it. On the fresh-prefill flavor the - # rope slice is left uninitialized instead — that call - # overwrites it in place from latent_cache. - copy_( - k_pe_dst, - expand( - reshape(k_pe_full, [tkv, 1, self.rope]), - [tkv, self.heads, self.rope], - ), - ) - thop_attention( - q=q_ctx, - k=reshape(k, [tkv, self.heads * self.qk_dim]), - v=v_view, - output=attn_ctx, - latent_cache=latent_arg, - q_pe=None, - local_layer_idx=i, - is_fused_qkv=False, - attention_input_type=1, - num_heads=self.heads, - num_kv_heads=self.heads, - head_size=self.qk_dim, - q_lora_rank=self.q_lora, - kv_lora_rank=self.kv_lora, - qk_nope_head_dim=self.nope, - qk_rope_head_dim=self.rope, - v_head_dim=self.v_dim, - cu_q_seqlens=None, - cu_kv_seqlens=None, - fmha_scheduler_counter=None, - # The context call's row count is `num_ctx_tokens` - # whatever the generation phase carries, and the entry - # certifies the MLA context flavors at 1 only. - predicted_tokens_per_seq=1, - **step, - **rope, - **self._call, + for i in range(num_layers): + q = self._qb(self._q_lora_norm(self._qa(x, layer=i), layer=i), layer=i) + kva = self._kva(x, layer=i) + ckv_raw, k_pe = torch.split(kva, [kv_lora, rope_dim], -1) + ckv = self._kv_lora_norm(ckv_raw, layer=i) + latent = torch.cat([ckv, k_pe], -1) + q_ctx, q_gen = torch.split(q, [tc, gen], 0) + latent_ctx, latent_gen = torch.split(latent, [tc, gen], 0) + + if core._cached_ctx: + q_ctx = self._mla_append(q=q_ctx, latent_cache=latent_ctx, layer=i) + ckv_full, k_pe_full = self._mla_load(layer=i) + latent_arg = None + else: + ckv_full, _ = torch.split(ckv, [tc, gen], 0) + k_pe_full = None + latent_arg = latent_ctx + # The context FMHA hard-codes V's row stride as the full packed [k_nope | v] + # width and reads the column block, so V must stay this split view. + kv = self._kvb(ckv_full, layer=i) + k_nope, v_view = torch.split(kv, [heads * nope, heads * v_dim], -1) + k_nope_dst.copy_(torch.reshape(k_nope, [tkv, heads, nope])) + if k_pe_full is not None: + k_pe_dst.copy_( + torch.reshape(k_pe_full, [tkv, 1, rope_dim]).expand([tkv, heads, rope_dim]) ) + attn_ctx = self._attn_ctx( + q=q_ctx, + k=torch.reshape(k, [tkv, heads * qk_dim]), + v=v_view, + output=attn_ctx, + latent_cache=latent_arg, + layer=i, + ) if gen: - q3 = reshape(q_gen, [gen, self.heads, self.qk_dim]) - q_nope, q_pe = split(q3, [self.nope, self.rope], -1) - fused_q = empty([gen, self.heads, self.lat_dim], dt, dev) - fq_nope, _ = split(fused_q, [self.kv_lora, self.rope], -1) - # Absorbed q: (q_nope @ W_k_nope) is what the latent-space - # dot product needs. Over an fp8 pool the next call **reads** - # this half to build the quantized query, so this BMM must - # have finished first — the two are issued in this order on - # one stream, which is what makes that safe. (On a bf16 pool - # they write disjoint halves and may overlap; assembling from - # that reading and switching to an fp8 cache is a silent race.) - bmm_out(transpose(q_nope, 0, 1), k_b, transpose(fq_nope, 0, 1)) - cu_q = empty([gen + 1], torch.int32, dev) - cu_kv = empty([gen + 1], torch.int32, dev) - counter = empty([1], torch.uint32, dev) - # The fp8 decode triple: the quantized query the FMHA reads - # instead of fused_q, and the two folded softmax/output scales - # it takes instead of either kv scale tensor. All three are - # written by the call below from q_scaling, the MLA dims and - # the read-side factor (None = 1.0). - quant_q = empty([gen, self.heads, self.lat_dim], torch.float8_e4m3fn, dev) - bmm1_scale = empty([2], torch.float32, dev) - bmm2_scale = empty([1], torch.float32, dev) - mla_rope_generation( - fused_q, - q_pe, - latent_gen, - rope["rotary_cos_sin"], - cu_q, - cu_kv, - counter, - bmm1_scale, - bmm2_scale, - quant_q, - md.kv_lens_cuda_runtime, - md.kv_lens_runtime, - md.prompt_lens_cpu_runtime, - num_ctx, - block_offsets, - pool_ptrs, - pool_map, - None, - None, - None, - None, - None, - [None, None], - gen_p, - i, - self.heads, - 1, - self.lat_dim, - _KV_RESIDUAL_DIM, - tokens_per_block, - md.max_seq_len, - 1, - self.quant_mode, - self.q_scaling, - self.q_lora, - self.kv_lora, - self.nope, - self.rope, - self.v_dim, - True, + q3 = torch.reshape(q_gen, [gen, heads, qk_dim]) + q_nope, q_pe = torch.split(q3, [nope, rope_dim], -1) + # Over an fp8 pool `_mla_gen` reads this half to build the quantized query, so + # this BMM must be issued first on the same stream; overlapping them is a silent + # race (they only write disjoint halves on a bf16 pool). + fq_nope_t = self._absorb(a=torch.transpose(q_nope, 0, 1), out=fq_nope_t, layer=i) + fused_q = self._mla_gen( + fused_q=fused_q, q_pe=q_pe, latent_cache=latent_gen, layer=i ) - lat_out = empty([gen, self.heads * self.kv_lora], dt, dev) - thop_attention( - q=reshape(fused_q, [gen, self.heads * self.lat_dim]), - k=None, - v=None, + lat_out = self._attn_gen( + q=torch.reshape(fused_q, [gen, heads * lat_dim]), output=lat_out, latent_cache=latent_gen, q_pe=q_pe, - local_layer_idx=i, - is_fused_qkv=True, - attention_input_type=2, - num_heads=self.heads, - num_kv_heads=1, - head_size=self.lat_dim, - q_lora_rank=self.q_lora, - kv_lora_rank=self.kv_lora, - qk_nope_head_dim=self.nope, - qk_rope_head_dim=self.rope, - v_head_dim=self.kv_lora, - cu_q_seqlens=cu_q, - cu_kv_seqlens=cu_kv, - fmha_scheduler_counter=counter, - mla_bmm1_scale=bmm1_scale, - mla_bmm2_scale=bmm2_scale, - quant_q_buffer=quant_q, - # The whole of drafting, as far as this op is concerned: - # the query block is `gen_p` rows per generation sequence, - # token-major, and `gen_p` is what produces the - # bottom-right-aligned within-block causal mask (draft row - # t sees [0, L_g - gen_p + t] and none of its later - # siblings). Certified at 1..4 over this fp8 cell. - predicted_tokens_per_seq=gen_p, - **step, - **rope, - **self._call, - ) - bmm_out( - transpose(reshape(lat_out, [gen, self.heads, self.kv_lora]), 0, 1), - v_b_t, - transpose(reshape(attn_gen, [gen, self.heads, self.v_dim]), 0, 1), + layer=i, ) + attn_gen_t = self._expand(a=lat_out_t, out=attn_gen_t, layer=i) - # Replicated o_proj over this rank's own tokens: complete as it - # stands, so the residual stream is updated with no collective. - o = cublas_mm(attn_out, w_o) - flashinfer_fused_add_rmsnorm(o, residual, w_n2, self.eps) - if moe_w[i] is None: - mlp_out = self._dense_mlp(o, mlp_w[i], dt) + o = self._o(attn_out, layer=i) + o, residual = self._norm2(o, residual, layer=i) + if i < dense_layers: + mlp_out = self._dense_mlp(o, i, dt) else: - # The shared expert and the routed round trip read the same `o` - # and meet only at the add below, so they are independent — but - # on one stream the shared expert's five kernels sit in front of - # a gather that is 94% exclusive on the device. Forking it onto - # a side stream lets it run inside that window. Both collective - # contracts certify a side stream joined to the current one on - # both ends, which is exactly the shape here; the fork/join pair - # is also what propagates a CUDA-graph capture into the branch - # and back, so a decode capture records both streams. - side = self._side_stream + side = core._side_stream main = torch.cuda.current_stream() - assert side is not None, "derive_after_load must run before forward" side.wait_stream(main) with torch.cuda.stream(side): - shared = self._dense_mlp(o, mlp_w[i], dt) - ( - w_router, - router_bias, - fc1_w, - fc1_s, - fc2_w, - fc2_s, - o1, - o1_gate, - o2, - routed_g, - ) = moe_w[i] - # The expert-parallel round trip. Gathering *before* the - # router is what keeps the four windows tiling the routing - # space exactly once: every rank routes the identical full - # token set, so a token's top-8 ids agree across ranks and - # each id falls in exactly one window. - o_pad = o if not pad_rows else pad(o, [0, 0, 0, pad_rows]) - o_all = allgather(o_pad, None, self.dp_group) - # The expert call is chunked to `_MOE_MAX_T` rows: the - # gathered set reaches 4 * max_num_tokens and both the runner's - # and the routing op's certified token columns stop at 8192. - # Every token is independent through routing, quantization and - # the expert GEMMs, so the chunks are the whole call, re-joined. + shared = self._dense_mlp(o, i, dt) + # Gathering before the router is what keeps the four expert windows tiling the + # routing space exactly once: every rank routes the identical token set, so a + # token's top-8 ids agree across ranks and each id falls in exactly one window. + o_pad = ( + o + if not pad_rows + else nn.functional.pad(o, [0, 0, 0, pad_rows], mode="constant", value=0.0) + ) + o_all = self._allgather(o_pad) parts = [] - for chunk in split(o_all, _moe_chunk_sizes(o_all.shape[0]), 0): - # Raw bf16 logits: noaux_tc_op applies the sigmoid - # itself, and its weight dtype follows the logits, so - # bf16 in means the MoE runner's bf16 topk_weights need - # no cast (the fp32 correction bias does not change that). - logits = cublas_mm(chunk, w_router) - topk_w, topk_ids = noaux_tc_op( - logits, - router_bias, - self.n_group, - self.topk_group, - self.topk, - self.routed_scale, - ) - # A second quantization of the same hidden states: the - # MoE runner reads the linear scale buffer as - # float8_e4m3fn, never the swizzled one the shared GEMM - # above consumed. - xq, xsf = fp4_quantize(chunk, routed_g, _SF_VEC, False, _SF_LINEAR) + for chunk in torch.split(o_all, _moe_chunk_sizes(o_all.shape[0]), 0): + logits = self._router(chunk, layer=i) + topk_w, topk_ids = self._noaux(logits, layer=i) + xq, xsf = self._moe_quant(chunk, layer=i) parts.append( - fp4_block_scale_moe_runner( - None, - None, - xq, - view_dtype(xsf, torch.float8_e4m3fn), - fc1_w, - fc1_s, - None, - None, - None, - None, - fc2_w, - fc2_s, - None, - o1, - o1_gate, - o2, - self.num_experts, - self.topk, - None, - None, - self.moe_inter, - self.expert_offset, - self.local_experts, - None, - _ROUTING_METHOD_INERT, - True, - _ACT_TYPE_SWIGLU, - topk_w, - topk_ids, + self._moe_runner( + hidden_states=xq, + hidden_states_scale=xsf.view(torch.float8_e4m3fn), + topk_weights=topk_w, + topk_ids=topk_ids, + layer=i, )[0] ) - routed_all = parts[0] if len(parts) == 1 else concat(parts, 0) - # The four windows' partials over the whole token set sum to - # the layer's routed output; the scatter hands this rank back - # exactly its own rows. bf16 in: the op sums, and it sums - # float8 as raw bytes. - routed_pad = reducescatter(routed_all, None, self.dp_group) + routed_all = parts[0] if len(parts) == 1 else torch.cat(parts, 0) + routed_pad = self._reducescatter(routed_all) routed = ( - routed_pad if not pad_rows else split(routed_pad, [num_tokens, pad_rows], 0)[0] + routed_pad + if not pad_rows + else torch.split(routed_pad, [num_tokens, pad_rows], 0)[0] ) - # Join: the add is the first reader of `shared` on this stream, - # and `shared` stays referenced until then, so the branch's - # allocations cannot be recycled underneath it. main.wait_stream(side) - mlp_out = add(routed, shared) - flashinfer_fused_add_rmsnorm(mlp_out, residual, next_norm[i], self.eps) + mlp_out = torch.add(routed, shared) + mlp_out, residual = self._next_norm(mlp_out, residual, layer=i) x = mlp_out return x + def _bind_mtp(self) -> None: + """Bind every op this target's draft step calls, once, against `core._mtp`. -class MTPLayer: - """The checkpoint's multi-token-prediction module, at layer index - `num_hidden_layers`, replayed once per draft step. - - Structurally one more decoder layer with a front end bolted on that mixes - in the embedding of the token being predicted: - - e = embed_tokens(input_ids) # the NEXT token - x = eh_proj(concat(enorm(e), hnorm(h))) # [T, 2H] -> [T, H] - x = x + MLA(input_layernorm(x)) - x = x + MoE(post_attention_layernorm(x)) - logits = lm_head(shared_head.norm(x)) # `shared_head`, below - - `h` is the hidden state the runtime hands in: the trunk's own output at - draft step 0, and this layer's own output at every step after. The layer - returns `x` **un-normalized** — `shared_head` is where the module's own - norm is applied — and the runtime slices that with its `gather_ids` for the - logits and feeds it forward as the next step's `h`. - - A deliberate copy of the trunk's layer body rather than a shared helper: - the trunk's loop is inlined and tuner-specialized, and three things differ - here anyway — the front end, the layer index (`num_hidden_layers`, the - extra pool layer the engine adds under a one-model MTP mode), and the MoE - dtype. The checkpoint excludes `model.layers.61*` from NVFP4 wholesale, so - this module is bf16 throughout and its routed experts go to `fused_moe` - where the trunk's go to `fp4_block_scale_moe_runner`. - - Not an `nn.Module`, and `mtp_layers` is a plain list: every weight this - layer reads lives in the core's ParameterDict, so registering the layer - again would put the trunk's parameters on a second path through the shell's - module tree. The runtime never inspects the container beyond `mtp_layers`, - `embed_tokens`, `lm_head` and an absent `model.d2t`. - - Two arguments of the documented calling convention are accepted and unused, - for the same reasons the trunk ignores them: `position_ids` (the MLA ops - take every position from `sequence_length`, not from this tensor) and - `spec_metadata` (the one thing the layer needs off it, the DP padding - basis, arrives as the `all_rank_num_tokens` keyword instead).""" - - def __init__(self, core: ModelingV2Core, logits_processor) -> None: - self.core = core - self.logits_processor = logits_processor + The same mirror of the trunk the MTP module is: its eleven attention-side + weights are layer 61's, its experts are bf16. `shared_head` and its norm stay + on `MTPLayer` -- the worker calls that entry point directly and it reads no + phase. + """ + core = self.core + cfg = core.model_config.pretrained_config + mw = core._mtp + eps = cfg.rms_norm_eps + heads = cfg.num_attention_heads + nope = cfg.qk_nope_head_dim + rope_dim = cfg.qk_rope_head_dim + v_dim = cfg.v_head_dim + kv_lora = cfg.kv_lora_rank + q_lora = cfg.q_lora_rank + qk_dim = nope + rope_dim + lat_dim = kv_lora + rope_dim + layer_idx = cfg.num_hidden_layers + self._mtp_enorm = FlashinferRmsnorm() + self._mtp_enorm.bind_const(weight=mw["enorm"], eps=eps) + self._mtp_hnorm = FlashinferRmsnorm() + self._mtp_hnorm.bind_const(weight=mw["hnorm"], eps=eps) + self._mtp_norm1 = FlashinferRmsnorm() + self._mtp_norm1.bind_const(weight=mw["norm1"], eps=eps) + self._mtp_q_norm = FlashinferRmsnorm() + self._mtp_q_norm.bind_const(weight=mw["q_norm"], eps=eps) + self._mtp_kv_norm = FlashinferRmsnorm() + self._mtp_kv_norm.bind_const(weight=mw["kv_norm"], eps=eps) + self._mtp_norm2 = FlashinferFusedAddRmsnorm() + self._mtp_norm2.bind_const(weight=mw["norm2"], eps=eps) + self._mtp_eh = CublasMm() + self._mtp_eh.bind_const(mat_b=mw["eh"]) + self._mtp_qa = CublasMm() + self._mtp_qa.bind_const(mat_b=mw["qa"]) + self._mtp_qb = CublasMm() + self._mtp_qb.bind_const(mat_b=mw["qb"]) + self._mtp_kva = CublasMm() + self._mtp_kva.bind_const(mat_b=mw["kva"]) + self._mtp_kvb = CublasMm() + self._mtp_kvb.bind_const(mat_b=mw["kvb"]) + self._mtp_o = CublasMm() + self._mtp_o.bind_const(mat_b=mw["o"]) + self._mtp_sh_gu = CublasMm() + self._mtp_sh_gu.bind_const(mat_b=mw["sh_gu"]) + self._mtp_sh_dn = CublasMm() + self._mtp_sh_dn.bind_const(mat_b=mw["sh_dn"]) + self._mtp_silu = FlashinferSiluAndMul() + self._mtp_absorb = BmmOut() + self._mtp_absorb.bind_const(b=mw["k_b"]) + self._mtp_expand = BmmOut() + self._mtp_expand.bind_const(b=mw["v_b_t"]) + self._mtp_allgather = Allgather() + self._mtp_allgather.bind_const(sizes=None, group=core.dp_group) + self._mtp_reducescatter = Reducescatter() + self._mtp_reducescatter.bind_const(sizes=None, group=core.dp_group) + self._mtp_router = CublasMm() + self._mtp_router.bind_const(mat_b=mw["router"], bias=None, out_dtype=torch.float32) + self._mtp_noaux = NoauxTcOp() + self._mtp_noaux.bind_const( + bias=mw["router_bias"], + n_group=cfg.n_group, + topk_group=cfg.topk_group, + topk=cfg.num_experts_per_tok, + routed_scaling_factor=float(cfg.routed_scaling_factor), + ) + self._mtp_moe = FusedMoe() + self._mtp_moe.bind_const( + fc1_expert_weights=mw["fc1"], + fc1_expert_biases=None, + fc2_expert_weights=mw["fc2"], + fc2_expert_biases=None, + quant_scales=[], + ep_size=core.ep_size, + ep_rank=core.ep_rank, + activation_type=_FUSED_MOE_ACT_SWIGLU, + ) + self._mtp_attn_ctx = ThopAttention() + self._mtp_attn_ctx.bind_const( + local_layer_idx=layer_idx, + num_heads=heads, + num_kv_heads=heads, + head_size=qk_dim, + q_lora_rank=q_lora, + kv_lora_rank=kv_lora, + qk_nope_head_dim=nope, + qk_rope_head_dim=rope_dim, + v_head_dim=v_dim, + is_fused_qkv=False, + attention_input_type=1, + q_pe=None, + cu_q_seqlens=None, + cu_kv_seqlens=None, + fmha_scheduler_counter=None, + predicted_tokens_per_seq=1, + **_CALL_INERT, + quant_mode=core.quant_mode, + q_scaling=core.q_scaling, + ) + self._mtp_attn_gen = ThopAttention() + self._mtp_attn_gen.bind_const( + local_layer_idx=layer_idx, + num_heads=heads, + num_kv_heads=1, + head_size=lat_dim, + q_lora_rank=q_lora, + kv_lora_rank=kv_lora, + qk_nope_head_dim=nope, + qk_rope_head_dim=rope_dim, + v_head_dim=kv_lora, + is_fused_qkv=True, + attention_input_type=2, + k=None, + v=None, + **_CALL_INERT, + quant_mode=core.quant_mode, + q_scaling=core.q_scaling, + ) + self._mtp_mla_append = MlaRopeAppendPagedKvAssignQ() + self._mtp_mla_append.bind_const( + layer_idx=layer_idx, + head_num=heads, + nope_size=nope, + rope_size=rope_dim, + lora_size=kv_lora, + kv_scale_orig_quant=None, + residual_dim=_KV_RESIDUAL_DIM, + beam_width=1, + quant_mode=core.quant_mode, + ) + self._mtp_mla_load = LoadPagedKvCacheForMla() + self._mtp_mla_load.bind_const( + layer_idx=layer_idx, + kv_scale_quant_orig=None, + kv_lora_rank=kv_lora, + qk_rope_head_dim=rope_dim, + beam_width=1, + quant_mode=core.quant_mode, + ) + self._mtp_mla_gen = MlaRopeGeneration() + self._mtp_mla_gen.bind_const( + layer_idx=layer_idx, + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + kv_cache_scale_orig_quant=None, + out_scale=None, + block_ids_per_seq=None, + helix_tensor_params=[None, None], + num_heads=heads, + num_kv_heads=1, + head_size=lat_dim, + residual_dim=_KV_RESIDUAL_DIM, + beam_width=1, + quant_mode=core.quant_mode, + q_scaling=core.q_scaling, + q_lora_rank=q_lora, + kv_lora_rank=kv_lora, + qk_nope_head_dim=nope, + qk_rope_head_dim=rope_dim, + v_head_dim=v_dim, + rope_append=True, + ) - def __call__(self, *args, **kwargs) -> torch.Tensor: - return self.forward(*args, **kwargs) + def _mtp_shared_mlp(self, x): + """The module's shared expert: one bf16 SwiGLU MLP over this rank's own tokens.""" + return self._mtp_sh_dn(self._mtp_silu(self._mtp_sh_gu(x))) - def shared_head( - self, hidden_states, lm_head, attn_metadata, return_context_logits - ) -> torch.Tensor: - """The module's own output head: its own RMS norm — a **distinct** - parameter from the trunk's `model.norm` — then the trunk's `lm_head`, - which the checkpoint's `shared_head.head` is a bitwise copy of. The - projection itself is the inherited shell's `logits_processor`, exactly - as on the non-speculative path.""" - mw = self.core._mtp - assert mw is not None, "load_weights must run before the draft loop" - return self.logits_processor.forward( - flashinfer_rmsnorm(hidden_states, mw["head_norm"], self.core.eps), - lm_head, - attn_metadata, - return_context_logits, - ) + def _mtp_routed_experts(self, x, all_rank_num_tokens, rows, dt): + """The module's expert-parallel round trip, the trunk's shape at a different dtype. - def _shared_mlp(self, x, mw): - """The module's shared expert: one bf16 SwiGLU MLP, replicated, over - this rank's own tokens — the same structure as the trunk's shared - expert at the same width, without the NVFP4 quantize/dequantize pair. - `sh_gu` holds the gate rows first, the half order silu_and_mul reads.""" - return cublas_mm(flashinfer_silu_and_mul(cublas_mm(x, mw["sh_gu"])), mw["sh_dn"]) - - def _routed_experts(self, x, mw, all_rank_num_tokens, rows, dt): - """The module's expert-parallel round trip, the trunk's shape at a - different dtype: pad to the group-wide row count, gather **before** the - router so all four ranks route the identical token set, run this rank's - 64-expert window, reduce-scatter the partials back. - - Two things differ from the trunk's. The router GEMM emits **fp32** - logits: `fused_moe` demands fp32 `token_final_scales` where the - trtllm-gen runner demands bf16, and `noaux_tc_op`'s weight dtype - follows its logits, so the dtype is chosen here rather than cast later. - And the expert call takes the global expert ids directly with - `ep_size`/`ep_rank` shifting this rank's window, where the trtllm-gen - runner takes an explicit offset/count pair. - - `fused_moe` requires a token's `topk` ids to be **distinct** above 256 - tokens — a repeat reads out of bounds in `finalizeMoeRoutingKernel`, - faulting or returning ~200-460 ulp of silent garbage — and this call - drives `T` to 8192. `noaux_tc_op` returns the indices of the top-k - largest corrected scores, and a top-k over expert indices cannot repeat - one, so the precondition holds structurally.""" + `fused_moe` wants fp32 `token_final_scales`, which is why the router GEMM emits fp32 + here and the trunk's emits bf16, and it requires a token's `topk` ids to be distinct + above 256 tokens — a repeat reads out of bounds in `finalizeMoeRoutingKernel`, faulting + or returning silent garbage — which a top-k over expert indices gives structurally.""" core = self.core dp_rows = _mtp_dp_rows(all_rank_num_tokens, core.rank, core.dp_size, rows) pad_rows = dp_rows - rows - x_pad = x if not pad_rows else pad(x, [0, 0, 0, pad_rows]) - x_all = allgather(x_pad, None, core.dp_group) + x_pad = ( + x + if not pad_rows + else nn.functional.pad(x, [0, 0, 0, pad_rows], mode="constant", value=0.0) + ) + x_all = self._mtp_allgather(x_pad) parts = [] - for chunk in split(x_all, _moe_chunk_sizes(x_all.shape[0]), 0): - logits = cublas_mm(chunk, mw["router"], None, torch.float32) - topk_w, topk_ids = noaux_tc_op( - logits, - mw["router_bias"], - core.n_group, - core.topk_group, - core.topk, - core.routed_scale, - ) + for chunk in torch.split(x_all, _moe_chunk_sizes(x_all.shape[0]), 0): + topk_w, topk_ids = self._mtp_noaux(self._mtp_router(chunk)) parts.append( - fused_moe( - chunk, - topk_ids, - topk_w, - mw["fc1"], - None, - mw["fc2"], - None, - dt, - [], - ep_size=core.ep_size, - ep_rank=core.ep_rank, - activation_type=_FUSED_MOE_ACT_SWIGLU, + self._mtp_moe( + input=chunk, + token_selected_experts=topk_ids, + token_final_scales=topk_w, + output_dtype=dt, )[0] ) - routed_all = parts[0] if len(parts) == 1 else concat(parts, 0) - routed_pad = reducescatter(routed_all, None, core.dp_group) - return routed_pad if not pad_rows else split(routed_pad, [rows, pad_rows], 0)[0] + routed_all = parts[0] if len(parts) == 1 else torch.cat(parts, 0) + routed_pad = self._mtp_reducescatter(routed_all) + return routed_pad if not pad_rows else torch.split(routed_pad, [rows, pad_rows], 0)[0] - def forward( + def mtp_forward( self, embed_tokens: torch.Tensor, all_rank_num_tokens, @@ -1721,249 +1157,152 @@ def forward( **kwargs, ) -> torch.Tensor: core = self.core - mw, rope = core._mtp, core._rope - assert mw is not None and rope is not None, "load_weights must run before the draft loop" - assert isinstance(attn_metadata, TrtllmAttentionMetadata) - assert core._step_contract_checked, ( - "the trunk's first-forward contract check has not run; the shell " - "calls the core before the worker, so this cannot be reached first" - ) + cfg = core.model_config.pretrained_config + heads = cfg.num_attention_heads + nope = cfg.qk_nope_head_dim + rope_dim = cfg.qk_rope_head_dim + v_dim = cfg.v_head_dim + kv_lora = cfg.kv_lora_rank + qk_dim = nope + rope_dim + lat_dim = kv_lora + rope_dim + + rope = core._rope md = attn_metadata + # Rebound every call, not once: the worker rewrites `attn_metadata` in place + # between draft step 0 and step 1+, and `core._rope` may have grown on the + # engine's first step. step = _build_step_args(md) + self._mtp_attn_ctx.bind_const(**step, **rope, **core._call) + self._mtp_attn_gen.bind_const(**step, **rope, **core._call) rows = hidden_states.shape[0] - assert input_ids.shape[0] == rows, ( - f"the draft step's {input_ids.shape[0]} token ids and " - f"{rows} hidden-state rows must describe the same tokens" - ) dt = hidden_states.dtype dev = hidden_states.device num_ctx = md.num_contexts tc = md.num_ctx_tokens gen = rows - tc - # **Read the phase from the metadata on every call.** The worker - # rewrites `attn_metadata` in place between draft step 0 and step 1+ — - # `_seq_lens` filled with 1, `num_contexts` and `num_ctx_tokens` to 0, - # one token per generation request — and this layer is invoked N times - # inside one forward, so a value computed on the first call is wrong on - # the rest (and, under capture, would be frozen wrong at all N - # positions). The read is host-side and sync-free: `on_update()` - # recomputes these from the pinned-host `_seq_lens` and the loop - # triggers it at exactly that boundary. The quotient is - # `runtime_draft_len + 1` on step 0 and exactly 1 afterwards. + # Read the phase from the metadata on every call: the worker rewrites `attn_metadata` in + # place between draft step 0 and step 1+, so a value computed on the first of the N + # calls in one forward is wrong on the rest, and frozen wrong under capture. gen_seqs = md.num_seqs - num_ctx gen_p = gen // gen_seqs if gen_seqs else 1 - assert gen_p * gen_seqs == gen, ( - f"{gen} generation rows do not divide over {gen_seqs} generation " - "sequences; a ragged per-sequence draft length cannot be expressed " - "to the MLA generation call" - ) tokens_per_block = md.tokens_per_block block_offsets = md.kv_cache_block_offsets pool_ptrs = md.host_kv_cache_pool_pointers pool_map = md.host_kv_cache_pool_mapping - assert tokens_per_block is not None, "paged KV cache is required" - assert block_offsets is not None and pool_ptrs is not None, "paged KV cache is required" - assert pool_map is not None, "paged KV cache is required" - ctx_kv_tokens = int(md.host_total_kv_lens[0]) if tc else 0 - # The engine raises the KV pool's layer count by - # `num_nextn_predict_layers` under a one-model MTP mode, so this - # module's attention addresses layer index `num_hidden_layers` in the - # same single pool the trunk's 61 layers use. Nothing in the target's - # config stub or manifest declares that. - layer_idx = core.num_layers - - e = embedding(input_ids, embed_tokens) - en = flashinfer_rmsnorm(e, mw["enorm"], core.eps) - hn = flashinfer_rmsnorm(hidden_states, mw["hnorm"], core.eps) + ctx_kv_tokens = int(md.host_total_kv_lens[0]) + # Under a one-model MTP mode the engine raises the KV pool's layer count by + # `num_nextn_predict_layers`, so this module addresses layer `num_hidden_layers` in the + # same pool as the trunk's 61 layers. Nothing in the config stub or manifest says so. + + e = nn.functional.embedding(input_ids, embed_tokens) + en = self._mtp_enorm(e) + hn = self._mtp_hnorm(hidden_states) halves = [en, hn] if _MTP_EMBED_BLOCK_FIRST else [hn, en] - x = cublas_mm(concat(halves, -1), mw["eh"]) + x = self._mtp_eh(torch.cat(halves, -1)) residual = x - xn = flashinfer_rmsnorm(x, mw["norm1"], core.eps) - attn_out = empty([rows, core.heads * core.v_dim], dt, dev) - attn_ctx, attn_gen = split(attn_out, [tc, gen], 0) - q = cublas_mm( - flashinfer_rmsnorm(cublas_mm(xn, mw["qa"]), mw["q_norm"], core.eps), - mw["qb"], - ) - kva = cublas_mm(xn, mw["kva"]) - ckv_raw, k_pe = split(kva, [core.kv_lora, core.rope], -1) - ckv = flashinfer_rmsnorm(ckv_raw, mw["kv_norm"], core.eps) - latent = concat([ckv, k_pe], -1) - q_ctx, q_gen = split(q, [tc, gen], 0) - latent_ctx, latent_gen = split(latent, [tc, gen], 0) - - if tc: - # A context request reaches the draft loop at step 0 only, fed - # `prompt[1:]` with its first accepted token written at the last - # position — the same row count and the same per-sequence lengths - # the trunk saw, so the context flavor the trunk settled at its - # first forward applies unchanged here. - if core._cached_ctx: - mla_rope_append_paged_kv_assign_q( - q_ctx, - latent_ctx, - num_ctx, - md.ctx_cached_token_indptr, - md.ctx_kv_indptr, - int(md.max_ctx_seq_len), - rope["rotary_cos_sin"], - core.heads, - core.nope, - core.rope, - core.kv_lora, - block_offsets, - pool_ptrs, - pool_map, - None, - _KV_RESIDUAL_DIM, - layer_idx, - tokens_per_block, - md.max_seq_len, - 1, - core.quant_mode, - ) - ckv_full, k_pe_full = load_paged_kv_cache_for_mla( - dt, - num_ctx, - ctx_kv_tokens, - int(md.max_ctx_kv_len), - md.ctx_kv_indptr, - block_offsets, - pool_ptrs, - pool_map, - None, - layer_idx, - core.kv_lora, - core.rope, - tokens_per_block, - md.max_seq_len, - 1, - core.quant_mode, - ) - latent_arg = None - else: - assert ctx_kv_tokens == tc, ( - "a context sequence arrived with a cached KV prefix while " - "the cached-KV metadata surface is off" - ) - ckv_full, _ = split(ckv, [tc, gen], 0) - k_pe_full = None - latent_arg = latent_ctx - tkv = ctx_kv_tokens - kv = cublas_mm(ckv_full, mw["kvb"]) - k_nope, v_view = split(kv, [core.heads * core.nope, core.heads * core.v_dim], -1) - k = empty([tkv, core.heads, core.qk_dim], dt, dev) - k_nope_dst, k_pe_dst = split(k, [core.nope, core.rope], -1) - copy_(k_nope_dst, reshape(k_nope, [tkv, core.heads, core.nope])) - if k_pe_full is not None: - copy_( - k_pe_dst, - expand( - reshape(k_pe_full, [tkv, 1, core.rope]), - [tkv, core.heads, core.rope], - ), - ) - thop_attention( + xn = self._mtp_norm1(x) + attn_out = torch.empty([rows, heads * v_dim], dtype=dt, device=dev) + attn_ctx, attn_gen = torch.split(attn_out, [tc, gen], 0) + q = self._mtp_qb(self._mtp_q_norm(self._mtp_qa(xn))) + kva = self._mtp_kva(xn) + ckv_raw, k_pe = torch.split(kva, [kv_lora, rope_dim], -1) + ckv = self._mtp_kv_norm(ckv_raw) + latent = torch.cat([ckv, k_pe], -1) + q_ctx, q_gen = torch.split(q, [tc, gen], 0) + latent_ctx, latent_gen = torch.split(latent, [tc, gen], 0) + + if core._cached_ctx: + q_ctx = self._mtp_mla_append( q=q_ctx, - k=reshape(k, [tkv, core.heads * core.qk_dim]), - v=v_view, - output=attn_ctx, - latent_cache=latent_arg, - q_pe=None, - local_layer_idx=layer_idx, - is_fused_qkv=False, - attention_input_type=1, - num_heads=core.heads, - num_kv_heads=core.heads, - head_size=core.qk_dim, - q_lora_rank=core.q_lora, - kv_lora_rank=core.kv_lora, - qk_nope_head_dim=core.nope, - qk_rope_head_dim=core.rope, - v_head_dim=core.v_dim, - cu_q_seqlens=None, - cu_kv_seqlens=None, - fmha_scheduler_counter=None, - predicted_tokens_per_seq=1, - **step, - **rope, - **core._call, + latent_cache=latent_ctx, + num_contexts=num_ctx, + cu_ctx_cached_kv_lens=md.ctx_cached_token_indptr, + cu_seq_lens=md.ctx_kv_indptr, + max_input_uncached_seq_len=int(md.max_ctx_seq_len), + cos_sin_cache=rope["rotary_cos_sin"], + kv_cache_block_offsets=block_offsets, + host_kv_cache_pool_pointers=pool_ptrs, + host_kv_cache_pool_mapping=pool_map, + tokens_per_block=tokens_per_block, + attention_window_size=md.max_seq_len, + ) + ckv_full, k_pe_full = self._mtp_mla_load( + out_dtype=dt, + num_contexts=num_ctx, + num_ctx_kv_tokens=ctx_kv_tokens, + max_ctx_kv_len=int(md.max_ctx_kv_len), + cu_ctx_kv_lens=md.ctx_kv_indptr, + kv_cache_block_offsets=block_offsets, + host_kv_cache_pool_pointers=pool_ptrs, + host_kv_cache_pool_mapping=pool_map, + tokens_per_block=tokens_per_block, + attention_window_size=md.max_seq_len, + ) + latent_arg = None + else: + ckv_full, _ = torch.split(ckv, [tc, gen], 0) + k_pe_full = None + latent_arg = latent_ctx + tkv = ctx_kv_tokens + kv = self._mtp_kvb(ckv_full) + k_nope, v_view = torch.split(kv, [heads * nope, heads * v_dim], -1) + k = torch.empty([tkv, heads, qk_dim], dtype=dt, device=dev) + k_nope_dst, k_pe_dst = torch.split(k, [nope, rope_dim], -1) + k_nope_dst.copy_(torch.reshape(k_nope, [tkv, heads, nope])) + if k_pe_full is not None: + k_pe_dst.copy_( + torch.reshape(k_pe_full, [tkv, 1, rope_dim]).expand([tkv, heads, rope_dim]) ) + attn_ctx = self._mtp_attn_ctx( + q=q_ctx, + k=torch.reshape(k, [tkv, heads * qk_dim]), + v=v_view, + output=attn_ctx, + latent_cache=latent_arg, + ) if gen: - q3 = reshape(q_gen, [gen, core.heads, core.qk_dim]) - q_nope, q_pe = split(q3, [core.nope, core.rope], -1) - fused_q = empty([gen, core.heads, core.lat_dim], dt, dev) - fq_nope, _ = split(fused_q, [core.kv_lora, core.rope], -1) - bmm_out(transpose(q_nope, 0, 1), mw["k_b"], transpose(fq_nope, 0, 1)) - cu_q = empty([gen + 1], torch.int32, dev) - cu_kv = empty([gen + 1], torch.int32, dev) - counter = empty([1], torch.uint32, dev) - quant_q = empty([gen, core.heads, core.lat_dim], torch.float8_e4m3fn, dev) - bmm1_scale = empty([2], torch.float32, dev) - bmm2_scale = empty([1], torch.float32, dev) - mla_rope_generation( - fused_q, - q_pe, - latent_gen, - rope["rotary_cos_sin"], - cu_q, - cu_kv, - counter, - bmm1_scale, - bmm2_scale, - quant_q, - md.kv_lens_cuda_runtime, - md.kv_lens_runtime, - md.prompt_lens_cpu_runtime, - num_ctx, - block_offsets, - pool_ptrs, - pool_map, - None, - None, - None, - None, - None, - [None, None], - gen_p, - layer_idx, - core.heads, - 1, - core.lat_dim, - _KV_RESIDUAL_DIM, - tokens_per_block, - md.max_seq_len, - 1, - core.quant_mode, - core.q_scaling, - core.q_lora, - core.kv_lora, - core.nope, - core.rope, - core.v_dim, - True, + q3 = torch.reshape(q_gen, [gen, heads, qk_dim]) + q_nope, q_pe = torch.split(q3, [nope, rope_dim], -1) + fused_q = torch.empty([gen, heads, lat_dim], dtype=dt, device=dev) + fq_nope, _ = torch.split(fused_q, [kv_lora, rope_dim], -1) + fq_nope_t = torch.transpose(fq_nope, 0, 1) + fq_nope_t = self._mtp_absorb(a=torch.transpose(q_nope, 0, 1), out=fq_nope_t) + cu_q = torch.empty([gen + 1], dtype=torch.int32, device=dev) + cu_kv = torch.empty([gen + 1], dtype=torch.int32, device=dev) + counter = torch.empty([1], dtype=torch.uint32, device=dev) + quant_q = torch.empty([gen, heads, lat_dim], dtype=torch.float8_e4m3fn, device=dev) + bmm1_scale = torch.empty([2], dtype=torch.float32, device=dev) + bmm2_scale = torch.empty([1], dtype=torch.float32, device=dev) + fused_q = self._mtp_mla_gen( + fused_q=fused_q, + q_pe=q_pe, + latent_cache=latent_gen, + rotary_cos_sin=rope["rotary_cos_sin"], + cu_q_seqlens=cu_q, + cu_kv_seqlens=cu_kv, + fmha_scheduler_counter=counter, + mla_bmm1_scale=bmm1_scale, + mla_bmm2_scale=bmm2_scale, + quant_q_buffer=quant_q, + sequence_length=md.kv_lens_cuda_runtime, + host_past_key_value_lengths=md.kv_lens_runtime, + host_context_lengths=md.prompt_lens_cpu_runtime, + num_contexts=num_ctx, + kv_cache_block_offsets=block_offsets, + host_kv_cache_pool_pointers=pool_ptrs, + host_kv_cache_pool_mapping=pool_map, + predicted_tokens_per_seq=gen_p, + tokens_per_block=tokens_per_block, + attention_window_size=md.max_seq_len, ) - lat_out = empty([gen, core.heads * core.kv_lora], dt, dev) - thop_attention( - q=reshape(fused_q, [gen, core.heads * core.lat_dim]), - k=None, - v=None, + lat_out = torch.empty([gen, heads * kv_lora], dtype=dt, device=dev) + lat_out = self._mtp_attn_gen( + q=torch.reshape(fused_q, [gen, heads * lat_dim]), output=lat_out, latent_cache=latent_gen, q_pe=q_pe, - local_layer_idx=layer_idx, - is_fused_qkv=True, - attention_input_type=2, - num_heads=core.heads, - num_kv_heads=1, - head_size=core.lat_dim, - q_lora_rank=core.q_lora, - kv_lora_rank=core.kv_lora, - qk_nope_head_dim=core.nope, - qk_rope_head_dim=core.rope, - v_head_dim=core.kv_lora, cu_q_seqlens=cu_q, cu_kv_seqlens=cu_kv, fmha_scheduler_counter=counter, @@ -1971,48 +1310,788 @@ def forward( mla_bmm2_scale=bmm2_scale, quant_q_buffer=quant_q, predicted_tokens_per_seq=gen_p, - **step, - **rope, - **core._call, ) - bmm_out( - transpose(reshape(lat_out, [gen, core.heads, core.kv_lora]), 0, 1), - mw["v_b_t"], - transpose(reshape(attn_gen, [gen, core.heads, core.v_dim]), 0, 1), + lat_out_t = torch.transpose(torch.reshape(lat_out, [gen, heads, kv_lora]), 0, 1) + attn_gen_t = torch.transpose(torch.reshape(attn_gen, [gen, heads, v_dim]), 0, 1) + attn_gen_t = self._mtp_expand(a=lat_out_t, out=attn_gen_t) + + o = self._mtp_o(attn_out) + o, residual = self._mtp_norm2(o, residual) + shared = self._mtp_shared_mlp(o) + routed = self._mtp_routed_experts(o, all_rank_num_tokens, rows, dt) + return torch.add(residual, torch.add(routed, shared)) + + +class DecodeTarget(Target): + """The specialization: routed to only when there are no context rows. + + Binds four fewer ops than the prefill target -- the context attention instance, the + kv_b projection and the two paged-KV helpers all serve a path this target's routing + guarantees never runs. + """ + + def __init__(self, core: "DeepseekV3Dep4ModelingV2") -> None: + """Bind every op this target calls, once, against the real weights.""" + super().__init__(core) + cfg = core.model_config.pretrained_config + + num_layers = cfg.num_hidden_layers + + heads = cfg.num_attention_heads + + nope = cfg.qk_nope_head_dim + + rope_dim = cfg.qk_rope_head_dim + + v_dim = cfg.v_head_dim + + kv_lora = cfg.kv_lora_rank + + q_lora = cfg.q_lora_rank + + lat_dim = kv_lora + rope_dim + + eps = cfg.rms_norm_eps + + dense_layers = cfg.first_k_dense_replace + + num_experts = cfg.n_routed_experts + + topk = cfg.num_experts_per_tok + + moe_inter = cfg.moe_intermediate_size + + w = core.w + + window = slice(core.expert_offset, core.expert_offset + core.local_experts) + + self._rms0 = FlashinferRmsnorm() + + self._rms0.bind_const(weight=w["l0_norm1"], eps=eps) + + self._qa = CublasMm() + + self._q_lora_norm = FlashinferRmsnorm() + + self._q_lora_norm.bind_const(eps=eps) + + self._qb = CublasMm() + + self._kva = CublasMm() + + self._kv_lora_norm = FlashinferRmsnorm() + + self._kv_lora_norm.bind_const(eps=eps) + + self._o = CublasMm() + + self._norm2 = FlashinferFusedAddRmsnorm() + + self._norm2.bind_const(eps=eps) + + self._next_norm = FlashinferFusedAddRmsnorm() + + self._next_norm.bind_const(eps=eps) + + self._absorb = BmmOut() + + self._expand = BmmOut() + + self._attn_gen = ThopAttention() + + self._attn_gen.bind_const( + num_heads=heads, + num_kv_heads=1, + head_size=lat_dim, + q_lora_rank=q_lora, + kv_lora_rank=kv_lora, + qk_nope_head_dim=nope, + qk_rope_head_dim=rope_dim, + v_head_dim=kv_lora, + is_fused_qkv=True, + attention_input_type=2, + k=None, + v=None, + **_CALL_INERT, + quant_mode=core.quant_mode, + q_scaling=core.q_scaling, + ) + + self._mla_gen = MlaRopeGeneration() + + self._mla_gen.bind_const( + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + kv_cache_scale_orig_quant=None, + out_scale=None, + block_ids_per_seq=None, + helix_tensor_params=[None, None], + num_heads=heads, + num_kv_heads=1, + head_size=lat_dim, + residual_dim=_KV_RESIDUAL_DIM, + beam_width=1, + quant_mode=core.quant_mode, + q_scaling=core.q_scaling, + q_lora_rank=q_lora, + kv_lora_rank=kv_lora, + qk_nope_head_dim=nope, + qk_rope_head_dim=rope_dim, + v_head_dim=kv_lora, + rope_append=True, + ) + + self._allgather = Allgather() + + self._allgather.bind_const(sizes=None, group=core.dp_group) + + self._reducescatter = Reducescatter() + + self._reducescatter.bind_const(sizes=None, group=core.dp_group) + + self._noaux = NoauxTcOp() + + self._noaux.bind_const( + n_group=cfg.n_group, + topk_group=cfg.topk_group, + topk=topk, + routed_scaling_factor=float(cfg.routed_scaling_factor), + ) + + self._router = CublasMm() + + self._mlp_gu_quant = Fp4Quantize() + + self._mlp_gu_quant.bind_const( + sf_vec_size=_SF_VEC, sf_use_ue8m0=False, is_sf_swizzled_layout=_SF_SWIZZLED + ) + + self._mlp_gu_gemm = Nvfp4Gemm() + + self._mlp_dn_quant = Fp4Quantize() + + self._mlp_dn_quant.bind_const( + sf_vec_size=_SF_VEC, sf_use_ue8m0=False, is_sf_swizzled_layout=_SF_SWIZZLED + ) + + self._mlp_dn_gemm = Nvfp4Gemm() + + self._moe_quant = Fp4Quantize() + + self._moe_quant.bind_const( + sf_vec_size=_SF_VEC, sf_use_ue8m0=False, is_sf_swizzled_layout=_SF_LINEAR + ) + + self._moe_runner = Fp4BlockScaleMoeRunner() + + self._moe_runner.bind_const( + routing_logits=None, + routing_bias=None, + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_bias=None, + num_experts=num_experts, + top_k=topk, + n_group=None, + topk_group=None, + intermediate_size=moe_inter, + local_expert_offset=core.expert_offset, + local_num_experts=core.local_experts, + routed_scaling_factor=None, + routing_method_type=_ROUTING_METHOD_INERT, + do_finalize=True, + act_type=_ACT_TYPE_SWIGLU, + ) + + hn = heads * nope + + for i in range(num_layers): + kvb = w[f"l{i}_kvb"] + self._qa.bind_layered(i, mat_b=w[f"l{i}_qa"].t()) + self._q_lora_norm.bind_layered(i, weight=w[f"l{i}_q_norm"]) + self._qb.bind_layered(i, mat_b=w[f"l{i}_qb"].t()) + self._kva.bind_layered(i, mat_b=w[f"l{i}_kva"].t()) + self._kv_lora_norm.bind_layered(i, weight=w[f"l{i}_kv_norm"]) + self._o.bind_layered(i, mat_b=w[f"l{i}_o"].t()) + self._norm2.bind_layered(i, weight=w[f"l{i}_norm2"]) + next_w = w[f"l{i + 1}_norm1"] if i + 1 < num_layers else w["final_norm"] + self._next_norm.bind_layered(i, weight=next_w) + self._absorb.bind_layered(i, b=kvb[:hn].reshape(heads, nope, kv_lora)) + self._expand.bind_layered( + i, b=torch.transpose(kvb[hn:].reshape(heads, v_dim, kv_lora), 1, 2) + ) + self._attn_gen.bind_layered(i, local_layer_idx=i) + self._mla_gen.bind_layered(i, layer_idx=i) + + # The checkpoint stores reciprocals (`input_scale` = 1/g_act, `weight_scale_2` = + # 1/g_w), so the quantizer's global scale is 1/input_scale and the GEMM alpha is + # their product — no further reciprocal anywhere. + isc1 = w[f"l{i}_mlp_isc1"] + isc2 = w[f"l{i}_mlp_isc2"] + self._mlp_gu_quant.bind_layered(i, global_scale=(1.0 / isc1).contiguous()) + self._mlp_gu_gemm.bind_layered( + i, + weight=w[f"l{i}_mlp_gu_w"], + weight_scale=w[f"l{i}_mlp_gu_s"], + alpha=(isc1 * w[f"l{i}_mlp_ws2_1"]).contiguous(), + ) + self._mlp_dn_quant.bind_layered(i, global_scale=(1.0 / isc2).contiguous()) + self._mlp_dn_gemm.bind_layered( + i, + weight=w[f"l{i}_mlp_dn_w"], + weight_scale=w[f"l{i}_mlp_dn_s"], + alpha=(isc2 * w[f"l{i}_mlp_ws2_2"]).contiguous(), + ) + + if i < dense_layers: + continue + self._router.bind_layered(i, mat_b=w[f"l{i}_router"].t()) + self._noaux.bind_layered(i, bias=w[f"l{i}_router_bias"]) + # One quantization of the gathered hidden states feeds every expert on every rank, + # so the routed FC1 activation scale has to be a single value and the same value on + # all four ranks, or the windows no longer sum to the whole layer; the shared + # expert's `input_scale` is the max over all routed ones, so it is the safe one. + e_isc2 = w[f"l{i}_e_isc2"][window] + gate1 = (w[f"l{i}_mlp_isc1"][0] * w[f"l{i}_e_ws2_1"][window]).contiguous() + self._moe_quant.bind_layered(i, global_scale=(1.0 / w[f"l{i}_mlp_isc1"]).contiguous()) + self._moe_runner.bind_layered( + i, + gemm1_weights=w[f"l{i}_fc1_w"], + gemm1_weights_scale=w[f"l{i}_fc1_s"].view(torch.float8_e4m3fn), + gemm2_weights=w[f"l{i}_fc2_w"], + gemm2_weights_scale=w[f"l{i}_fc2_s"].view(torch.float8_e4m3fn), + output1_scale_scalar=(gate1 / e_isc2).contiguous(), + output1_scale_gate_scalar=gate1, + output2_scale_scalar=(e_isc2 * w[f"l{i}_e_ws2_2"][window]).contiguous(), ) + if core.mtp_enabled: + self._bind_mtp() + + def step_args(self, md: TrtllmAttentionMetadata) -> dict: + """The batch state both MLA calls of one step share. - o = cublas_mm(attn_out, mw["o"]) - flashinfer_fused_add_rmsnorm(o, residual, mw["norm2"], core.eps) - shared = self._shared_mlp(o, mw) - routed = self._routed_experts(o, mw, all_rank_num_tokens, rows, dt) - # The layer's output is the residual stream itself, un-normalized: - # `shared_head` owns the module's norm, and the runtime feeds this - # tensor straight back in as the next draft step's `h`. - return add(residual, add(routed, shared)) + States the two context counts instead of reading them back: this target runs + only when `num_contexts` is zero, so reading them would ask the metadata to + confirm what the routing already decided. + """ + step = _build_step_args(md) + step["num_contexts"] = 0 + step["num_ctx_tokens"] = 0 + return step + + def _dense_mlp(self, x, layer: int, dt): + """One NVFP4 SwiGLU MLP over this rank's own tokens, used for both the dense layers' MLP + and every MoE layer's shared expert.""" + xq, xsf = self._mlp_gu_quant(x, layer=layer) + gu = self._mlp_gu_gemm(act_fp4=xq, act_sf=xsf, output_dtype=dt, layer=layer) + act = flashinfer_silu_and_mul(gu) + aq, asf = self._mlp_dn_quant(act, layer=layer) + return self._mlp_dn_gemm(act_fp4=aq, act_sf=asf, output_dtype=dt, layer=layer) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None = None, + position_ids: torch.IntTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + lora_params: dict | None = None, + **kwargs, + ) -> torch.Tensor: + core = self.core + self._check_step_contract(attn_metadata) + + cfg = core.model_config.pretrained_config + heads = cfg.num_attention_heads + nope = cfg.qk_nope_head_dim + rope_dim = cfg.qk_rope_head_dim + v_dim = cfg.v_head_dim + kv_lora = cfg.kv_lora_rank + qk_dim = nope + rope_dim + lat_dim = kv_lora + rope_dim + num_layers = cfg.num_hidden_layers + dense_layers = cfg.first_k_dense_replace + + md = attn_metadata + rope = core._rope + + step = self.step_args(md) + self._attn_gen.bind_const(**step, **rope) + # An explicit subset rather than **step: mla_rope_generation's raw_call + # does not take the projection's engine-bookkeeping keys. + self._mla_gen.bind_const( + rotary_cos_sin=rope["rotary_cos_sin"], + tokens_per_block=md.tokens_per_block, + attention_window_size=md.max_seq_len, + sequence_length=md.kv_lens_cuda_runtime, + host_past_key_value_lengths=md.kv_lens_runtime, + host_context_lengths=md.prompt_lens_cpu_runtime, + num_contexts=0, + kv_cache_block_offsets=md.kv_cache_block_offsets, + host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, + ) + + if inputs_embeds is None: + h = nn.functional.embedding(input_ids, core.w["embed"]) + else: + h = inputs_embeds + num_tokens = h.shape[0] + gen = num_tokens + dt = h.dtype + dev = h.device + # Query rows per generation sequence (the whole draft chain under MTP): the only thing + # that tells the attention ops the query block is taller, and what aligns the + # bottom-right within-block causal mask. + gen_seqs = md.num_seqs + gen_p = gen // gen_seqs if gen_seqs else 1 + dp_rows = core._dp_rows(md, num_tokens) + pad_rows = dp_rows - num_tokens + + attn_out = torch.empty([num_tokens, heads * v_dim], dtype=dt, device=dev) + attn_gen = attn_out + # One set of attention scratch for all layers: every tensor here is fully + # rewritten by layer i's _mla_gen before layer i's _attn_gen reads it, and + # the whole attention path runs on the main stream, so layer i+1 cannot + # overwrite what layer i still reads. 8 allocations per step, not 8 x 61. + fused_q = torch.empty([gen, heads, lat_dim], dtype=dt, device=dev) + fq_nope, _ = torch.split(fused_q, [kv_lora, rope_dim], -1) + fq_nope_t = torch.transpose(fq_nope, 0, 1) + lat_out = torch.empty([gen, heads * kv_lora], dtype=dt, device=dev) + lat_out_t = torch.transpose(torch.reshape(lat_out, [gen, heads, kv_lora]), 0, 1) + attn_gen_t = torch.transpose(torch.reshape(attn_gen, [gen, heads, v_dim]), 0, 1) + cu_q = torch.empty([gen + 1], dtype=torch.int32, device=dev) + cu_kv = torch.empty([gen + 1], dtype=torch.int32, device=dev) + counter = torch.empty([1], dtype=torch.uint32, device=dev) + quant_q = torch.empty([gen, heads, lat_dim], dtype=torch.float8_e4m3fn, device=dev) + bmm1_scale = torch.empty([2], dtype=torch.float32, device=dev) + bmm2_scale = torch.empty([1], dtype=torch.float32, device=dev) + scratch = dict( + cu_q_seqlens=cu_q, + cu_kv_seqlens=cu_kv, + fmha_scheduler_counter=counter, + mla_bmm1_scale=bmm1_scale, + mla_bmm2_scale=bmm2_scale, + quant_q_buffer=quant_q, + predicted_tokens_per_seq=gen_p, + ) + self._mla_gen.bind_const(**scratch) + self._attn_gen.bind_const(**scratch) + x = self._rms0(h) + residual = h + for i in range(num_layers): + q = self._qb(self._q_lora_norm(self._qa(x, layer=i), layer=i), layer=i) + kva = self._kva(x, layer=i) + ckv_raw, k_pe = torch.split(kva, [kv_lora, rope_dim], -1) + ckv = self._kv_lora_norm(ckv_raw, layer=i) + latent = torch.cat([ckv, k_pe], -1) + q_gen = q + latent_gen = latent + + q3 = torch.reshape(q_gen, [gen, heads, qk_dim]) + q_nope, q_pe = torch.split(q3, [nope, rope_dim], -1) + # Over an fp8 pool `_mla_gen` reads this half to build the quantized query, so + # this BMM must be issued first on the same stream; overlapping them is a silent + # race (they only write disjoint halves on a bf16 pool). + fq_nope_t = self._absorb(a=torch.transpose(q_nope, 0, 1), out=fq_nope_t, layer=i) + fused_q = self._mla_gen(fused_q=fused_q, q_pe=q_pe, latent_cache=latent_gen, layer=i) + lat_out = self._attn_gen( + q=torch.reshape(fused_q, [gen, heads * lat_dim]), + output=lat_out, + latent_cache=latent_gen, + q_pe=q_pe, + layer=i, + ) + attn_gen_t = self._expand(a=lat_out_t, out=attn_gen_t, layer=i) + + o = self._o(attn_out, layer=i) + o, residual = self._norm2(o, residual, layer=i) + if i < dense_layers: + mlp_out = self._dense_mlp(o, i, dt) + else: + side = core._side_stream + main = torch.cuda.current_stream() + side.wait_stream(main) + with torch.cuda.stream(side): + shared = self._dense_mlp(o, i, dt) + # Gathering before the router is what keeps the four expert windows tiling the + # routing space exactly once: every rank routes the identical token set, so a + # token's top-8 ids agree across ranks and each id falls in exactly one window. + o_pad = ( + o + if not pad_rows + else nn.functional.pad(o, [0, 0, 0, pad_rows], mode="constant", value=0.0) + ) + o_all = self._allgather(o_pad) + parts = [] + for chunk in torch.split(o_all, _moe_chunk_sizes(o_all.shape[0]), 0): + logits = self._router(chunk, layer=i) + topk_w, topk_ids = self._noaux(logits, layer=i) + xq, xsf = self._moe_quant(chunk, layer=i) + parts.append( + self._moe_runner( + hidden_states=xq, + hidden_states_scale=xsf.view(torch.float8_e4m3fn), + topk_weights=topk_w, + topk_ids=topk_ids, + layer=i, + )[0] + ) + routed_all = parts[0] if len(parts) == 1 else torch.cat(parts, 0) + routed_pad = self._reducescatter(routed_all) + routed = ( + routed_pad + if not pad_rows + else torch.split(routed_pad, [num_tokens, pad_rows], 0)[0] + ) + main.wait_stream(side) + mlp_out = torch.add(routed, shared) + mlp_out, residual = self._next_norm(mlp_out, residual, layer=i) + x = mlp_out + return x + + def _bind_mtp(self) -> None: + """Bind every op this target's draft step calls, once, against `core._mtp`. + + The same mirror of the trunk the MTP module is: its eleven attention-side + weights are layer 61's, its experts are bf16. `shared_head` and its norm stay + on `MTPLayer` -- the worker calls that entry point directly and it reads no + phase. + """ + core = self.core + cfg = core.model_config.pretrained_config + mw = core._mtp + eps = cfg.rms_norm_eps + heads = cfg.num_attention_heads + nope = cfg.qk_nope_head_dim + rope_dim = cfg.qk_rope_head_dim + v_dim = cfg.v_head_dim + kv_lora = cfg.kv_lora_rank + q_lora = cfg.q_lora_rank + lat_dim = kv_lora + rope_dim + layer_idx = cfg.num_hidden_layers + self._mtp_enorm = FlashinferRmsnorm() + self._mtp_enorm.bind_const(weight=mw["enorm"], eps=eps) + self._mtp_hnorm = FlashinferRmsnorm() + self._mtp_hnorm.bind_const(weight=mw["hnorm"], eps=eps) + self._mtp_norm1 = FlashinferRmsnorm() + self._mtp_norm1.bind_const(weight=mw["norm1"], eps=eps) + self._mtp_q_norm = FlashinferRmsnorm() + self._mtp_q_norm.bind_const(weight=mw["q_norm"], eps=eps) + self._mtp_kv_norm = FlashinferRmsnorm() + self._mtp_kv_norm.bind_const(weight=mw["kv_norm"], eps=eps) + self._mtp_norm2 = FlashinferFusedAddRmsnorm() + self._mtp_norm2.bind_const(weight=mw["norm2"], eps=eps) + self._mtp_eh = CublasMm() + self._mtp_eh.bind_const(mat_b=mw["eh"]) + self._mtp_qa = CublasMm() + self._mtp_qa.bind_const(mat_b=mw["qa"]) + self._mtp_qb = CublasMm() + self._mtp_qb.bind_const(mat_b=mw["qb"]) + self._mtp_kva = CublasMm() + self._mtp_kva.bind_const(mat_b=mw["kva"]) + self._mtp_o = CublasMm() + self._mtp_o.bind_const(mat_b=mw["o"]) + self._mtp_sh_gu = CublasMm() + self._mtp_sh_gu.bind_const(mat_b=mw["sh_gu"]) + self._mtp_sh_dn = CublasMm() + self._mtp_sh_dn.bind_const(mat_b=mw["sh_dn"]) + self._mtp_silu = FlashinferSiluAndMul() + self._mtp_absorb = BmmOut() + self._mtp_absorb.bind_const(b=mw["k_b"]) + self._mtp_expand = BmmOut() + self._mtp_expand.bind_const(b=mw["v_b_t"]) + self._mtp_allgather = Allgather() + self._mtp_allgather.bind_const(sizes=None, group=core.dp_group) + self._mtp_reducescatter = Reducescatter() + self._mtp_reducescatter.bind_const(sizes=None, group=core.dp_group) + self._mtp_router = CublasMm() + self._mtp_router.bind_const(mat_b=mw["router"], bias=None, out_dtype=torch.float32) + self._mtp_noaux = NoauxTcOp() + self._mtp_noaux.bind_const( + bias=mw["router_bias"], + n_group=cfg.n_group, + topk_group=cfg.topk_group, + topk=cfg.num_experts_per_tok, + routed_scaling_factor=float(cfg.routed_scaling_factor), + ) + self._mtp_moe = FusedMoe() + self._mtp_moe.bind_const( + fc1_expert_weights=mw["fc1"], + fc1_expert_biases=None, + fc2_expert_weights=mw["fc2"], + fc2_expert_biases=None, + quant_scales=[], + ep_size=core.ep_size, + ep_rank=core.ep_rank, + activation_type=_FUSED_MOE_ACT_SWIGLU, + ) + self._mtp_attn_gen = ThopAttention() + self._mtp_attn_gen.bind_const( + local_layer_idx=layer_idx, + num_heads=heads, + num_kv_heads=1, + head_size=lat_dim, + q_lora_rank=q_lora, + kv_lora_rank=kv_lora, + qk_nope_head_dim=nope, + qk_rope_head_dim=rope_dim, + v_head_dim=kv_lora, + is_fused_qkv=True, + attention_input_type=2, + k=None, + v=None, + **_CALL_INERT, + quant_mode=core.quant_mode, + q_scaling=core.q_scaling, + ) + self._mtp_mla_gen = MlaRopeGeneration() + self._mtp_mla_gen.bind_const( + layer_idx=layer_idx, + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + kv_cache_scale_orig_quant=None, + out_scale=None, + block_ids_per_seq=None, + helix_tensor_params=[None, None], + num_heads=heads, + num_kv_heads=1, + head_size=lat_dim, + residual_dim=_KV_RESIDUAL_DIM, + beam_width=1, + quant_mode=core.quant_mode, + q_scaling=core.q_scaling, + q_lora_rank=q_lora, + kv_lora_rank=kv_lora, + qk_nope_head_dim=nope, + qk_rope_head_dim=rope_dim, + v_head_dim=v_dim, + rope_append=True, + ) + + def _mtp_shared_mlp(self, x): + """The module's shared expert: one bf16 SwiGLU MLP over this rank's own tokens.""" + return self._mtp_sh_dn(self._mtp_silu(self._mtp_sh_gu(x))) + + def _mtp_routed_experts(self, x, all_rank_num_tokens, rows, dt): + """The module's expert-parallel round trip, the trunk's shape at a different dtype. + + `fused_moe` wants fp32 `token_final_scales`, which is why the router GEMM emits fp32 + here and the trunk's emits bf16, and it requires a token's `topk` ids to be distinct + above 256 tokens — a repeat reads out of bounds in `finalizeMoeRoutingKernel`, faulting + or returning silent garbage — which a top-k over expert indices gives structurally.""" + core = self.core + dp_rows = _mtp_dp_rows(all_rank_num_tokens, core.rank, core.dp_size, rows) + pad_rows = dp_rows - rows + x_pad = ( + x + if not pad_rows + else nn.functional.pad(x, [0, 0, 0, pad_rows], mode="constant", value=0.0) + ) + x_all = self._mtp_allgather(x_pad) + parts = [] + for chunk in torch.split(x_all, _moe_chunk_sizes(x_all.shape[0]), 0): + topk_w, topk_ids = self._mtp_noaux(self._mtp_router(chunk)) + parts.append( + self._mtp_moe( + input=chunk, + token_selected_experts=topk_ids, + token_final_scales=topk_w, + output_dtype=dt, + )[0] + ) + routed_all = parts[0] if len(parts) == 1 else torch.cat(parts, 0) + routed_pad = self._mtp_reducescatter(routed_all) + return routed_pad if not pad_rows else torch.split(routed_pad, [rows, pad_rows], 0)[0] + + def mtp_forward( + self, + embed_tokens: torch.Tensor, + all_rank_num_tokens, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> torch.Tensor: + core = self.core + cfg = core.model_config.pretrained_config + heads = cfg.num_attention_heads + nope = cfg.qk_nope_head_dim + rope_dim = cfg.qk_rope_head_dim + v_dim = cfg.v_head_dim + kv_lora = cfg.kv_lora_rank + qk_dim = nope + rope_dim + lat_dim = kv_lora + rope_dim + + rope = core._rope + md = attn_metadata + # Rebound every call, not once: the worker rewrites `attn_metadata` in place + # between draft step 0 and step 1+, and `core._rope` may have grown on the + # engine's first step. + step = _build_step_args(md) + self._mtp_attn_gen.bind_const(**step, **rope, **core._call) + rows = hidden_states.shape[0] + dt = hidden_states.dtype + dev = hidden_states.device + gen = rows + # Read the phase from the metadata on every call: the worker rewrites `attn_metadata` in + # place between draft step 0 and step 1+, so a value computed on the first of the N + # calls in one forward is wrong on the rest, and frozen wrong under capture. + gen_seqs = md.num_seqs + gen_p = gen // gen_seqs if gen_seqs else 1 + tokens_per_block = md.tokens_per_block + block_offsets = md.kv_cache_block_offsets + pool_ptrs = md.host_kv_cache_pool_pointers + pool_map = md.host_kv_cache_pool_mapping + # Under a one-model MTP mode the engine raises the KV pool's layer count by + # `num_nextn_predict_layers`, so this module addresses layer `num_hidden_layers` in the + # same pool as the trunk's 61 layers. Nothing in the config stub or manifest says so. + + e = nn.functional.embedding(input_ids, embed_tokens) + en = self._mtp_enorm(e) + hn = self._mtp_hnorm(hidden_states) + halves = [en, hn] if _MTP_EMBED_BLOCK_FIRST else [hn, en] + x = self._mtp_eh(torch.cat(halves, -1)) + + residual = x + xn = self._mtp_norm1(x) + attn_out = torch.empty([rows, heads * v_dim], dtype=dt, device=dev) + attn_gen = attn_out + q = self._mtp_qb(self._mtp_q_norm(self._mtp_qa(xn))) + kva = self._mtp_kva(xn) + ckv_raw, k_pe = torch.split(kva, [kv_lora, rope_dim], -1) + ckv = self._mtp_kv_norm(ckv_raw) + latent = torch.cat([ckv, k_pe], -1) + q_gen = q + latent_gen = latent + + q3 = torch.reshape(q_gen, [gen, heads, qk_dim]) + q_nope, q_pe = torch.split(q3, [nope, rope_dim], -1) + fused_q = torch.empty([gen, heads, lat_dim], dtype=dt, device=dev) + fq_nope, _ = torch.split(fused_q, [kv_lora, rope_dim], -1) + fq_nope_t = torch.transpose(fq_nope, 0, 1) + fq_nope_t = self._mtp_absorb(a=torch.transpose(q_nope, 0, 1), out=fq_nope_t) + cu_q = torch.empty([gen + 1], dtype=torch.int32, device=dev) + cu_kv = torch.empty([gen + 1], dtype=torch.int32, device=dev) + counter = torch.empty([1], dtype=torch.uint32, device=dev) + quant_q = torch.empty([gen, heads, lat_dim], dtype=torch.float8_e4m3fn, device=dev) + bmm1_scale = torch.empty([2], dtype=torch.float32, device=dev) + bmm2_scale = torch.empty([1], dtype=torch.float32, device=dev) + fused_q = self._mtp_mla_gen( + fused_q=fused_q, + q_pe=q_pe, + latent_cache=latent_gen, + rotary_cos_sin=rope["rotary_cos_sin"], + cu_q_seqlens=cu_q, + cu_kv_seqlens=cu_kv, + fmha_scheduler_counter=counter, + mla_bmm1_scale=bmm1_scale, + mla_bmm2_scale=bmm2_scale, + quant_q_buffer=quant_q, + sequence_length=md.kv_lens_cuda_runtime, + host_past_key_value_lengths=md.kv_lens_runtime, + host_context_lengths=md.prompt_lens_cpu_runtime, + num_contexts=0, + kv_cache_block_offsets=block_offsets, + host_kv_cache_pool_pointers=pool_ptrs, + host_kv_cache_pool_mapping=pool_map, + predicted_tokens_per_seq=gen_p, + tokens_per_block=tokens_per_block, + attention_window_size=md.max_seq_len, + ) + lat_out = torch.empty([gen, heads * kv_lora], dtype=dt, device=dev) + lat_out = self._mtp_attn_gen( + q=torch.reshape(fused_q, [gen, heads * lat_dim]), + output=lat_out, + latent_cache=latent_gen, + q_pe=q_pe, + cu_q_seqlens=cu_q, + cu_kv_seqlens=cu_kv, + fmha_scheduler_counter=counter, + mla_bmm1_scale=bmm1_scale, + mla_bmm2_scale=bmm2_scale, + quant_q_buffer=quant_q, + predicted_tokens_per_seq=gen_p, + ) + lat_out_t = torch.transpose(torch.reshape(lat_out, [gen, heads, kv_lora]), 0, 1) + attn_gen_t = torch.transpose(torch.reshape(attn_gen, [gen, heads, v_dim]), 0, 1) + attn_gen_t = self._mtp_expand(a=lat_out_t, out=attn_gen_t) + + o = self._mtp_o(attn_out) + o, residual = self._mtp_norm2(o, residual) + shared = self._mtp_shared_mlp(o) + routed = self._mtp_routed_experts(o, all_rank_num_tokens, rows, dt) + return torch.add(residual, torch.add(routed, shared)) + + +class MTPLayer: + """The checkpoint's multi-token-prediction module, at layer index `num_hidden_layers`, + replayed once per draft step. + + It differs from a trunk layer in three ways the code alone does not show: an `eh_proj` front + end that mixes the next token's embedding into `h`, bf16 experts (the checkpoint excludes + `model.layers.61*` from NVFP4, so they go to `fused_moe`, not the trtllm-gen runner), and an + output returned **un-normalized** — `shared_head` applies the module's own norm, and the + runtime feeds the un-normalized tensor back in as the next step's `h`.""" + + def __init__(self, core: DeepseekV3Dep4ModelingV2, logits_processor) -> None: + self.core = core + self.logits_processor = logits_processor + + def __call__(self, *args, **kwargs) -> torch.Tensor: + return self.forward(*args, **kwargs) + + def post_load_weights(self) -> None: + """Bind the output head's own norm; everything else the draft step calls now + lives on the targets, bound in their `__init__` from `core._mtp`.""" + core = self.core + self._head_norm = FlashinferRmsnorm() + self._head_norm.bind_const( + weight=core._mtp["head_norm"], + eps=core.model_config.pretrained_config.rms_norm_eps, + ) + + def shared_head( + self, hidden_states, lm_head, attn_metadata, return_context_logits + ) -> torch.Tensor: + """The module's own output head: its own RMS norm — a distinct parameter from the + trunk's `model.norm` — then the trunk's `lm_head`.""" + return self.logits_processor.forward( + self._head_norm(hidden_states), + lm_head, + attn_metadata, + return_context_logits, + ) + + def forward( + self, + embed_tokens: torch.Tensor, + all_rank_num_tokens, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> torch.Tensor: + """Route the draft step by the phase this call's metadata shows. + + Per call, not per engine forward: the worker rewrites `attn_metadata` in place + between draft step 0 and step 1+, so one engine step's N draft calls can cross + from prefill to decode. + """ + return self.core._select_target(attn_metadata).mtp_forward( + embed_tokens, all_rank_num_tokens, input_ids, hidden_states, attn_metadata, **kwargs + ) class DraftModel: """The container the spec worker reaches this target's drafter through. - The runtime never constructs it and inspects exactly four names on it: - `mtp_layers` (only `[0]` is ever indexed — MTP-Eagle replays one layer), - `embed_tokens` and `lm_head`, which it hands to the layer and to - `shared_head`, and `model.d2t`, read with a nested `getattr(..., None)` and - correctly absent here because draft and target share one vocabulary. - - `embed_tokens` is the **trunk's** embedding weight and `lm_head` the - trunk's head. The checkpoint's `model.layers.61.embed_tokens.weight` and - `.shared_head.head.weight` are `torch.equal` to those two, so pointing at - them is exact and saves 1.85 GB per rank; they stay a predicted non-load in - the weight manifest even with MTP on. - - `embed_tokens` resolves through the core on every read rather than being - snapshotted here: this container is built in the shell's `__init__`, where - every parameter is still a **meta** tensor, and the engine materializes the - registry by replacing those tensor objects. A reference captured at - construction stays on meta and fails at the first draft step.""" - - def __init__(self, core: ModelingV2Core, lm_head, logits_processor) -> None: + `embed_tokens` must resolve through the core on every read: this container is built while + every parameter is still a meta tensor, and the engine materializes the registry by + replacing the tensor objects, so a reference snapshotted here stays on meta.""" + + def __init__(self, core: DeepseekV3Dep4ModelingV2, lm_head, logits_processor) -> None: self.core = core self.mtp_layers = [MTPLayer(core, logits_processor)] self.lm_head = lm_head @@ -2024,43 +2103,20 @@ def embed_tokens(self) -> torch.Tensor: @register_auto_model("ModelingV2DeepseekR10528Nvfp4Sm103Dep4") class ModelingV2DeepseekR10528Nvfp4Sm103Dep4( - DecoderModelForCausalLM[ModelingV2Core, PretrainedConfig] + DecoderModelForCausalLM[DeepseekV3Dep4ModelingV2, PretrainedConfig] ): def __init__(self, model_config: ModelConfig): cfg = model_config.pretrained_config - assert cfg is not None super().__init__( - ModelingV2Core(model_config), + DeepseekV3Dep4ModelingV2(model_config), config=model_config, hidden_size=cfg.hidden_size, vocab_size=cfg.vocab_size, ) - # The speculative branch, built only when `speculative_config` asked for - # it. `spec_config` is already resolved on the model config by the time - # the model is built, and the checkpoint — not the target — picked the - # mode: `num_nextn_predict_layers: 1` gives MTP-Eagle one-model, where - # `max_draft_len` is a serving knob rather than a checkpoint property. - # - # The branch is written out rather than inherited from the in-tree - # one-engine shell on purpose: that shell builds its drafter through a - # module-level function with no override point, which dispatches on the - # config's `model_type` — still the upstream family name, since the - # stub config patches `architectures` only — and would construct - # trtllm's own MTP layer instead of this target's. self.spec_config = getattr(model_config, "spec_config", None) self.draft_model = None self.spec_worker = None if self.spec_config is not None: - mode = self.spec_config.spec_dec_mode - assert mode.is_mtp_eagle_one_model(), ( - f"this target implements the one-model MTP-Eagle draft loop; " - f"the engine resolved spec_dec_mode {mode!r}" - ) - assert self.model.mtp_enabled, "core built without the MTP module" - assert hasattr(self, "logits_processor"), ( - "the inherited shell no longer exposes logits_processor; the " - "draft head and the shell's own gather both project through it" - ) self.draft_model = DraftModel(self.model, self.lm_head, self.logits_processor) self.spec_worker = get_spec_worker(self.spec_config, model_config, model_config.mapping) @@ -2077,30 +2133,12 @@ def forward( ): """One engine step. - Without a spec worker this **delegates to the inherited base forward** - with every argument it was handed, rather than reimplementing it: the - release criterion was measured on that base forward, and calling it is - the only way to keep this path bit-identical. The parameter list - mirrors the base's exactly for that reason, and `resource_manager` - deliberately stays inside `**kwargs` — naming it here would drop it - from what the base receives. - - With one, the shell owns the logits gather: every one-model mode sets - `without_logits`, so the engine applies no second gather and the worker - needs the trunk's hidden states **ungathered** beside the gathered - logits. `position_ids` is handed on in the engine's `[1, T]` shape — - the worker squeezes it itself, and flattening here would produce a - silently wrong draft position sequence.""" + `resource_manager` must stay inside `**kwargs` — naming it here drops it from what the + base forward receives — and `position_ids` must be handed on in the engine's `[1, T]` + shape, since the worker squeezes it itself and flattening it here is silently wrong.""" if self.spec_worker is None: - assert spec_metadata is None, ( - "spec_metadata arrived without a spec worker; this engine was " - "built without a speculative_config" - ) return super().forward( attn_metadata, - # A typing no-op: the base declares `input_ids: torch.IntTensor - # = None`, a non-Optional annotation with a None default, so - # the value is forwarded exactly as received — None included. cast(torch.IntTensor, input_ids), position_ids, inputs_embeds, @@ -2109,15 +2147,6 @@ def forward( lora_params, **kwargs, ) - assert spec_metadata is not None, ( - "the spec worker is built but the engine passed no spec_metadata" - ) - # `spec_metadata` is not forwarded into the core: the trunk genuinely - # does not read it — on MTP_EAGLE_ONE_MODEL the runtime sets - # `layers_to_capture = ()`, so `is_layer_capture()` is False everywhere - # and no hidden-state capture hook is owed (that is Eagle3's - # requirement, not this mode's) — and the core's assert that it is - # absent is a stronger guarantee than deleting the assert would be. hidden = self.model( attn_metadata=attn_metadata, input_ids=input_ids, @@ -2126,12 +2155,8 @@ def forward( lora_params=lora_params, **kwargs, ) - # `gather_ids` holds one row per context request (its last token) and - # `runtime_draft_len + 1` per generation request. `embedding` is the - # catalog's row-lookup entry (torch.nn.functional.embedding), used here - # for what it is — `hidden[gather_ids]`. logits = self.logits_processor.forward( - embedding(spec_metadata.gather_ids, hidden), + nn.functional.embedding(spec_metadata.gather_ids, hidden), self.lm_head, attn_metadata, True, @@ -2148,8 +2173,11 @@ def forward( ) def load_weights(self, weights, *args, **kwargs): - _weights.load(self, weights) + _weights.MODEL_WEIGHTS.load(self, weights) def post_load_weights(self): super().post_load_weights() - self.model.derive_after_load() + self.model.post_load_weights() + if self.draft_model is not None: + for layer in self.draft_model.mtp_layers: + layer.post_load_weights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/__init__.py new file mode 100644 index 000000000000..1059445198db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/attention/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/attention/load_paged_kv_cache_for_mla.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/attention/load_paged_kv_cache_for_mla.py new file mode 100644 index 000000000000..419cf43e6de4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/attention/load_paged_kv_cache_for_mla.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gather context rows out of the paged MLA latent cache, via the trtllm op.""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: quant_mode bit for an fp8 KV cache (tensorrt_llm.quantization.mode.QuantMode). +QM_FP8_KV_CACHE = 128 + + +class LoadPagedKvCacheForMla(OpWrapper): + """Copy positions `[0, L_s)` of every context sequence out of the paged MLA + latent cache into two new contiguous tensors, `(compressed_kv, k_pe)`. + + A pure gather: the pool is read and never written, and the slot one past + each sequence's range is not touched. Nothing is recomputed -- when the pool + holds e4m3 the op dequantizes, and that is the whole of its arithmetic. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # deepseek-r1's MLA latent geometry, which is the only caller: a 512-wide + # compressed_kv and a 64-wide k_pe sharing one 576-wide pool row. The page + # sizes are the engine default and the value a tuned MLA target opts into; + # both are covered because the gather addresses pages, not tokens. + _R1 = dict(kv_lora_rank=512, qk_rope_head_dim=64, quant_mode=QM_FP8_KV_CACHE) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="production context call: fp8 pool, scale omitted, bf16 out, page 32", + spec=dict( + cached=[0, 32, 100, 511], + new=[37, 64, 300, 1], + tokens_per_block=32, + out_dtype=torch.bfloat16, + read_scale=None, + **_R1, + ), + ), + Cell( + why="page 64, the size a tuned MLA target opts into", + spec=dict( + cached=[0, 32, 100, 511], + new=[37, 64, 300, 1], + tokens_per_block=64, + out_dtype=torch.bfloat16, + read_scale=None, + **_R1, + ), + ), + Cell( + why="an explicit read scale, which the omitted form has to equal at 1.0", + spec=dict( + cached=[0, 32], + new=[37, 64], + tokens_per_block=64, + out_dtype=torch.bfloat16, + read_scale=1.0, + **_R1, + ), + ), + Cell( + why="a non-unit read scale: the dequantization is the op's only arithmetic", + spec=dict( + cached=[0, 32], + new=[37, 64], + tokens_per_block=64, + out_dtype=torch.bfloat16, + read_scale=0.375, + **_R1, + ), + ), + ) + + note = """ + No `is_valid`. Everything this op could get wrong is either rejected by the + op itself or is a property of what some earlier call *wrote* -- a read scale + that does not match the write scale returns plausible numbers, but nothing + visible at this call site can tell. The test drives that case directly + instead, which is the only place it can be seen. + + A bf16 pool (`quant_mode=0`) works and is a straight copy, but every shipped + target runs the fp8 pool, so it is not a cell. + """ + + def raw_call( + self, + out_dtype: torch.dtype, + num_contexts: int, + num_ctx_kv_tokens: int, + max_ctx_kv_len: int, + cu_ctx_kv_lens: torch.Tensor, + kv_cache_block_offsets: torch.Tensor, + host_kv_cache_pool_pointers: torch.Tensor, + host_kv_cache_pool_mapping: torch.Tensor, + kv_scale_quant_orig: Optional[torch.Tensor], + layer_idx: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + tokens_per_block: int, + attention_window_size: int, + beam_width: int, + quant_mode: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.load_paged_kv_cache_for_mla( + out_dtype, + num_contexts, + num_ctx_kv_tokens, + max_ctx_kv_len, + cu_ctx_kv_lens, + kv_cache_block_offsets, + host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping, + kv_scale_quant_orig, + layer_idx, + kv_lora_rank, + qk_rope_head_dim, + tokens_per_block, + attention_window_size, + beam_width, + quant_mode, + ) + + def reference( + self, + stored: torch.Tensor, + kv_scale_quant_orig: Optional[torch.Tensor], + out_dtype: torch.dtype, + ) -> torch.Tensor: + """`float(e4m3 byte) * scale`, in fp32, rounded once to `out_dtype`. + + The second documented exception to mirroring `raw_call`, after the + collectives: this op's output is a function of what the paged cache + holds, and its arguments only say where to look. `stored` is the rows + the gather was pointed at; the test reads them out of the pool, which is + the only place they exist. + + An omitted scale behaves exactly as 1.0 -- stated here rather than in + prose because a cell drives both forms against this one line. + """ + wide = stored.float() + if kv_scale_quant_orig is not None: + wide = wide * kv_scale_quant_orig[0] + return wide.to(out_dtype) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """A gate of exactly zero: one multiply and one rounding, on both sides. + + The op dequantizes in fp32 and rounds once; so does `reference`. There + is no accumulation to disagree about, so any difference is a wrong + gather, a wrong scale, or a wrong row -- never drift. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0) + + +load_paged_kv_cache_for_mla = LoadPagedKvCacheForMla() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/attention/mla_rope_append_paged_kv_assign_q.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/attention/mla_rope_append_paged_kv_assign_q.py new file mode 100644 index 000000000000..ddb1546d9bb2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/attention/mla_rope_append_paged_kv_assign_q.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""RoPE the new context tokens and append their latent rows to the paged MLA cache.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: quant_mode bit for an fp8 KV cache. +QM_FP8_KV_CACHE = 128 + +#: Half an e4m3 ulp. e4m3 carries 3 explicit mantissa bits, so a +#: round-to-nearest lands within 2**-3 relative. +E4M3_ULP_RTOL = 2.0**-3 + +#: How much of the appended output may miss bit-exactness before the gate calls +#: it drift rather than rounding. See `compare` for why this is not zero and why +#: it has never been approached. +E4M3_MAX_INEXACT_FRACTION = 1e-3 + + +class MlaRopeAppendPagedKvAssignQ(OpWrapper): + """RoPE each new context token's `q_pe` (in `q`) and `k_pe` (in + `latent_cache`) in place at its absolute position, then append the token's + latent row `[compressed_kv | rope(k_pe)]` to the paged MLA KV cache. + + Two buffers are written and one is the paged pool itself, which is why the + cells below check what landed in the cache rather than a return value: the + op returns None. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # deepseek-r1's context call is the only caller. Its latent geometry is + # fixed by the checkpoint -- 512-wide compressed_kv, 64-wide k_pe, 128 + # heads -- so the cells vary what a step actually varies: how much of each + # sequence was already cached, how many tokens are new, and the page size. + _R1 = dict( + head_num=128, + nope_size=128, + rope_size=64, + lora_size=512, + residual_dim=0, + quant_mode=QM_FP8_KV_CACHE, + ) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="production context call: fp8 pool, page 32, mixed cached lengths", + spec=dict(cached=[0, 32, 100], new=[37, 64, 300], tokens_per_block=32, **_R1), + ), + Cell( + why="page 64, the size a tuned MLA target opts into", + spec=dict(cached=[0, 32, 100], new=[37, 64, 300], tokens_per_block=64, **_R1), + ), + Cell( + why="a single new token on top of a long cached prefix -- the decode-adjacent shape", + spec=dict(cached=[511], new=[1], tokens_per_block=64, **_R1), + ), + ) + + note = """ + No `is_valid` beyond `residual_dim`. What else could go wrong here is either + rejected by the op or is a property of what an earlier call wrote -- a write + scale that disagrees with the read scale appends plausible bytes, and + nothing visible at this call site can tell. The test drives that directly. + """ + + def raw_call( + self, + q: torch.Tensor, + latent_cache: torch.Tensor, + num_contexts: int, + cu_ctx_cached_kv_lens: torch.Tensor, + cu_seq_lens: torch.Tensor, + max_input_uncached_seq_len: int, + cos_sin_cache: torch.Tensor, + head_num: int, + nope_size: int, + rope_size: int, + lora_size: int, + kv_cache_block_offsets: torch.Tensor, + host_kv_cache_pool_pointers: torch.Tensor, + host_kv_cache_pool_mapping: torch.Tensor, + kv_scale_orig_quant: Optional[torch.Tensor], + residual_dim: int, + layer_idx: int, + tokens_per_block: int, + attention_window_size: int, + beam_width: int, + quant_mode: int, + ) -> torch.Tensor: + torch.ops.trtllm.mla_rope_append_paged_kv_assign_q( + q, + latent_cache, + num_contexts, + cu_ctx_cached_kv_lens, + cu_seq_lens, + max_input_uncached_seq_len, + cos_sin_cache, + head_num, + nope_size, + rope_size, + lora_size, + kv_cache_block_offsets, + host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping, + kv_scale_orig_quant, + residual_dim, + layer_idx, + tokens_per_block, + attention_window_size, + beam_width, + quant_mode, + ) + return q + + def reference( + self, + latent_row: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + rope_size: int, + kv_scale_orig_quant: Optional[torch.Tensor] = None, + store_dtype: torch.dtype = torch.float8_e4m3fn, + ) -> torch.Tensor: + """The latent row this op is supposed to append, fp32 then stored once. + + The same documented exception to mirroring as the other paged-cache + entries: this op's effect is on the pool, and its arguments only say + where. `latent_row` is `[compressed_kv | k_pe]` before the rotation, and + `cos` / `sin` are the halves of the rope table already selected for each + row's absolute position; the test reads both out of the environment, + which is the only place they exist. + + The pairing is **GPT-J interleaved** -- consecutive elements form a pair + -- not the split-half neox form. The two are indistinguishable by shape + and give different answers, which is what makes stating it here worth + more than stating it in prose. + + `compressed_kv` passes through untouched. The scale is applied once, + before the store, so the stored bytes are a function of the row rather + than of the order the kernel visited it in. + """ + row = latent_row.float() + lora = row[..., : row.shape[-1] - rope_size] + pe = row[..., row.shape[-1] - rope_size :] + pairs = pe.reshape(*pe.shape[:-1], rope_size // 2, 2) + roped = torch.empty_like(pairs) + roped[..., 0] = pairs[..., 0] * cos - pairs[..., 1] * sin + roped[..., 1] = pairs[..., 0] * sin + pairs[..., 1] * cos + out = torch.cat([lora, roped.reshape(pe.shape)], dim=-1) + if kv_scale_orig_quant is not None: + out = out * kv_scale_orig_quant[0] + return out.to(store_dtype) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Half an e4m3 ulp, *and* a floor on how much may miss bit-exactness. + + The rotation is where kernel and reference can disagree: both evaluate + `x*cos -+ y*sin`, and they may do it in different orders. e4m3's coarse + rounding absorbs that completely in practice -- every appended byte + measured on this architecture, across every scale and every cell, came + back bit-exact, so neither half of this gate has ever been approached. + + Both halves are kept anyway, and both are measured rather than assumed: + the test runs four deliberately wrong mirrors of the same bytes and + asserts each blows the half it is able to blow. A tolerance alone would + pass a mirror that is uniformly a little wrong; a bit-exact demand alone + would fail the first time a reordering did survive the rounding. + """ + assert out.shape == ref.shape, (out.shape, ref.shape) + torch.testing.assert_close(out.float(), ref.float(), rtol=E4M3_ULP_RTOL, atol=0.0) + inexact = int( + (out.reshape(-1).view(torch.uint8) != ref.reshape(-1).view(torch.uint8)).sum() + ) + allowed = max(1, int(E4M3_MAX_INEXACT_FRACTION * out.numel())) + assert inexact <= allowed, ( + f"not bit-exact enough: {inexact}/{out.numel()} bytes differ (allowed {allowed})" + ) + + def is_valid( + self, + q: torch.Tensor, + latent_cache: torch.Tensor, + num_contexts: int, + cu_ctx_cached_kv_lens: torch.Tensor, + cu_seq_lens: torch.Tensor, + max_input_uncached_seq_len: int, + cos_sin_cache: torch.Tensor, + head_num: int, + nope_size: int, + rope_size: int, + lora_size: int, + kv_cache_block_offsets: torch.Tensor, + host_kv_cache_pool_pointers: torch.Tensor, + host_kv_cache_pool_mapping: torch.Tensor, + kv_scale_orig_quant: Optional[torch.Tensor], + residual_dim: int, + *unused_args: object, + **unused_kwargs: object, + ) -> None: + # `residual_dim` must be 0 or `rope_size`, and the op rejects non-zero + # unless the KV pool is FP4. Every caller here runs a bf16 or fp8-e4m3 + # pool, so 0 is the only legal value and a non-zero one would be read + # against a pool layout that does not exist. + assert residual_dim == 0, ( + f"residual_dim must be 0 for a bf16 or fp8 KV pool; got {residual_dim}. " + "A non-zero value is only legal against an FP4 pool" + ) + + +mla_rope_append_paged_kv_assign_q = MlaRopeAppendPagedKvAssignQ() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/attention/mla_rope_generation.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/attention/mla_rope_generation.py new file mode 100644 index 000000000000..66ec226d28e8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/attention/mla_rope_generation.py @@ -0,0 +1,446 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MLA generation-phase RoPE + latent KV-cache append + FMHA scheduler-buffer fill.""" + +from __future__ import annotations + +import math +from typing import List, Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: quant_mode bit for an fp8 KV cache. +QM_FP8_KV_CACHE = 128 + +#: Half an e4m3 ulp. e4m3 carries 3 explicit mantissa bits, so a +#: round-to-nearest lands within 2**-3 relative. +E4M3_ULP_RTOL = 2.0**-3 + +#: How much of a quantized output may miss bit-exactness before the gate calls +#: it drift rather than rounding. See `compare`. +E4M3_MAX_INEXACT_FRACTION = 1e-3 + +#: DeepSeek-R1's YaRN attention temperature, carried as `q_scaling = 1/mscale^2` +#: with `mscale = 0.1 * ln(40) + 1`. The target passes this, which is why a +#: `q_scaling != 1` cell exists: it lands in `mla_bmm1_scale` and nowhere else. +_R1_MSCALE = 0.1 * math.log(40.0) + 1.0 +Q_SCALING_R1 = 1.0 / (_R1_MSCALE * _R1_MSCALE) + + +class MlaRopeGeneration(OpWrapper): + """One decode step's MLA preprocessing, over the batch's generation + sequences only. + + Each generation sequence carries `P = predicted_tokens_per_seq` new query + tokens -- 1 is an ordinary decode, `P > 1` the MTP shape. The per-token + tensors hold `G * P` rows token-major within a sequence, and row `(g, t)` + sits at absolute position `sequence_length[num_contexts + g] - P + t`, so + each row is rotated at its own angle. + + Over an fp8 latent pool the call reads `fused_q`, writes the quantized + query into `quant_q_buffer`, appends the quantized latent row to the pool, + and fills `mla_bmm1_scale` / `mla_bmm2_scale` plus the two scheduler + buffers. It returns None, so the cells below read the buffers and the pool + rather than a return value. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The DeepSeek-R1-0528 cell: 128 heads, 512/64 latent geometry, fp8 pool, + # single-layer. Fixed by the checkpoint, so the cells vary what a step + # varies -- how many new tokens, the page size, and whether the scales are + # passed explicitly. + _R1 = dict( + num_heads=128, + kv_lora_rank=512, + qk_rope_head_dim=64, + qk_nope_head_dim=128, + quant_mode=QM_FP8_KV_CACHE, + ) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="production decode: P=1, page 32, scales omitted, R1's YaRN q_scaling", + spec=dict( + predicted_tokens_per_seq=1, + tokens_per_block=32, + q_scaling=Q_SCALING_R1, + write_scale=None, + read_scale=None, + **_R1, + ), + ), + Cell( + why="the MTP shape: P=3 consecutive positions per sequence, page 32", + spec=dict( + predicted_tokens_per_seq=3, + tokens_per_block=32, + q_scaling=Q_SCALING_R1, + write_scale=None, + read_scale=None, + **_R1, + ), + ), + Cell( + why="explicit reciprocal scales: w and r reach different outputs, so both must be driven", + spec=dict( + predicted_tokens_per_seq=1, + tokens_per_block=32, + q_scaling=1.0, + write_scale=0.5, + read_scale=2.0, + **_R1, + ), + ), + ) + + note = """ + Only the fp8-e4m3 latent pool (`quant_mode=128`) is certified. The op also + has a bf16-pool path, where `quant_q_buffer` is unused and the roped q lands + in `fused_q`'s tail instead; nothing here drives it, so it is not certified + -- the prose contract this class replaced claimed it was, with fifteen tests + behind the claim and all fifteen on the fp8 path. + + The schema annotation is wrong in both directions: it marks `fused_q` and + `q_pe` mutable, but `q_pe` is never modified and on this path neither is + `fused_q`, while `cu_q_seqlens`, `cu_kv_seqlens`, `fmha_scheduler_counter`, + `quant_q_buffer` and both scale buffers -- all filled by the call -- carry no + annotation at all. Do not rely on its alias info under functionalization. + + Re-running a prepared step is idempotent rather than double-appending: the + positions come from `sequence_length`, so the same slots are rewritten. + + Not certified: `P` above 4, `block_ids_per_seq` (the flash-MLA layout), + helix, `out_scale`, an `attention_window_size` below the KV length, and + multi-layer pools. + """ + + def raw_call( + self, + fused_q: torch.Tensor, + q_pe: torch.Tensor, + latent_cache: torch.Tensor, + rotary_cos_sin: Optional[torch.Tensor], + cu_q_seqlens: torch.Tensor, + cu_kv_seqlens: torch.Tensor, + fmha_scheduler_counter: torch.Tensor, + mla_bmm1_scale: Optional[torch.Tensor], + mla_bmm2_scale: Optional[torch.Tensor], + quant_q_buffer: Optional[torch.Tensor], + sequence_length: torch.Tensor, + host_past_key_value_lengths: torch.Tensor, + host_context_lengths: torch.Tensor, + num_contexts: int, + kv_cache_block_offsets: Optional[torch.Tensor], + host_kv_cache_pool_pointers: Optional[torch.Tensor], + host_kv_cache_pool_mapping: Optional[torch.Tensor], + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + kv_cache_scale_orig_quant: Optional[torch.Tensor], + out_scale: Optional[torch.Tensor], + block_ids_per_seq: Optional[torch.Tensor], + helix_tensor_params: List[Optional[torch.Tensor]], + predicted_tokens_per_seq: int, + layer_idx: int, + num_heads: int, + num_kv_heads: int, + head_size: int, + residual_dim: int, + tokens_per_block: int, + attention_window_size: int, + beam_width: int, + quant_mode: int, + q_scaling: float, + q_lora_rank: int, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + rope_append: bool, + kv_norm_weight: Optional[torch.Tensor] = None, + kv_norm_eps: float = 1e-6, + precomputed_cu_seqlens: bool = False, + precomputed_fmha_scheduler: bool = False, + kv_only: bool = False, + kv_done_elsewhere: bool = False, + quant_scale_qkv: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + torch.ops.trtllm.mla_rope_generation( + fused_q, + q_pe, + latent_cache, + rotary_cos_sin, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + sequence_length, + host_past_key_value_lengths, + host_context_lengths, + num_contexts, + kv_cache_block_offsets, + host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping, + kv_scale_orig_quant, + kv_scale_quant_orig, + kv_cache_scale_orig_quant, + out_scale, + block_ids_per_seq, + helix_tensor_params, + predicted_tokens_per_seq, + layer_idx, + num_heads, + num_kv_heads, + head_size, + residual_dim, + tokens_per_block, + attention_window_size, + beam_width, + quant_mode, + q_scaling, + q_lora_rank, + kv_lora_rank, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, + rope_append, + kv_norm_weight, + kv_norm_eps, + precomputed_cu_seqlens, + precomputed_fmha_scheduler, + kv_only, + kv_done_elsewhere, + quant_scale_qkv, + ) + return fused_q + + def reference( + self, + row: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + rope_size: int, + kv_scale_orig_quant: Optional[torch.Tensor] = None, + store_dtype: torch.dtype = torch.float8_e4m3fn, + ) -> torch.Tensor: + """One quantized `[.. , C+R]` row this op is supposed to write. + + The same documented exception to mirroring as the other paged-cache + entries: the op's effect is on buffers and its arguments only say where. + One function covers both quantized outputs, because they are the same + computation on different inputs -- `[fused_q[..., :C] | q_pe]` for a + `quant_q_buffer` row, `[compressed_kv | k_pe]` for an appended cache + row. `cos` / `sin` are the rope table's halves already selected for the + row's absolute position; the test reads them out of the environment, + which is the only place they exist. + + Two details decide whether a mirror agrees with the kernel: + + The pairing is **GPT-J interleaved** -- consecutive elements form a pair + -- not the split-half neox form. The two are indistinguishable by shape + and give different answers. + + The rotation is rounded to the activation dtype **before** quantizing, + which is what the kernel does. Quantizing the fp32 rotation directly + disagrees on ~3% of the roped elements: two roundings, not one. + + The passthrough half is not re-rounded -- it arrives in `row`'s dtype + and is only scaled. The scale is applied once, so the stored bytes are a + function of the row rather than of the order the kernel visited it in. + """ + pe = row[..., row.shape[-1] - rope_size :] + lora = row[..., : row.shape[-1] - rope_size] + pairs = pe.float().reshape(*pe.shape[:-1], rope_size // 2, 2) + roped = torch.empty_like(pairs) + roped[..., 0] = pairs[..., 0] * cos - pairs[..., 1] * sin + roped[..., 1] = pairs[..., 0] * sin + pairs[..., 1] * cos + # Round to the activation dtype first: the kernel's own intermediate. + out = torch.cat([lora, roped.reshape(pe.shape).to(row.dtype)], dim=-1).float() + if kv_scale_orig_quant is not None: + out = out * kv_scale_orig_quant[0] + return out.to(store_dtype) + + def reference_bmm_scales( + self, + kv_scale_quant_orig: Optional[torch.Tensor], + q_scaling: float, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """The `(mla_bmm1_scale, mla_bmm2_scale)` this call is supposed to fill. + + A per-batch scalar pair, not per sequence, and `P` does not enter it. + The read-side factor `r` appears squared in one and plain in the other + because the consuming fp8 FMHA dequantizes both operands of BMM1 and one + of BMM2 -- see `is_valid` for what that costs when `r != 1/w`. + + Computed in double and rounded once, so the comparison measures the + kernel: every way of getting this scale wrong misses by a factor rather + than by ulps. + """ + r = 1.0 if kv_scale_quant_orig is None else float(kv_scale_quant_orig[0]) + x = r * r / (q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim)) + return ( + torch.tensor([x, x * math.log2(math.e)], dtype=torch.float32), + torch.tensor([r], dtype=torch.float32), + ) + + def reference_scheduler_buffers( + self, + kv_lens: List[int], + num_heads: int, + predicted_tokens_per_seq: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """The `(cu_q_seqlens, cu_kv_seqlens)` this call is supposed to fill. + + `kv_lens` is the generation sequences' total KV lengths, in batch order, + already including all `P` of this step's tokens. The q buffer is in + units of q rows -- `num_heads` rows per generation *token*, so + `num_heads * P` per sequence -- and the kv buffer in tokens. Exactly + `len(kv_lens) + 1` entries are written whatever `P` is; anything past + that is left alone. + """ + p = predicted_tokens_per_seq + cu_q = [i * num_heads * p for i in range(len(kv_lens) + 1)] + cu_kv = [0] + for length in kv_lens: + cu_kv.append(cu_kv[-1] + length) + return ( + torch.tensor(cu_q, dtype=torch.int32), + torch.tensor(cu_kv, dtype=torch.int32), + ) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Half an e4m3 ulp, *and* a floor on how much may miss bit-exactness. + + The rotation is where kernel and reference can disagree: both evaluate + `x*cos -+ y*sin` and may do it in different orders. e4m3's coarse + rounding absorbs that in practice -- every quantized element measured on + this architecture, at every scale and every cell, came back bit-exact, + so neither half of this gate has been approached. + + Both halves are kept because they fail differently: a tolerance alone + passes a mirror that is uniformly a little wrong, and a bit-exact demand + alone fails the first time a reordering does survive the rounding. + """ + assert out.shape == ref.shape, (out.shape, ref.shape) + torch.testing.assert_close(out.float(), ref.float(), rtol=E4M3_ULP_RTOL, atol=0.0) + inexact = int( + (out.reshape(-1).view(torch.uint8) != ref.reshape(-1).view(torch.uint8)).sum() + ) + allowed = max(1, int(E4M3_MAX_INEXACT_FRACTION * out.numel())) + assert inexact <= allowed, ( + f"not bit-exact enough: {inexact}/{out.numel()} bytes differ (allowed {allowed})" + ) + + def is_valid( + self, + fused_q: torch.Tensor, + q_pe: torch.Tensor, + latent_cache: torch.Tensor, + rotary_cos_sin: Optional[torch.Tensor], + cu_q_seqlens: torch.Tensor, + cu_kv_seqlens: torch.Tensor, + fmha_scheduler_counter: torch.Tensor, + mla_bmm1_scale: Optional[torch.Tensor], + mla_bmm2_scale: Optional[torch.Tensor], + quant_q_buffer: Optional[torch.Tensor], + sequence_length: torch.Tensor, + host_past_key_value_lengths: torch.Tensor, + host_context_lengths: torch.Tensor, + num_contexts: int, + kv_cache_block_offsets: Optional[torch.Tensor], + host_kv_cache_pool_pointers: Optional[torch.Tensor], + host_kv_cache_pool_mapping: Optional[torch.Tensor], + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + kv_cache_scale_orig_quant: Optional[torch.Tensor], + out_scale: Optional[torch.Tensor], + block_ids_per_seq: Optional[torch.Tensor], + helix_tensor_params: List[Optional[torch.Tensor]], + predicted_tokens_per_seq: int, + layer_idx: int, + num_heads: int, + num_kv_heads: int, + head_size: int, + residual_dim: int, + tokens_per_block: int, + attention_window_size: int, + beam_width: int, + quant_mode: int, + q_scaling: float, + q_lora_rank: int, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + rope_append: bool, + kv_norm_weight: Optional[torch.Tensor] = None, + kv_norm_eps: float = 1e-6, + precomputed_cu_seqlens: bool = False, + precomputed_fmha_scheduler: bool = False, + kv_only: bool = False, + kv_done_elsewhere: bool = False, + quant_scale_qkv: Optional[torch.Tensor] = None, + ) -> None: + """The four inputs this op takes and answers wrongly. + + The signature mirrors `raw_call` position for position because + `validating` forwards a call's arguments verbatim and every caller here + passes them positionally: a narrower signature would collect + `residual_dim` and `quant_mode` into a catch-all and check nothing, + silently. Only `quant_q_buffer`, `kv_scale_orig_quant`, + `kv_scale_quant_orig`, `residual_dim`, `quant_mode` and + `kv_norm_weight` are read. + """ + # `residual_dim` must be 0 or `rope_size`, and non-zero requires an FP4 + # KV pool. This entry certifies an fp8-e4m3 pool, against which a + # non-zero value reads a layout that does not exist. + assert residual_dim == 0, ( + f"residual_dim must be 0 for an fp8-e4m3 KV pool; got {residual_dim}. " + "A non-zero value is only legal against an FP4 pool" + ) + + # Over an fp8 pool the quantized query has nowhere to go without this + # buffer, and its absence is not checked: the kernel launches anyway, + # the run dies with an illegal memory access, the CUDA context is lost + # and the process aborts with nothing pointing here. + if quant_mode & QM_FP8_KV_CACHE: + assert quant_q_buffer is not None, ( + "an fp8 KV cache (quant_mode & 128) requires quant_q_buffer; the op does " + "not presence-check it and launches into a null pointer, killing the " + "CUDA context" + ) + + # A non-None weight folds the kv_a_layernorm into this kernel, which + # then reads latent_cache RAW. A caller that already normalized -- every + # caller here -- would be normalizing twice, and the result is plausible. + assert kv_norm_weight is None, ( + "kv_norm_weight folds kv_a_layernorm into this kernel, which then reads " + "latent_cache unnormalized; a caller that already normalized would " + "normalize twice" + ) + + # The two scales are used independently and are not derived from each + # other: `w` quantizes what this call writes, `r` is folded into the + # scales the consuming fp8 FMHA dequantizes with. That cancellation is + # only correct when `r = 1/w`, and nothing downstream can detect a + # non-reciprocal pair -- the decode is silently mis-scaled. Driving them + # apart on purpose is a characterization of the op, not a call a target + # makes, so such a case belongs outside `validating()`. + if kv_scale_orig_quant is not None or kv_scale_quant_orig is not None: + w = 1.0 if kv_scale_orig_quant is None else float(kv_scale_orig_quant[0]) + r = 1.0 if kv_scale_quant_orig is None else float(kv_scale_quant_orig[0]) + assert w != 0.0 and math.isclose(r, 1.0 / w, rel_tol=1e-6), ( + f"kv_scale_quant_orig ({r}) must be the reciprocal of kv_scale_orig_quant " + f"({w}): the consuming FMHA undoes the write scale with it, and a " + "non-reciprocal pair mis-scales the decode with nothing to show for it" + ) + + +mla_rope_generation = MlaRopeGeneration() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/moe/fused_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/moe/fused_moe.py new file mode 100644 index 000000000000..375b3a721de5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/moe/fused_moe.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused mixture-of-experts layer: expert permutation + grouped FC1/FC2 GEMMs +with a gated activation between them + routing-weighted combine, in one call. + +Everything a target author needs is below and all of it is checked: what the +op computes (`reference`), what it refuses (`is_valid`), how close the two +must be and why (`compare`), and the range CI drives (`CELLS`). +""" + +from __future__ import annotations + +from typing import List, Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The one geometry this entry is certified at: the R1-0528 MTP layer's MoE, +#: which is the only thing that calls this op. Repeated into every cell so a +#: cell prints as a complete configuration rather than a delta. +_R1 = dict(hidden=7168, inter=2048, experts_local=64, ep_size=4, topk=8) + + +class FusedMoe(OpWrapper): + """One MoE layer over pre-routed tokens; returns a fresh `[T, hidden]`. + + Fusion boundary: permutation, both GEMMs, the gated activation between + them, and the routing-weighted combine. Routing itself happens before the + call (`noaux_tc_op`), and the op\'s own `n_group`/`topk_group` stay inert. + + Of this op\'s ~50 parameters the shipped call sets twelve. The rest -- + biases, the quantized paths, LoRA, alltoall, min-latency, the swiglu + alpha/beta/limit triple -- are left at their defaults and are not + certified here. `activation_type=5` is Swiglu in *this* op\'s enum, which + is not the trtllm-gen runner\'s enum of the same name. + """ + + #: What CI drives. The geometry is fixed -- one op, one caller, one shape + #: family -- so the cells vary what the caller actually varies: how many + #: tokens arrive in a chunk, and which expert window this rank owns. + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode: a single token, the shape that dominates serving", + spec=dict(**_R1, num_tokens=1, ep_rank=0), + ), + Cell( + why="a mid-sized batch between the decode and chunk-bound extremes", + spec=dict(**_R1, num_tokens=256, ep_rank=0), + ), + Cell( + why="the chunk bound itself: the target splits the gathered token " + "set at 8192 because that is where this op stops being certified", + spec=dict(**_R1, num_tokens=8192, ep_rank=0), + ), + Cell( + why="a non-zero expert window -- rank 3 owns global ids [192, 256), " + "so an off-by-one in the window lands on the wrong experts", + spec=dict(**_R1, num_tokens=256, ep_rank=3), + ), + Cell( + why="token count not a multiple of anything: the tail path", + spec=dict(**_R1, num_tokens=37, ep_rank=1), + ), + ) + + note = """ + Tokens whose expert id falls outside this rank\'s window contribute + nothing; a token routed entirely outside it comes back as zeros. That is + the mechanism expert parallelism rests on, not an error, and the + reference reproduces it. + + Repeated ids in one token\'s top-k are summed, not deduplicated. + + The autotuner\'s hot path is exercised by test_r1_mtp_autotuned_tactics + rather than by a cell: what a tuning pass leaves in its cache is the + tuner\'s behaviour, and `misc/test_autotuner.py` owns it. + """ + + def raw_call( + self, + input: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: Optional[torch.Tensor], + fc1_expert_weights: torch.Tensor, + fc1_expert_biases: Optional[torch.Tensor], + fc2_expert_weights: torch.Tensor, + fc2_expert_biases: Optional[torch.Tensor], + output_dtype: torch.dtype, + quant_scales: List[torch.Tensor], + input_sf: Optional[torch.Tensor] = None, + swizzled_input_sf: bool = True, + swiglu_alpha: Optional[torch.Tensor] = None, + swiglu_beta: Optional[torch.Tensor] = None, + swiglu_limit: Optional[torch.Tensor] = None, + tp_size: int = 1, + tp_rank: int = 0, + ep_size: int = 1, + ep_rank: int = 0, + cluster_size: int = 1, + cluster_rank: int = 0, + enable_alltoall: bool = False, + use_deepseek_fp8_block_scale: bool = False, + use_w4_group_scaling: bool = False, + use_int8_woq_per_channel: bool = False, + use_mxfp8_act_scaling: bool = False, + min_latency_mode: bool = False, + use_fused_finalize: bool = True, + tune_max_num_tokens: int = 8192, + tuner_num_tokens: Optional[int] = None, + tuner_top_k: Optional[int] = None, + activation_type: int = 5, # ActivationType.Swiglu + unpadded_hidden_size: Optional[int] = None, + out_tensor: Optional[torch.Tensor] = None, + use_dynamic_fc2_scale: bool = False, + use_mxfp8_weight_scaling: bool = False, + fc1_lora_ranks: Optional[torch.Tensor] = None, + fc1_lora_weight_ptrs: Optional[torch.Tensor] = None, + fc2_lora_ranks: Optional[torch.Tensor] = None, + fc2_lora_weight_ptrs: Optional[torch.Tensor] = None, + gated_lora_ranks: Optional[torch.Tensor] = None, + gated_lora_weight_ptrs: Optional[torch.Tensor] = None, + host_request_types: Optional[torch.Tensor] = None, + host_context_lengths: Optional[torch.Tensor] = None, + lora_max_low_rank: int = 0, + fc1_slot_lora_ranks: Optional[torch.Tensor] = None, + fc1_slot_lora_weight_ptrs: Optional[torch.Tensor] = None, + fc2_slot_lora_ranks: Optional[torch.Tensor] = None, + fc2_slot_lora_weight_ptrs: Optional[torch.Tensor] = None, + gated_slot_lora_ranks: Optional[torch.Tensor] = None, + gated_slot_lora_weight_ptrs: Optional[torch.Tensor] = None, + token_to_slot: Optional[torch.Tensor] = None, + ) -> List[torch.Tensor]: + """Run one MoE layer over pre-routed tokens. + + Returns `[out]` with `out` a fresh `[num_tokens, hidden_size]` tensor in + `output_dtype`, or `[]` when `out_tensor` is given (written in place). + """ + return torch.ops.trtllm.fused_moe( + input, + token_selected_experts, + token_final_scales, + fc1_expert_weights, + fc1_expert_biases, + fc2_expert_weights, + fc2_expert_biases, + output_dtype, + quant_scales, + input_sf=input_sf, + swizzled_input_sf=swizzled_input_sf, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + swiglu_limit=swiglu_limit, + tp_size=tp_size, + tp_rank=tp_rank, + ep_size=ep_size, + ep_rank=ep_rank, + cluster_size=cluster_size, + cluster_rank=cluster_rank, + enable_alltoall=enable_alltoall, + use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale, + use_w4_group_scaling=use_w4_group_scaling, + use_int8_woq_per_channel=use_int8_woq_per_channel, + use_mxfp8_act_scaling=use_mxfp8_act_scaling, + min_latency_mode=min_latency_mode, + use_fused_finalize=use_fused_finalize, + tune_max_num_tokens=tune_max_num_tokens, + tuner_num_tokens=tuner_num_tokens, + tuner_top_k=tuner_top_k, + activation_type=activation_type, + unpadded_hidden_size=unpadded_hidden_size, + out_tensor=out_tensor, + use_dynamic_fc2_scale=use_dynamic_fc2_scale, + use_mxfp8_weight_scaling=use_mxfp8_weight_scaling, + fc1_lora_ranks=fc1_lora_ranks, + fc1_lora_weight_ptrs=fc1_lora_weight_ptrs, + fc2_lora_ranks=fc2_lora_ranks, + fc2_lora_weight_ptrs=fc2_lora_weight_ptrs, + gated_lora_ranks=gated_lora_ranks, + gated_lora_weight_ptrs=gated_lora_weight_ptrs, + host_request_types=host_request_types, + host_context_lengths=host_context_lengths, + lora_max_low_rank=lora_max_low_rank, + fc1_slot_lora_ranks=fc1_slot_lora_ranks, + fc1_slot_lora_weight_ptrs=fc1_slot_lora_weight_ptrs, + fc2_slot_lora_ranks=fc2_slot_lora_ranks, + fc2_slot_lora_weight_ptrs=fc2_slot_lora_weight_ptrs, + gated_slot_lora_ranks=gated_slot_lora_ranks, + gated_slot_lora_weight_ptrs=gated_slot_lora_weight_ptrs, + token_to_slot=token_to_slot, + ) + + def is_valid( + self, + input: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: Optional[torch.Tensor], + fc1_expert_weights: torch.Tensor, + fc1_expert_biases: Optional[torch.Tensor], + fc2_expert_weights: torch.Tensor, + fc2_expert_biases: Optional[torch.Tensor], + output_dtype: torch.dtype, + *unused_args: object, + **unused_kwargs: object, + ) -> None: + """The one input the op takes and answers wrongly. + + On the unquantized high-precision path the store is done in the + activation dtype while the output buffer is allocated with + output_dtype, so a mismatch reinterprets the written bits. Observed to + return plausible-looking wrong values, never to raise -- which is why + this is a guard and not a comment. + """ + unquantized = input.dtype in (torch.bfloat16, torch.float16) + if unquantized and fc1_expert_weights.dtype == input.dtype: + assert output_dtype == input.dtype, ( + f"output_dtype ({output_dtype}) must equal input.dtype ({input.dtype}) " + "on the unquantized path; a mismatch is written as raw activation-dtype " + "bits into an output_dtype buffer and silently gives wrong values" + ) + + def reference( + self, + input: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: Optional[torch.Tensor], + fc1_expert_weights: torch.Tensor, + fc1_expert_biases: Optional[torch.Tensor] = None, + fc2_expert_weights: torch.Tensor = None, + fc2_expert_biases: Optional[torch.Tensor] = None, + output_dtype: torch.dtype = torch.bfloat16, + quant_scales: Optional[List[torch.Tensor]] = None, + *unused_args: object, + ep_size: int = 1, + ep_rank: int = 0, + **unused_kwargs: object, + ) -> torch.Tensor: + """What the op computes, in plain torch, fp32 through both GEMMs. + + Takes the arguments that change the result rather than mirroring + `raw_call`. The other ~40 are inert on the certified path -- a + reference that accepted them would have to ignore them, which reads + as coverage it does not have. + + `fc1_expert_weights[e]` is `[up | gate]` stacked on dim 0. Ids are + global; this rank owns `[ep_rank * E_local, (ep_rank + 1) * E_local)` + and everything outside contributes nothing, which is how expert + parallelism tiles the routing space. + """ + num_tokens, hidden = input.shape + num_local = fc1_expert_weights.shape[0] + out = torch.zeros(num_tokens, hidden, dtype=torch.float32, device=input.device) + xf = input.float() + for local_e in range(num_local): + mask = token_selected_experts == ep_rank * num_local + local_e + tok, slot = mask.nonzero(as_tuple=True) + if tok.numel() == 0: + continue + w_up, w_gate = fc1_expert_weights[local_e].float().chunk(2, dim=0) + h_up = xf[tok] @ w_up.t() + h_gate = xf[tok] @ w_gate.t() + gated = h_gate * torch.sigmoid(h_gate) + # the kernel materializes the FC1 activation in the input dtype + inter = (gated * h_up).to(input.dtype).float() + y = inter @ fc2_expert_weights[local_e].float().t() + weight = ( + torch.ones(tok.numel(), device=input.device) + if token_final_scales is None + else token_final_scales[tok, slot].float() + ) + out.index_add_(0, tok, y * weight.unsqueeze(1)) + return out.to(input.dtype) + + #: Relative distance between neighbouring representable values + #: (1 + 2^-m for m stored mantissa bits): bf16 has 7, fp16 has 10. + _ULP = {torch.bfloat16: 2.0**-8, torch.float16: 2.0**-11} + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Wider than the base default, with the derivation rather than a number. + + Kernel and reference consume bit-identical operands and differ only in + GEMM accumulation order and in which side of a rounding boundary each + intermediate lands -- one flipped intermediate moves an output element + by about one ulp of that row\'s scale. torch\'s default cannot express + that: its bf16 `atol=1e-5` sits three orders of magnitude below one + output ulp of a two-GEMM chain, and per-element `rtol` is meaningless + where cancellation drives `|ref|` to zero. + + So: 8 ulp of the row\'s largest magnitude element-wise, 4 ulp of + relative RMS in aggregate. Measured worst case across everything this + entry drives is 3.98 / 1.43 ulp once the whole tactic space is walked + -- 2.0x and 2.8x margin. Both gates were sized against deliberately + wrong computations, which land two orders of magnitude out. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + ulp = self._ULP[ref.dtype] + row_scale = ref.float().abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + torch.testing.assert_close( + out.float() / row_scale, ref.float() / row_scale, rtol=0.0, atol=8 * ulp + ) + rel_rms = ( + (out.float() - ref.float()).pow(2).mean().sqrt() + / ref.float().pow(2).mean().sqrt().clamp_min(1e-9) + ).item() + assert rel_rms <= 4 * ulp, f"relative RMS {rel_rms:.3e} > {4 * ulp:.3e}" + + +fused_moe = FusedMoe() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/moe/noaux_tc_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/moe/noaux_tc_op.py new file mode 100644 index 000000000000..1d69933d0c54 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/moe/noaux_tc_op.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Grouped top-k expert routing by biased sigmoid score, via the trtllm op.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class NoauxTcOp(OpWrapper): + """Route each token to `topk` experts by biased sigmoid score. + + Selection uses `sigmoid(router_logits) + bias`; the combine weights are + gathered from the *unbiased* `sigmoid(router_logits)`, renormalized to sum + to 1 and scaled by `routed_scaling_factor`. The bias steers which experts + are chosen and never reaches the weights -- getting that backwards is the + single easiest way to misread this op. + + Returns `(topk_weights [T, topk] in router_logits.dtype, topk_ids [T, topk] + int32)`, both newly allocated and contiguous. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # deepseek-r1's router is the only caller: 256 experts in 8 groups, 4 groups + # kept, 8 experts per token, scaled by 2.5. The token counts are the ones + # the target produces -- one per decode step, and the chunk it splits a + # prefill into. + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode: one token through the full 256-expert router", + spec=dict(tokens=1, experts=256, n_group=8, topk_group=4, topk=8, scale=2.5), + ), + Cell( + why="prefill chunk: the batch the target routes at once", + spec=dict(tokens=2048, experts=256, n_group=8, topk_group=4, topk=8, scale=2.5), + ), + Cell( + why="a token count that is not a multiple of any tile", + spec=dict(tokens=37, experts=256, n_group=8, topk_group=4, topk=8, scale=2.5), + ), + ) + + note = """ + The kernel evaluates sigmoid as `0.5 * tanh(x/2) + 0.5`, which saturates to + exactly 1.0 at x >= ~17 and exactly 0.0 at x <= ~-18.5. `reference` uses the + same form; the 1/(1+exp(-x)) form disagrees in the tails. + + Ties in the selection score are broken toward the lower expert id, which + `reference` reproduces with a stable descending sort. + """ + + def raw_call( + self, + router_logits: torch.Tensor, + bias: torch.Tensor, + n_group: int, + topk_group: int, + topk: int, + routed_scaling_factor: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.noaux_tc_op( + router_logits, bias, n_group, topk_group, topk, routed_scaling_factor + ) + + def reference( + self, + router_logits: torch.Tensor, + bias: torch.Tensor, + n_group: int, + topk_group: int, + topk: int, + routed_scaling_factor: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Native torch, in the kernel's own arithmetic. + + Two details are the kernel's rather than the obvious choice, and both + are load-bearing: sigmoid is the tanh form, and the renormalization is + evaluated in fp64. An all-fp32 normalization disagrees on a third of the + elements at `routed_scaling_factor = 2.5`; the fp64 one is bit-exact. + """ + num_experts = router_logits.shape[-1] + scores = 0.5 * torch.tanh(0.5 * router_logits.float()) + 0.5 + choice = scores + bias.float() + if n_group > 1: + grouped = choice.view(-1, n_group, num_experts // n_group) + group_score = torch.topk(grouped, k=2, dim=-1).values.sum(-1) + keep = torch.sort(group_score, dim=-1, descending=True, stable=True).indices[ + :, :topk_group + ] + mask = torch.zeros_like(group_score).scatter_(-1, keep, 1.0) + mask = mask.unsqueeze(-1).expand_as(grouped).reshape(choice.shape) + choice = torch.where( + mask.bool(), choice, torch.tensor(float("-inf"), device=choice.device) + ) + ids = torch.sort(choice, dim=-1, descending=True, stable=True).indices[:, :topk] + weights = torch.gather(scores, 1, ids) + total = weights.sum(-1, keepdim=True).double() + weights = weights.double() / (total + 1e-20) * routed_scaling_factor + return weights.to(router_logits.dtype), ids.to(torch.int32) + + def is_valid( + self, + router_logits: torch.Tensor, + bias: torch.Tensor, + n_group: int, + topk_group: int, + topk: int, + routed_scaling_factor: float, + ) -> None: + # Three domains the op does not police. All three were observed to + # return a plausible-looking wrong answer instead of raising: the kernel + # addresses both tensors from data_ptr() as dense buffers and ignores + # strides, and topk past the end of a row emits expert ids outside the + # expert range at non-zero weight. + assert router_logits.is_contiguous(), ( + "router_logits must be contiguous; a strided view is read as a dense " + "[num_tokens, num_experts] buffer and silently gives wrong routing" + ) + assert bias.is_contiguous(), ( + "bias must be contiguous; a strided view is read as a dense " + "[num_experts] buffer and silently gives wrong routing" + ) + assert topk <= router_logits.shape[-1], ( + f"topk ({topk}) exceeds num_experts ({router_logits.shape[-1]}); the " + "kernel reads out of bounds and emits out-of-range expert ids" + ) + + +noaux_tc_op = NoauxTcOp() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/weights.py index ed36ef94bc0e..112e4fa1ee47 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/weights.py @@ -90,8 +90,12 @@ key either consumed or in the two predicted non-load sets above. """ +from types import SimpleNamespace + import torch +from tensorrt_llm._torch._experimental.modeling_v2._weights import ModelWeights, W + SF_BLOCK = 16 # NVFP4: one e4m3 scale per 16 elements along K @@ -131,7 +135,8 @@ def _swizzle(x: torch.Tensor) -> torch.Tensor: def _reorder_kv_b(t: torch.Tensor, core) -> torch.Tensor: """`[H*(nope+v), C]` per-head-interleaved -> `[H*nope ; H*v]` blocks.""" - h, n, v, c = core.heads, core.nope, core.v_dim, core.kv_lora + cfg = core.model_config.pretrained_config + h, n, v, c = cfg.num_attention_heads, cfg.qk_nope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank t3 = t.reshape(h, n + v, c) return torch.cat( [t3[:, :n, :].reshape(h * n, c), t3[:, n:, :].reshape(h * v, c)], dim=0 @@ -189,65 +194,6 @@ def _mtp_fc1(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: return torch.cat(t, dim=0).contiguous() -def _mtp_rows(rows: dict, core) -> None: - """The MTP module's manifest rows, at layer index `num_hidden_layers`. - Everything here is bf16 and stored HF `[out, in]`, so the only transforms - are the same `kv_b_proj` row regroup the trunk's attention needs and two - plain row concatenations.""" - p = f"model.layers.{core.num_layers}" - rows["mtp_enorm"] = [(f"{p}.enorm.weight", None, None)] - rows["mtp_hnorm"] = [(f"{p}.hnorm.weight", None, None)] - rows["mtp_eh"] = [(f"{p}.eh_proj.weight", None, None)] - rows["mtp_norm1"] = [(f"{p}.input_layernorm.weight", None, None)] - rows["mtp_qa"] = [(f"{p}.self_attn.q_a_proj.weight", None, None)] - rows["mtp_q_norm"] = [(f"{p}.self_attn.q_a_layernorm.weight", None, None)] - rows["mtp_qb"] = [(f"{p}.self_attn.q_b_proj.weight", None, None)] - rows["mtp_kva"] = [(f"{p}.self_attn.kv_a_proj_with_mqa.weight", None, None)] - rows["mtp_kv_norm"] = [(f"{p}.self_attn.kv_a_layernorm.weight", None, None)] - rows["mtp_kvb"] = [(f"{p}.self_attn.kv_b_proj.weight", None, _reorder_kv_b)] - rows["mtp_o"] = [(f"{p}.self_attn.o_proj.weight", None, None)] - rows["mtp_k_scale"] = [(f"{p}.self_attn.k_proj.k_scale", (0,), None)] - rows["mtp_v_scale"] = [(f"{p}.self_attn.v_proj.v_scale", (0,), None)] - rows["mtp_norm2"] = [(f"{p}.post_attention_layernorm.weight", None, None)] - rows["mtp_head_norm"] = [(f"{p}.shared_head.norm.weight", None, None)] - rows["mtp_router"] = [(f"{p}.mlp.gate.weight", None, None)] - rows["mtp_router_bias"] = [(f"{p}.mlp.gate.e_score_correction_bias", None, None)] - s = f"{p}.mlp.shared_experts" - rows["mtp_sh_gu"] = [((f"{s}.gate_proj.weight", f"{s}.up_proj.weight"), None, _cat_rows)] - rows["mtp_sh_dn"] = [(f"{s}.down_proj.weight", None, None)] - fc1, fc2 = [], [] - for local in range(core.local_experts): - q = f"{p}.mlp.experts.{core.expert_offset + local}" - fc1.append(((f"{q}.up_proj.weight", f"{q}.gate_proj.weight"), (local,), _mtp_fc1)) - fc2.append((f"{q}.down_proj.weight", (local,), None)) - rows["mtp_fc1"] = fc1 - rows["mtp_fc2"] = fc2 - - -def _dense_mlp_rows(rows: dict, key: str, prefix: str) -> None: - """Manifest rows shared by the dense MLPs of layers 0-2 and the shared - experts: one fused gate_up NVFP4 linear plus one down NVFP4 linear. Both - are replicated and run over this rank's own tokens.""" - rows[f"{key}_gu_w"] = [ - ((f"{prefix}.gate_proj.weight", f"{prefix}.up_proj.weight"), None, _cat_rows) - ] - rows[f"{key}_gu_s"] = [ - ( - (f"{prefix}.gate_proj.weight_scale", f"{prefix}.up_proj.weight_scale"), - None, - _cat_rows_swizzle, - ) - ] - rows[f"{key}_dn_w"] = [(f"{prefix}.down_proj.weight", None, None)] - rows[f"{key}_dn_s"] = [(f"{prefix}.down_proj.weight_scale", None, _swizzle_only)] - rows[f"{key}_isc1"] = [(f"{prefix}.gate_proj.input_scale", (0,), None)] - rows[f"{key}_isc1_up"] = [(f"{prefix}.up_proj.input_scale", (0,), None)] - rows[f"{key}_ws2_1"] = [(f"{prefix}.gate_proj.weight_scale_2", (0,), None)] - rows[f"{key}_ws2_1_up"] = [(f"{prefix}.up_proj.weight_scale_2", (0,), None)] - rows[f"{key}_isc2"] = [(f"{prefix}.down_proj.input_scale", (0,), None)] - rows[f"{key}_ws2_2"] = [(f"{prefix}.down_proj.weight_scale_2", (0,), None)] - - def _materialize(entry) -> torch.Tensor: """Realize one checkpoint entry: `[:]` on a lazy safetensors slice, but a 0-dim tensor (every NVFP4 per-tensor scalar and every KV scale) rejects @@ -255,77 +201,471 @@ def _materialize(entry) -> torch.Tensor: return entry if getattr(entry, "ndim", 1) == 0 else entry[:] -def _manifest(core) -> dict: - """target param key -> list of (ckpt key or key tuple, index into the - param | None, source transform | None).""" - rows: dict = {} - for i in range(core.num_layers): - p = f"model.layers.{i}" - rows[f"l{i}_norm1"] = [(f"{p}.input_layernorm.weight", None, None)] - rows[f"l{i}_qa"] = [(f"{p}.self_attn.q_a_proj.weight", None, None)] - rows[f"l{i}_q_norm"] = [(f"{p}.self_attn.q_a_layernorm.weight", None, None)] - rows[f"l{i}_qb"] = [(f"{p}.self_attn.q_b_proj.weight", None, None)] - rows[f"l{i}_kva"] = [(f"{p}.self_attn.kv_a_proj_with_mqa.weight", None, None)] - rows[f"l{i}_kv_norm"] = [(f"{p}.self_attn.kv_a_layernorm.weight", None, None)] - rows[f"l{i}_kvb"] = [(f"{p}.self_attn.kv_b_proj.weight", None, _reorder_kv_b)] - rows[f"l{i}_o"] = [(f"{p}.self_attn.o_proj.weight", None, None)] - # The fp8 KV-cache scales sit under the k/v projections the MLA - # checkpoint does not otherwise have. - rows[f"l{i}_k_scale"] = [(f"{p}.self_attn.k_proj.k_scale", (0,), None)] - rows[f"l{i}_v_scale"] = [(f"{p}.self_attn.v_proj.v_scale", (0,), None)] - rows[f"l{i}_norm2"] = [(f"{p}.post_attention_layernorm.weight", None, None)] - if i < core.dense_layers: - _dense_mlp_rows(rows, f"l{i}_mlp", f"{p}.mlp") - continue - _dense_mlp_rows(rows, f"l{i}_mlp", f"{p}.mlp.shared_experts") - rows[f"l{i}_router"] = [(f"{p}.mlp.gate.weight", None, None)] - rows[f"l{i}_router_bias"] = [(f"{p}.mlp.gate.e_score_correction_bias", None, None)] - fc1_w, fc1_s, fc2_w, fc2_s = [], [], [], [] - for local in range(core.local_experts): - q = f"{p}.mlp.experts.{core.expert_offset + local}" - fc1_w.append( - ( - (f"{q}.up_proj.weight", f"{q}.gate_proj.weight"), - (local,), - _expert_fc1_w, - ) +def _trunk_layers(d): + """Every trunk layer, plus the MTP module when it is live. + + The MTP module is layer `num_hidden_layers` of the same checkpoint and of the same + KV pool, and eleven of its weights are a trunk layer's exactly -- same shape, same + key template, same transform. They are those weights, one layer further along, not + a second family that happens to look alike. + """ + return range(d.num_layers + (1 if d.mtp_enabled else 0)) + + +def _moe_layers(d): + """Layers with a router: the trunk's MoE layers, plus MTP when live.""" + return range(d.dense_layers, d.num_layers + (1 if d.mtp_enabled else 0)) + + +def _mtp_only(d): + """The MTP module alone, and nothing at all when it is not live.""" + return [d.num_layers] if d.mtp_enabled else [] + + +def _dense_mlp_layers(d): + """The leading dense layers, whose MLP is the checkpoint's own `mlp`.""" + return range(d.dense_layers) + + +def _shared_expert_layers(d): + """MoE trunk layers, whose `mlp` of the same shape is the shared expert. + + MTP is excluded: its shared expert is bf16, not NVFP4, so it is a different weight + rather than this one at another layer. + """ + return range(d.dense_layers, d.num_layers) + + +def _one(template): + """A single checkpoint key, no index into the destination.""" + return lambda d, i, t=template: [(t.format(i=i), None)] + + +def _scalar(template): + """A per-tensor scalar, written to element 0 of a one-element parameter.""" + return lambda d, i, t=template: [(t.format(i=i), (0,))] + + +def _pair(a, b): + """Two checkpoint tensors handed to the transform together.""" + return lambda d, i, x=a, y=b: [((x.format(i=i), y.format(i=i)), None)] + + +def _expert_window(suffix, *, pair_with=None): + """This rank's slice of the routed experts, one source per local expert. + + The checkpoint key depends on the rank -- `expert_offset` is this rank's first + expert -- which is why the source is computed rather than templated. + """ + + def build(d, i): + out = [] + for local in range(d.local_experts): + q = f"model.layers.{i}.mlp.experts.{d.expert_offset + local}" + if pair_with is None: + out.append((f"{q}.{suffix}", (local,))) + else: + out.append(((f"{q}.{pair_with}", f"{q}.{suffix}"), (local,))) + return out + + return build + + +def _expert_scalars(suffix): + """A per-tensor scalar replicated over the whole routing space, not windowed. + + `derive_after_load` asserts the shared expert's activation scale against the max + over all routed experts, so every one of them has to be present on every rank. + """ + return lambda d, i: [ + (f"model.layers.{i}.mlp.experts.{e}.{suffix}", (e,)) for e in range(d.num_experts) + ] + + +#: The six per-tensor NVFP4 scalars every MLP and every routed expert carries, as +#: (parameter suffix, checkpoint suffix). The checkpoint stores reciprocals. +_MLP_SCALARS = ( + ("isc1", "gate_proj.input_scale"), + ("isc1_up", "up_proj.input_scale"), + ("ws2_1", "gate_proj.weight_scale_2"), + ("ws2_1_up", "up_proj.weight_scale_2"), + ("isc2", "down_proj.input_scale"), + ("ws2_2", "down_proj.weight_scale_2"), +) + +_EXPERT_SCALARS = _MLP_SCALARS + + +def _nvfp4_mlp(key, prefix, inter, layers): + """The ten rows of one NVFP4 SwiGLU MLP: a fused gate_up, a down, and six scalars. + + Used twice over disjoint layer ranges -- the dense layers' own `mlp` and the MoE + layers' shared expert -- because the two have the same shape in terms of their own + intermediate width, and different widths. + """ + return ( + W( + f"{key}_gu_w", + shape=lambda d: (2 * inter(d), d.hidden // 2), + dtype=torch.uint8, + src=_pair(f"{prefix}.gate_proj.weight", f"{prefix}.up_proj.weight"), + transform=_cat_rows, + layers=layers, + ), + W( + f"{key}_gu_s", + shape=lambda d: (2 * inter(d) * (d.hidden // SF_BLOCK),), + dtype=torch.uint8, + src=_pair(f"{prefix}.gate_proj.weight_scale", f"{prefix}.up_proj.weight_scale"), + transform=_cat_rows_swizzle, + layers=layers, + ), + W( + f"{key}_dn_w", + shape=lambda d: (d.hidden, inter(d) // 2), + dtype=torch.uint8, + src=_one(f"{prefix}.down_proj.weight"), + layers=layers, + ), + W( + f"{key}_dn_s", + shape=lambda d: (d.hidden * (inter(d) // SF_BLOCK),), + dtype=torch.uint8, + src=_one(f"{prefix}.down_proj.weight_scale"), + transform=_swizzle_only, + layers=layers, + ), + *( + W( + f"{key}_{n}", + shape=lambda d: (1,), + dtype=torch.float32, + src=_scalar(f"{prefix}.{s}"), + layers=layers, ) - fc1_s.append( - ( - (f"{q}.up_proj.weight_scale", f"{q}.gate_proj.weight_scale"), - (local,), - _expert_fc1_s, - ) + for n, s in _MLP_SCALARS + ), + ) + + +class DeepseekV3Dep4Weights(ModelWeights): + """deepseek-r1-0528-nvfp4 / sm_103 / dep4.""" + + RELEASE_AFTER = "_fc2_s" + + WEIGHTS: tuple[W, ...] = ( + # --- attention and the two layer norms: every trunk layer, and MTP --- + W( + "norm1", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.input_layernorm.weight"), + layers=_trunk_layers, + ), + W( + "qa", + shape=lambda d: (d.q_lora, d.hidden), + src=_one("model.layers.{i}.self_attn.q_a_proj.weight"), + layers=_trunk_layers, + ), + W( + "q_norm", + shape=lambda d: (d.q_lora,), + src=_one("model.layers.{i}.self_attn.q_a_layernorm.weight"), + layers=_trunk_layers, + ), + W( + "qb", + shape=lambda d: (d.heads * d.qk_dim, d.q_lora), + src=_one("model.layers.{i}.self_attn.q_b_proj.weight"), + layers=_trunk_layers, + ), + W( + "kva", + shape=lambda d: (d.lat_dim, d.hidden), + src=_one("model.layers.{i}.self_attn.kv_a_proj_with_mqa.weight"), + layers=_trunk_layers, + ), + W( + "kv_norm", + shape=lambda d: (d.kv_lora,), + src=_one("model.layers.{i}.self_attn.kv_a_layernorm.weight"), + layers=_trunk_layers, + ), + W( + "kvb", + shape=lambda d: (d.heads * (d.nope + d.v_dim), d.kv_lora), + src=_one("model.layers.{i}.self_attn.kv_b_proj.weight"), + transform=_reorder_kv_b, + layers=_trunk_layers, + ), + W( + "o", + shape=lambda d: (d.hidden, d.heads * d.v_dim), + src=_one("model.layers.{i}.self_attn.o_proj.weight"), + layers=_trunk_layers, + ), + # The fp8 KV-cache scales sit under the k/v projections this MLA checkpoint does + # not otherwise have. + W( + "k_scale", + shape=lambda d: (1,), + dtype=torch.float32, + src=_scalar("model.layers.{i}.self_attn.k_proj.k_scale"), + layers=_trunk_layers, + ), + W( + "v_scale", + shape=lambda d: (1,), + dtype=torch.float32, + src=_scalar("model.layers.{i}.self_attn.v_proj.v_scale"), + layers=_trunk_layers, + ), + W( + "norm2", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.post_attention_layernorm.weight"), + layers=_trunk_layers, + ), + # --- the NVFP4 MLP: the same parameter key over two layer families --- + # Layers 0..first_k_dense_replace-1 have the checkpoint's own `mlp`; the MoE + # layers have a shared expert of the same shape under `mlp.shared_experts`. + # Two entries over disjoint layers rather than one: the intermediate width + # differs, so these are different weights that share a key. + *_nvfp4_mlp("mlp", "model.layers.{i}.mlp", lambda d: d.dense_inter, _dense_mlp_layers), + *_nvfp4_mlp( + "mlp", + "model.layers.{i}.mlp.shared_experts", + lambda d: d.shared_inter, + _shared_expert_layers, + ), + # --- routing: MoE trunk layers, and MTP --- + W( + "router", + shape=lambda d: (d.num_experts, d.hidden), + src=_one("model.layers.{i}.mlp.gate.weight"), + layers=_moe_layers, + ), + W( + "router_bias", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_one("model.layers.{i}.mlp.gate.e_score_correction_bias"), + layers=_moe_layers, + ), + # --- this rank's routed experts, NVFP4 --- + W( + "fc1_w", + shape=lambda d: (d.local_experts, 2 * d.moe_inter, d.hidden // 2), + dtype=torch.uint8, + src=_expert_window("gate_proj.weight", pair_with="up_proj.weight"), + transform=_expert_fc1_w, + layers=_shared_expert_layers, + ), + W( + "fc1_s", + shape=lambda d: (d.local_experts, 2 * d.moe_inter, d.hidden // SF_BLOCK), + dtype=torch.uint8, + src=_expert_window("gate_proj.weight_scale", pair_with="up_proj.weight_scale"), + transform=_expert_fc1_s, + layers=_shared_expert_layers, + ), + W( + "fc2_w", + shape=lambda d: (d.local_experts, d.hidden, d.moe_inter // 2), + dtype=torch.uint8, + src=_expert_window("down_proj.weight"), + transform=_expert_fc2_w, + layers=_shared_expert_layers, + ), + W( + "fc2_s", + shape=lambda d: (d.local_experts, d.hidden, d.moe_inter // SF_BLOCK), + dtype=torch.uint8, + src=_expert_window("down_proj.weight_scale"), + transform=_expert_fc2_s, + layers=_shared_expert_layers, + ), + # --- the routed experts' per-tensor scalars, over the whole routing space --- + *( + W( + f"e_{n}", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_expert_scalars(s), + layers=_shared_expert_layers, ) - fc2_w.append((f"{q}.down_proj.weight", (local,), _expert_fc2_w)) - fc2_s.append((f"{q}.down_proj.weight_scale", (local,), _expert_fc2_s)) - rows[f"l{i}_fc1_w"] = fc1_w - rows[f"l{i}_fc1_s"] = fc1_s - rows[f"l{i}_fc2_w"] = fc2_w - rows[f"l{i}_fc2_s"] = fc2_s - # The per-tensor scalars are replicated over the whole routing space, - # not windowed: derive_after_load asserts the shared-expert - # activation scale against the max over all 256 routed experts. - isc1, isc1_up, ws2_1, ws2_1_up, isc2, ws2_2 = [], [], [], [], [], [] - for e in range(core.num_experts): - q = f"{p}.mlp.experts.{e}" - isc1.append((f"{q}.gate_proj.input_scale", (e,), None)) - isc1_up.append((f"{q}.up_proj.input_scale", (e,), None)) - ws2_1.append((f"{q}.gate_proj.weight_scale_2", (e,), None)) - ws2_1_up.append((f"{q}.up_proj.weight_scale_2", (e,), None)) - isc2.append((f"{q}.down_proj.input_scale", (e,), None)) - ws2_2.append((f"{q}.down_proj.weight_scale_2", (e,), None)) - rows[f"l{i}_e_isc1"] = isc1 - rows[f"l{i}_e_isc1_up"] = isc1_up - rows[f"l{i}_e_ws2_1"] = ws2_1 - rows[f"l{i}_e_ws2_1_up"] = ws2_1_up - rows[f"l{i}_e_isc2"] = isc2 - rows[f"l{i}_e_ws2_2"] = ws2_2 - rows["final_norm"] = [("model.norm.weight", None, None)] - rows["embed"] = [("model.embed_tokens.weight", None, None)] - if core.mtp_enabled: - _mtp_rows(rows, core) - return rows + for n, s in _EXPERT_SCALARS + ), + # --- the MTP module's own weights: the ones a trunk layer has no counterpart for --- + W( + "enorm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.enorm.weight"), + layers=_mtp_only, + ), + W( + "hnorm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.hnorm.weight"), + layers=_mtp_only, + ), + W( + "eh", + shape=lambda d: (d.hidden, 2 * d.hidden), + src=_one("model.layers.{i}.eh_proj.weight"), + layers=_mtp_only, + ), + W( + "head_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.shared_head.norm.weight"), + layers=_mtp_only, + ), + # bf16, not NVFP4: the checkpoint excludes `model.layers.61*` from the export, + # which is also why these go to `fused_moe` rather than the trtllm-gen runner. + W( + "sh_gu", + shape=lambda d: (2 * d.shared_inter, d.hidden), + src=_pair( + "model.layers.{i}.mlp.shared_experts.gate_proj.weight", + "model.layers.{i}.mlp.shared_experts.up_proj.weight", + ), + transform=_cat_rows, + layers=_mtp_only, + ), + W( + "sh_dn", + shape=lambda d: (d.hidden, d.shared_inter), + src=_one("model.layers.{i}.mlp.shared_experts.down_proj.weight"), + layers=_mtp_only, + ), + W( + "fc1", + shape=lambda d: (d.local_experts, 2 * d.moe_inter, d.hidden), + src=_expert_window("gate_proj.weight", pair_with="up_proj.weight"), + transform=_mtp_fc1, + layers=_mtp_only, + ), + W( + "fc2", + shape=lambda d: (d.local_experts, d.hidden, d.moe_inter), + src=_expert_window("down_proj.weight"), + layers=_mtp_only, + ), + # --- not per-layer --- + W("final_norm", shape=lambda d: (d.hidden,), src="model.norm.weight", layers=None), + W( + "embed", + shape=lambda d: (d.vocab, d.hidden), + src="model.embed_tokens.weight", + layers=None, + ), + ) + + def expected_unconsumed(self, core, weights) -> set: + return _offwindow_expert_keys(core) | _mtp_keys(core, weights) + + def load(self, model, weights) -> None: + """Overrides the base: multi-key rows feed their transform several tensors at + once, the per-tensor scalars are 0-dim and reject `[:]`, and this target asserts + coverage in both directions rather than only checkpoint-side.""" + core = model.model + manifest = self.manifest(core) + consumed: set = set() + + def fill(param: torch.nn.Parameter, ckpt_key, index, transform) -> None: + keys = ckpt_key if isinstance(ckpt_key, tuple) else (ckpt_key,) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + dst = param.data if index is None else param.data[index] + # Materialize the checkpoint entries where the destination lives: the + # expert relayouts are row gathers over ~15 MB per expert and the + # concatenations allocate scratch of the same order. + src_all = tuple( + _materialize(weights[key]).to(dst.device, non_blocking=True) for key in keys + ) + if transform is not None: + src = transform(src_all if isinstance(ckpt_key, tuple) else src_all[0], core) + else: + assert len(src_all) == 1, "a multi-key row needs a transform" + src = src_all[0] + assert dst.shape == src.shape, (ckpt_key, tuple(dst.shape), tuple(src.shape)) + assert src.dtype == dst.dtype, (ckpt_key, src.dtype, dst.dtype) + dst.copy_(src, non_blocking=True) + consumed.update(keys) + + assert set(manifest.keys()) == set(core.w.keys()), ( + "manifest/parameter drift", + set(manifest.keys()) ^ set(core.w.keys()), + ) + for param_key, sources in manifest.items(): + for ckpt_key, index, transform in sources: + fill(core.w[param_key], ckpt_key, index, transform) + # The expert relayouts allocate per-expert scratch 64 times per layer; + # release it before the next parameter so peak load memory stays one + # layer deep. + if param_key.endswith(self.RELEASE_AFTER): + torch.cuda.empty_cache() + + # Shell-registered exception: the base class owns lm_head (untied), and + # under attention DP it builds it **whole** -- `[vocab, hidden]` on every + # rank. That is the opposite of a tensor-parallel target, where the same + # shell builds a vocab-parallel `[vocab/tp, hidden]`, and it is the only + # consistent choice here: each rank holds different tokens, so a vocab + # shard could not be completed by a collective over rows no other rank + # computed. Asserted rather than adapted -- a shape change is a different + # logits path. + cfg = core.model_config.pretrained_config + expected_lm_head_shape = (cfg.vocab_size, cfg.hidden_size) + assert tuple(model.lm_head.weight.shape) == expected_lm_head_shape, ( + f"attention DP replicates lm_head; the shell built " + f"{tuple(model.lm_head.weight.shape)} instead of {expected_lm_head_shape}" + ) + fill(model.lm_head.weight, "lm_head.weight", None, None) + # Parameter-side coverage, the other direction: the shell's construction + # is topology-dependent, so name every registered parameter outside the + # target's own ParameterDict instead of assuming lm_head is the only one. + shell = {n for n, _ in model.named_parameters() if not n.startswith("model.w.")} + assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" + + torch.cuda.synchronize() + leftover = set(weights.keys()) - consumed + expected = self.expected_unconsumed(core, weights) + assert leftover == expected, ( + "checkpoint coverage: leftover keys are not exactly this rank's " + f"off-window experts plus the MTP layers; unexpected " + f"{sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" + ) + + def dims(self, core) -> SimpleNamespace: + cfg = core.model_config.pretrained_config + mapping = core.model_config.mapping + local_experts = cfg.n_routed_experts // mapping.moe_ep_size + return SimpleNamespace( + num_layers=cfg.num_hidden_layers, + hidden=cfg.hidden_size, + heads=cfg.num_attention_heads, + nope=cfg.qk_nope_head_dim, + rope_dim=cfg.qk_rope_head_dim, + v_dim=cfg.v_head_dim, + kv_lora=cfg.kv_lora_rank, + q_lora=cfg.q_lora_rank, + qk_dim=cfg.qk_nope_head_dim + cfg.qk_rope_head_dim, + lat_dim=cfg.kv_lora_rank + cfg.qk_rope_head_dim, + vocab=cfg.vocab_size, + dense_layers=cfg.first_k_dense_replace, + num_experts=cfg.n_routed_experts, + moe_inter=cfg.moe_intermediate_size, + shared_inter=cfg.moe_intermediate_size * cfg.n_shared_experts, + dense_inter=cfg.intermediate_size, + local_experts=local_experts, + expert_offset=local_experts * mapping.moe_ep_rank, + mtp_enabled=getattr(core.model_config, "spec_config", None) is not None, + dtype=core.model_config.torch_dtype, + ) + + +MODEL_WEIGHTS = DeepseekV3Dep4Weights() def _offwindow_expert_keys(core) -> set: @@ -336,14 +676,16 @@ def _offwindow_expert_keys(core) -> set: The trunk's MoE layers store six such tensors per expert (NVFP4 data plus block scales); the MTP module's are bf16, so its off-window experts leave three each and no `weight_scale` at all.""" + cfg = core.model_config.pretrained_config + num_layers = cfg.num_hidden_layers lo, hi = core.expert_offset, core.expert_offset + core.local_experts - layers = list(range(core.dense_layers, core.num_layers)) + layers = list(range(cfg.first_k_dense_replace, num_layers)) if core.mtp_enabled: - layers += list(range(core.num_layers, core.num_layers + core.mtp_layers)) + layers += list(range(num_layers, num_layers + core.mtp_layers)) keys = set() for i in layers: - quantized = i < core.num_layers - for e in range(core.num_experts): + quantized = i < num_layers + for e in range(cfg.n_routed_experts): if lo <= e < hi: continue q = f"model.layers.{i}.mlp.experts.{e}" @@ -367,9 +709,8 @@ def _mtp_keys(core, weights) -> set: `lm_head`, which the draft-model container points at instead of loading them twice. (The off-window experts are left over too, but they belong to the family above and are named there.)""" - prefixes = tuple( - f"model.layers.{i}." for i in range(core.num_layers, core.num_layers + core.mtp_layers) - ) + num_layers = core.model_config.pretrained_config.num_hidden_layers + prefixes = tuple(f"model.layers.{i}." for i in range(num_layers, num_layers + core.mtp_layers)) if not prefixes: return set() keys = {k for k in weights if k.startswith(prefixes)} @@ -377,72 +718,3 @@ def _mtp_keys(core, weights) -> set: return keys aliased = (".embed_tokens.weight", ".shared_head.head.weight") return {k for k in keys if k.endswith(aliased)} - - -def load(model, weights) -> None: - core = model.model - manifest = _manifest(core) - consumed: set = set() - - def fill(param: torch.nn.Parameter, ckpt_key, index, transform) -> None: - keys = ckpt_key if isinstance(ckpt_key, tuple) else (ckpt_key,) - for key in keys: - assert key in weights, f"checkpoint key missing: {key}" - dst = param.data if index is None else param.data[index] - # Materialize the checkpoint entries where the destination lives: the - # expert relayouts are row gathers over ~15 MB per expert and the - # concatenations allocate scratch of the same order. The per-tensor - # scalars are stored 0-dim, which `[:]` rejects. - src_all = tuple( - _materialize(weights[key]).to(dst.device, non_blocking=True) for key in keys - ) - if transform is not None: - src = transform(src_all if isinstance(ckpt_key, tuple) else src_all[0], core) - else: - assert len(src_all) == 1, "a multi-key row needs a transform" - src = src_all[0] - assert dst.shape == src.shape, (ckpt_key, tuple(dst.shape), tuple(src.shape)) - assert src.dtype == dst.dtype, (ckpt_key, src.dtype, dst.dtype) - dst.copy_(src, non_blocking=True) - consumed.update(keys) - - assert set(manifest.keys()) == set(core.w.keys()), ( - "manifest/parameter drift", - set(manifest.keys()) ^ set(core.w.keys()), - ) - for param_key, sources in manifest.items(): - for ckpt_key, index, transform in sources: - fill(core.w[param_key], ckpt_key, index, transform) - # The expert relayouts allocate per-expert scratch 64 times per - # layer; release it before the next parameter so peak load memory - # stays one layer deep. - if param_key.endswith("_fc2_s"): - torch.cuda.empty_cache() - - # Shell-registered exception: the base class owns lm_head (untied), and - # under attention DP it builds it **whole** — `[vocab, hidden]` on every - # rank. That is the opposite of a tensor-parallel target, where the same - # shell builds a vocab-parallel `[vocab/tp, hidden]`, and it is the only - # consistent choice here: each rank holds different tokens, so a vocab - # shard could not be completed by a collective over rows no other rank - # computed. Asserted rather than adapted — a shape change is a different - # logits path. - assert tuple(model.lm_head.weight.shape) == (core.vocab, core.hidden), ( - f"attention DP replicates lm_head; the shell built " - f"{tuple(model.lm_head.weight.shape)} instead of {(core.vocab, core.hidden)}" - ) - fill(model.lm_head.weight, "lm_head.weight", None, None) - # Parameter-side coverage, the other direction: the shell's construction - # is topology-dependent, so name every registered parameter outside the - # target's own ParameterDict instead of assuming lm_head is the only one. - shell = {n for n, _ in model.named_parameters() if not n.startswith("model.w.")} - assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" - - torch.cuda.synchronize() - leftover = set(weights.keys()) - consumed - expected = _offwindow_expert_keys(core) | _mtp_keys(core, weights) - assert leftover == expected, ( - "checkpoint coverage: leftover keys are not exactly this rank's " - f"off-window experts plus the MTP layers; unexpected " - f"{sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/__init__.py new file mode 100644 index 000000000000..43d736a01ea5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DeepseekV41ForCausalLM targets.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/__init__.py new file mode 100644 index 000000000000..bac8b83ea30a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DeepSeek-V4.1-Flash-NVFP4 / sm_103 / unified {tep4, dep4, tp4}.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/engram_host_table.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/engram_host_table.py new file mode 100644 index 000000000000..07176eed4811 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/engram_host_table.py @@ -0,0 +1,2314 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One shared, pinned host copy per node of DeepSeek-V4.1's Engram embedding tables. + +The Engram tables (``layers.{1,14}.engram.embed.weight`` ``[~384M, 256]`` +float8_e4m3fn plus ``.scale`` ``[~384M, 8]`` float8_e8m0fnu, 202.76 GB in all) +are larger than the rest of the checkpoint. Under ``host_offload`` they stay in +host memory and each GPU gathers the rows it needs through a mapped pointer +(the target's ``ops/embedding/engram_gather``). This module owns the host side, with two +backends selected by ``$TRTLLM_DSV41_ENGRAM_HOST_TABLE_BACKEND`` +(``host_table_backend``): + +* ``vmm`` (the default): CUDA VMM host-NUMA allocations (``cuMemCreate`` with + ``CU_MEM_LOCATION_TYPE_HOST_NUMA``, 1 GiB slabs alternating between the CPU + sockets the creating process may allocate on: both when it sees GPUs of both + sockets, else its GPUs' socket), shared between processes as POSIX file descriptors over a + node-local Unix socket. The GPU maps them in its own page tables at the + allocation granularity, so a random row gather is not bound by the CPU page + tables' 64 KiB pages. See ``VmmEngramHostTable``. +* ``shm``: the ``/dev/shm`` file described below, registered with + ``cudaHostRegister``. + +The rest of this docstring describes the ``shm`` backend. + +* **One copy per node.** Local rank 0 fills a file on a tmpfs (``/dev/shm`` by + default) and every local rank maps that file ``MAP_SHARED``, so the node holds + the bytes once however many ranks map them. A named file, + unlike a ``memfd`` handed between ranks, needs no common PID namespace and + outlives the engine, which is what lets the next boot skip the fill. +* **Pinned per rank.** Each rank ``cudaHostRegister``s the table regions of its + own mapping and checks that the tensor views report ``is_pinned()``; + ``cudaHostGetDevicePointer`` alone also succeeds on pageable memory. +* **A fast, checked fill.** Parallel ``pread`` of the safetensors byte ranges + straight into the shared mapping, with no intermediate tensors, and a CRC32 + per 64 MiB chunk recorded in a manifest at the head of the file. +* **Reuse.** If the file exists, is complete, and its manifest names the same + checkpoint bytes (shard path, size, mtime, tensor offsets and shapes), the + fill is skipped and a sample of the chunk CRCs is re-checked. + +File layout: a 1 MiB header (magic, version, state, JSON manifest) followed by +the tensors, each starting at a multiple of 512 MiB, in a file whose size is a +multiple of 512 MiB. That is a PMD huge page on a 64 KiB-page arm64 kernel and a +multiple of 2 MiB elsewhere, so a tmpfs that allows huge pages (``huge=`` mount +option, or ``shmem_enabled`` not ``never``) or a hugetlbfs file +(``$TRTLLM_DSV41_ENGRAM_HOST_TABLE_PATH``) can back the tensors with them. The +gaps are never written and occupy no memory on a tmpfs. + +Page size decides the gather rate. The GPU translates every row address through +the process's page tables. With 64 KiB pages and a table past roughly 10 GB, +random row gathers become translation-bound. Measured on GB300 at 4096 tokens +x 24 lookups: 3.1 ms per call over the full 101 GB layer table, against 0.14 ms +for the same table on 512 MiB pages or on 2 MiB-granular CUDA host allocations. +``fill_shared_engram_table`` asks for huge pages (``MADV_HUGEPAGE``) and records in +the manifest whether it got them. + +Placement: ``numa="interleave"`` (the default) spreads the pages over both +sockets. Every GPU then reads at its local rate, whereas memory on the remote +socket alone was 1.7x slower. + +Lifetime: every attached process holds a shared ``flock`` on ``.lock`` +for as long as it is attached, and the kernel drops that lock when the process +exits or dies. The file is kept for reuse by default. +``remove_shared_engram_table`` unlinks it only when no process holds that lock, +unless ``force`` is given. A process that dies leaves nothing pinned, and the +survivors' mappings are unaffected. After an unlink, tmpfs frees the pages when +the last mapping goes away. A fill with ``owned=True`` (what an engine does) +also writes ``.owner`` naming the new file; ``close(remove_if_owned=True)`` +then removes that file once its last holder detaches. A file filled from the +command line has no owner and stays until ``remove``, so consecutive engines +reuse it. + +Residency switch: ``TRTLLM_DSV41_ENGRAM_HOST_OFFLOAD`` (``host_offload_enabled``) +decides whether the DeepSeek-V4.1 target keeps its Engram tables here (unset or +a true value) or row-sharded on the devices (a false value). + +The module imports only the standard library and torch at load time, so it also +runs as a script without importing ``tensorrt_llm``: +``python engram_host_table.py {fill,status,verify,remove} ...``. +""" + +from __future__ import annotations + +import argparse +import ctypes +import fcntl +import json +import logging +import math +import mmap +import os +import platform +import re +import socket +import struct +import threading +import time +import warnings +import zlib +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict, dataclass +from typing import Callable, Dict, List, Optional, Sequence, Tuple + +import torch + +__all__ = [ + "EngramTableSource", + "FillReport", + "SharedEngramHostTable", + "VmmEngramHostTable", + "attach_shared_engram_table", + "default_table_path", + "describe_shared_engram_table", + "fill_shared_engram_table", + "find_engram_table_sources", + "host_offload_enabled", + "host_table_backend", + "open_shared_engram_table", + "open_vmm_engram_table", + "remove_shared_engram_table", + "verify_shared_engram_table", +] + +_logger = logging.getLogger(__name__) + +PATH_ENV = "TRTLLM_DSV41_ENGRAM_HOST_TABLE_PATH" +"""str: environment variable overriding the shared file's path.""" + +HOST_OFFLOAD_ENV = "TRTLLM_DSV41_ENGRAM_HOST_OFFLOAD" +"""str: environment variable selecting the Engram tables' residency (see ``host_offload_enabled``).""" + + +def host_offload_enabled() -> bool: + """Whether the Engram tables live in this module's shared pinned host copy. + + ``$TRTLLM_DSV41_ENGRAM_HOST_OFFLOAD`` unset: yes. ``1``/``true``/``yes``/``on``: + yes. ``0``/``false``/``no``/``off``: no, the tables are row-sharded on the + devices. Anything else, the empty string included, raises: a typo must not + silently select a residency. + """ + value = os.environ.get(HOST_OFFLOAD_ENV) + if value is None: + return True + token = value.strip().lower() + if token in ("1", "true", "yes", "on"): + return True + if token in ("0", "false", "no", "off"): + return False + raise ValueError(f"{HOST_OFFLOAD_ENV} must be a boolean value, got {value!r}") + + +BACKEND_ENV = "TRTLLM_DSV41_ENGRAM_HOST_TABLE_BACKEND" +"""str: environment variable selecting the host table's backend (see ``host_table_backend``).""" + + +def host_table_backend() -> str: + """``$TRTLLM_DSV41_ENGRAM_HOST_TABLE_BACKEND``: ``vmm`` (unset: the default) or ``shm``. + + Anything else raises, so a typo cannot silently select a backend. + """ + value = os.environ.get(BACKEND_ENV) + if value is None: + return "vmm" + token = value.strip().lower() + if token in ("vmm", "shm"): + return token + raise ValueError(f"{BACKEND_ENV} must be 'vmm' or 'shm', got {value!r}") + + +_MAGIC = b"DSV41ENG" +_VERSION = 1 +_STATE_FILLING = 0 +_STATE_READY = 1 +_HEADER = struct.Struct("<8sIIQQ") # magic, version, state, manifest bytes, file bytes +_MANIFEST_AT = 64 +_HEADER_BYTES = 1 << 20 +_ALIGN = 1 << 29 +_CHUNK_BYTES = 1 << 26 +_DIRECT_ALIGN = 1 << 20 +_SAMPLE_CHUNKS = 8 + +_ENGRAM_KEY = re.compile(r"layers\.(\d+)\.engram\.embed\.(weight|scale)") +_DTYPES = { + "F8_E4M3": torch.float8_e4m3fn, + "F8_E8M0": torch.float8_e8m0fnu, + "U8": torch.uint8, +} +_MPOL_BIND = 2 +_MPOL_INTERLEAVE = 3 +_SYS_MBIND = {"aarch64": 235, "x86_64": 237} + + +def default_table_path() -> str: + """Path of the node's shared table file: ``$TRTLLM_DSV41_ENGRAM_HOST_TABLE_PATH``, else + ``/dev/shm/trtllm-dsv41-engram-.bin``. + + A fixed name per user means a node can hold at most one table per user: a + fill for different checkpoint bytes replaces the file rather than adding a + second one. + """ + return os.environ.get(PATH_ENV) or f"/dev/shm/trtllm-dsv41-engram-{os.getuid()}.bin" # nosec B108 + + +@dataclass(frozen=True) +class EngramTableSource: + """One Engram table tensor as stored in a safetensors shard. + + Attributes: + key: checkpoint name, e.g. ``layers.1.engram.embed.weight``. + layer: the decoder layer the table belongs to. + shard: absolute path of the shard that holds it. + offset: absolute byte offset of the tensor's first byte in ``shard``. + nbytes: tensor size in bytes (every element is one byte). + shape: ``(rows, 256)`` for a weight, ``(rows, 8)`` for a scale. + dtype: safetensors dtype string, ``F8_E4M3`` or ``F8_E8M0``. + """ + + key: str + layer: int + shard: str + offset: int + nbytes: int + shape: Tuple[int, ...] + dtype: str + + +@dataclass(frozen=True) +class FillReport: + """What ``fill_shared_engram_table`` did. + + Attributes: + path: the shared file. + reused: True if a matching complete file was found and the fill was skipped. + nbytes: table bytes (the sum of the tensors, not counting the header or gaps). + seconds: wall time of the call. + threads: reader threads used (or that would have been used). + method: ``buffered`` or ``direct``. + numa: the page placement policy applied to a new file. + verified_chunks: chunk CRCs re-checked on reuse. + numa_bytes: resident table bytes per NUMA node after a fill (empty on reuse). + """ + + path: str + reused: bool + nbytes: int + seconds: float + threads: int + method: str + numa: str + verified_chunks: int + numa_bytes: Dict[int, int] + + @property + def gbps(self) -> float: + """Table bytes per second of wall time, in GB/s.""" + return self.nbytes / self.seconds / 1e9 if self.seconds > 0 else math.inf + + +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +def _read_safetensors_header(path: str) -> Tuple[int, dict]: + with open(path, "rb") as f: + (length,) = struct.unpack(" List[EngramTableSource]: + """Locate every ``layers..engram.embed.{weight,scale}`` tensor of a checkpoint. + + Reads ``model.safetensors.index.json`` and the header of each shard that + holds a table. Returns the tensors ordered by layer, weight before scale. + """ + index_path = os.path.join(checkpoint_dir, "model.safetensors.index.json") + with open(index_path) as f: + weight_map = json.load(f)["weight_map"] + headers: Dict[str, Tuple[int, dict]] = {} + found = [] + for key, shard_name in weight_map.items(): + match = _ENGRAM_KEY.fullmatch(key) + if match is None: + continue + shard = os.path.realpath(os.path.join(checkpoint_dir, shard_name)) + if shard not in headers: + headers[shard] = _read_safetensors_header(shard) + data_start, header = headers[shard] + entry = header[key] + begin, end = entry["data_offsets"] + shape = tuple(int(d) for d in entry["shape"]) + if entry["dtype"] not in _DTYPES: + raise ValueError(f"{key}: unsupported dtype {entry['dtype']}") + if end - begin != math.prod(shape): + raise ValueError( + f"{key}: {end - begin} bytes for shape {shape}; expected one byte per element" + ) + found.append( + EngramTableSource( + key, + int(match.group(1)), + shard, + data_start + begin, + end - begin, + shape, + entry["dtype"], + ) + ) + if not found: + raise ValueError(f"no layers..engram.embed.{{weight,scale}} tensors in {index_path}") + found.sort(key=lambda s: (s.layer, s.key.endswith(".scale"))) + by_layer: Dict[int, List[EngramTableSource]] = {} + for source in found: + by_layer.setdefault(source.layer, []).append(source) + for layer, pair in by_layer.items(): + if len(pair) != 2 or pair[0].shape[0] != pair[1].shape[0]: + raise ValueError( + f"layer {layer}: expected a weight and a scale with equal row counts, got {pair}" + ) + return found + + +def _layout(sources: Sequence[EngramTableSource]) -> dict: + """Where each tensor goes in the shared file, and the checkpoint bytes it comes from.""" + offset = _round_up(_HEADER_BYTES, _ALIGN) + tensors = [] + for source in sources: + st = os.stat(source.shard) + tensors.append( + { + "key": source.key, + "layer": source.layer, + "shape": list(source.shape), + "dtype": source.dtype, + "offset": offset, + "nbytes": source.nbytes, + "source": { + "shard": source.shard, + "shard_bytes": st.st_size, + "shard_mtime_ns": st.st_mtime_ns, + "offset": source.offset, + }, + } + ) + offset = _round_up(offset + source.nbytes, _ALIGN) + last = tensors[-1] + return { + "format": _VERSION, + "align": _ALIGN, + "chunk_bytes": _CHUNK_BYTES, + "total_bytes": _round_up(last["offset"] + last["nbytes"], _ALIGN), + "tensors": tensors, + } + + +def _identity(manifest: dict) -> dict: + return { + k: manifest.get(k) for k in ("format", "align", "chunk_bytes", "total_bytes", "tensors") + } + + +def _read_header(path: str) -> Tuple[Optional[int], Optional[dict], int]: + """(state, manifest, file bytes); state and manifest are None if the file has no valid header.""" + fd = os.open(path, os.O_RDONLY) + try: + size = os.fstat(fd).st_size + raw = os.pread(fd, _HEADER_BYTES, 0) + finally: + os.close(fd) + if len(raw) < _MANIFEST_AT: + return None, None, size + magic, version, state, manifest_bytes, _ = _HEADER.unpack_from(raw) + if magic != _MAGIC or version != _VERSION or manifest_bytes > _HEADER_BYTES - _MANIFEST_AT: + return None, None, size + try: + manifest = json.loads(raw[_MANIFEST_AT : _MANIFEST_AT + manifest_bytes]) + except ValueError: + return None, None, size + return state, manifest, size + + +def _write_header(mm: mmap.mmap, state: int, manifest: dict) -> None: + body = json.dumps(manifest, separators=(",", ":")).encode() + if len(body) > _HEADER_BYTES - _MANIFEST_AT: + raise ValueError( + f"manifest of {len(body)} bytes does not fit the {_HEADER_BYTES}-byte header" + ) + mm[_MANIFEST_AT : _MANIFEST_AT + len(body)] = body + mm[: _HEADER.size] = _HEADER.pack(_MAGIC, _VERSION, state, len(body), manifest["total_bytes"]) + + +def _check(path: str, layout: Optional[dict]) -> Tuple[Optional[dict], str]: + """The manifest of a complete file matching ``layout`` (any layout if None), else (None, why not).""" + try: + state, manifest, size = _read_header(path) + except FileNotFoundError: + return None, "absent" + if manifest is None: + return None, "no valid header" + if state != _STATE_READY: + return None, "incomplete (a fill started and did not finish)" + if size != manifest["total_bytes"]: + return None, f"{size} bytes, manifest says {manifest['total_bytes']}" + if layout is not None and _identity(manifest) != _identity(layout): + return None, "made from different checkpoint bytes" + return manifest, "ready" + + +def _lock_path(path: str) -> str: + return path + ".lock" + + +def _open_lock(path: str) -> int: + return os.open(_lock_path(path), os.O_RDWR | os.O_CREAT, 0o600) + + +def _owner_path(path: str) -> str: + return path + ".owner" + + +def _write_owner(path: str, owned: bool) -> None: + """Record (or clear) that the file now at ``path`` was created by an engine.""" + marker = _owner_path(path) + if owned: + with open(marker, "w") as f: + f.write(str(os.stat(path).st_ino)) + return + try: + os.unlink(marker) + except FileNotFoundError: + pass + + +def _owned_inode(path: str) -> Optional[int]: + """The inode an engine recorded as its own at ``path``, or None.""" + try: + with open(_owner_path(path)) as f: + return int(f.read().strip()) + except (FileNotFoundError, ValueError): + return None + + +def _flock_within(fd: int, operation: int, timeout_s: float, what: str) -> None: + deadline = time.monotonic() + timeout_s + while True: + try: + fcntl.flock(fd, operation | fcntl.LOCK_NB) + return + except BlockingIOError: + if time.monotonic() >= deadline: + raise TimeoutError(what) from None + time.sleep(0.2) + + +def _lock_holders(path: str) -> List[Tuple[int, str]]: + """(pid, READ|WRITE) of every flock on ``.lock``, from /proc/locks.""" + try: + inode = os.stat(_lock_path(path)).st_ino + with open("/proc/locks") as f: + lines = f.read().splitlines() + except OSError: + return [] + holders = [] + for line in lines: + fields = line.split() + if "->" in fields: + fields.remove("->") + if len(fields) >= 6 and fields[1] == "FLOCK" and int(fields[5].split(":")[-1]) == inode: + holders.append((int(fields[4]), fields[3])) + return holders + + +def _meminfo() -> Dict[str, int]: + out = {} + with open("/proc/meminfo") as f: + for line in f: + name, value = line.split(":", 1) + parts = value.split() + out[name] = int(parts[0]) * (1024 if len(parts) > 1 else 1) + return out + + +def _cgroup_available() -> Optional[int]: + """Bytes this cgroup (v2) can still charge before its memory.max, counting clean page cache + as reclaimable; None without a limit.""" + try: + with open("/proc/self/cgroup") as f: + rel = next(line.split("::", 1)[1].strip() for line in f if line.startswith("0::")) + root = os.path.join("/sys/fs/cgroup", rel.lstrip("/")) + with open(os.path.join(root, "memory.max")) as f: + limit = f.read().strip() + if limit == "max": + return None + with open(os.path.join(root, "memory.current")) as f: + current = int(f.read()) + stat = {} + with open(os.path.join(root, "memory.stat")) as f: + for line in f: + name, value = line.split() + stat[name] = int(value) + except (OSError, StopIteration, ValueError): + return None + return int(limit) - current + max(0, stat.get("file", 0) - stat.get("shmem", 0)) + + +def _check_capacity(path: str, nbytes: int, headroom_bytes: int) -> None: + st = os.statvfs(os.path.dirname(os.path.abspath(path))) + fs_free = st.f_bavail * st.f_frsize + if fs_free < nbytes: + raise RuntimeError( + f"{os.path.dirname(path)} has {fs_free / 1e9:.1f} GB free; the table needs {nbytes / 1e9:.1f} GB" + ) + available = _meminfo().get("MemAvailable", 0) + cgroup = _cgroup_available() + limit = available if cgroup is None else min(available, cgroup) + if limit < nbytes + headroom_bytes: + raise RuntimeError( + f"the table needs {nbytes / 1e9:.1f} GB plus {headroom_bytes / 1e9:.1f} GB headroom; " + f"{available / 1e9:.1f} GB MemAvailable, " + f"cgroup {'unlimited' if cgroup is None else f'{cgroup / 1e9:.1f} GB'}" + ) + + +def _node_list(text: str) -> List[int]: + nodes: List[int] = [] + for part in text.strip().split(","): + if part: + lo, _, hi = part.partition("-") + nodes.extend(range(int(lo), int(hi or lo) + 1)) + return nodes + + +def _device_socket_cpus(device: torch.device) -> Optional[set]: + """CPUs of the NUMA node the device's PCI function reports, or None if it reports none.""" + from cuda.bindings import runtime as cudart + + err, bus = cudart.cudaDeviceGetPCIBusId(64, device.index) + if err != cudart.cudaError_t.cudaSuccess: + return None + domain, bus_no, rest = bytes(bus).split(b"\x00", 1)[0].decode().lower().split(":") + function = f"{int(domain, 16):04x}:{int(bus_no, 16):02x}:{rest}" + try: + with open(f"/sys/bus/pci/devices/{function}/numa_node") as f: + node = int(f.read()) + if node < 0: + return None + with open(f"/sys/devices/system/node/node{node}/cpulist") as f: + return set(_node_list(f.read())) + except OSError: + return None + + +def _cpu_memory_nodes() -> List[int]: + """NUMA nodes with both CPUs and memory (on GB300: the two Grace sockets, not the GPU nodes).""" + with open("/sys/devices/system/node/has_cpu") as f: + cpu = set(_node_list(f.read())) + with open("/sys/devices/system/node/has_memory") as f: + mem = set(_node_list(f.read())) + return sorted(cpu & mem) + + +def _mbind(addr: int, length: int, mode: int, nodes: Sequence[int]) -> None: + number = _SYS_MBIND.get(platform.machine()) + if number is None: + raise RuntimeError(f"no mbind syscall number for {platform.machine()}") + words = max(nodes) // 64 + 1 + mask = (ctypes.c_ulong * words)() + for node in nodes: + mask[node // 64] |= 1 << (node % 64) + libc = ctypes.CDLL(None, use_errno=True) + libc.syscall.restype = ctypes.c_long + rc = libc.syscall( + ctypes.c_long(number), + ctypes.c_void_p(addr), + ctypes.c_ulong(length), + ctypes.c_int(mode), + mask, + ctypes.c_ulong(words * 64 + 1), + ctypes.c_uint(0), + ) + if rc != 0: + err = ctypes.get_errno() + raise OSError(err, f"mbind: {os.strerror(err)}") + + +def _parse_numa_policy(numa: str) -> Tuple[Optional[int], List[int]]: + """``interleave`` -> all CPU+memory nodes; ``local`` -> first touch; ```` -> bind to node n.""" + if numa == "local": + return None, [] + if numa == "interleave": + return _MPOL_INTERLEAVE, _cpu_memory_nodes() + return _MPOL_BIND, [int(numa)] + + +def _mapping_numa_bytes(start: int, end: int) -> Dict[int, int]: + """Resident bytes per NUMA node of this process's mappings inside [start, end).""" + out: Dict[int, int] = {} + with open("/proc/self/numa_maps") as f: + for line in f: + fields = line.split() + if not start <= int(fields[0], 16) < end: + continue + page = next( + (int(x.split("=")[1]) * 1024 for x in fields if x.startswith("kernelpagesize_kB=")), + mmap.PAGESIZE, + ) + for x in fields: + if x[0] == "N" and "=" in x: + node, pages = x[1:].split("=") + out[int(node)] = out.get(int(node), 0) + int(pages) * page + return out + + +def _mapping_smaps(start: int, end: int) -> Dict[str, int]: + """Sum of this process's smaps fields (in bytes) over the mappings inside [start, end).""" + wanted = ( + "Rss", + "Pss", + "Shared_Clean", + "Shared_Dirty", + "Private_Clean", + "Private_Dirty", + "AnonHugePages", + "ShmemPmdMapped", + "FilePmdMapped", + "Locked", + ) + out = dict.fromkeys(wanted, 0) + inside = False + with open("/proc/self/smaps") as f: + for line in f: + head = line.split(None, 1)[0] + if "-" in head and not head.endswith(":"): + inside = start <= int(head.split("-")[0], 16) < end + elif inside and head[:-1] in out: + out[head[:-1]] += int(line.split()[1]) * 1024 + return out + + +def _mmap_address(mm: mmap.mmap) -> int: + anchor = ctypes.c_char.from_buffer(mm) + try: + return ctypes.addressof(anchor) + finally: + del anchor + + +def _read_buffered(fd: int, dst: memoryview, src_offset: int) -> None: + done = 0 + while done < len(dst): + n = os.preadv(fd, [dst[done:]], src_offset + done) + if n <= 0: + raise OSError(f"short read at byte {src_offset + done} ({done} of {len(dst)})") + done += n + + +class _BounceBuffer: + """A page-aligned buffer that one reader thread uses for its O_DIRECT reads.""" + + def __init__(self, nbytes: int) -> None: + self.buffer = mmap.mmap(-1, nbytes) + self.view = memoryview(self.buffer) + self._anchor = ctypes.c_char.from_buffer(self.buffer) + self.address = ctypes.addressof(self._anchor) + + def release(self) -> None: + self.view.release() + del self._anchor + self.buffer.close() + + +class _PerThreadBounce(threading.local): + """Gives each thread its own ``_BounceBuffer`` and records it for release by the caller.""" + + def __init__( + self, nbytes: int, registry: List[_BounceBuffer], registry_lock: threading.Lock + ) -> None: + super().__init__() + self.buffer = _BounceBuffer(nbytes) + with registry_lock: + registry.append(self.buffer) + + +def _parallel_fill(mm: mmap.mmap, layout: dict, threads: int, method: str) -> Dict[str, List[int]]: + """Copy every tensor's checkpoint bytes into ``mm``; returns the CRC32 of each chunk.""" + if method not in ("buffered", "direct"): + raise ValueError(f"method must be 'buffered' or 'direct', got {method!r}") + chunk = layout["chunk_bytes"] + jobs = [] + crcs: Dict[str, List[int]] = {} + for t in layout["tensors"]: + count = -(-t["nbytes"] // chunk) + crcs[t["key"]] = [0] * count + for i in range(count): + n = min(chunk, t["nbytes"] - i * chunk) + jobs.append( + ( + i, + t["key"], + t["source"]["shard"], + t["source"]["offset"] + i * chunk, + t["offset"] + i * chunk, + n, + ) + ) + jobs.sort() # chunk-major, so every shard is read concurrently + flags = os.O_RDONLY | (os.O_DIRECT if method == "direct" else 0) + fds = {shard: os.open(shard, flags) for shard in {j[2] for j in jobs}} + view = memoryview(mm) + bounces: List[_BounceBuffer] = [] + base = _mmap_address(mm) if method == "direct" else 0 + local = ( + _PerThreadBounce(chunk + 2 * _DIRECT_ALIGN, bounces, threading.Lock()) + if method == "direct" + else None + ) + + def run(job: tuple) -> None: + idx, key, shard, src, dst, n = job + fd = fds[shard] + if method == "direct": + bounce = local.buffer + lo = src - src % _DIRECT_ALIGN + need = src + n - lo + span = _round_up(need, _DIRECT_ALIGN) + got = 0 + while got < need: + r = os.preadv(fd, [bounce.view[got:span]], lo + got) + if r <= 0: + raise OSError(f"short O_DIRECT read at byte {lo + got} of {shard}") + got += r + ctypes.memmove(base + dst, bounce.address + (src - lo), n) + else: + _read_buffered(fd, view[dst : dst + n], src) + os.posix_fadvise(fd, src, n, os.POSIX_FADV_DONTNEED) + crcs[key][idx] = zlib.crc32(view[dst : dst + n]) + + try: + with ThreadPoolExecutor(max_workers=threads) as pool: + for _ in pool.map(run, jobs): + pass + finally: + for fd in fds.values(): + os.close(fd) + for bounce in bounces: + bounce.release() + view.release() + return crcs + + +def _chunks_to_verify(manifest: dict, mode: str) -> List[Tuple[dict, int]]: + if mode not in ("none", "sample", "full"): + raise ValueError(f"verify must be 'none', 'sample' or 'full', got {mode!r}") + picks = [] + for t in manifest["tensors"]: + count = len(manifest["crc32"][t["key"]]) + if mode == "full": + idxs = range(count) + elif mode == "sample": + idxs = sorted( + { + round(i * (count - 1) / max(1, _SAMPLE_CHUNKS - 1)) + for i in range(min(count, _SAMPLE_CHUNKS)) + } + ) + else: + idxs = [] + picks.extend((t, i) for i in idxs) + return picks + + +def _verify(path: str, manifest: dict, mode: str, threads: int) -> Tuple[int, List[str]]: + """(chunks checked, chunks whose CRC differs from the manifest).""" + picks = _chunks_to_verify(manifest, mode) + if not picks: + return 0, [] + chunk = manifest["chunk_bytes"] + fd = os.open(path, os.O_RDONLY) + try: + mm = mmap.mmap(fd, manifest["total_bytes"], flags=mmap.MAP_SHARED, prot=mmap.PROT_READ) + finally: + os.close(fd) + view = memoryview(mm) + + def crc_differs(pick: Tuple[dict, int]) -> Optional[str]: + t, i = pick + lo = t["offset"] + i * chunk + hi = t["offset"] + min((i + 1) * chunk, t["nbytes"]) + return ( + None if zlib.crc32(view[lo:hi]) == manifest["crc32"][t["key"]][i] else f"{t['key']}#{i}" + ) + + try: + with ThreadPoolExecutor(max_workers=threads) as pool: + bad = [b for b in pool.map(crc_differs, picks) if b is not None] + finally: + view.release() + mm.close() + return len(picks), bad + + +def _create_and_fill( + path: str, layout: dict, threads: int, method: str, numa: str +) -> Dict[int, int]: + total = layout["total_bytes"] + fd = os.open(path, os.O_RDWR | os.O_CREAT | os.O_EXCL, 0o600) + try: + os.ftruncate(fd, total) + mm = mmap.mmap(fd, total, flags=mmap.MAP_SHARED, prot=mmap.PROT_READ | mmap.PROT_WRITE) + except BaseException: + os.close(fd) + os.unlink(path) + raise + os.close(fd) + try: + base = _mmap_address(mm) + mode, nodes = _parse_numa_policy(numa) + for t in layout["tensors"]: + length = _round_up(t["nbytes"], mmap.PAGESIZE) + try: + mm.madvise(mmap.MADV_HUGEPAGE, t["offset"], length) + except OSError: + pass # a kernel without THP; the table is served from base pages + if mode is not None: + _mbind(base + t["offset"], length, mode, nodes) + _write_header(mm, _STATE_FILLING, layout) + start = time.monotonic() + crcs = _parallel_fill(mm, layout, threads, method) + seconds = time.monotonic() - start + numa_bytes = _mapping_numa_bytes(base, base + total) + manifest = dict(layout) + manifest["crc32"] = crcs + manifest["fill"] = { + "host": socket.gethostname(), + "pid": os.getpid(), + "unix_time": time.time(), + "seconds": seconds, + "threads": threads, + "method": method, + "numa": numa, + "numa_bytes": {str(k): v for k, v in numa_bytes.items()}, + "huge_page_bytes": _mapping_smaps(base, base + total)["ShmemPmdMapped"], + } + _write_header(mm, _STATE_READY, manifest) + except BaseException: + mm.close() + try: + os.unlink(path) + except FileNotFoundError: + pass + raise + mm.close() + return numa_bytes + + +def fill_shared_engram_table( + checkpoint_dir: str, + path: Optional[str] = None, + *, + threads: Optional[int] = None, + method: str = "buffered", + numa: str = "interleave", + verify: str = "sample", + headroom_bytes: int = 64 << 30, + lock_timeout_s: float = 1800.0, + owned: bool = False, +) -> FillReport: + """Make ``path`` hold the checkpoint's Engram tables, filling it only if it does not already. + + Call on one process per node (local rank 0). A complete file whose + manifest names the same checkpoint bytes, and whose sampled chunk CRCs + still match, is reused. Otherwise the file is removed and refilled under + an exclusive lock, so the node never holds two copies. That waits for + processes still attached to the old file to detach, up to + ``lock_timeout_s``. A fill that fails removes its partial file. One that + is killed leaves the file marked incomplete, and the next call refills it. + + Args: + checkpoint_dir: the HF checkpoint directory (``model.safetensors.index.json``). + path: the shared file; ``default_table_path()`` if None. + threads: reader threads; ``min(64, cpu_count)`` if None. + method: ``buffered`` (``preadv`` into the mapping, then drop the shard's + page cache) or ``direct`` (``O_DIRECT`` into a per-thread buffer, then copy). + numa: page placement for a new file: ``interleave`` over the NUMA nodes + that have CPUs and memory, ``local`` (where the reader threads run), + or a node number. + verify: on reuse, re-check ``none``, a ``sample`` of 8 chunks per tensor, + or the ``full`` set of chunk CRCs. + headroom_bytes: host memory that must remain available after a new fill; + the fill refuses to start otherwise. + lock_timeout_s: how long to wait for the exclusive lock. + owned: when this call creates the file, mark it as the caller's (``.owner``) + so that ``close(remove_if_owned=True)`` removes it after its last holder + detaches; without it a new file carries no owner. A reused file keeps + whatever owner it had. + """ + path = path or default_table_path() + layout = _layout(find_engram_table_sources(checkpoint_dir)) + table_bytes = sum(t["nbytes"] for t in layout["tensors"]) + threads = threads or min(64, os.cpu_count() or 8) + start = time.monotonic() + lock_fd = _open_lock(path) + try: + _flock_within( + lock_fd, fcntl.LOCK_SH, lock_timeout_s, f"{path} is being filled by another process" + ) + for exclusive in (False, True): + manifest, reason = _check(path, layout) + if manifest is not None: + checked, bad = _verify(path, manifest, verify, threads) + if not bad: + return FillReport( + path, + True, + table_bytes, + time.monotonic() - start, + threads, + method, + numa, + checked, + {}, + ) + reason = f"CRC mismatch in chunks {bad[:4]}" + if exclusive: + break + fcntl.flock(lock_fd, fcntl.LOCK_UN) + _flock_within( + lock_fd, + fcntl.LOCK_EX, + lock_timeout_s, + f"{path} needs a refill ({reason}) but is still attached by {_lock_holders(path)}", + ) + if os.path.exists(path): + _logger.warning("replacing %s: %s", path, reason) + os.unlink(path) + _check_capacity(path, layout["total_bytes"], headroom_bytes) + numa_bytes = _create_and_fill(path, layout, threads, method, numa) + _write_owner(path, owned) + return FillReport( + path, False, table_bytes, time.monotonic() - start, threads, method, numa, 0, numa_bytes + ) + finally: + os.close(lock_fd) + + +def verify_shared_engram_table( + path: Optional[str] = None, *, full: bool = True, threads: Optional[int] = None +) -> Tuple[int, List[str]]: + """Recompute chunk CRCs of a complete table file: (chunks checked, chunks that differ).""" + path = path or default_table_path() + manifest, reason = _check(path, None) + if manifest is None: + raise RuntimeError(f"{path}: {reason}") + return _verify( + path, manifest, "full" if full else "sample", threads or min(64, os.cpu_count() or 8) + ) + + +class SharedEngramHostTable: + """A process's mapping of the node's shared Engram table file. + + Obtain one with ``attach_shared_engram_table`` or ``open_shared_engram_table``. + ``weight(layer)`` / ``scale(layer)`` are CPU tensor views of the mapping, + ``[rows, 256]`` float8_e4m3fn and ``[rows, 8]`` float8_e8m0fnu, that + ``engram_gather`` reads directly with ``row_start=0, row_end=rows``. Never + write through them: every rank on the node shares these bytes. The mapping + is read-only (a write faults) unless registering it needs write access, which + is the case where the device lacks ``cudaDevAttrHostRegisterReadOnlySupported`` + (GB300 with driver 580 does). The views stay valid until ``close()``. A CUDA + graph that captured a gather on them must be released before then. + + Register on the device's socket (``attach_shared_engram_table`` binds the + registering thread there by default). The GPU walks this process's page + tables, and with the process unbound the socket-0 GPUs of a GB300 node + gathered 1.45x slower than with it bound (measured). + """ + + def __init__( + self, + path: str, + manifest: dict, + mm: mmap.mmap, + lock_fd: int, + read_only: bool, + inode: int = 0, + ) -> None: + self.path = path + """str: the shared file.""" + self.manifest = manifest + """dict: the file's manifest (layout, source identity, chunk CRCs, fill record).""" + self.read_only = read_only + """bool: whether this process mapped the file read-only.""" + self.inode = inode + """int: inode of the file this process mapped.""" + self.register_seconds = 0.0 + """float: wall time of this process's ``cudaHostRegister`` calls.""" + self.register_cpus: Optional[set] = None + """set: CPUs the registering thread was bound to (its device's socket), or None if left unbound.""" + self.fill_report: Optional[FillReport] = None + """FillReport: what ``open_shared_engram_table`` did on local rank 0; None on the other ranks.""" + self._mm = mm + self._lock_fd = lock_fd + self._registered: List[int] = [] + self._closed = False + # One storage per tensor, starting at the tensor: Tensor.is_pinned() asks CUDA about the + # storage's first byte, and only the tensor regions are registered, not the header. + self._views: Dict[str, torch.Tensor] = {} + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) # torch warns when the buffer is read-only + for t in manifest["tensors"]: + region = torch.frombuffer( + mm, dtype=torch.uint8, count=t["nbytes"], offset=t["offset"] + ) + self._views[t["key"]] = region.view(_DTYPES[t["dtype"]]).view(*t["shape"]) + first = manifest["tensors"][0] + self._base = self._views[first["key"]].data_ptr() - first["offset"] + + @property + def layers(self) -> List[int]: + """Decoder layers that have a table, ascending.""" + return sorted({t["layer"] for t in self.manifest["tensors"]}) + + @property + def nbytes(self) -> int: + """Table bytes, all layers, weights and scales.""" + return sum(t["nbytes"] for t in self.manifest["tensors"]) + + def weight(self, layer: int) -> torch.Tensor: + """``[rows, 256]`` float8_e4m3fn host view of ``layers..engram.embed.weight``.""" + return self._views[f"layers.{layer}.engram.embed.weight"] + + def scale(self, layer: int) -> torch.Tensor: + """``[rows, 8]`` float8_e8m0fnu host view of ``layers..engram.embed.scale``.""" + return self._views[f"layers.{layer}.engram.embed.scale"] + + def register(self, device: Optional[torch.device] = None) -> None: + """``cudaHostRegister`` every table region of this mapping and check ``is_pinned()``. + + ``Portable | Mapped``, plus ``ReadOnly`` for a read-only mapping. The + CUDA context of ``device`` (the current device if None) performs the + registration; ``Portable`` makes it valid in every context of the process. + Registration pins (and so populates) every page of the table in this process. + """ + if self._registered: + return + from cuda.bindings import runtime as cudart + + flags = cudart.cudaHostRegisterPortable | cudart.cudaHostRegisterMapped + if self.read_only: + flags |= cudart.cudaHostRegisterReadOnly + device = _cuda_device(device) + start = time.monotonic() + with torch.cuda.device(device): + for key, view in self._views.items(): + (err,) = cudart.cudaHostRegister(view.data_ptr(), view.numel(), flags) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError( + f"cudaHostRegister({key}, {view.numel()} bytes, read_only={self.read_only}) failed: {err}" + ) + self._registered.append(view.data_ptr()) + self.register_seconds = time.monotonic() - start + unpinned = [k for k, v in self._views.items() if not v.is_pinned()] + if unpinned: + raise RuntimeError(f"registered but not reported pinned: {unpinned}") + + def numa_bytes(self) -> Dict[int, int]: + """Resident bytes of this mapping per NUMA node (from /proc/self/numa_maps).""" + return _mapping_numa_bytes(self._base, self._base + self.manifest["total_bytes"]) + + def smaps(self) -> Dict[str, int]: + """This process's smaps totals for the mapping, in bytes (Rss, Pss, ShmemPmdMapped, ...).""" + return _mapping_smaps(self._base, self._base + self.manifest["total_bytes"]) + + def check_gather_views(self, device: torch.device) -> None: + """Raise unless every table view is page-locked and has a mapped device pointer on ``device``.""" + from cuda.bindings import runtime as cudart + + for key, view in self._views.items(): + if not view.is_pinned(): + raise RuntimeError(f"{key}: the shared table view is not page-locked") + with torch.cuda.device(device): + err, _ = cudart.cudaHostGetDevicePointer(view.data_ptr(), 0) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"{key}: cudaHostGetDevicePointer {err}") + + def describe(self) -> str: + """One log line: backend, file, fill or reuse, registration.""" + report, cpus = self.fill_report, self.register_cpus + return ( + f"backend=shm {self.path} {self.nbytes / 1e9:.2f} GB" + + ( + f", {'reused' if report.reused else 'filled'} in {report.seconds:.2f} s" + if report is not None + else "" + ) + + f", registered in {self.register_seconds:.2f} s on CPUs " + + (f"{min(cpus)}-{max(cpus)}" if cpus else "unbound") + + f", read_only={self.read_only}" + ) + + def close(self, remove_if_unused: bool = False, remove_if_owned: bool = False) -> None: + """Unregister, unmap and release this process's lock. + + Args: + remove_if_unused: also unlink the file if no other process is attached. + remove_if_owned: the same, but only when an engine's fill created the file + (``fill_shared_engram_table(owned=True)``); a file filled from the command + line is left for the next attach. + """ + if self._closed: + return + self._closed = True + if self._registered: + from cuda.bindings import runtime as cudart + + for ptr in self._registered: + cudart.cudaHostUnregister(ptr) + self._registered.clear() + self._views.clear() + try: + self._mm.close() + except BufferError: + _logger.warning( + "%s: a caller still holds a table view; the mapping ends with it", self.path + ) + fcntl.flock(self._lock_fd, fcntl.LOCK_UN) + os.close(self._lock_fd) + if remove_if_unused or (remove_if_owned and _owned_inode(self.path) == self.inode): + remove_shared_engram_table(self.path) + + def __enter__(self) -> "SharedEngramHostTable": + return self + + def __exit__(self, *exc) -> None: + self.close() + + +def _cuda_device(device: Optional[torch.device]) -> torch.device: + device = torch.device("cuda") if device is None else torch.device(device) + return device if device.index is not None else torch.device("cuda", torch.cuda.current_device()) + + +def _read_only_registration_supported(device: torch.device) -> bool: + from cuda.bindings import runtime as cudart + + err, value = cudart.cudaDeviceGetAttribute( + cudart.cudaDeviceAttr.cudaDevAttrHostRegisterReadOnlySupported, device.index + ) + return err == cudart.cudaError_t.cudaSuccess and bool(value) + + +def attach_shared_engram_table( + path: Optional[str] = None, + *, + checkpoint_dir: Optional[str] = None, + device: Optional[torch.device] = None, + register: bool = True, + writable: Optional[bool] = None, + wait_s: float = 0.0, + bind_to_device_socket: bool = True, +) -> SharedEngramHostTable: + """Map a complete table file in this process and (by default) register it. + + Args: + path: the shared file; ``default_table_path()`` if None. + checkpoint_dir: if given, the file must have been filled from exactly + this checkpoint's bytes. + device: the device whose context registers the mapping (current if None). + register: ``cudaHostRegister`` the table regions (see ``SharedEngramHostTable.register``). + writable: map read-write. None: only when ``register`` is set and the + device cannot register read-only memory. + wait_s: keep retrying this long while the file is absent, being filled or incomplete. + bind_to_device_socket: run the registration on the CPUs of ``device``'s NUMA + node that the calling thread may already use (that thread only; its previous + affinity is restored). The registration populates this process's page tables + for the whole table, and the GPU translates every gathered row through them. + """ + path = path or default_table_path() + if register: + device = _cuda_device(device) + if writable is None: + writable = register and not _read_only_registration_supported(device) + layout = _layout(find_engram_table_sources(checkpoint_dir)) if checkpoint_dir else None + lock_fd = _open_lock(path) + deadline = time.monotonic() + wait_s + while True: + try: + fcntl.flock(lock_fd, fcntl.LOCK_SH | fcntl.LOCK_NB) + except BlockingIOError: + reason = "being filled" + else: + manifest, reason = _check(path, layout) + if manifest is not None: + break + fcntl.flock(lock_fd, fcntl.LOCK_UN) + if time.monotonic() >= deadline: + os.close(lock_fd) + raise RuntimeError(f"shared Engram table {path} is not usable: {reason}") + time.sleep(0.2) + try: + fd = os.open(path, os.O_RDWR if writable else os.O_RDONLY) + try: + inode = os.fstat(fd).st_ino + prot = mmap.PROT_READ | (mmap.PROT_WRITE if writable else 0) + mm = mmap.mmap(fd, manifest["total_bytes"], flags=mmap.MAP_SHARED, prot=prot) + finally: + os.close(fd) + except BaseException: + os.close(lock_fd) + raise + table = SharedEngramHostTable(path, manifest, mm, lock_fd, read_only=not writable, inode=inode) + if register: + previous = os.sched_getaffinity(0) + # Narrow, never widen: a thread its launcher confined elsewhere stays where it is. + cpus = (_device_socket_cpus(device) or set()) & previous if bind_to_device_socket else None + try: + if cpus: + os.sched_setaffinity(0, cpus) + table.register_cpus = cpus + table.register(device) + except BaseException: + table.close() + raise + finally: + if cpus: + os.sched_setaffinity(0, previous) + return table + + +def open_shared_engram_table( + checkpoint_dir: str, + *, + local_rank: int, + barrier: Optional[Callable[[], None]] = None, + path: Optional[str] = None, + device: Optional[torch.device] = None, + wait_s: float = 1800.0, + **fill_kwargs, +) -> SharedEngramHostTable: + """The weight-loading entry point: local rank 0 fills (or reuses), every local rank attaches. + + Dispatches on ``host_table_backend()``: ``vmm`` goes to ``open_vmm_engram_table`` + (``path`` and the fill keywords then serve only its shm fallback, which a VMM + create failure takes unless ``$TRTLLM_DSV41_ENGRAM_HOST_TABLE_BACKEND`` names + ``vmm`` explicitly); ``shm`` fills and attaches the ``/dev/shm`` file as described here. + + Args: + checkpoint_dir: the HF checkpoint directory. + local_rank: this process's rank on the node; rank 0 fills. + barrier: a node-local (or global) barrier, e.g. ``tensorrt_llm._utils.local_mpi_barrier``. + With it, the other ranks attach after rank 0 returns. Without it, + they poll the file for up to ``wait_s``. + path: the shared file; ``default_table_path()`` if None. + device: this rank's device (current if None). + wait_s: how long a non-filling rank waits without a barrier. + **fill_kwargs: forwarded to ``fill_shared_engram_table`` on rank 0. + """ + if host_table_backend() == "vmm": + shm_kwargs = {k: v for k, v in fill_kwargs.items() if k != "threads"} + return open_vmm_engram_table( + checkpoint_dir, + local_rank=local_rank, + barrier=barrier, + device=device, + wait_s=wait_s, + threads=fill_kwargs.get("threads"), + headroom_bytes=fill_kwargs.get("headroom_bytes", 64 << 30), + shm_fallback_path=None if _vmm_forced() else (path or default_table_path()), + shm_fill_kwargs=shm_kwargs, + ) + path = path or default_table_path() + report = None + if local_rank == 0: + report = fill_shared_engram_table(checkpoint_dir, path, **fill_kwargs) + _logger.info( + "Engram host table %s: %s %.1f GB in %.1f s", + path, + "reused" if report.reused else "filled", + report.nbytes / 1e9, + report.seconds, + ) + if barrier is not None: + barrier() + table = attach_shared_engram_table( + path, + checkpoint_dir=checkpoint_dir, + device=device, + wait_s=0.0 if barrier is not None else wait_s, + ) + table.fill_report = report + return table + + +# --------------------------------------------------------------------------- +# The vmm backend: CUDA VMM host-NUMA allocations shared by file descriptor. +# +# Why not the shm file: a gather through a registered tmpfs mapping is translated +# through the process's CPU page tables, and on a 64 KiB-page kernel whose shmem +# gets no huge pages (shmem_enabled=never, where MADV_COLLAPSE on shmem fails with +# EINVAL) a random-row gather over the 101 GB layer table is translation-bound: +# 3.1 ms at 4096 tokens x 24 lookups, against 0.14 ms for host-NUMA VMM memory, +# which the GPU maps in its own page tables at the allocation granularity. +# +# Sharing: the process that creates the table exports each slab's handle as a +# POSIX fd, and every process holding the table serves the fds over one abstract +# Unix socket per node, user and checkpoint layout (the first holder to bind the +# name serves; the others retry the bind every second, so the table stays +# reachable while any holder lives). A process imports the fds, maps them into its +# own address range and grants its device access. The physical memory lives until +# the last process that maps it exits. + +_VMM_SLAB_BYTES = 1 << 30 +_VMM_FD_BATCH = 64 +_VMM_PROTOCOL = 1 +_VMM_MIN_GRANULARITY = 1 << 21 +_VMM_RETRY_MAX_SLEEP_S = 10.0 + +CREATE_RETRY_ENV = "TRTLLM_DSV41_ENGRAM_VMM_CREATE_RETRY_S" +"""str: environment variable bounding how long a VMM table create retries on host out-of-memory (seconds).""" + + +class VmmCreateError(RuntimeError): + """The VMM table's host allocations could not be made (after the out-of-memory retries).""" + + +def vmm_create_retry_seconds() -> float: + """``$TRTLLM_DSV41_ENGRAM_VMM_CREATE_RETRY_S`` (default 120): how long ``_vmm_create`` keeps retrying + while the host is short of memory, for example while the previous engine's table is still being + released. 0 tries once.""" + value = os.environ.get(CREATE_RETRY_ENV) + if value is None: + return 120.0 + seconds = float(value) + if not math.isfinite(seconds) or seconds < 0: + raise ValueError( + f"{CREATE_RETRY_ENV} must be a non-negative number of seconds, got {value!r}" + ) + return seconds + + +def _vmm_forced() -> bool: + """Whether ``$TRTLLM_DSV41_ENGRAM_HOST_TABLE_BACKEND`` names ``vmm`` explicitly (no shm fallback then).""" + return os.environ.get(BACKEND_ENV, "").strip().lower() == "vmm" + + +def _numa_memory(nodes: Sequence[int]) -> Dict[int, Dict[str, int]]: + """``{node: {"MemFree": bytes, "FilePages": bytes}}`` from ``/sys/devices/system/node/node*/meminfo``.""" + out: Dict[int, Dict[str, int]] = {} + for node in nodes: + fields: Dict[str, int] = {} + try: + with open(f"/sys/devices/system/node/node{node}/meminfo") as f: + for line in f: + parts = line.split() # "Node 0 MemFree: 123456 kB" + if len(parts) >= 4 and parts[2].rstrip(":") in ("MemFree", "FilePages"): + fields[parts[2].rstrip(":")] = int(parts[3]) * 1024 + except OSError: + pass + out[node] = fields + return out + + +def _cu_mem_create(size: int, node: int): + """``cuMemCreate`` of ``size`` host bytes on NUMA ``node``: the raw (error, handle) result.""" + from cuda.bindings import driver as cu + + return cu.cuMemCreate(size, _vmm_prop(node), 0) + + +def default_vmm_socket_name(layout: dict) -> str: + """Abstract Unix socket name serving the node's VMM table for ``layout``. + + ``$TRTLLM_DSV41_ENGRAM_HOST_TABLE_PATH`` (else ``trtllm-dsv41-engram-``) + plus a hash of the layout identity, so two checkpoints never share a name. + """ + base = os.environ.get(PATH_ENV) or f"trtllm-dsv41-engram-{os.getuid()}" + ident = json.dumps(_identity(layout), sort_keys=True).encode() + return f"{base}.vmm.{zlib.crc32(ident):08x}" + + +def _cu_call(result, what: str): + """Unwrap a cuda-python driver result tuple, raising on an error code.""" + from cuda.bindings import driver as cu + + err = result[0] + if err != cu.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"{what} failed: {err}") + if len(result) == 1: + return None + return result[1] if len(result) == 2 else result[1:] + + +def _vmm_prop(node: int): + from cuda.bindings import driver as cu + + prop = cu.CUmemAllocationProp() + prop.type = cu.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.location.type = cu.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA + prop.location.id = node + prop.requestedHandleTypes = ( + cu.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + ) + return prop + + +def _visible_device_host_nodes() -> Dict[int, int]: + """``{device ordinal: host NUMA node}`` of this process's visible devices (``CU_DEVICE_ATTRIBUTE_HOST_NUMA_ID``; + -1 where the device reports none).""" + from cuda.bindings import driver as cu + + attr = getattr(cu.CUdevice_attribute, "CU_DEVICE_ATTRIBUTE_HOST_NUMA_ID", None) + out: Dict[int, int] = {} + for i in range(int(_cu_call(cu.cuDeviceGetCount(), "cuDeviceGetCount"))): + dev = _cu_call(cu.cuDeviceGet(i), "cuDeviceGet") + out[i] = ( + -1 + if attr is None + else int(_cu_call(cu.cuDeviceGetAttribute(attr, dev), "cuDeviceGetAttribute")) + ) + return out + + +def _vmm_usable_nodes(nodes: Sequence[int]) -> Tuple[List[int], Dict[int, str]]: + """The subset of ``nodes`` this process can place host-NUMA VMM memory on, and why each other node was refused. + + The driver refuses ``cuMemCreate(HOST_NUMA, node)`` with ``CUDA_ERROR_INVALID_VALUE`` for a + CPU socket that none of the process's visible devices is attached to (for example socket 1 + when only GPUs 0 and 1 are visible). Each node is probed with a one-granule allocation, which + is released at once. Needs a current CUDA context. + """ + from cuda.bindings import driver as cu + + flag = cu.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED + usable: List[int] = [] + refused: Dict[int, str] = {} + for node in nodes: + prop = _vmm_prop(node) + err, gran = cu.cuMemGetAllocationGranularity(prop, flag) + if err != cu.CUresult.CUDA_SUCCESS: + refused[node] = f"cuMemGetAllocationGranularity: {getattr(err, 'name', err)}" + continue + err, handle = cu.cuMemCreate(max(int(gran), _VMM_MIN_GRANULARITY), prop, 0) + if err != cu.CUresult.CUDA_SUCCESS: + refused[node] = f"cuMemCreate: {getattr(err, 'name', err)}" + continue + _cu_call(cu.cuMemRelease(handle), "cuMemRelease") + usable.append(node) + return usable, refused + + +def _vmm_granularity(nodes: Sequence[int]) -> int: + from cuda.bindings import driver as cu + + flag = cu.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED + gran = max( + int( + _cu_call( + cu.cuMemGetAllocationGranularity(_vmm_prop(n), flag), + "cuMemGetAllocationGranularity", + ) + ) + for n in nodes + ) + gran = max(gran, _VMM_MIN_GRANULARITY) + if _ALIGN % gran or _VMM_SLAB_BYTES % gran: + raise RuntimeError(f"VMM granularity {gran} does not divide the table alignment {_ALIGN}") + return gran + + +@dataclass(frozen=True) +class _VmmSlab: + """One ``cuMemCreate`` allocation: ``size`` bytes at ``offset`` of the table's address range, on NUMA ``node``.""" + + offset: int + size: int + node: int + + +def _vmm_plan(layout: dict, granularity: int, nodes: Sequence[int]) -> List[_VmmSlab]: + """Cover every tensor of ``layout`` with slabs of at most 1 GiB, alternating over ``nodes``. + + With both sockets usable, alternating 1 GiB slabs spread each table evenly over them, + which is what page interleaving does for the shm file: every GPU then reads half its + rows from its own socket. With one usable socket every slab lands on it. + """ + slabs: List[_VmmSlab] = [] + for t in layout["tensors"]: + start = t["offset"] + end = start + _round_up(t["nbytes"], granularity) + while start < end: + size = min(_VMM_SLAB_BYTES, end - start) + slabs.append(_VmmSlab(start, size, nodes[len(slabs) % len(nodes)])) + start += size + return slabs + + +def _access_desc(location_type: str, location_id: int, flags: str): + from cuda.bindings import driver as cu + + desc = cu.CUmemAccessDesc() + desc.location.type = getattr(cu.CUmemLocationType, f"CU_MEM_LOCATION_TYPE_{location_type}") + desc.location.id = location_id + desc.flags = getattr(cu.CUmemAccess_flags, f"CU_MEM_ACCESS_FLAGS_PROT_{flags}") + return desc + + +def _device_bytes(ptr: int, nbytes: int, device: torch.device) -> torch.Tensor: + """A ``[nbytes]`` uint8 CUDA tensor on ``device`` over memory this module owns (no allocation, no copy). + + Built from the pointer and the device as given: ``torch.as_tensor`` on a + ``__cuda_array_interface__`` would ask ``cudaPointerGetAttributes`` for the + device, which names no device for host-NUMA memory. + """ + storage = torch._C._construct_storage_from_data_pointer(ptr, device, nbytes) + view = torch.empty(0, dtype=torch.uint8, device=device).set_(storage, 0, (nbytes,), (1,)) + if view.data_ptr() != ptr or view.device != device: + raise RuntimeError( + f"view of {ptr:#x} landed at {view.data_ptr():#x} on {view.device}, expected {device}" + ) + return view + + +def _recv_exact(sock: socket.socket, n: int) -> bytes: + data = sock.recv(n) + if len(data) != n: + raise RuntimeError(f"Engram VMM table server sent {len(data)} bytes, expected {n}") + return data + + +def _vmm_fetch(name: str, deadline: float) -> Optional[Tuple[dict, List[int]]]: + """(message, fds) from the node's server for ``name``, or None if nobody serves it before ``deadline``.""" + while True: + sock = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET) + try: + sock.connect("\0" + name) + except (ConnectionRefusedError, FileNotFoundError): + sock.close() + if time.monotonic() >= deadline: + return None + time.sleep(0.2) + continue + fds: List[int] = [] + try: + sock.settimeout(120) + (length,) = struct.unpack(" None: + super().__init__(name="dsv41-engram-vmm-server", daemon=True) + self._table = table + self._stop_event = threading.Event() + self.attempted = threading.Event() + self.serving = False + + def run(self) -> None: + while not self._stop_event.is_set(): + sock = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET) + try: + sock.bind("\0" + self._table.path) + except OSError: + sock.close() + self.attempted.set() + self._stop_event.wait(1.0) + continue + try: + sock.listen(64) + sock.settimeout(0.5) + self.serving = True + self.attempted.set() + while not self._stop_event.is_set(): + try: + conn, _ = sock.accept() + except socket.timeout: + continue + with conn: + try: + self._table._serve(conn) + except OSError as error: + _logger.warning("Engram VMM table: serving a client failed: %s", error) + finally: + self.serving = False + sock.close() + + def stop(self) -> None: + self._stop_event.set() + self.join(timeout=5) + + +class VmmEngramHostTable: + """This process's mapping of the node's Engram tables in CUDA VMM host-NUMA memory. + + Obtain one with ``open_vmm_engram_table``. ``weight(layer)`` / ``scale(layer)`` + are CUDA tensor views of host memory (``[rows, 256]`` float8_e4m3fn and ``[rows, + 8]`` float8_e8m0fnu) that ``engram_gather`` reads with ``row_start=0, + row_end=rows``; no device memory backs them. Never write through them: every + process on the node shares these bytes. The views stay valid until ``close()``, + and a CUDA graph that captured a gather on them must be released before then. + """ + + def __init__( + self, + name: str, + manifest: dict, + slabs: List[_VmmSlab], + granularity: int, + handles: list, + fds: List[int], + va: int, + device: torch.device, + created: bool, + ) -> None: + self.path = name + """str: the abstract Unix socket name the table is served on.""" + self.manifest = manifest + """dict: layout, source identity, chunk CRCs and the creator's fill record.""" + self.read_only = not created + """bool: whether this process mapped the table without host write access (every importer).""" + self.created = created + """bool: whether this process created and filled the table.""" + self.register_seconds = 0.0 + """float: wall time of this process's import, map and access calls.""" + self.register_cpus: Optional[set] = None + """None: the vmm backend binds no CPUs (kept for interface parity with the shm backend).""" + self.fill_report: Optional[FillReport] = None + """FillReport: what ``open_vmm_engram_table`` did on local rank 0; None on the other ranks.""" + self.device = device + self._slabs = slabs + self._granularity = granularity + self._handles = handles + self._fds = fds + self._va = va + self._closed = False + self._server: Optional[_VmmServer] = None + self._views: Dict[str, torch.Tensor] = {} + with torch.cuda.device(device): + for t in manifest["tensors"]: + raw = _device_bytes(va + t["offset"], t["nbytes"], device) + self._views[t["key"]] = raw.view(_DTYPES[t["dtype"]]).view(*t["shape"]) + + @property + def layers(self) -> List[int]: + """Decoder layers that have a table, ascending.""" + return sorted({t["layer"] for t in self.manifest["tensors"]}) + + @property + def nbytes(self) -> int: + """Table bytes, all layers, weights and scales.""" + return sum(t["nbytes"] for t in self.manifest["tensors"]) + + @property + def allocated_bytes(self) -> int: + """Host bytes of the slabs (the tensors rounded up to the allocation granularity).""" + return sum(s.size for s in self._slabs) + + def weight(self, layer: int) -> torch.Tensor: + """``[rows, 256]`` float8_e4m3fn CUDA view of ``layers..engram.embed.weight``.""" + return self._views[f"layers.{layer}.engram.embed.weight"] + + def scale(self, layer: int) -> torch.Tensor: + """``[rows, 8]`` float8_e8m0fnu CUDA view of ``layers..engram.embed.scale``.""" + return self._views[f"layers.{layer}.engram.embed.scale"] + + def host_buffer(self, key: str) -> memoryview: + """Read-only CPU bytes of one tensor (this process has host read access to every slab).""" + t = next(t for t in self.manifest["tensors"] if t["key"] == key) + return ( + memoryview((ctypes.c_char * t["nbytes"]).from_address(self._va + t["offset"])) + .cast("B") + .toreadonly() + ) + + def numa_bytes(self) -> Dict[int, int]: + """Allocated host bytes per NUMA node.""" + out: Dict[int, int] = {} + for s in self._slabs: + out[s.node] = out.get(s.node, 0) + s.size + return out + + def verify(self, full: bool = True, threads: Optional[int] = None) -> Tuple[int, List[str]]: + """Recompute chunk CRCs from this process's host view: (chunks checked, chunks that differ).""" + picks = _chunks_to_verify(self.manifest, "full" if full else "sample") + chunk = self.manifest["chunk_bytes"] + base = memoryview( + (ctypes.c_char * self.manifest["total_bytes"]).from_address(self._va) + ).cast("B") + + def crc_differs(pick: Tuple[dict, int]) -> Optional[str]: + t, i = pick + lo = t["offset"] + i * chunk + hi = t["offset"] + min((i + 1) * chunk, t["nbytes"]) + return ( + None + if zlib.crc32(base[lo:hi]) == self.manifest["crc32"][t["key"]][i] + else f"{t['key']}#{i}" + ) + + try: + with ThreadPoolExecutor(max_workers=threads or min(64, os.cpu_count() or 8)) as pool: + bad = [b for b in pool.map(crc_differs, picks) if b is not None] + finally: + base.release() + return len(picks), bad + + def check_gather_views(self, device: torch.device) -> None: + """Raise unless every table view is a CUDA view on ``device`` (what ``engram_gather`` must be handed: + a CPU view of VMM memory passes ``cudaHostGetDevicePointer`` and then gathers wrong rows).""" + for key, view in self._views.items(): + if not view.is_cuda or view.device != torch.device(device): + raise RuntimeError(f"{key}: expected a CUDA view on {device}, got {view.device}") + + def describe(self) -> str: + """One log line: backend, socket name, create or import, NUMA split, serving state.""" + report = self.fill_report + numa = ",".join(f"node{n}={b / 1e9:.1f}GB" for n, b in sorted(self.numa_bytes().items())) + return ( + f"backend=vmm {self.path} {self.nbytes / 1e9:.2f} GB in {len(self._slabs)} slabs ({numa})" + + ( + f", {'imported' if report.reused else 'created and filled'} in {report.seconds:.2f} s" + if report is not None + else "" + ) + + f", mapped in {self.register_seconds:.2f} s, {'creator' if self.created else 'importer'}" + + f", serving={bool(self._server and self._server.serving)}" + ) + + def _serve(self, conn: socket.socket) -> None: + conn.settimeout(60) + _, uid, _ = struct.unpack( + "3i", conn.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, struct.calcsize("3i")) + ) + if uid != os.getuid(): + return + body = json.dumps( + { + "protocol": _VMM_PROTOCOL, + "manifest": self.manifest, + "granularity": self._granularity, + "slabs": [[s.offset, s.size, s.node] for s in self._slabs], + }, + separators=(",", ":"), + ).encode() + conn.send(struct.pack(" None: + """Start (or keep) this process's server thread and wait for its first bind attempt.""" + if self._server is None: + self._server = _VmmServer(self) + self._server.start() + self._server.attempted.wait(wait_s) + + def close(self, remove_if_unused: bool = False, remove_if_owned: bool = False) -> None: + """Stop serving, unmap, release this process's handles and fds. + + The physical memory is freed when the last process that maps it closes or + exits. ``remove_if_unused`` / ``remove_if_owned`` are accepted for interface + parity with the shm backend and have nothing to remove. + """ + del remove_if_unused, remove_if_owned + if self._closed: + return + self._closed = True + if self._server is not None: + self._server.stop() + self._views.clear() + from cuda.bindings import driver as cu + + for s in self._slabs: + cu.cuMemUnmap(self._va + s.offset, s.size) + cu.cuMemAddressFree(self._va, self.manifest["total_bytes"]) + for h in self._handles: + cu.cuMemRelease(h) + self._handles = [] + for fd in self._fds: + os.close(fd) + self._fds = [] + + def __enter__(self) -> "VmmEngramHostTable": + return self + + def __exit__(self, *exc) -> None: + self.close() + + +def _vmm_map( + slabs: List[_VmmSlab], handles: list, total_bytes: int, device: torch.device, host_write: bool +) -> int: + """Reserve one address range, map every slab at its offset and grant ``device`` and the host access.""" + from cuda.bindings import driver as cu + + va = int(_cu_call(cu.cuMemAddressReserve(total_bytes, _ALIGN, 0, 0), "cuMemAddressReserve")) + mapped = [] + try: + for s, h in zip(slabs, handles): + _cu_call(cu.cuMemMap(va + s.offset, s.size, 0, h, 0), "cuMemMap") + mapped.append(s) + # The creator writes through the host mapping; an importer asks for read-only access and + # takes the first combination the driver accepts. + if host_write: + choices = [("READWRITE", "READWRITE")] + else: + choices = [("READ", "READ"), ("READWRITE", "READ"), ("READWRITE", "READWRITE")] + descs = [ + [_access_desc("DEVICE", device.index, d), _access_desc("HOST", 0, h)] + for d, h in choices + ] + for s in slabs: + for i, pair in enumerate(descs): + result = cu.cuMemSetAccess(va + s.offset, s.size, pair, len(pair)) + if result[0] == cu.CUresult.CUDA_SUCCESS: + break + if i == len(descs) - 1: + _cu_call(result, f"cuMemSetAccess(device {device.index}, host)") + except BaseException: + for s in mapped: + cu.cuMemUnmap(va + s.offset, s.size) + cu.cuMemAddressFree(va, total_bytes) + raise + return va + + +def _vmm_allocate( + slabs: List[_VmmSlab], ctx, threads: int, headroom_bytes: int +) -> Tuple[list, int]: + """``cuMemCreate`` every slab; (handles, attempts). + + Host memory can be short for a while after another process released a table + (its 203 GB are returned asynchronously, and VMM host memory is not charged to + the cgroup), so a shortfall -- ``MemAvailable`` below the table plus + ``headroom_bytes``, or ``CUDA_ERROR_OUT_OF_MEMORY`` from any slab -- releases + whatever was allocated and retries with backoff for ``vmm_create_retry_seconds()``. + Free memory per NUMA node is logged before the first attempt and after every + shortfall; it is not a gate, since clean page cache counts against ``MemFree`` + and ``cuMemCreate`` reclaims it. Raises ``VmmCreateError`` when the time is up + or on any other ``cuMemCreate`` error. + """ + from cuda.bindings import driver as cu + + allocated = sum(s.size for s in slabs) + need = {n: sum(s.size for s in slabs if s.node == n) for n in sorted({s.node for s in slabs})} + retry_s = vmm_create_retry_seconds() + deadline = time.monotonic() + retry_s + + def memory_line() -> str: + numa = _numa_memory(list(need)) + cgroup = _cgroup_available() + return ( + "need " + + ", ".join(f"node{n} {b / 1e9:.1f} GB" for n, b in need.items()) + + "; MemFree " + + ", ".join(f"node{n} {m.get('MemFree', 0) / 1e9:.1f} GB" for n, m in numa.items()) + + " (FilePages " + + ", ".join(f"node{n} {m.get('FilePages', 0) / 1e9:.1f} GB" for n, m in numa.items()) + + f"); MemAvailable {_meminfo().get('MemAvailable', 0) / 1e9:.1f} GB" + + f"; cgroup {'unlimited' if cgroup is None else f'{cgroup / 1e9:.1f} GB'}" + ) + + _logger.info("Engram VMM table create: %s", memory_line()) + attempts = 0 + while True: + attempts += 1 + available = _meminfo().get("MemAvailable", 0) + cgroup = _cgroup_available() + limit = available if cgroup is None else min(available, cgroup) + handles: list = [None] * len(slabs) + failure: List[str] = [] + stop = threading.Event() + + def create(i: int) -> None: + if stop.is_set(): + return + _cu_call(cu.cuCtxSetCurrent(ctx), "cuCtxSetCurrent") + err, handle = _cu_mem_create(slabs[i].size, slabs[i].node) + if err == cu.CUresult.CUDA_SUCCESS: + handles[i] = handle + return + stop.set() + failure.append( + f"cuMemCreate({slabs[i].size}, node {slabs[i].node}): {getattr(err, 'name', err)}" + ) + if err != cu.CUresult.CUDA_ERROR_OUT_OF_MEMORY: + raise VmmCreateError(failure[-1]) + + short = None + if limit < allocated + headroom_bytes: + short = ( + f"{allocated / 1e9:.1f} GB plus {headroom_bytes / 1e9:.1f} GB headroom needed, " + f"{limit / 1e9:.1f} GB available" + ) + else: + try: + with ThreadPoolExecutor(max_workers=min(16, threads)) as pool: + for _ in pool.map(create, range(len(slabs))): + pass + except BaseException: + for h in handles: + if h is not None: + cu.cuMemRelease(h) + raise + if not failure: + return handles, attempts + for h in handles: + if h is not None: + cu.cuMemRelease(h) + short = failure[0] + remaining = deadline - time.monotonic() + if remaining <= 0: + raise VmmCreateError( + f"the Engram VMM table could not be allocated after {attempts} attempt(s) over {retry_s:.0f} s " + f"(${CREATE_RETRY_ENV}): {short}; {memory_line()}" + ) + sleep_s = min(_VMM_RETRY_MAX_SLEEP_S, float(2 ** (attempts - 1)), remaining) + _logger.warning( + "Engram VMM table create attempt %d: host memory short (%s); %s; retrying in %.0f s (%.0f s left)", + attempts, + short, + memory_line(), + sleep_s, + remaining, + ) + time.sleep(sleep_s) + + +def _vmm_create( + name: str, layout: dict, device: torch.device, threads: int, headroom_bytes: int +) -> Tuple[VmmEngramHostTable, float]: + """Allocate every slab, map it read-write for this process and the host, fill it from the checkpoint. + + Slabs alternate over the CPU sockets this process may allocate on (``_vmm_usable_nodes``): + both when it sees GPUs of both sockets, else only the socket of its GPUs. An importer on + the other socket maps the same slabs and reads them over the socket link. + """ + from cuda.bindings import driver as cu + + start = time.monotonic() + with torch.cuda.device(device): + torch.cuda.synchronize() # makes the device's primary context current on this thread + ctx = _cu_call(cu.cuCtxGetCurrent(), "cuCtxGetCurrent") + if not int(ctx): + ctx = _cu_call(cu.cuDevicePrimaryCtxRetain(device.index), "cuDevicePrimaryCtxRetain") + candidates = _cpu_memory_nodes() + _cu_call(cu.cuCtxPushCurrent(ctx), "cuCtxPushCurrent") + try: + nodes, refused = _vmm_usable_nodes(candidates) + finally: + _cu_call(cu.cuCtxPopCurrent(), "cuCtxPopCurrent") + device_nodes = _visible_device_host_nodes() + if not nodes: + raise VmmCreateError( + f"no CPU NUMA node of {candidates} accepts a host-NUMA VMM allocation from this process " + f"(visible devices' host nodes {device_nodes}; {refused}); set {BACKEND_ENV}=shm" + ) + if refused: + _logger.info( + "Engram VMM table: NUMA node(s) %s refuse host-NUMA allocations from this process (visible devices' " + "host nodes %s: %s); placing every slab on %s", + sorted(refused), + device_nodes, + refused, + nodes, + ) + granularity = _vmm_granularity(nodes) + slabs = _vmm_plan(layout, granularity, nodes) + handles, attempts = _vmm_allocate(slabs, ctx, threads, headroom_bytes) + try: + va = _vmm_map(slabs, handles, layout["total_bytes"], device, host_write=True) + except BaseException: + for h in handles: + cu.cuMemRelease(h) + raise + alloc_seconds = time.monotonic() - start + fds: List[int] = [] + try: + buf = memoryview((ctypes.c_char * layout["total_bytes"]).from_address(va)).cast("B") + fill_start = time.monotonic() + crcs = _parallel_fill(buf, layout, threads, "buffered") + fill_seconds = time.monotonic() - fill_start + buf.release() + posix_fd = cu.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + for h in handles: + fds.append( + int( + _cu_call( + cu.cuMemExportToShareableHandle(h, posix_fd, 0), + "cuMemExportToShareableHandle", + ) + ) + ) + except BaseException: + for fd in fds: + os.close(fd) + for s in slabs: + cu.cuMemUnmap(va + s.offset, s.size) + cu.cuMemAddressFree(va, layout["total_bytes"]) + for h in handles: + cu.cuMemRelease(h) + raise + manifest = dict(layout) + manifest["crc32"] = crcs + manifest["fill"] = { + "host": socket.gethostname(), + "pid": os.getpid(), + "unix_time": time.time(), + "backend": "vmm", + "alloc_seconds": alloc_seconds, + "seconds": fill_seconds, + "threads": threads, + "granularity": granularity, + "slab_bytes": _VMM_SLAB_BYTES, + "numa_bytes": {str(n): sum(s.size for s in slabs if s.node == n) for n in nodes}, + "numa_nodes_refused": {str(n): why for n, why in refused.items()}, + "device_host_nodes": {str(d): n for d, n in device_nodes.items()}, + "create_attempts": attempts, + } + table = VmmEngramHostTable( + name, manifest, slabs, granularity, handles, fds, va, device, created=True + ) + table.register_seconds = alloc_seconds + return table, alloc_seconds + fill_seconds + + +def _vmm_import( + name: str, message: dict, fds: List[int], device: torch.device +) -> VmmEngramHostTable: + """Import the served fds and map them for this process's device (read) and the host (read).""" + from cuda.bindings import driver as cu + + if message.get("protocol") != _VMM_PROTOCOL: + raise RuntimeError( + f"Engram VMM table server speaks protocol {message.get('protocol')}, expected {_VMM_PROTOCOL}" + ) + start = time.monotonic() + manifest = message["manifest"] + slabs = [_VmmSlab(*s) for s in message["slabs"]] + posix_fd = cu.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + handles = [] + try: + for fd in fds: + handles.append( + _cu_call( + cu.cuMemImportFromShareableHandle(fd, posix_fd), + "cuMemImportFromShareableHandle", + ) + ) + va = _vmm_map(slabs, handles, manifest["total_bytes"], device, host_write=False) + except BaseException: + for h in handles: + cu.cuMemRelease(h) + raise + table = VmmEngramHostTable( + name, manifest, slabs, message["granularity"], handles, fds, va, device, created=False + ) + table.register_seconds = time.monotonic() - start + return table + + +def open_vmm_engram_table( + checkpoint_dir: str, + *, + local_rank: int, + barrier: Optional[Callable[[], None]] = None, + device: Optional[torch.device] = None, + wait_s: float = 1800.0, + threads: Optional[int] = None, + headroom_bytes: int = 64 << 30, + name: Optional[str] = None, + shm_fallback_path: Optional[str] = None, + shm_fill_kwargs: Optional[dict] = None, +) -> "VmmEngramHostTable | SharedEngramHostTable": + """The vmm backend's weight-loading entry point; the same collective contract as ``open_shared_engram_table``. + + Local rank 0, under an exclusive lock per node, imports the table from a + process on this node that already serves it (another engine, for one), or + else creates and fills it; either way it serves the table from then on. + After ``barrier()`` every other local rank imports it from a server and + stands by to serve it too. Every rank maps the whole table. + + With ``shm_fallback_path``, a create that fails (``VmmCreateError``: host + memory still short after the retries, or no usable NUMA node) makes local + rank 0 warn and fill the ``shm`` table at that path instead, and every rank + then returns that ``SharedEngramHostTable``. If rank 0 fails outright it still + joins the barrier before raising, so the other ranks fail instead of hanging. + + Args: + checkpoint_dir: the HF checkpoint directory. + local_rank: this process's rank on the node. + barrier: a node-local (or global) barrier. Without it, the other ranks wait + up to ``wait_s`` for a server. + device: this rank's device (current if None). + wait_s: how long to wait for the per-node lock, and (without a barrier) for a server. + threads: fill threads; ``min(64, cpu_count)`` if None. + headroom_bytes: host memory that must remain available after a new table is allocated. + name: the abstract socket name; ``default_vmm_socket_name(layout)`` if None. + shm_fallback_path: the ``shm`` table file to fall back to; None: no fallback. + shm_fill_kwargs: keywords for ``fill_shared_engram_table`` on a fallback. + """ + device = _cuda_device(device) + layout = _layout(find_engram_table_sources(checkpoint_dir)) + name = name or default_vmm_socket_name(layout) + threads = threads or min(64, os.cpu_count() or 8) + table: Optional[VmmEngramHostTable] = None + shm_report: Optional[FillReport] = None + failure: Optional[Exception] = None + start = time.monotonic() + if local_rank == 0: + try: + lock_fd = _open_lock(default_table_path() + ".vmm") + try: + _flock_within( + lock_fd, + fcntl.LOCK_EX, + wait_s, + f"the Engram VMM table lock of {name} is held elsewhere", + ) + got = _vmm_fetch(name, time.monotonic() + 3.0) + if got is not None and _identity(got[0]["manifest"]) == _identity(layout): + table = _vmm_import(name, got[0], got[1], device) + reused = True + else: + if got is not None: + for fd in got[1]: + os.close(fd) + raise RuntimeError(f"{name} serves a table of different checkpoint bytes") + try: + table, _ = _vmm_create(name, layout, device, threads, headroom_bytes) + except VmmCreateError as error: + if shm_fallback_path is None: + raise + message = ( + f"Engram VMM table create failed ({error}); falling back to the shm backend at " + f"{shm_fallback_path}. Set {BACKEND_ENV}=vmm to fail instead." + ) + _logger.warning(message) + warnings.warn(message, RuntimeWarning, stacklevel=2) + reused = False + if table is not None: + table.start_serving() + finally: + os.close(lock_fd) + if table is None: + shm_report = fill_shared_engram_table( + checkpoint_dir, shm_fallback_path, threads=threads, **(shm_fill_kwargs or {}) + ) + else: + table.fill_report = FillReport( + name, + reused, + table.nbytes, + time.monotonic() - start, + threads, + "vmm", + "slab-interleave", + 0, + table.numa_bytes(), + ) + _logger.info("Engram host table %s", table.describe()) + except Exception as error: # joined with the barrier below, then re-raised + failure = error + if barrier is not None: + barrier() + if failure is not None: + raise failure + if local_rank == 0 and table is None: + shm = attach_shared_engram_table( + shm_fallback_path, checkpoint_dir=checkpoint_dir, device=device, wait_s=0.0 + ) + shm.fill_report = shm_report + return shm + if table is None: + got = None + if barrier is not None and shm_fallback_path is not None: + # Rank 0 serves before the barrier whenever it holds a VMM table, so one + # refused connect plus a complete shm table for these bytes means it fell back. + got = _vmm_fetch(name, time.monotonic()) + if got is None and _check(shm_fallback_path, layout)[0] is not None: + _logger.warning( + "Engram VMM table %s is not served; attaching the shm table %s", + name, + shm_fallback_path, + ) + return attach_shared_engram_table( + shm_fallback_path, checkpoint_dir=checkpoint_dir, device=device, wait_s=0.0 + ) + if got is None: + got = _vmm_fetch(name, time.monotonic() + (30.0 if barrier is not None else wait_s)) + if got is None: + raise RuntimeError(f"no process on this node serves the Engram VMM table {name}") + if _identity(got[0]["manifest"]) != _identity(layout): + for fd in got[1]: + os.close(fd) + raise RuntimeError(f"{name} serves a table of different checkpoint bytes") + table = _vmm_import(name, got[0], got[1], device) + table.start_serving(wait_s=0.0) + return table + + +def remove_shared_engram_table(path: Optional[str] = None, *, force: bool = False) -> bool: + """Unlink the table file if no process is attached (or regardless, with ``force``). + + Processes still attached keep a valid mapping after a forced unlink; the + memory is freed when the last of them detaches. Returns True if the file was unlinked. + """ + path = path or default_table_path() + lock_fd = _open_lock(path) + try: + try: + fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + if not force: + return False + try: + os.unlink(path) + except FileNotFoundError: + return False + try: + os.unlink(_owner_path(path)) + except FileNotFoundError: + pass + return True + finally: + os.close(lock_fd) + + +def describe_shared_engram_table(path: Optional[str] = None) -> dict: + """State of the table file: header state, sizes, resident bytes, fill record, attached lock holders.""" + path = path or default_table_path() + try: + state, manifest, size = _read_header(path) + st = os.stat(path) + except FileNotFoundError: + return {"path": path, "state": "absent"} + out = { + "path": path, + "state": { + None: "invalid", + _STATE_FILLING: "filling-or-abandoned", + _STATE_READY: "ready", + }.get(state, "unknown"), + "file_bytes": size, + "resident_bytes": st.st_blocks * 512, + "attached": _lock_holders(path), + "engine_owned": _owned_inode(path) == st.st_ino, + } + if manifest is not None: + out["tensors"] = [ + {k: t[k] for k in ("key", "shape", "dtype", "offset", "nbytes")} + for t in manifest["tensors"] + ] + out["fill"] = manifest.get("fill") + return out + + +def main(argv: Optional[Sequence[str]] = None) -> int: + """Command line: ``fill``, ``status``, ``verify``, ``remove``.""" + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + sub = parser.add_subparsers(dest="command", required=True) + fill = sub.add_parser("fill", help="fill (or reuse) the shared table") + fill.add_argument("--checkpoint", required=True) + fill.add_argument("--threads", type=int) + fill.add_argument("--method", default="buffered", choices=("buffered", "direct")) + fill.add_argument("--numa", default="interleave") + fill.add_argument("--verify", default="sample", choices=("none", "sample", "full")) + for name in ("status", "verify", "remove"): + sub.add_parser(name) + sub.choices["verify"].add_argument("--sample", action="store_true") + sub.choices["remove"].add_argument("--force", action="store_true") + for p in sub.choices.values(): + p.add_argument("--path") + args = parser.parse_args(argv) + if args.command == "fill": + report = fill_shared_engram_table( + args.checkpoint, + args.path, + threads=args.threads, + method=args.method, + numa=args.numa, + verify=args.verify, + ) + print(json.dumps(dict(asdict(report), gbps=report.gbps))) + elif args.command == "status": + print(json.dumps(describe_shared_engram_table(args.path), indent=1)) + elif args.command == "verify": + start = time.monotonic() + checked, bad = verify_shared_engram_table(args.path, full=not args.sample) + print(json.dumps({"checked": checked, "bad": bad, "seconds": time.monotonic() - start})) + return 1 if bad else 0 + else: + print(json.dumps({"removed": remove_shared_engram_table(args.path, force=args.force)})) + return 0 + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + raise SystemExit(main()) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/input_processor.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/input_processor.py new file mode 100644 index 000000000000..64c4cfd409da --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/input_processor.py @@ -0,0 +1,678 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DeepSeek-V4.1 multimodal input processor: image preprocessing and image-span expansion. + +This is the checkpoint's ``inference/image_processor.py`` (``prepare_vl_inputs``) as a TensorRT-LLM +input processor, plus the host-side tables the model's vision path indexes with. An image becomes +an ``n_vit_h x n_vit_w`` grid of 14x14 patches for the ViT and an ``n_llm_h x n_llm_w`` grid after +the aligner's 3x3 downsample, which the language model sees as the span + + [START] + ([IMAGE] * n_llm_w + [NEW_LINE]) * n_llm_h + [END] + +of ``n_llm_h * (n_llm_w + 1) + 2`` positions. Every span position carries ``image_token_id`` +(``<｜deepseek_image｜>``, 129264) in ``input_ids``; only the per-position type tells them apart. + +Preprocessing, exactly as the checkpoint does it: + +* an image under ``min_pixels`` (544 x 544) is scaled up to about that area, and both sides are + rounded up to multiples of the patch size; +* if the resulting span would exceed ``max_image_tokens`` (1024), the size shrinks to the largest + aspect-preserving one that fits; +* the image is then fitted into that size with ``ImageOps.pad`` (aspect-preserving resize, then + gray (127, 127, 127) padding). A plain resize to the target size is used only when + ``max_wh_ratio`` is set and the image is at least that wide; this checkpoint sets it to None, so + it never happens; +* pixels become ``(x / 255 - 0.5) / 0.5`` in fp32, then bf16, then row-major patches with the + channels first inside each patch. + +Text is never rendered here. A prompt arrives already rendered (the ``custom_tokenizer: +deepseek_v4`` chat renderer produces the V4.1 prompt byte-for-byte, and the placeholder registry +puts one ``<｜deepseek_image｜>`` per image before the text, joined by a blank line); this module +tokenizes it and expands each placeholder into its span. + +What the model receives: ``multimodal_data`` of a request with images +(``P`` = patches over all images of the request, ``C`` = aligner cells = IMAGE slots over all +images, ``S`` = span positions over all images, ``L`` = prompt length, ``N`` = images) + +``multimodal_data["image"]``, the vision path's inputs, in prompt order: + +* ``pixel_values``: bf16 ``[P, 3, 14, 14]``. ``patch_embed`` reads ``pixel_values.flatten(1)``. +* ``vit_grid_hw``: ``list[list[int]]``, ``[n_vit_h, n_vit_w]`` per image. +* ``llm_grid_hw``: ``list[list[int]]``, ``[n_llm_h, n_llm_w]`` per image. +* ``vit_cu_seqlens``: ``list[int]`` of length ``N + 1``, the patch offset of each image. ViT + attention is bidirectional within one image and never across images. +* ``vit_rope_cos`` / ``vit_rope_sin``: fp32 ``[P, 32]``, the ViT's 2D rotary table: per patch the + angles ``[row * inv_freq(16), col * inv_freq(16)]``, ``inv_freq = 1 / theta ** (arange(0, 32, + 2) / 32)``, applied rotate-half style to each 64-wide head. +* ``aligner_index``: int64 ``[C, 9]``. Row ``k`` lists, in ``(kh, kw)`` order, the ViT output rows + of the 3x3 block that aligner cell ``k`` reads; the value ``P`` is the zero padding the aligner + adds on the bottom and right edges. With ``vit_out`` ``[P, 1024]``, + ``cat([vit_out, zeros(1, 1024)])[aligner_index].transpose(1, 2).reshape(C, 9216)`` is exactly + the aligner's ``F.unfold`` input ``(channel, kh, kw)``, for all images at once. +* ``span_types``: int64 ``[S]``, the types of every span position (``IMAGE_START`` 0, + ``IMAGE`` 1, ``IMAGE_NEW_LINE`` 2, ``IMAGE_END`` 3). +* ``span_source_index``: int64 ``[S]``. With ``aligner_out`` ``[C, hidden]``, + ``cat([aligner_out, image_start[None], image_end[None], image_newline[None]])[span_source_index]`` + is the ``[S, hidden]`` span embedding: IMAGE slots take aligner rows in reading order and the + delimiters take the learned vectors (rows ``C``, ``C + 1``, ``C + 2``). +* ``span_starts`` / ``span_lengths``: ``list[int]``, each span's first prompt position and length. +* ``aligner_cu_cells``: ``list[int]`` of length ``N + 1``, the aligner-cell offset of each image. +* ``aligner_index_local`` / ``span_source_local``: the same two tables with every index local to its + own image, so any run of whole images can be assembled without the others: row ``k`` of image + ``i``'s cells indexes ``cat([vit_out_i, zeros(1, 1024)])`` (``P_i`` = that image's patch count is + the zero pad), and its span positions index ``cat([aligner_out_i, image_start, image_end, + image_newline])`` (delimiters at ``C_i``, ``C_i + 1``, ``C_i + 2``). + +``multimodal_data["token_layout"]``, per prompt position (index by absolute prompt position): + +* ``token_types``: int8 ``[L]``, ``TEXT`` (-1) outside spans, the span type inside. +* ``image_mask``: bool ``[L]``, True on every span position; it selects the MoE gate's ``bias_vl``. +* ``engram_dead_mask``: bool ``[L]``, True on every span position; Engram stores those positions as + DEAD (no n-gram crosses them, and their gate is 0). The reference's Engram ``token_mask`` is its + negation. + +Every ``image_token_id`` in a processed prompt is a span position and every span position is one, +so ``image_mask == engram_dead_mask == (input_ids == image_token_id)`` on the prompt. + +The span embedding is one row per span position, which is what the runtime's multimodal splice +expects: ``get_mm_token_ids`` is ``[image_token_id]`` and there are no id-level special tokens, so +``multimodal_embed_mask_cumsum`` counts every span position. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, ClassVar, Dict, List, Optional, Sequence, Tuple + +import numpy as np +import torch +from PIL import Image, ImageOps +from transformers import AutoTokenizer, PretrainedConfig + +from tensorrt_llm.inputs.content_format import ContentFormat +from tensorrt_llm.inputs.data import TextPrompt +from tensorrt_llm.inputs.registry import ( + BaseMultimodalInputProcessor, + DefaultInputProcessor, + ExtraProcessedInputs, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, +) +from tensorrt_llm.sampling_params import SamplingParams + +__all__ = [ + "IMAGE", + "IMAGE_END", + "IMAGE_NEW_LINE", + "IMAGE_PLACEHOLDER", + "IMAGE_START", + "INPUT_PROCESSOR", + "MODEL_TYPE", + "PLACEHOLDER_METADATA", + "PAD_COLOR", + "TEXT", + "DeepseekV41InputProcessor", + "ImageGrid", + "PreprocessedImage", + "VisionGeometry", + "build_vision_inputs", + "expand_image_placeholders", + "image_span_types", + "num_image_tokens", + "plan_image_grid", + "preprocess_image", + "to_rgb_image", + "vit_rope_table", +] + +# Per-position type codes, the checkpoint's own values. +TEXT = -1 +IMAGE_START, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(4) + +IMAGE_PLACEHOLDER = "<｜deepseek_image｜>" +PAD_COLOR = (127, 127, 127) + + +@dataclass(frozen=True) +class VisionGeometry: + """The vision numbers preprocessing and the tables depend on, read from the checkpoint config. + + Attributes: + patch_size: ViT patch side in pixels (``vision_config.patch_size``). + downsample_ratio: aligner block side in patches (``vision_config.downsample_ratio``). + max_image_tokens: largest span an image may cost (``vision_config.max_image_tokens``). + min_pixels: smaller images are scaled up to about this area (``vision_config.min_pixels``). + max_wh_ratio: plain-resize threshold; None pads every image (``vision_config.max_wh_ratio``). + hidden_size: ViT width (``vision_config.hidden_size``). + num_heads: ViT heads (``vision_config.num_attention_heads``). + num_layers: ViT depth; 0 means the checkpoint has no vision tower. + rope_theta: base of the ViT's 2D rotary table (``vision_config.rope_theta``). + image_token_id: the id every span position carries (top-level ``image_token_id``). + """ + + patch_size: int + downsample_ratio: int + max_image_tokens: int + min_pixels: int + max_wh_ratio: Optional[float] + hidden_size: int + num_heads: int + num_layers: int + rope_theta: float + image_token_id: int + + @classmethod + def from_config(cls, config: PretrainedConfig) -> "VisionGeometry": + """Read the geometry from a ``deepseek_v41`` composite config.""" + vision = config.vision_config + if isinstance(vision, dict): + vision = PretrainedConfig(**vision) + return cls( + patch_size=int(vision.patch_size), + downsample_ratio=int(vision.downsample_ratio), + max_image_tokens=int(vision.max_image_tokens), + min_pixels=int(vision.min_pixels), + max_wh_ratio=vision.max_wh_ratio, + hidden_size=int(vision.hidden_size), + num_heads=int(vision.num_attention_heads), + num_layers=int(vision.num_hidden_layers), + rope_theta=vision.rope_theta, + image_token_id=int(config.image_token_id), + ) + + @property + def rope_dim(self) -> int: + """Rotary angles per patch: half of one head's width.""" + return self.hidden_size // self.num_heads // 2 + + +@dataclass(frozen=True) +class ImageGrid: + """Where an image of a given size lands: its padded pixel size and both grids.""" + + n_llm_h: int + n_llm_w: int + best_height: int + best_width: int + n_vit_h: int + n_vit_w: int + + @property + def num_span_tokens(self) -> int: + return num_image_tokens(self.n_llm_h, self.n_llm_w) + + +@dataclass(frozen=True) +class PreprocessedImage: + """One image, ready for the ViT. + + Attributes: + patches: bf16 ``[n_vit_h * n_vit_w, 3, patch, patch]``, row-major, channels first. + grid: the image's grid. + """ + + patches: torch.Tensor + grid: ImageGrid + + +def num_image_tokens(n_llm_h: int, n_llm_w: int) -> int: + """Span length of an ``n_llm_h x n_llm_w`` token grid, delimiters included.""" + return n_llm_h * (n_llm_w + 1) + 2 + + +def _llm_grid(best_height: int, best_width: int, patch_size: int, downsample_ratio: int): + return math.ceil((best_height // patch_size) / downsample_ratio), math.ceil( + (best_width // patch_size) / downsample_ratio + ) + + +def _solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token): + """Largest aspect-preserving pixel size whose span still fits in ``max_n_token``.""" + r = height / width + max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5 + max_h_float = max_w_float * r + cell = patch_size * downsample_ratio + if max_w_float < 1.0: # very tall: a single column + return (max_n_token - 2) // 2 * cell, cell + if max_h_float < 1.0: # very wide: a single row + return cell, (max_n_token - 3) * cell + beta = min(math.floor(max_w_float) * cell / width, math.floor(max_h_float) * cell / height) + return ( + math.floor(height * beta / patch_size) * patch_size, + math.floor(width * beta / patch_size) * patch_size, + ) + + +def _safe_resize(height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token): + n_llm_h, n_llm_w = _llm_grid(best_height, best_width, patch_size, downsample_ratio) + if num_image_tokens(n_llm_h, n_llm_w) > max_n_token: + best_height, best_width = _solve_resize_ratio( + height, width, patch_size, downsample_ratio, max_n_token + ) + n_llm_h, n_llm_w = _llm_grid(best_height, best_width, patch_size, downsample_ratio) + if num_image_tokens(n_llm_h, n_llm_w) > max_n_token: + raise ValueError( + f"image grid {n_llm_h}x{n_llm_w} still exceeds {max_n_token} tokens after shrinking" + ) + return n_llm_h, n_llm_w, best_height, best_width + + +def plan_image_grid(width: int, height: int, geometry: VisionGeometry) -> ImageGrid: + """The grid an image of this original pixel size gets; depends on the size alone.""" + p = geometry.patch_size + if geometry.max_wh_ratio is not None and width > height * geometry.max_wh_ratio: + width = height * geometry.max_wh_ratio + if 0 < width * height < geometry.min_pixels: + ratio = (geometry.min_pixels / (width * height)) ** 0.5 + width = int(width * ratio) + height = int(height * ratio) + best_width = math.ceil(width / p) * p + best_height = math.ceil(height / p) * p + n_llm_h, n_llm_w, best_height, best_width = _safe_resize( + height, + width, + best_height, + best_width, + p, + geometry.downsample_ratio, + geometry.max_image_tokens, + ) + return ImageGrid( + n_llm_h=n_llm_h, + n_llm_w=n_llm_w, + best_height=best_height, + best_width=best_width, + n_vit_h=best_height // p, + n_vit_w=best_width // p, + ) + + +def to_rgb_image(image: Any) -> Image.Image: + """Bring an image from any runtime media format to an RGB ``PIL.Image``. + + Accepts a ``PIL.Image`` (converted with ``convert("RGB")``, as the checkpoint does), a + ``torch.Tensor`` ``[C, H, W]`` (float in ``[0, 1]``, the runtime's ``"pt"`` format, or uint8), or + a ``numpy.ndarray`` ``[H, W, C]`` of uint8 (the runtime's ``"np"`` format). Float tensors are + mapped back to the uint8 values they were made from, so the pixels are the decoder's. + """ + if isinstance(image, Image.Image): + return image.convert("RGB") + if isinstance(image, torch.Tensor): + tensor = image.detach().to("cpu") + if tensor.dim() != 3: + raise ValueError(f"image tensor must be [C, H, W], got shape {tuple(tensor.shape)}") + if tensor.is_floating_point(): + tensor = (tensor.float().clamp(0.0, 1.0) * 255.0).round().to(torch.uint8) + elif tensor.dtype != torch.uint8: + raise ValueError(f"integer image tensors must be uint8, got {tensor.dtype}") + array = tensor.permute(1, 2, 0).contiguous().numpy() + elif isinstance(image, np.ndarray): + if image.dtype != np.uint8 or image.ndim not in (2, 3): + raise ValueError( + f"numpy images must be uint8 [H, W] or [H, W, C], got {image.dtype} {image.shape}" + ) + array = image + else: + raise TypeError(f"unsupported image type {type(image).__name__}") + if array.ndim == 3 and array.shape[2] == 1: + array = array[:, :, 0] + return Image.fromarray(np.ascontiguousarray(array)).convert("RGB") + + +def _image_size(image: Any) -> Tuple[int, int]: + """(width, height) without decoding pixels.""" + if isinstance(image, Image.Image): + return image.width, image.height + if isinstance(image, torch.Tensor): + return int(image.shape[-1]), int(image.shape[-2]) + if isinstance(image, np.ndarray): + return int(image.shape[1]), int(image.shape[0]) + raise TypeError(f"unsupported image type {type(image).__name__}") + + +def preprocess_image(image: Any, geometry: VisionGeometry) -> PreprocessedImage: + """Fit one image into its grid and cut it into normalized bf16 patches.""" + image = to_rgb_image(image) + grid = plan_image_grid(image.width, image.height, geometry) + p = geometry.patch_size + target = (grid.best_width, grid.best_height) + if geometry.max_wh_ratio is not None and image.width >= geometry.max_wh_ratio * image.height: + image = image.resize(target) + else: + image = ImageOps.pad(image, target, color=PAD_COLOR) + x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255 + x = ((x - 0.5) / 0.5).to(torch.bfloat16) + patches = ( + x.reshape(3, grid.n_vit_h, p, grid.n_vit_w, p) + .permute(1, 3, 0, 2, 4) + .reshape(grid.n_vit_h * grid.n_vit_w, 3, p, p) + ) + return PreprocessedImage(patches=patches, grid=grid) + + +def image_span_types(n_llm_h: int, n_llm_w: int) -> torch.Tensor: + """int64 ``[num_image_tokens]``: the span's per-position types in reading order.""" + types = [IMAGE_START] + types += ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h + types.append(IMAGE_END) + return torch.tensor(types, dtype=torch.int64) + + +def vit_rope_table(n_h: int, n_w: int, dim: int, theta: float) -> Tuple[torch.Tensor, torch.Tensor]: + """fp32 ``[n_h * n_w, dim]`` cos and sin of the ViT's 2D rotary angles, patches row-major.""" + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w) + wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w) + freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float() * inv_freq + freqs = freqs.flatten(1) + return freqs.cos(), freqs.sin() + + +def _aligner_index(grid: ImageGrid, ratio: int, patch_base: int) -> torch.Tensor: + """int64 ``[n_llm_h * n_llm_w, ratio**2]`` patch rows per aligner cell; -1 marks padding.""" + rows = torch.arange(grid.n_llm_h * ratio).view(grid.n_llm_h, 1, ratio, 1) + cols = torch.arange(grid.n_llm_w * ratio).view(1, grid.n_llm_w, 1, ratio) + index = patch_base + rows * grid.n_vit_w + cols + valid = (rows < grid.n_vit_h) & (cols < grid.n_vit_w) + index = torch.where(valid, index, torch.full_like(index, -1)) + return index.reshape(grid.n_llm_h * grid.n_llm_w, ratio * ratio) + + +def build_vision_inputs( + images: Sequence[PreprocessedImage], span_starts: Sequence[int], geometry: VisionGeometry +) -> Dict[str, Any]: + """The ``multimodal_data["image"]`` payload for a request's images, in prompt order.""" + pixel_values, cos, sin, aligner, span_types, span_image_rows = [], [], [], [], [], [] + aligner_local, span_local = [], [] + vit_cu_seqlens = [0] + aligner_cu_cells = [0] + num_cells = 0 + for image in images: + grid = image.grid + pixel_values.append(image.patches) + c, s = vit_rope_table(grid.n_vit_h, grid.n_vit_w, geometry.rope_dim, geometry.rope_theta) + cos.append(c) + sin.append(s) + aligner.append(_aligner_index(grid, geometry.downsample_ratio, vit_cu_seqlens[-1])) + num_image_patches = grid.n_vit_h * grid.n_vit_w + local = _aligner_index(grid, geometry.downsample_ratio, 0) + local[local < 0] = num_image_patches + aligner_local.append(local) + types = image_span_types(grid.n_llm_h, grid.n_llm_w) + span_types.append(types) + rows = torch.full_like(types, -1) + cells = grid.n_llm_h * grid.n_llm_w + rows[types == IMAGE] = torch.arange(num_cells, num_cells + cells) + span_image_rows.append(rows) + rows_local = torch.full_like(types, -1) + rows_local[types == IMAGE] = torch.arange(cells) + rows_local[types == IMAGE_START] = cells + rows_local[types == IMAGE_END] = cells + 1 + rows_local[types == IMAGE_NEW_LINE] = cells + 2 + span_local.append(rows_local) + num_cells += cells + vit_cu_seqlens.append(vit_cu_seqlens[-1] + num_image_patches) + aligner_cu_cells.append(num_cells) + + num_patches = vit_cu_seqlens[-1] + aligner_index = torch.cat(aligner) + aligner_index[aligner_index < 0] = num_patches + span_types_all = torch.cat(span_types) + span_source_index = torch.cat(span_image_rows) + span_source_index[span_types_all == IMAGE_START] = num_cells + span_source_index[span_types_all == IMAGE_END] = num_cells + 1 + span_source_index[span_types_all == IMAGE_NEW_LINE] = num_cells + 2 + return { + "pixel_values": torch.cat(pixel_values), + "vit_grid_hw": [[image.grid.n_vit_h, image.grid.n_vit_w] for image in images], + "llm_grid_hw": [[image.grid.n_llm_h, image.grid.n_llm_w] for image in images], + "vit_cu_seqlens": vit_cu_seqlens, + "vit_rope_cos": torch.cat(cos), + "vit_rope_sin": torch.cat(sin), + "aligner_index": aligner_index, + "span_types": span_types_all, + "span_source_index": span_source_index, + "span_starts": list(span_starts), + "span_lengths": [image.grid.num_span_tokens for image in images], + "aligner_cu_cells": aligner_cu_cells, + "aligner_index_local": torch.cat(aligner_local), + "span_source_local": torch.cat(span_local), + } + + +def expand_image_placeholders( + prompt_token_ids: Sequence[int], images: Sequence[Any], geometry: VisionGeometry +) -> Tuple[List[int], List[int], List[PreprocessedImage], List[int]]: + """Expand every placeholder token into its image's span, preprocessing images in order. + + Returns: + ``(token_ids, token_types, images, span_starts)``: the expanded ids, the per-position types + (``TEXT`` outside spans), the preprocessed images and each span's first position. + + Raises: + ValueError: the number of placeholder tokens is not the number of images, or the + checkpoint has no vision tower and images were given. + """ + image_token_id = geometry.image_token_id + num_placeholders = sum(token == image_token_id for token in prompt_token_ids) + if num_placeholders != len(images): + raise ValueError( + f"Found {num_placeholders} image placeholder tokens ({IMAGE_PLACEHOLDER}, id " + f"{image_token_id}) but got {len(images)} images" + ) + if num_placeholders and geometry.num_layers == 0: + raise ValueError( + "The model config has no vision tower (vision_config.num_hidden_layers == 0) but the " + "prompt contains images" + ) + tokens: List[int] = [] + token_types: List[int] = [] + processed: List[PreprocessedImage] = [] + span_starts: List[int] = [] + image_iter = iter(images) + for token in prompt_token_ids: + if token != image_token_id: + tokens.append(token) + token_types.append(TEXT) + continue + image = preprocess_image(next(image_iter), geometry) + types = image_span_types(image.grid.n_llm_h, image.grid.n_llm_w) + span_starts.append(len(tokens)) + processed.append(image) + tokens += [image_token_id] * types.numel() + token_types += types.tolist() + return tokens, token_types, processed, span_starts + + +class DeepseekV41InputProcessor(BaseMultimodalInputProcessor): + """Tokenizes a rendered DeepSeek-V4.1 prompt and expands its image placeholders. + + A prompt without images is tokenized exactly as ``DefaultInputProcessor`` does, so text-only + requests get the token ids they would get with no multimodal registration at all. + + Args: + model_path: checkpoint directory; the tokenizer is loaded from it when none is given. + config: the checkpoint's ``deepseek_v41`` config (``vision_config`` and ``image_token_id`` + are read). + tokenizer: the engine's tokenizer, returned unchanged by ``tokenizer`` so the engine keeps + its chat renderer; None (per-rank construction) loads the checkpoint's on first use. + trust_remote_code: passed to ``AutoTokenizer`` when the tokenizer is loaded here. + """ + + # The language model attends causally over image spans, so a span may be split across + # prefill chunks. + mm_bidirectional_blocks: ClassVar[bool] = False + # Expansion works on token ids, so prompt_token_ids + images never needs a detokenize. + supports_token_id_mm_expansion: ClassVar[bool] = True + + def __init__( + self, + model_path: str, + config: PretrainedConfig, + tokenizer: Any = None, + trust_remote_code: bool = True, + **kwargs, + ) -> None: + super().__init__( + model_path=model_path, + config=config, + tokenizer=tokenizer, + trust_remote_code=trust_remote_code, + **kwargs, + ) + self._geometry = VisionGeometry.from_config(config) + self._text_processor: Optional[DefaultInputProcessor] = None + self._placeholder_checked = False + + @property + def geometry(self) -> VisionGeometry: + return self._geometry + + @property + def config(self) -> PretrainedConfig: + return self._config + + @property + def tokenizer(self) -> Any: + if self._tokenizer is None: + self._tokenizer = AutoTokenizer.from_pretrained( + self._model_path, trust_remote_code=self._trust_remote_code + ) + return self._tokenizer + + @property + def model_path(self) -> str: + return self._model_path + + @property + def processor(self) -> None: + """The checkpoint ships no HF processor; this class is the whole preprocessing path.""" + return None + + @property + def dtype(self) -> torch.dtype: + """The checkpoint's dtype, and the dtype ``pixel_values`` are produced in.""" + return torch.bfloat16 + + def get_mm_token_ids(self) -> torch.Tensor: + return torch.tensor([self._geometry.image_token_id]) + + def get_mm_special_token_ids(self) -> None: + # START, NEW_LINE and END share the image id, so none is distinguishable by id; each takes + # an embedding row like an IMAGE slot does. + return None + + def get_num_tokens_per_image(self, *, image: Any, **kwargs) -> int: + """Span length of an image, delimiters included; depends only on its size.""" + width, height = _image_size(image) + return plan_image_grid(width, height, self._geometry).num_span_tokens + + def get_preferred_media_io_kwargs(self) -> Dict[str, Dict[str, Any]]: + """Per-modality media-IO decode defaults: none, so the server's own settings apply. + + ``OpenAIServer`` asks every ``BaseMultimodalInputProcessor`` for these when it starts. + The registry's default lives on ``BaseMultimodalDummyInputsBuilder``, which this + processor does not implement, so it answers here with that same empty default: + ``to_rgb_image`` takes the runtime's PIL, ``"pt"`` and ``"np"`` image formats alike. + """ + return {} + + def call_with_text_prompt( + self, inputs: TextPrompt, sampling_params: SamplingParams + ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: + images = self._images(inputs) + if not images: + if self._text_processor is None: + self._text_processor = DefaultInputProcessor(None, None, self.tokenizer) + token_ids, _ = self._text_processor(inputs, sampling_params) + self._check_no_placeholders(token_ids) + return token_ids, None + self._check_placeholder_token() + prompt_token_ids = self.tokenizer.encode( + inputs["prompt"], add_special_tokens=sampling_params.add_special_tokens + ) + return self._expand(prompt_token_ids, images, sampling_params) + + def call_with_token_ids( + self, inputs: TextPrompt, sampling_params: SamplingParams + ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: + return self._expand(list(inputs["prompt_token_ids"]), self._images(inputs), sampling_params) + + def _expand( + self, + prompt_token_ids: List[int], + images: List[Any], + sampling_params: SamplingParams, + ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: + tokens, token_types, processed, span_starts = expand_image_placeholders( + prompt_token_ids, images, self._geometry + ) + if not processed: + return tokens, None + limit = sampling_params.truncate_prompt_tokens + if limit is not None and len(tokens) > limit: + raise ValueError( + f"The prompt is {len(tokens)} tokens with its image spans expanded, more than " + f"truncate_prompt_tokens={limit}; truncating it would cut an image span" + ) + types = torch.tensor(token_types, dtype=torch.int8) + image_mask = types >= 0 + multimodal_data = { + "image": build_vision_inputs(processed, span_starts, self._geometry), + "token_layout": { + "token_types": types, + "image_mask": image_mask, + "engram_dead_mask": image_mask.clone(), + }, + } + return tokens, {"multimodal_data": multimodal_data} + + @staticmethod + def _images(inputs: TextPrompt) -> List[Any]: + mm_data = inputs.get("multi_modal_data") or {} + unsupported = sorted(k for k, v in mm_data.items() if k != "image" and v is not None) + if unsupported: + raise ValueError(f"DeepSeek-V4.1 takes images only, got {unsupported}") + images = mm_data.get("image") + if images is None: + return [] + return list(images) if isinstance(images, (list, tuple)) else [images] + + def _check_no_placeholders(self, token_ids: Sequence[int]) -> None: + count = sum(token == self._geometry.image_token_id for token in token_ids) + if count: + raise ValueError( + f"Found {count} image placeholder tokens ({IMAGE_PLACEHOLDER}, id " + f"{self._geometry.image_token_id}) but got 0 images" + ) + + def _check_placeholder_token(self) -> None: + """The placeholder string must tokenize to the config's image id.""" + if self._placeholder_checked: + return + placeholder_id = self.tokenizer.convert_tokens_to_ids(IMAGE_PLACEHOLDER) + unk_id = getattr(self.tokenizer, "unk_token_id", None) + if placeholder_id is not None and placeholder_id != unk_id: + if placeholder_id != self._geometry.image_token_id: + raise ValueError( + f"{IMAGE_PLACEHOLDER} is token {placeholder_id} in this tokenizer but the " + f"config's image_token_id is {self._geometry.image_token_id}" + ) + self._placeholder_checked = True + + +# What ``modeling_v2_input_processor`` registers when a target claims the +# checkpoint: the processor, the multimodal model type the placeholder registry +# keys it by, and the placement, which follows the checkpoint's own prompt +# encoder (``encoding/encoding.py``): one ``<｜deepseek_image｜>`` per image, all +# before the text, blocks joined by a blank line, the content handed to the chat +# renderer as a plain string. +MODEL_TYPE = "deepseek_v41" +PLACEHOLDER_METADATA = MultimodalPlaceholderMetadata( + placeholder_map={"image": IMAGE_PLACEHOLDER}, + placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, + placeholders_separator="\n\n", + content_format=ContentFormat.STRING, +) +INPUT_PROCESSOR = DeepseekV41InputProcessor diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/modeling.py new file mode 100644 index 000000000000..32239fed5ba2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/modeling.py @@ -0,0 +1,6150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: deepseek-v4.1-flash-nvfp4 / sm_103 / unified {tep4, dep4, tp4, tep8, dep8}. + +The core holds what the model statically is: geometry, weights, the pool views +and every table derived after load. A ``PrefillTarget`` and a ``DecodeTarget`` +each run one step over it (``_target.py``), and the core routes every step by +``phase_of``: decode when the step has no context request, prefill otherwise +(a mixed batch is prefill's). Every trtllm op a target calls is an entry -- the +shared ``catalog/`` or this target's own ``ops/`` -- instantiated by the target +and bound once to its weights; native torch is called directly. + +DeepSeek-V4.1-Flash (``DeepseekV41ForCausalLM``), text path: 40 blocks of +hidden 5120 over a 4-copy hyper-connection stream (``hc_mult`` 4), each block +an attention sublayer and a MoE sublayer (384 routed NVFP4 experts at top-6 +plus one FP8 shared expert), with Engram n-gram memory added into the stream +before blocks 1 and 14. + +**Images (the checkpoint's ``vision.py`` / ``model.py`` VL path).** A context +chunk of a request that carries images (``multimodal_params``, built by the +``deepseek_v41`` input processor) takes its image-span rows from +``VisionTower``: the 32-block DeepSeek-ViT over each image's 14x14 patches +(bidirectional attention within one image), the 3x3 aligner and the learned +span delimiters. Those rows replace the token embedding at the span positions +(every span position carries id 129264). On the same rows the MoE gate selects +experts with ``gate.bias_vl`` instead of ``gate.bias`` (two routing calls and a +per-token select), and Engram treats them as DEAD: no n-gram hashes across +them and their Engram gate is 0. Generated tokens are text. The tower is +replicated on every rank and runs one image per call, only for the images a +chunk's rows reach; the context step itself is never captured, but each tower +call replays a CUDA graph captured at load for its padded patch count +(``TRTLLM_DSV41_VIT_CUDA_GRAPH``, unset = on; see ``VisionTower``). + +**DSpark (``speculative_config`` DSpark, block 5).** The +checkpoint's three ``mtp.*`` stages are a block drafter, not chained MTP heads. +With it on, the trunk hands the runtime's ``DSparkSpecMetadata`` the INPUT +stream of layers 37-39 (the metadata averages the 4 hc copies), generation +rows carry 1 + 5 tokens, and the shell drives ``DSv4DSparkWorker`` with +``DSparkDraftModel`` below: per step it writes the bonus position's latent into +each stage's 128-row ring (the worker owns the rings), runs the three stages +over the 5-row block ``[bonus, noise x 4]`` -- each stage's attention reads its +ring plus the whole block, non-causal -- projects through the shell's +``lm_head`` and applies the sequential Markov correction with greedy feedback. +The drafter is replicated on every rank of every member and has no collective +of its own; its logits leave the vocab-parallel ``lm_head`` gathered, so every +TP rank drafts the same tokens. Its one autotuned op, the routed MoE, runs in +the autotuner's ``fallback_tactic()`` scope, so every rank takes the same +untuned tactic whatever its tuner cache holds. Verification is full-block +(static): the confidence head is not computed. + +**Encoder SWA bounded replay (switch ``TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY``, +resolved by the V4.1 cache manager: unset means on under attention TP -- +tep4, tp4 -- aggregated and without a KV connector, off elsewhere).** The +prefix cache keeps main KV and index K only. A hit of ``reused`` tokens (a +multiple of 2, the compress ratio) starts the request's first chunk +``R = min(128, reused)`` positions early: those R replay rows run every layer +and rebuild the sliding windows (a query at ``p`` attends window keys +``[max(reused - R, p - 127), p]`` and nothing below the replay start: the +context attention takes its window counts from the metadata's +``v41_swa_seq_lens_cuda``, kv length minus the SWA floor, since the kernel +reads a -1 inside the count as a zero K/V row, not as a mask), and they write +nothing else. The +compressors of the KV sources run over the other rows only, packed as a step +that starts at the reuse boundary, so no cached main-KV or index-K entry is +recomputed or overwritten; the Engram history of the three positions before +the replay start is stored first, from the prompt ids. A context request's +last row is always a suffix row, so replay rows produce no logits. Under +attention DP (dep4) each rank replays the hits of the requests it holds; its +replay rows are rows of its step like any other. Disaggregated serving never +replays (the manager refuses an explicit 1). + +**Decoder SWA bounded replay (switch ``TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY``).** +Layers ``0 .. L`` -- ``L`` the last KV-source layer, 20, whose ratio-1 +compressor and index K are the decoder's global KV -- run over every prompt +token. Layers ``L+1 .. 39`` run, in a request's final context chunk only, +over the last ``min(N, 128)`` prompt positions, their sliding window truncated +at ``r = N - min(N, 128)``: a query at ``p`` attends window keys +``[max(r, p - 127), p]`` and nothing below ``r`` (the bounded layers' metadata +view carries the tail's own floor ``r`` and window seq_lens ``N - r``). Main KV +and index K stay exact; the decoder layers' +window KV is the replay's. The layer-``L+1`` input of tail positions that +arrive in earlier chunks (hc stream, the lagged pre mixes, layer ``L``'s +top-k and index lengths, and layer 20's index query) is kept per request +between chunks, and dropped with the request once the cache manager no longer +holds it (a request cancelled or evicted between two of its chunks never +reaches its last one). Past 2048 x 8 visible positions the re-index layers +mask their scores with layer 20's candidate blocks: the decoder pass scores +its tail rows' masks from layer 20's carried query against the step's index +K, which is the mask a full prefill gives those rows; so any prompt length +the engine admits runs bounded. Only tail rows +exist past layer ``L``, so the step returns hidden rows for them and for the +generation rows (every other row is zero) and DSpark captures only tail rows. +Prompt logprobs and context logits are refused, as is a captured context +forward. Under attention DP (dep4) each rank bounds its own requests, and the +bounded layers' MoE is a collective: in every eager step the ranks exchange +their passes over the bounded layers (planned on the host before the layers +run: a bounded rank's slices, or its step rows as one pass) and every rank +runs as many passes as the rank with the most, each padded to that pass's +largest row count. A rank with fewer passes, or no decoder row, joins the +remaining ones with zero rows, and each pass gathers the router's image-span +mask again from its own rows. +Generation rows run every layer. On a prefix-cache hit with encoder replay on +the replay start lies at or below the tail's first position, so every tail row +is forwarded and the decoder layers' window KV is this request's own. With +encoder replay off, a hit inside the tail would keep its cached rows' decoder +window KV, which no forward of the request wrote; so with block reuse on, +decoder replay runs only beside encoder replay. The V4.1 cache manager resolves +that (an unset switch resolves off there with one line naming why; an explicit +1 raises at construction), and the model takes decoder replay as the manager +resolved it. + +**Attention.** 64 query heads over ONE shared 512-wide latent (K == V), RoPE +on the last 64 channels (adjacent pairs), inverse RoPE on the output, a +grouped low-rank output projection (8 groups x 1024) and fp32 per-head sinks. +Every layer attends a 128-token sliding window; layers 2-19 add ratio-2 +compressed KV (sources 2/8/14), layers 20-39 ratio-1 compressed KV (source +20). Compressed positions are chosen per query by an indexer (32 heads x 128, +MXFP4 q/k, top-512); layer 20 additionally publishes a candidate-block mask +(top-2048 blocks of 8) that the re-index layers 24/28/32/36 apply. The +top-512 is deterministic: the same set on every call and every rank, a tie at +the 512th score going to the lower position, in ascending position order, so +the sparse attention sums a row's columns in a fixed order. The per-layer role +table is asserted against the checkpoint at construction. + +**Hyper-connections (Single-Pass mHC).** ``post``/``comb`` are used by the +sublayer that computed them, but ``pre`` is one sublayer late: attention +collapses the stream with the previous FFN's mixes, the FFN with this +attention's, the final norm with layer 39's FFN mixes, and layer 0 with a +one-hot on copy 0 (which is the embedding itself). + +**Cache writes.** No V4.1 attention +op appends K/V -- the sparse kernel only reads. So this forward writes its +pools itself through scatter entries, at slots the runtime assigned +(``DeepseekV4TrtllmAttentionMetadata``'s V4.1 address book): the SWA window KV +(FP8-e4m3 x UE8M0/32), the compressed main KV (E2M1 x E4M3/16, no global +scale), the index K (native MXFP4), the ratio-2 compressor's pending pair +(fp32) and the Engram token history. It never allocates, moves or frees a page; +a negative slot means skip. + +**KV storage.** The window KV and the main KV are stored NATIVELY +by default: 528-byte rows (512 E4M3 codes + 16 UE8M0 bytes) and 288-byte rows +(256 bytes of packed E2M1 + 32 E4M3 scales), so the persistent KV is 890 bytes +per token (3 x (288 + 68) / 2 + (288 + 68)). The sparse-attention kernel reads +bf16 rows only, so each step dequantizes exactly the rows it reads into two +bf16 buffers and hands the kernel those, with its index table renumbered onto +them: per layer, the 128 window rows of every query (row ``t * 128 + c`` is +column ``c`` of query ``t``; a -1 column becomes a zero row, which the kernel +treats as it treats -1), and the compressed rows either per KV source, every +compressed entry of every request of the step (row ``b * L + e``), or per +top-k set, every query's top-k entries (row ``t * 512 + k``) -- whichever +buffer is smaller for the engine's batch, ``max_seq_len`` and +``max_num_tokens``. Dequantizing a native row gives the +bf16 value its bf16-storage container held, bit for bit, so native and bf16 +storage produce the same tokens. ``TRTLLM_DSV41_KV_STORAGE=bf16`` keeps both +pools in bf16 containers of the same value sets (2730 bytes per token) and +the kernel reads the pools directly. + +**FP8 attention (opt-in, lossy).** ``TRTLLM_DSV41_FP8_ATTENTION=1`` (native +storage only) runs the sparse attention with FP8 tensor-core math: the read +buffers hold E4M3 rows with one power-of-two scale per layer (per compressed +source group; bounded from the norm weights so nothing saturates), the query is +quantized to E4M3 with a power-of-two scale from the step's max|q|, and the +same trtllm-gen kernel family runs its ``QkvE4m3OBfloat16`` variants, which +also round the probabilities to E4M3 for the PV product. The window's MXFP8 +values stay exact in E4M3 in its normal range; the query, the compressed +NVFP4 values (up to 6 significant bits) and the probabilities do not, so the +output departs from the reference's bf16 attention. The DSpark drafter's +attention stays bf16. + +**Parallel members, folded from the mapping at construction:** + +* ``tep4`` -- attention TP4 (16 heads, groups 2r/2r+1 per rank, fp32 + all-reduce after ``wo_b``), routed experts EP4 (96 per rank), fp32 + all-reduce of the routed partials; +* ``dep4`` -- attention DP4 (all 64 heads, each rank its own requests, no + attention collective), EP4 with the token all-gather ahead of the router and + an fp32 reduce-scatter after the experts; +* ``tp4`` -- attention TP4, routed experts MoE-TP4 (all 384, I = 576 per rank), + fp32 all-reduce of the routed partials; +* ``tep8`` / ``dep8`` -- the same over two NVL72 trays (8 ranks): attention TP8 + (8 heads, output group r per rank) or DP8, routed experts EP8 (48 per rank). + The IPC all-reduce workspace is single-node, so a TP group spanning the two + trays sums on NCCL, which runs over the NVLink fabric (MNNVL) between them. + +The indexer, compressor, hyper-connections, shared expert, embedding and every +KV pool are replicated. + +**Pipeline parallelism (``pp_size`` > 1).** Each stage runs its layers +(``mapping.pp_layers``) with any member above as its TP layout, on its own +pools: a stage holds every layer its layers read from (a compressed layer's KV +source and index source, the candidate source of the re-index layers, the +Engram history owner), which leaves this checkpoint's stage boundaries at +layers 1 and 20 (PP2's even split 0-19 | 20-39; construction raises +otherwise). A later stage receives the hc stream ``[T, 4, 5120]`` and the +previous stage's last FFN mixes (its first attention's lagged pre; that +layer's ``hc_ffn`` triple is loaded there too) through ``pp_recv_tensors``, and +an earlier one sends them with ``pp_send_tensors``. The embedding and Engram +run on the stage that holds them, the final norm, the lm_head, DSpark and the +decoder tail on the last stage. + +**Engram residency (host_offload).** By default the two Engram tables (202.76 +GB) stay in the node's one shared host copy +(the sibling ``engram_host_table.py``: CUDA VMM host-NUMA allocations +shared by file descriptor, or with ``TRTLLM_DSV41_ENGRAM_HOST_TABLE_BACKEND=shm`` +one registered ``/dev/shm`` file): every rank maps the whole table and gathers +its own tokens' rows through the mapping, so no rank holds any of the table in +device memory and no collective follows the gather. ``TRTLLM_DSV41_ENGRAM_HOST_OFFLOAD=0`` +row-shards the tables 4 ways on the devices instead: a bf16 all-reduce of the +zero-filled lookups (tep4/tp4), or all-gather of the hash ids + reduce-scatter +(dep4). The two give the same bits -- each looked-up value is a sum with one +non-zero term -- except that the collective turns a -0 row element into +0. +""" + +import atexit +import math +import os +import sys +import time +import weakref +from types import SimpleNamespace + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._target import Phase, Target, phase_of +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import advance_step_generation +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( # noqa: E501 + FlashinferSiluAndMul, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + ThopAttention, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import Allgather +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.reducescatter import Reducescatter +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.bmm_out import BmmOut +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import CublasMm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( # noqa: E501 + Fp4BlockScaleMoeRunner, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.mxe4m3_mxe2m1_block_scale_moe_runner import ( # noqa: E501 + Mxe4m3Mxe2m1BlockScaleMoeRunner, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + FlashinferRmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + Fp4Quantize, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.mxfp8_quantize import ( + Mxfp8Quantize, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import ( + DeepseekV4TrtllmAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import DeepseekV4AttentionType +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.v41_layout import ( + DEEPSEEK_V41_DECODER_REPLAY_ENV, + DEEPSEEK_V41_KV_STORAGE_BF16, + DEEPSEEK_V41_KV_STORAGE_NATIVE, + DeepseekV41Layout, + use_deepseek_v41_layout_mode, +) +from tensorrt_llm._torch.autotuner import fallback_tactic +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_utils import ( + DecoderModel, + DecoderModelForCausalLM, + register_auto_model, +) +from tensorrt_llm._torch.speculative import get_spec_worker +from tensorrt_llm._utils import local_mpi_barrier, local_mpi_comm, local_mpi_rank + +from . import engram_host_table as _engram_host +from . import weights as _weights +from .ops.activation.moe_gelu import MoeGelu +from .ops.activation.silu_and_mul import SiluAndMul +from .ops.attention.candidate_block_select import CandidateBlockSelect +from .ops.attention.compressor_paged_kv_compress import CompressorPagedKvCompress +from .ops.attention.compressor_prefill_reduction import CompressorPrefillReduction +from .ops.attention.deepgemm_fp8_fp4_mqa_logits import DeepgemmFp8Fp4MqaLogits +from .ops.attention.deepgemm_fp8_fp4_paged_mqa_logits import ( + DeepgemmFp8Fp4PagedMqaLogits, + DeepgemmPagedMqaLogitsMetadata, +) +from .ops.attention.dsv4_local_to_global_indices import Dsv4LocalToGlobalIndices +from .ops.attention.flashinfer_top_k_ragged_deterministic import FlashinferTopKRaggedDeterministic +from .ops.attention.indexer_k_cache_gather import IndexerKCacheGather +from .ops.attention.indexer_k_cache_scatter import IndexerKCacheScatter +from .ops.attention.mla_rope_inplace import MlaRopeInplace +from .ops.attention.trtllm_sparse_mla_dsv4 import TrtllmSparseMlaDsv4 +from .ops.attention.trtllm_sparse_mla_dsv4_fp8 import TrtllmSparseMlaDsv4Fp8 +from .ops.comm.allreduce import Allreduce +from .ops.comm.pp_recv_tensors import PpRecvTensors +from .ops.comm.pp_send_tensors import PpSendTensors +from .ops.embedding.engram_gather import EngramGather +from .ops.embedding.engram_hash import EngramHash +from .ops.gemm.mhc_gemm_sqrsum_fma import MhcGemmSqrsumFma +from .ops.gemm.mxfp8_mxfp8_gemm import Mxfp8Mxfp8Gemm +from .ops.moe.gate_forward import GateForward +from .ops.norm.engram_gate import EngramGate +from .ops.norm.mhc_big_fuse import MhcBigFuse +from .ops.norm.mhc_hc_head_apply import MhcHcHeadApply +from .ops.norm.mhc_post_mapping import MhcPostMapping +from .ops.quantization.e4m3_quantize_pow2_dynamic import E4m3QuantizePow2Dynamic +from .ops.quantization.mxfp8_dequantize import Mxfp8Dequantize +from .ops.quantization.mxfp8_fake_quantize import Mxfp8FakeQuantize +from .ops.quantization.mxfp8_to_e4m3 import Mxfp8ToE4m3 +from .ops.quantization.nvfp4_fake_quantize import Nvfp4FakeQuantize +from .ops.quantization.nvfp4_kv_to_bf16 import Nvfp4KvToBf16 +from .ops.quantization.nvfp4_kv_to_e4m3 import Nvfp4KvToE4m3 + +_HC = _weights.HC # hyper-connection copies +_MIX = _weights.MIX # 24 = 4 pre + 4 post + 16 comb +_WINDOW = 128 +_TOPK = 512 +_ROPE = 64 +_IDX_HEADS = _weights.IDX_HEADS +_IDX_DIM = _weights.IDX_DIM +_PAGE = 128 # tokens_per_block of every V4.1 pool +_CAND_BLOCK = 8 +_CAND_TOPK_BLOCKS = 2048 +_SWIGLU_LIMIT = 10.0 +# NVFP4 block (one e4m3 scale per 16 along K) / MXFP4 index-K block (UE8M0 per 32). +_SF16 = 16 +_SF32 = 32 +# The MoE runner reads its activation scales LINEAR; the swizzled buffer has +# the same byte count at T % 128 == 0 (every graph batch) and is silently wrong. +_SF_LINEAR = False +# Pre-routed path of the trtllm-gen block-scale runner: routing inert. +_ROUTING_METHOD_INERT = 1 +_ACT_SWIGLU = 0 +# The runner's and the routing op's token columns stop at 8192. +_MOE_MAX_T = 8192 # capacity: fixed: the MoE runner's token columns; larger steps chunk +# Rows the index top-k's zero offsets cover at load. The first forward grows +# them to the engine's step capacity (max_num_tokens): a step with more rows +# than offsets trips the top-k wrapper's one-entry-per-row assert. +_TOPK_MAX_ROWS = 8192 # capacity: max_num_tokens (grown on the first step) +# fp32 logits elements (1 GiB) one prefill index call may produce: longer +# contexts run the query rows in tiles (``_index_row_tiles``). +_INDEX_TILE_ELEMS = 1 << 28 # capacity: fixed: per-call logits budget; longer contexts tile +# Rows one index-K gather call may write: row 2^25 x 64 bytes is past int32 +# (``_gather_index_k``). +_GATHER_MAX_ROWS = 1 << 25 # capacity: fixed: per-call gather rows; longer gathers chunk +# Bytes of candidate-block masks (one bool per context row x gathered key) a +# prefill step keeps for its re-index layers; past this each re-index layer +# re-scores its tiles' masks instead (``_index_prefill``). +_CAND_MASK_BYTES = 1 << 31 # capacity: fixed: kept-mask budget; past it masks are re-scored +# comm/allreduce strategies: fp32 sums on AUTO (workspace), bf16 on NCCL -- +# mixing dtypes on one workspace's one-shot kernel is wrong or hangs. +_AR_NCCL = 0 +_AR_AUTO = 3 +# Tensor-parallel (or attention-DP) sizes this target folds: heads 64/tp and +# output groups 8/tp per rank, experts 384/tp per rank under EP. +_TP_SIZES = (2, 4, 8) +# Sparse-attention workspace (the certified size), zero-initialised once. +_ATTN_WS_BYTES = 128 * 1024 * 1024 # capacity: fixed: the certified workspace size + + +# -------------------------------------------------------------------------- +# Runtime-value dump: metadata of every entry call (shapes, dtypes, the +# constants passed), a set of tuples appended to in Python only. CUDA-graph +# capture executes Python, so every captured batch size lands here; replay +# executes none of it. Off unless TRTLLM_MODELING_V2_VALUE_DUMP=1: the dump is +# for auditing a run against the entries' certified cells, and a served engine +# must not print it at every exit. +_VALUE_DUMP_ENV = "TRTLLM_MODELING_V2_VALUE_DUMP" +_VALUE_DUMP = os.environ.get(_VALUE_DUMP_ENV, "").strip() == "1" +_RUNTIME_VALUES: set = set() + + +def _rv(entry: str, *meta) -> None: + if _VALUE_DUMP: + _RUNTIME_VALUES.add((entry,) + tuple(meta)) + + +def _emit_runtime_values() -> None: + if not _RUNTIME_VALUES: + return + rank = os.environ.get("OMPI_COMM_WORLD_RANK", os.environ.get("RANK", "?")) + lines = [f"=== MODELINGV2 RUNTIME VALUE DUMP (deepseek_v41 unified, rank {rank}) ==="] + lines += [repr(v) for v in sorted(_RUNTIME_VALUES, key=repr)] + lines.append("=== END MODELINGV2 RUNTIME VALUE DUMP ===") + sys.stdout.write("\n".join(lines) + "\n") + sys.stdout.flush() + + +if _VALUE_DUMP: + atexit.register(_emit_runtime_values) + + +def _sh(t: torch.Tensor) -> tuple: + return tuple(t.shape) + + +# -------------------------------------------------------------------------- +# Load-time helpers (not on the forward path). + + +def _rope_table(base: float, positions: int, yarn: dict | None, device) -> torch.Tensor: + """``[positions, 2, 32]`` fp32 cos/sin table, the checkpoint reference's + ``precompute_freqs_cis`` recipe: frequency interpolation only (YaRN ramp + between ``beta_fast`` and ``beta_slow``), no mscale, built in fp32 on the + device the reference builds it on.""" + dim = _ROPE + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)) + if yarn is not None: + orig, factor = yarn["original"], yarn["factor"] + + def corrected_dim(rotations): + return dim * math.log(orig / (rotations * 2 * math.pi)) / (2 * math.log(base)) + + low = max(math.floor(corrected_dim(yarn["beta_fast"])), 0) + high = min(math.ceil(corrected_dim(yarn["beta_slow"])), dim - 1) + ramp = ( + (torch.arange(dim // 2, dtype=torch.float32, device=device) - low) + / max(high - low, 1e-3) + ).clamp(0, 1) + smooth = 1 - ramp + freqs = freqs / factor * (1 - smooth) + freqs * smooth + angles = torch.outer(torch.arange(positions, dtype=torch.float32, device=device), freqs) + return torch.stack([angles.cos(), angles.sin()], dim=1).contiguous() + + +def _is_prime(n: int) -> bool: + if n < 2: + return False + if n % 2 == 0: + return n == 2 + f = 3 + while f * f <= n: + if n % f == 0: + return False + f += 2 + return True + + +def _engram_layout(cfg, tokenizer_json: str, device) -> dict: + """The Engram hash constants of the checkpoint's ``inference/engram.py``: + the compressed-vocab token map (normalizer chain over each decoded token, + raw token string for partial UTF-8), per-(layer, shift) odd multipliers + from numpy PCG64 seeded ``10007 * layer_id``, the 48 distinct primes above + ``engram_vocab_size - 1`` drawn layer -> n-gram -> head, and the per-layer + bucket offsets (exclusive cumsum of the layer's primes).""" + import numpy as np + from tokenizers import Regex, Tokenizer, normalizers + + layer_ids = tuple(cfg.engram_layer_ids) + n_max = int(cfg.engram_max_ngram_size) + heads = int(cfg.engram_n_heads) + tok = Tokenizer.from_file(tokenizer_json) + sentinel = "\ue000" + normalizer = normalizers.Sequence( + [ + normalizers.NFKC(), + normalizers.NFD(), + normalizers.StripAccents(), + normalizers.Lowercase(), + normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), + normalizers.Replace(Regex(r"^ $"), sentinel), + normalizers.Strip(), + normalizers.Replace(sentinel, " "), + ] + ) + vocab = tok.get_vocab_size(with_added_tokens=True) + assert vocab == cfg.vocab_size, (vocab, cfg.vocab_size) + key_to_new: dict = {} + lookup = [0] * vocab + for token_id in range(vocab): + text = tok.decode([token_id], skip_special_tokens=False) + if "\ufffd" in text: + key = tok.id_to_token(token_id) + else: + normalized = normalizer.normalize_str(text) + key = normalized if normalized else text + new_id = key_to_new.get(key) + if new_id is None: + new_id = len(key_to_new) + key_to_new[key] = new_id + lookup[token_id] = new_id + classes = len(key_to_new) + assert classes == cfg.engram_compressed_vocab_size, ( + f"Engram compressed vocab {classes} != checkpoint {cfg.engram_compressed_vocab_size}; " + "every hash multiplier derives from it" + ) + bound = max(1, (np.iinfo(np.int64).max // classes) // 2) + mults = [] + for layer_id in layer_ids: + gen = np.random.default_rng(10007 * layer_id) + vals = gen.integers(low=0, high=bound, size=(n_max,), dtype=np.int64) + mults.append([int(v) * 2 + 1 for v in vals]) + primes, seen = [], set() + engram_vocab = _weights._engram_vocab_size(cfg) + for _ in layer_ids: + per_ngram = [] + for _ in range(n_max - 1): + sizes, cur = [], engram_vocab - 1 + for _ in range(heads): + cur += 1 + while not _is_prime(cur) or cur in seen: + cur += 1 + seen.add(cur) + sizes.append(cur) + per_ngram.append(sizes) + primes.append(per_ngram) + offsets = [] + for li, layer in enumerate(primes): + flat = [p for per in layer for p in per] + assert sum(flat) == int(cfg.engram_num_embeddings[li]), ( + f"Engram layer {layer_ids[li]}: primes sum to {sum(flat)}, table has " + f"{cfg.engram_num_embeddings[li]} rows" + ) + offsets.append([int(x) for x in np.cumsum([0] + flat[:-1])]) + return { + "token_map": torch.tensor(lookup, dtype=torch.int32, device=device), + "multipliers": torch.tensor(mults, dtype=torch.int64, device=device), + "primes": torch.tensor(primes, dtype=torch.int64, device=device), + "offsets": torch.tensor(offsets, dtype=torch.int64, device=device), + "pad_id": int(lookup[int(cfg.engram_pad_token_id)]), + } + + +def _gemm_backend(m: int) -> str: + """mxfp8_mxfp8_gemm backend by row count: CuTe DSL is fastest at M <= 32, + cuDNN above (both certified).""" + return "cute-dsl" if m <= 32 else "cudnn" + + +def _tile_n(m: int) -> int: + """mhc_gemm_sqrsum_fma's measured tile: 1 at M <= 128, 6 to 2048, 24 above.""" + if m <= 128: + return 1 + if m <= 2048: + return 6 + return 24 + + +def _moe_chunks(total: int) -> list: + full, rest = divmod(total, _MOE_MAX_T) + return [_MOE_MAX_T] * full + ([rest] if rest else []) + + +def _index_row_tiles(rows: int, s_len: int) -> list: + """Row ranges ``[a, b)`` of one prefill index pass over ``s_len`` gathered + keys, so that no range's fp32 logits exceed ``_INDEX_TILE_ELEMS`` (the op + pads each row to ``align(s_len + 256, 256)`` columns). One range when they + fit; otherwise ranges of a multiple of 64 rows where possible. Every op of + the pass (logits, candidate select, mask, top-k) is per row, so the tiling + does not change any row's result.""" + width = -(-(s_len + 256) // 256) * 256 + per = max(1, _INDEX_TILE_ELEMS // width) + if per >= rows: + return [(0, rows)] + if per >= 64: + per -= per % 64 + return [(a, min(rows, a + per)) for a in range(0, rows, per)] + + +# -------------------------------------------------------------------------- +# DSpark drafter geometry (``mtp.*``; the checkpoint's ``inference/model.py`` +# DSparkBlock / DSparkAttention / forward_spec). Its weights are the +# ``DRAFT_WEIGHTS`` table. + +_DS_BLOCK = 5 +_DS_STAGES = 3 +_DS_EXPERTS = 128 +_DS_TOPK = 3 +# Index row of a block query: 128 ring columns, the 5 block rows, 3 x -1 so +# the width (136) is a multiple of 4; 133 columns are live. +_DS_BLK_COLS = 8 +_DS_IDX_W = _WINDOW + _DS_BLK_COLS +_DS_IDX_LEN = _WINDOW + _DS_BLOCK +# The MXFP4 runner reads its MXFP8 activation linear, padded to 512 columns. +_DS_MOE_ALIGN = 512 +# trtllm_sparse_mla_dsv4 Precondition 8: a pool the call loads a -1 index +# from must span >= 128 mapped rows from its first row. The block K/V rows sit +# at the head of a persistent buffer this many rows longer than 5 x ring rows. +_DS_BLK_PAD_ROWS = 128 + + +# -------------------------------------------------------------------------- +# Vision tower geometry (``vision_config``; the checkpoint's +# ``inference/vision.py`` ViT + Aligner) and the per-call row envelope the +# entries certify for it. Its weights are the ``vis*`` rows of the table. + +_IMAGE_TOKEN = 129264 # every image-span position carries this id +_VIT_LAYERS = _weights.VIT_LAYERS +_VIT_DIM = _weights.VIT_DIM +_VIT_HEADS = _weights.VIT_HEADS +_VIT_HEAD_DIM = 64 +_VIT_INTER = _weights.VIT_INTER +_VIT_PATCH_IN = _weights.VIT_PATCH_IN # a patch flattened (channel, kh, kw) +_VIT_RATIO = _weights.VIT_RATIO # the aligner's 3x3 block +_VIT_EPS = 1e-6 +_VIT_ROPE_HALF = _VIT_HEAD_DIM // 2 # 16 row + 16 column angles per patch +# thop_attention no-pool: bidirectional within each sequence (= image). +_VIT_MASK_BIDIRECTIONAL = 0 +# thop's max_seq_len / max_context_length / attention_window_size: inert +# without a pool, but must cover the longest image (9189 patches). +_VIT_ATTN_MAX_LEN = 16384 # capacity: fixed: the largest image's patch count fits +# thop_attention's reference-checked per-call envelope (cublas_mm and the +# RoPE are certified beyond it); it sizes the load-time ``arange`` tables. The +# tower runs one image, at most 9189 patches, per call. +_VIT_MAX_PATCHES = 17037 # capacity: fixed: the certified per-call patch envelope +# moe_gelu and the aligner's GEMMs are certified to this many cells per call. +_AL_MAX_ROWS = 2048 # capacity: fixed: certified cells per call; more cells chunk +# The patch embed's cublas_mm (bf16 [M, 588] x [588, 1024]) picks a split-K +# kernel with bf16 partial sums at these row counts only, and the entry's +# is_valid refuses them. The patch rows are zero-padded past the band instead +# (``_patch_embed_rows``) and the padded rows' output is dropped. +_VIT_SPLITK_ROWS = (2305, 2432) +_VIT_SPLITK_PAD = 2433 +# The largest image: the span limit (max_image_tokens 1024) allows 9 * 1021 +# patches (a 3 x 3063 grid); the tower runs one image per call. +_VIT_MAX_IMAGE_PATCHES = 9216 # capacity: fixed: max_image_tokens 1024 of the checkpoint +# ViT CUDA graphs: one captured tower call per padded patch count, 256-patch +# steps to 4096 and 512-patch steps to 9216. No bucket falls in the split-K +# band above, so a replay never needs the pad. +_VIT_GRAPH_ENV = "TRTLLM_DSV41_VIT_CUDA_GRAPH" +_VIT_GRAPH_BUCKETS = tuple(range(1024, 4096 + 1, 256)) + tuple( + range(4608, _VIT_MAX_IMAGE_PATCHES + 1, 512) +) +# Sequence slots of a captured call: the image, then the pad sequence over +# rows [P, B) (empty when P == B). +_VIT_GRAPH_SLOTS = 2 +assert not any(_VIT_SPLITK_ROWS[0] <= b <= _VIT_SPLITK_ROWS[1] for b in _VIT_GRAPH_BUCKETS) +assert _VIT_GRAPH_BUCKETS[-1] == _VIT_MAX_IMAGE_PATCHES <= _VIT_MAX_PATCHES + + +def _patch_embed_rows(patches: int) -> int: + """Rows the patch-embed GEMM runs for ``patches`` patch rows: the rows + themselves, or ``_VIT_SPLITK_PAD`` when they fall in the split-K band.""" + lo, hi = _VIT_SPLITK_ROWS + return _VIT_SPLITK_PAD if lo <= patches <= hi else patches + + +def _vision_geometry(cfg) -> None: + """Assert the vision geometry this assembly is written for.""" + vc = getattr(cfg, "vision_config", None) + assert vc is not None, "the checkpoint config has no vision_config" + if isinstance(vc, dict): + vc = PretrainedConfig(**vc) + got = ( + int(vc.num_hidden_layers), + int(vc.hidden_size), + int(vc.num_attention_heads), + int(vc.intermediate_size), + int(vc.patch_size), + int(vc.downsample_ratio), + float(vc.rope_theta), + ) + assert got == (_VIT_LAYERS, _VIT_DIM, _VIT_HEADS, _VIT_INTER, 14, _VIT_RATIO, 10000.0), got + assert int(cfg.image_token_id) == _IMAGE_TOKEN, cfg.image_token_id + + +def _dspark_geometry(tc) -> SimpleNamespace: + """The drafter geometry this assembly is written for, read off + ``text_config`` -- the drafter's own ``dspark_*`` keys, not the trunk's + 384 experts at top-6.""" + g = SimpleNamespace( + hidden=int(tc.hidden_size), + vocab=int(tc.vocab_size), + eps=float(tc.rms_norm_eps), + heads=int(tc.num_attention_heads), + head_dim=int(tc.head_dim), + q_lora=int(tc.q_lora_rank), + o_groups=int(tc.o_groups), + o_lora=int(tc.o_lora_rank), + window=int(tc.sliding_window), + theta=float(tc.rope_theta), + inter=int(tc.moe_intermediate_size), + stages=int(tc.num_nextn_predict_layers), + block=int(tc.dspark_block_size), + noise=int(tc.dspark_noise_token_id), + markov_rank=int(tc.dspark_markov_rank), + experts=int(tc.dspark_n_routed_experts), + topk=int(tc.dspark_num_experts_per_tok), + targets=tuple(tc.dspark_target_layer_ids), + route_scale=float(tc.routed_scaling_factor), + swiglu=float(tc.swiglu_limit), + ) + n = int(tc.num_hidden_layers) + ratios = list(tc.compress_ratios) + assert ratios[n:] == [0] * g.stages, ( + f"drafter stages are window-only (theta 10000, no YaRN); got {ratios[n:]}" + ) + assert (g.hidden, g.vocab, g.eps) == (5120, 129280, 1e-20), (g.hidden, g.vocab, g.eps) + assert (g.heads, g.head_dim, g.q_lora, g.o_groups, g.o_lora) == (64, 512, 1280, 8, 1024) + assert (g.window, g.theta) == (_WINDOW, 10000.0) + assert (g.stages, g.block, g.experts, g.topk) == ( + _DS_STAGES, + _DS_BLOCK, + _DS_EXPERTS, + _DS_TOPK, + ), ( + g.stages, + g.block, + g.experts, + g.topk, + ) + assert (g.inter, g.markov_rank, g.noise) == (2304, 256, 128799) + assert g.targets == (37, 38, 39), g.targets + assert (g.route_scale, g.swiglu) == (1.5, _SWIGLU_LIMIT) + return g + + +# -------------------------------------------------------------------------- +# Bounded replays (tech report 3.2.2): the decoder +# replay's switch and tail plan, and the witness counters of both replays +# (``TRTLLM_MODELING_V2_WITNESS=1``). + +#: The decoder replay's switch, read once at construction: 1/true/on/yes, 0/false/off/no, +#: or unset (on, under attention TP and attention DP alike). The engine's cache manager +#: reads it too and turns an unset switch off where block reuse is on while encoder +#: replay is off (``DeepseekV4CacheManager.v41_decoder_replay``). +_REPLAY_ENV = DEEPSEEK_V41_DECODER_REPLAY_ENV +#: Set to 1 by the replay witness: one counter line per context request, when its +#: context phase completes. +_WITNESS_ENV = "TRTLLM_MODELING_V2_WITNESS" +# FP8 attention: E4M3 query and E4M3 pools (one power-of-two scale each) through +# the FP8-MMA sparse-MLA kernel. Off by default: it is lossy against the +# reference's bf16 attention (see ``_FP8_KV_BOUND``). +_FP8_ATTN_ENV = "TRTLLM_DSV41_FP8_ATTENTION" +# |kv| <= sqrt(512) * max|norm weight| for an RMS-normed 512-vector (RoPE keeps +# each rotated pair's norm), and the stored MXFP8 / NVFP4 rounding raises a +# block's amax by at most 2^-4: the per-tensor KV scale 2^ceil(log2(bound/448)) +# never saturates. +_FP8_KV_BOUND = 512**0.5 * (1.0 + 2.0**-4) + + +def _env_switch(name: str): + """The explicit value of boolean switch ``name``, or None when unset.""" + raw = os.environ.get(name, "").strip().lower() + if raw == "": + return None + if raw in ("1", "true", "on", "yes"): + return True + if raw in ("0", "false", "off", "no"): + return False + raise ValueError(f"{name}={os.environ[name]!r}: expected 1/true/on/yes or 0/false/off/no") + + +def _witness_line(**fields) -> None: + sys.stdout.write( + "MODELING_V2_WITNESS " + " ".join(f"{k}={int(v)}" for k, v in fields.items()) + "\n" + ) + sys.stdout.flush() + + +def _dp_pass_rows(gathered: list) -> list: + """dep4: the padded row count of each pass over the bounded layers, from + every rank's ``[count, rows_0, rows_1, ...]``: as many passes as the rank + with the most, each the max over the ranks that hold one.""" + passes = max(r[0] for r in gathered) + return [max(r[1 + k] if k < r[0] else 0 for r in gathered) for k in range(passes)] + + +class _TailRecord: + """A context request between two of its chunks: the position its next + chunk must start at, its counters, and (under replay) the stash -- the + layer-(L+1) input rows of its last ``window`` forwarded positions + ``[start, end)``, keyed like the step's per-row tensors. A first chunk + after a prefix hit under encoder replay starts ``replayed`` positions + before the reuse boundary: the prefix cache served ``first_position + + replayed`` tokens.""" + + __slots__ = ("end", "reused", "replayed", "fwd", "dec", "start", "stash") + + def __init__(self, first_position: int, replayed: int = 0): + self.end = first_position + self.reused = first_position + replayed + self.replayed = replayed + self.fwd = 0 + self.dec = 0 + self.start = first_position + self.stash = None + + +class _TailMetadata: + """The step's attention metadata as the bounded layers see it: the per-row + and per-request fields of the decoder row set; every other attribute is + the engine's.""" + + def __init__(self, md, fields: dict): + self.__dict__["_md"] = md + self.__dict__.update(fields) + + def __getattr__(self, name): + return getattr(self.__dict__["_md"], name) + + +class _TailSlots: + """``v41_swa_slot_mapping`` of the decoder row set, per local layer: a tail + row's slot through the layer's SWA block table at the row's position (a + stash row's position lies in an earlier chunk, outside this step's address + book); generation rows as the metadata has them. ``l2g`` is the prefill + target's ``dsv4_local_to_global_indices`` entry.""" + + def __init__(self, md, req, self_local, T, tc, Td, tcd, dev, l2g=None): + self.md, self.req, self.self_local, self.dev = md, req, self_local, dev + self.T, self.tc, self.Td, self.tcd = T, tc, Td, tcd + self.l2g = l2g + self._slots = {} + + def __getitem__(self, local): + s = self._slots.get(local) + if s is not None: + return s + md = self.md + s = torch.empty([self.Td], dtype=torch.int64, device=self.dev) + if self.tcd: + # Column 0 of a [row position, -1, ...] table is the row's own slot. + bt = md.sliding_block_tables[local, DeepseekV4AttentionType.SWA.value] + _rv("dsv4_local_to_global_indices", self.tcd, _WINDOW) + g = self.l2g(self.req, bt, self.self_local, _PAGE, 0) + s[: self.tcd].copy_(torch.reshape(torch.split(g, [1, _WINDOW - 1], 1)[0], [self.tcd])) + if self.Td > self.tcd: + s[self.tcd :].copy_(md.v41_swa_slot_mapping[local][self.tc : self.T]) + self._slots[local] = s + return s + + +class _TailCapture: + """``spec_metadata`` for the bounded layers. DSpark reads a captured + layer's input over the step's rows, so the tail rows are put back at their + rows of the step first; the tail rows of earlier chunks have no row there + and are captured on their own for the ring (``_DecoderTail.stage``).""" + + def __init__(self, spec, tail): + self.__dict__["_spec"] = spec + self.__dict__["_tail"] = tail + + def is_layer_capture(self, layer_id: int) -> bool: + return self._spec.is_layer_capture(layer_id) + + def maybe_capture_hidden_states(self, layer_id: int, hidden_states: torch.Tensor) -> None: + self._spec.maybe_capture_hidden_states( + layer_id, self._tail.stage(layer_id, hidden_states, self._spec) + ) + + def __getattr__(self, name): + return getattr(self.__dict__["_spec"], name) + + +class _DecoderTail: + """Decoder SWA bounded replay for one step that has context requests. + + Planned before the layers run, from host metadata only (request ids, chunk + and cached lengths, the engine's last-chunk flags). A request whose chunk + is not its last contributes no decoder row, and its last ``window`` + forwarded positions become its stash. A request in its last chunk + contributes its tail ``[s, N)``, ``s = max(N - min(N, window), first + forwarded position)``: its stash rows below this chunk, then this chunk's. + The decoder row set is those tails, request after request, then the + generation rows; it runs the bounded layers in slices of at most the + engine's step capacity, cut at request boundaries (``enter``). ``bound`` + is False when no context row would change (every request in its last + chunk, ``N <= window``, nothing stashed): the step then runs every layer + over every row. With replay off the plan only counts rows for the + witness. + + Prefill-only: a decode step has no context request to plan. ``l2g`` is the + prefill target's ``dsv4_local_to_global_indices`` entry, which the tail + rows' window slots are resolved through.""" + + def __init__(self, core, md, st, dev, l2g=None): + self.core, self.md, self.dev = core, md, dev + self.l2g = l2g + T, tc, nc = st["T"], st["tc"], st["nc"] + self.T, self.tc, self.nc = T, tc, nc + W = core.replay_window + last = getattr(md, "context_last_chunk", None) + assert last is not None and len(last) == nc, ( + "decoder replay and its witness counter need the engine's per-request " + f"context_last_chunk flags; got {last!r} for {nc} context requests" + ) + if core.replay: + flags = getattr(md, "context_return_logits", None) + assert flags is not None and len(flags) == nc, flags + if any(flags): + raise ValueError( + "decoder SWA bounded replay runs the layers after the last KV-source layer " + "over each prompt's last 128 positions only, so prompt logprobs and context " + f"logits do not exist; serve them with {_REPLAY_ENV}=0" + ) + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "decoder SWA bounded replay selects a context step's decoder rows on the " + "host; a captured context forward (prefill CUDA graphs) cannot run it" + ) + records = core._tail_records + rids = list(md.request_ids) + for rid in rids[nc:]: + records.pop(rid, None) # a generation row: its context phase is over + seq = md.seq_lens[:nc].tolist() + cached = md.cached_token_lens_cpu[:nc].tolist() + # Encoder replay rows at the head of a hit's first chunk: they are + # forwarded rows like any other (a tail starts at or after them). + replayed = md.v41_num_replay_tokens_cpu[:nc].tolist() + reqs, off, bound, need = [], 0, False, W + 1 + for b in range(nc): + L, c, R, rid = int(seq[b]), int(cached[b]), int(replayed[b]), rids[b] + rec = records.get(rid) + if rec is None or rec.end != c or R: + rec = _TailRecord(c, R) # the first chunk of this context phase + r = SimpleNamespace( + b=b, rid=rid, off=off, L=L, c=c, R=R, N=c + L, last=bool(last[b]), rec=rec + ) + if core.replay: + stashed = rec.stash is not None + if not r.last: + bound = True + r.keep = max(r.N - W, rec.start if stashed else c) + else: + r.t0 = r.N - min(r.N, W) + r.s = max(r.t0, rec.start if stashed else c) + bound = bound or r.N > W or stashed + need = max(need, r.N + 1) + reqs.append(r) + off += L + assert off == tc, (off, tc) + # Decoder rows the step needs (tails, then generation rows). + self.rows = sum(r.N - r.s for r in reqs if r.last) + st["gen"] if core.replay else 0 + core._ensure_ar(max(need, nc * W + 1), dev) + self.reqs = reqs + self.bound = bound + self.counted = False + self.slices, self.cur, self.h_full, self.out = [], None, None, None + self.h_layer = {} + + def count_forwarded(self, x) -> None: + """Layer 0: each request's rows of its input ``x``.""" + for r in self.reqs: + r.rec.fwd += x[r.off : r.off + r.L].shape[0] + + def count_full(self, h) -> None: + """Unbounded step: layer L+1 runs over each request's rows of ``h``.""" + for r in self.reqs: + r.rec.dec += h[r.off : r.off + r.L].shape[0] + self.counted = True + + def enter(self, rows: dict, st: dict, cache: dict, max_rows: int) -> list: + """At the first bounded layer. Stash the rows of requests whose chunk + is not their last, and return the decoder passes: the decoder row set + cut at request boundaries into slices of at most ``max_rows`` rows, + each as (its rows of every per-row tensor in ``rows`` -- the step's + layer-(L+1) inputs --, its step dict, its per-step cache). The tail + requests come in step order; the generation rows go in the last slice. + No decoder row: an empty list. + + ``max_rows`` is the engine's step capacity (max_num_tokens), which + every per-row buffer and every certified op shape is sized for. A + finishing request's tail holds min(N, 128) rows, the rows of its + earlier chunks included, so the row set outgrows the step itself when + more than max_rows / 128 requests finish in it. Every bounded op is + per row or per request, so the slices compute what one pass would.""" + core, md = self.core, self.md + W = core.replay_window + T, gen = self.T, st["gen"] + assert max_rows >= W, (max_rows, W) + rows = dict(rows) + # Per-row inputs of the bounded layers besides the stream: the last + # KV-source layer's top-k (layers 21-23 read it), the window + top-k + # counts and the visible compressed length (index row ends). + rows["topk"] = cache["topk"][1] + rows["lens"] = md.sparse_mla_topk_lens[1][:T] + rows["vis"] = md.compressed_visible_lens_cuda[1][:T] + # The candidate source's index query (MXFP4 q, its UE8M0 words, the + # head weights): a tail row's candidate-block mask is scored from it in + # the decoder pass, against the index K of the step that runs it. + rows["cq"], rows["cqs"], rows["cw"] = cache["cand_q"] + names = list(rows) + for r in self.reqs: + if r.last: + continue + rec = r.rec + parts = {n: [] for n in names} + if rec.stash is not None and r.keep < r.c: + for n in names: + parts[n].append(rec.stash[n][r.keep - rec.start :]) + a = max(r.keep, r.c) - r.c + for n in names: + parts[n].append(rows[n][r.off + a : r.off + r.L]) + rec.stash = {n: torch.cat(parts[n], 0) for n in names} + rec.start = r.keep + fin = [r for r in self.reqs if r.last] + for r in fin: + r.n_old = r.c - r.s if (r.rec.stash is not None and r.s < r.c) else 0 + r.a = max(r.t0, r.c) - r.c + r.n = r.N - r.s + r.rec.dec += r.n + self.h_full = rows["h"] + self.counted = True + groups = self._groups(max_rows) + last = len(groups) - 1 + self.slices = [ + self._slice(g, gen if k == last else 0, rows, names, st, cache) + for k, g in enumerate(groups) + ] + # The passes are the ones pass_rows planned (attention DP pads to them). + assert [sl.Td for sl in self.slices] == self.pass_rows(max_rows), ( + [sl.Td for sl in self.slices], + self.pass_rows(max_rows), + ) + if self.slices: + self.out = torch.empty([T, core.hidden], dtype=torch.bfloat16, device=self.dev) + self.out.copy_(core._const["zero_bf16"].expand([T, core.hidden])) + return [(sl.dec, sl.st, sl.cache) for sl in self.slices] + + def _groups(self, max_rows: int) -> list: + """The finishing requests cut at request boundaries into slices of at + most ``max_rows`` rows (a tail is ``N - s`` rows); the generation rows + go in the last. Host-only, so it is known before the layers run.""" + gen = self.T - self.tc + groups, cur, n_cur = [], [], 0 + for r in (r for r in self.reqs if r.last): + n = r.N - r.s + if cur and n_cur + n > max_rows: + groups.append(cur) + cur, n_cur = [], 0 + cur.append(r) + n_cur += n + if gen and cur and n_cur + gen > max_rows: + groups.append(cur) + cur = [] + if cur or gen: + groups.append(cur) + return groups + + def pass_rows(self, max_rows: int) -> list: + """The row count of each decoder pass ``enter`` will return.""" + groups = self._groups(max_rows) + gen = self.T - self.tc + return [ + sum(r.N - r.s for r in g) + (gen if k == len(groups) - 1 else 0) + for k, g in enumerate(groups) + ] + + def _slice(self, fin, gen, rows, names, st, cache): + """One decoder pass: the tails of ``fin`` (in order), then the step's + generation rows when ``gen``; its rows, step dict and per-step cache.""" + core, md, dev = self.core, self.md, self.dev + W = core.replay_window + T, tc, nc = self.T, self.tc, self.nc + G = st["G"] if gen else 0 + B = st["B"] + ar = core._ar32 + parts = {n: [] for n in names} + pos, req, win, kvl, flo, wsl, ks, cu = [], [], [], [], [], [], [], [0] + q = 0 + for r in fin: + rec = r.rec + for n in names: + if r.n_old: + parts[n].append(rec.stash[n][r.s - rec.start :]) + parts[n].append(rows[n][r.off + r.a : r.off + r.L]) + r.q = q + got = sum(p.shape[0] for p in parts["h"][-(2 if r.n_old else 1) :]) + assert got == r.n, (got, r.N, r.s) + pos.append(ar[r.s : r.N]) + req.append(ar[r.b : r.b + 1].expand([r.n])) + j0 = r.s - r.t0 + win.append( + torch.where( + core._const["tri"][j0 : j0 + r.n], + torch.reshape(ar[r.t0 : r.t0 + W], [1, W]).expand([r.n, W]), + -1, + ) + ) + kvl.append(md.kv_lens_cuda_runtime[r.b : r.b + 1]) + # The tail's SWA floor is t0, not the engine's floor for request b: + # its window count comes from kv length N minus t0. + flo.append(ar[r.t0 : r.t0 + 1]) + wsl.append(ar[r.N - r.t0 : r.N - r.t0 + 1]) + ks.append(md.v41_ctx_index_ks[1][r.off + r.a : r.off + r.a + 1].expand([r.n])) + q += r.n + cu.append(q) + tcd, ncd = q, len(fin) + Td = tcd + gen + max_tail = max((r.n for r in fin), default=0) + _rv("decoder_replay.rows", T, tcd, ncd, gen, max_tail) + if gen: + for n in names: + parts[n].append(rows[n][tc:T]) + pos.append(st["pos"][tc:T]) + req.append(md.req_idx_per_token[tc:T]) + win.append(md.swa_local_indices_cuda[tc:T]) + kvl.append(md.kv_lens_cuda_runtime[nc:B]) + dec = {n: torch.cat(parts[n], 0) for n in names} + pos_d = torch.cat(pos, 0) + req_d = torch.cat(req, 0) + vis = dec.pop("vis") + fields = { + "req_idx_per_token": req_d, + "swa_local_indices_cuda": torch.cat(win, 0), + "sparse_mla_topk_lens": {1: dec.pop("lens")}, + "compressed_visible_lens_cuda": {1: vis}, + "kv_lens_cuda_runtime": torch.cat(kvl, 0), + # Generation rows read their tables at [ncd:ncd+G]; the engine's + # are at [nc:nc+G]. + "indexer_compress_block_tables": {1: md.indexer_compress_block_tables[1][nc - ncd :]}, + # Per tail request: its floor and the context attention's seq_lens + # (kv length - floor), in the decoder pass's request order. + "v41_swa_floor_cuda": torch.cat(flo, 0) if ncd else None, + "v41_swa_seq_lens_cuda": torch.cat(wsl, 0) if ncd else None, + "num_contexts": ncd, + "num_ctx_tokens": tcd, + "num_seqs": ncd + G, + "v41_max_ctx_q_len": max_tail, + } + self_local, ring = None, None + if ncd: + fields["cu_seq_lens_cuda"] = torch.cat([ar[x : x + 1] for x in cu], 0) + ks_d = torch.cat(ks, 0) + fields["v41_ctx_index_ks"] = {1: ks_d} + fields["v41_ctx_index_ke"] = {1: torch.add(ks_d, vis[:tcd])} + self_local = torch.where( + core._const["col0"].expand([tcd, W]), + torch.reshape(pos_d[:tcd], [tcd, 1]).expand([tcd, W]), + -1, + ) + old = [r for r in fin if r.n_old] + if old and core.ds is not None: + ring = ( + old, + torch.empty( + [sum(r.n_old for r in old), len(core.ds.targets) * core.hidden], + dtype=torch.bfloat16, + device=dev, + ), + ) + fields["v41_swa_slot_mapping"] = _TailSlots( + md, req_d[:tcd], self_local, T, tc, Td, tcd, dev, self.l2g + ) + # nc / B count the decoder pass's requests (its context tails, then the + # generation requests); request indices (req_idx_per_token, block-table + # rows) stay the engine's, whose request count is B_engine. + std = dict(st) + std.update( + md=_TailMetadata(md, fields), + T=Td, + tc=tcd, + nc=ncd, + G=G, + gen=gen, + B=ncd + G, + B_engine=B, + pos=pos_d, + max_ctx_q=max_tail, + ) + cached = { + "topk": {1: dec.pop("topk")}, + "ctx_k": cache["ctx_k"], + "sched": cache["sched"], + "cand_q": (dec.pop("cq"), dec.pop("cqs"), dec.pop("cw")), + } + if "cand_gen" in cache: + cached["cand_gen"] = cache["cand_gen"] + return SimpleNamespace( + fin=fin, + gen=gen, + tcd=tcd, + Td=Td, + ring=ring, + ring_done=set(), + dec=dec, + st=std, + cache=cached, + ) + + def activate(self, k: int) -> None: + """Slice k is the decoder pass that runs next.""" + self.cur = self.slices[k] + + def stage(self, layer_id: int, h_flat, spec): + """The step-row view of a bounded layer's input for DSpark's capture: + the post-layer-L stream with the running slice's decoder rows written + back to their rows of the step. With several slices each captured + layer keeps its own copy of the step rows, so an earlier slice's rows + still hold this layer's input when a later slice captures (a capture + overwrites the layer's rows; the last slice's holds every row). The + tail rows of earlier chunks go through the metadata's own capture + first (it averages the hc copies), and their rows are kept for the + shell to write into the request's ring.""" + core, sl = self.core, self.cur + hd = torch.reshape(h_flat, [sl.Td, _HC, core.hidden]) + if sl.ring is not None: + ring, buf = sl.ring + k = core.ds.targets.index(layer_id) + src = torch.cat([hd[r.q : r.q + r.n_old] for r in ring], 0) + n = src.shape[0] + spec.maybe_capture_hidden_states(layer_id, torch.reshape(src, [n, core.hc_k])) + cols = [core.hidden * k, core.hidden, core.hidden * (len(core.ds.targets) - k - 1)] + torch.split(buf, cols, 1)[1].copy_(torch.split(spec.get_hidden_states(n), cols, 1)[1]) + sl.ring_done.add(k) + hf = self.h_full + if len(self.slices) > 1: + hf = self.h_layer.get(layer_id) + if hf is None: + hf = self.h_layer[layer_id] = torch.empty( + list(self.h_full.shape), dtype=self.h_full.dtype, device=self.dev + ) + hf.copy_(self.h_full) + for r in sl.fin: + hf[r.off + r.a : r.off + r.L].copy_(hd[r.q + r.n_old : r.q + r.n]) + if sl.gen: + hf[self.tc : self.T].copy_(hd[sl.tcd : sl.Td]) + return torch.reshape(hf, [self.T, core.hc_k]) + + def scatter(self, y) -> None: + """The running slice's output rows into the step's [T, hidden] output + (zeros on rows that no bounded layer ran).""" + sl = self.cur + for r in sl.fin: + self.out[r.off + r.a : r.off + r.L].copy_(y[r.q + r.n_old : r.q + r.n]) + if sl.gen: + self.out[self.tc : self.T].copy_(y[sl.tcd : sl.Td]) + + def close(self, out, spec): + """The step's output rows, then ``finish``: ``out`` for an unbounded + step, the slices' scattered rows for a bounded one. With no decoder + row at all the captured layers still read the step's rows (their + post-layer-L stream) and every output row is zero.""" + if self.bound: + if not self.slices: + if spec is not None: + flat = torch.reshape(self.h_full, [self.T, self.core.hc_k]) + for i in range(self.core.late, self.core.num_layers): + if spec.is_layer_capture(i): + spec.maybe_capture_hidden_states(i, flat) + out = torch.empty([self.T, self.core.hidden], dtype=torch.bfloat16, device=self.dev) + out.copy_(self.core._const["zero_bf16"].expand([self.T, self.core.hidden])) + else: + out = self.out + self.finish() + return out + + def finish(self) -> None: + """Carry each request to its next chunk; print the witness counter of + each whose context phase completed.""" + core = self.core + assert self.counted, "the step never reached the first bounded layer" + for sl in self.slices: + if sl.ring is not None and len(sl.ring_done) == len(core.ds.targets): + ring, buf = sl.ring + ar, off = core._ar32, 0 + for r in ring: + frames = torch.empty([r.n_old], dtype=torch.int64, device=self.dev) + frames.copy_(ar[r.s + 1 : r.c + 1]) # frame = position + 1 + core._tail_ring.append((r.rid, buf[off : off + r.n_old], frames)) + off += r.n_old + for r in self.reqs: + rec = r.rec + rec.end = r.N + if r.last: + core._tail_records.pop(r.rid, None) + if core._witness: + _witness_line( + prompt_tokens=r.N, + reused_tokens=rec.reused, + replayed_tokens=rec.replayed, + forwarded_tokens=rec.fwd, + decoder_tokens=rec.dec, + ) + else: + core._tail_records[r.rid] = rec + + +class DeepseekV41UnifiedModelingV2(DecoderModel): + """The core: what DeepSeek-V4.1-Flash statically is on this rank. + + Geometry and the member read off the mapping, the weights (``w``, and the + drafter's ``w_ds`` under DSpark), the per-layer operand table and the rope, + Engram and constant tables derived after load, the engine's pool views, and + the host-side bookkeeping that outlives a step (decoder-replay records). It + runs no layer itself: ``forward`` routes each step to the prefill or the + decode target, which hold and bind every entry they call. + """ + + def __init__(self, model_config: ModelConfig): + super().__init__(model_config) + cfg = model_config.pretrained_config + assert cfg is not None + tc = getattr(cfg, "text_config", None) + tc = cfg if tc is None else tc + # The engine picks the KV cache manager and the attention metadata off + # the built model's sparse config: the layout mode gives it the + # deepseek_v4 runtime whose V4.1 pools and metadata this core reads. + use_deepseek_v41_layout_mode(model_config) + + # ---- identity: topology, dtype ---------------------------------------- + mapping = model_config.mapping + tp = mapping.tp_size + assert tp in _TP_SIZES and mapping.cp_size == 1, ( + f"unified target: tp in {_TP_SIZES}, no context parallelism; got tp={tp} cp={mapping.cp_size}" + ) + assert mapping.world_size == tp * mapping.pp_size, (mapping.world_size, tp, mapping.pp_size) + self.dp = bool(mapping.enable_attention_dp) + self.ep = mapping.moe_ep_size + self.moe_tp = mapping.moe_tp_size + if self.ep == tp: + assert self.moe_tp == 1, f"EP{tp} members split experts one way only" + self.member = f"dep{tp}" if self.dp else f"tep{tp}" + else: + assert self.ep == 1 and self.moe_tp == tp and not self.dp, ( + f"members are tep, dep, tp; got tp={tp} ep={self.ep} moe_tp={self.moe_tp} dp={self.dp}" + ) + self.member = f"tp{tp}" + # Pipeline stages: this rank runs `pp_layer_ids` (the engine's split, + # `mapping.pp_layers`), and checks below that none of them reads a pool + # or a per-step product of another stage. + self.pp = mapping.pp_size + self.pp_rank = mapping.pp_rank + self.first_pp = mapping.is_first_pp_rank() + self.last_pp = mapping.is_last_pp_rank() + if self.pp > 1: + self.member = f"pp{self.pp}x{self.member}" + self.rank = mapping.rank + self.tp = tp + self.tp_rank = mapping.tp_rank + self.world = mapping.world_size + # The collective group: this stage's TP (or attention-DP) ranks, as + # global MPI ranks. + self.group = list(mapping.tp_group) + self.gpus_per_node = mapping.gpus_per_node + # The IPC all-reduce workspace needs the whole group on one node; a + # group spanning nodes (tep8 over two NVL72 trays) sums on NCCL, which + # rides the NVLink fabric (MNNVL) between the trays. + self.ar_one_node = len({r // self.gpus_per_node for r in self.group}) == 1 + self.attn_tp = 1 if self.dp else tp + self.attn_rank = 0 if self.dp else mapping.tp_rank + + dt = model_config.torch_dtype + assert dt == torch.bfloat16, f"bf16 target, engine resolved {dt}" + assert not getattr(tc, "tie_word_embeddings", False), "untied head" + assert not getattr(tc, "attention_bias", False), "no attention bias" + assert tc.hidden_act == "silu", tc.hidden_act + # fp8 32x32 dense weights plus NVFP4 routed experts: the engine reads + # this export's global quant_algo as FP8_BLOCK_SCALES (its DeepSeek-V4.1 + # loader) or MIXED_PRECISION (a modelopt reading); the experts' NVFP4 + # comes from the checkpoint's per-layer quantization either way. + algo = getattr(model_config.quant_config, "quant_algo", None) + assert str(getattr(algo, "value", algo)) in ("FP8_BLOCK_SCALES", "MIXED_PRECISION"), ( + f"fp8 32x32 dense + NVFP4 experts checkpoint; engine read quant_algo {algo!r}" + ) + # The quantization shape is a routing criterion; a KV-cache quantization + # is an engine argument the routing context does not carry. The pools' + # formats are this core's own (see "KV storage" above). + assert getattr(model_config.quant_config, "kv_cache_quant_algo", None) is None, ( + "the core stores the checkpoint's own KV value sets (TRTLLM_DSV41_KV_STORAGE " + "picks native or bf16 rows); a declared KV-cache quantization is not supported" + ) + + # ---- geometry ------------------------------------------------------- + self.num_layers = tc.num_hidden_layers + self.hidden = tc.hidden_size + self.vocab = tc.vocab_size + self.eps = float(tc.rms_norm_eps) + assert (self.num_layers, self.hidden, self.vocab) == (40, 5120, 129280), ( + self.num_layers, + self.hidden, + self.vocab, + ) + assert self.eps == 1e-20, f"every RMS statistic uses eps 1e-20, got {self.eps}" + assert cfg.num_hidden_layers == self.num_layers and cfg.hidden_size == self.hidden, ( + "mirrored top-level config disagrees with text_config" + ) + self.heads = tc.num_attention_heads + self.head_dim = tc.head_dim + self.q_lora = tc.q_lora_rank + self.o_groups = tc.o_groups + self.o_lora = tc.o_lora_rank + assert (self.heads, self.head_dim, self.q_lora) == (64, 512, 1280) + assert tc.num_key_value_heads == 1, "one shared latent KV head" + assert tc.qk_rope_head_dim == _ROPE, "RoPE on the last 64 channels" + assert (self.o_groups, self.o_lora) == (8, 1024), "grouped low-rank output" + assert tc.sliding_window == _WINDOW + self.nope = self.head_dim - _ROPE # 448 + assert self.heads % self.attn_tp == 0 and self.o_groups % self.attn_tp == 0 + # Every per-rank window -- heads, output groups, experts, Engram rows, + # pipeline layers -- is the weight table's, derived from the mapping + # once (``weights.dims``). + d = _weights.MODEL_WEIGHTS.dims(self) + self.h_loc = d.h_loc + self.g_loc = d.g_loc + # TP heads [h r, h r + h) (h = 64 / tp) belong to groups [g r, g r + g) + # (g = 8 / tp): wq_b rows, the sinks and wo_a's groups split on the same + # boundary (tp4: 16 heads, groups 2r, 2r+1; tp8: 8 heads, group r). + assert self.h_loc * self.o_groups == self.g_loc * self.heads + self.group_in = d.group_in # 4096 + assert self.group_in == 4096 + + # Indexer, compressor, candidates. + assert (tc.index_n_heads, tc.index_head_dim, tc.index_topk) == ( + _IDX_HEADS, + _IDX_DIM, + _TOPK, + ) + assert (tc.candidate_block_size, tc.candidate_topk_blocks) == ( + _CAND_BLOCK, + _CAND_TOPK_BLOCKS, + ) + self.cand_src = int(tc.candidate_source_layer_id) + ratios = d.ratios + self.ratios = ratios + self.kv_sources = d.kv_sources + self.index_sources = d.index_sources + # The per-layer role table this assembly is written for: ratio 0 + # = SWA only; ratio 1 = SWA + layer 20's global KV. A checkpoint with + # another table is another assembly. + assert ratios == [0, 0] + [2] * 18 + [1] * 20, ratios + assert self.kv_sources == (2, 8, 14, 20), self.kv_sources + assert self.index_sources == (2, 8, 14, 20, 24, 28, 32, 36), self.index_sources + assert self.cand_src == 20 + self.kv_src = [None] * self.num_layers + self.idx_src = [None] * self.num_layers + for i in range(self.num_layers): + if ratios[i]: + self.kv_src[i] = max(s for s in self.kv_sources if s <= i) + self.idx_src[i] = max(s for s in self.index_sources if s <= i) + assert ratios[self.kv_src[i]] == ratios[i] == ratios[self.idx_src[i]] + # This stage's layers. A stage runs on its own pools only: every + # compressed layer's KV source and index source, the candidate source of + # the re-index layers and the Engram history owner must be on it. For + # this checkpoint that leaves stage boundaries at layers 1 and 20 (so + # PP2's even split, 0-19 | 20-39, and PP3 as 1 | 19 | 20). + self.pp_layer_ids = d.pp_layer_ids + self.local = d.local + assert self.pp_layer_ids == list(range(self.pp_layer_ids[0], self.pp_layer_ids[-1] + 1)) + layout = DeepseekV41Layout.from_config(cfg) + err = layout.stage_dependency_error(self.pp_layer_ids) + if err is not None: + raise ValueError( + f"pipeline stage {self.pp_rank} (layers {self.pp_layer_ids[0]}-{self.pp_layer_ids[-1]}): " + f"{err}; this checkpoint's stage boundaries can only be at layers " + f"{list(layout.pipeline_boundaries())} (pp_partition)" + ) + # The first attention of a later stage collapses its input with the + # previous stage's last FFN mixes: that layer's hc_ffn triple is held too. + self.pp_prev = d.pp_prev + rope = getattr(tc, "rope_scaling", None) or getattr(tc, "rope_parameters", None) + assert isinstance(rope, dict), rope + assert rope.get("rope_type", rope.get("type")) == "yarn", rope + assert "mscale" not in rope and "mscale_all_dim" not in rope, ( + "V4.1's YaRN is frequency interpolation only; a magnitude scale is another model" + ) + self.yarn = { + "factor": float(rope["factor"]), + "original": int(rope["original_max_position_embeddings"]), + "beta_fast": float(rope["beta_fast"]), + "beta_slow": float(rope["beta_slow"]), + } + assert self.yarn == { + "factor": 16.0, + "original": 65536, + "beta_fast": 32.0, + "beta_slow": 1.0, + }, self.yarn + self.theta = float(tc.rope_theta) + self.comp_theta = float(tc.compress_rope_theta) + assert (self.theta, self.comp_theta) == (10000.0, 160000.0) + + # Hyper-connections. + assert (tc.hc_mult, tc.hc_sinkhorn_iters, float(tc.hc_eps)) == (_HC, 20, 1e-6) + self.hc_eps = float(tc.hc_eps) + self.hc_k = d.hc_k # 20480 + + # MoE. + assert tc.scoring_func == "sqrtsoftplus" and tc.topk_method == "noaux_tc" + assert tc.norm_topk_prob and float(tc.routed_scaling_factor) == 1.5 + assert float(tc.swiglu_limit) == _SWIGLU_LIMIT + assert tc.n_shared_experts == 1 + self.num_experts = d.num_experts + self.topk = tc.num_experts_per_tok + self.moe_inter = d.moe_inter + assert (self.num_experts, self.topk, self.moe_inter) == (384, 6, 2304) + self.route_scale = float(tc.routed_scaling_factor) + if self.ep == tp: + assert self.num_experts % tp == 0 + self.e_loc, self.e_off = d.e_loc, d.e_off + self.i_loc, self.i_off = d.i_loc, d.i_off + # fp4_block_scale_moe_runner's layout rules at the per-rank width. + assert self.hidden % 256 == 0 and self.i_loc % 64 == 0 + + # Engram. + self.engram_layers = d.engram_layers + assert self.engram_layers == (1, 14), self.engram_layers + assert (tc.engram_max_ngram_size, tc.engram_n_heads, tc.engram_head_dim) == (4, 8, 256) + self.eng_cols = d.eng_cols # 24 + self.eng_dim = d.eng_dim + self.eng_rows = d.eng_rows + # host_offload: the tables are not parameters; every rank reads all + # rows of the node's shared pinned copy (attach_engram_host_table). + self.eng_host = d.eng_host + self.eng_part, self.eng_start, self.eng_end = d.eng_part, d.eng_start, d.eng_end + self._eng_table = None + self.eng_local = bool(self.local & set(self.engram_layers)) + # The node's ranks attach the shared table together (local rank 0 fills + # it, then a node-local barrier), so every rank of a node attaches when + # a rank of that node holds the Engram layers. + eng_stages = {mapping.pp_rank_of_layer(i, self.num_layers) for i in self.engram_layers} + node = self.rank // self.gpus_per_node + self.eng_attach = self.eng_host and any( + r // tp in eng_stages for r in range(self.world) if r // self.gpus_per_node == node + ) + if self.eng_host: + # The attach's node barrier runs on the node-local communicator, whose + # (lazy) split is collective over every rank: split it here, where every + # rank passes, or a node whose ranks do not attach (pp2 x tp4's second + # tray) would strand the attaching ranks in the split. + local_mpi_comm() + self.eng_wkv_n = d.eng_wkv_n # 25600 + self.eng_in = d.eng_in # 6144 + self._tc = tc + self._ckpt_dir = getattr(cfg, "_name_or_path", None) or getattr(cfg, "name_or_path", None) + + # Vision tower + aligner (always built: the engine serves images for + # every configuration of this checkpoint). + _vision_geometry(cfg) + + # DSpark: the drafter's parameters exist only when the engine was + # built with a speculative_config. + spec = getattr(model_config, "spec_config", None) + self.ds = None + if spec is not None: + assert spec.spec_dec_mode.is_dspark() and spec.draft_is_embedded_in_target, ( + "this target's mtp.* stages are the checkpoint's embedded DSpark drafter; " + f"the engine resolved spec_dec_mode {spec.spec_dec_mode!r}" + ) + self.ds = _dspark_geometry(tc) + assert int(spec.max_draft_len) == self.ds.block, ( + f"DSpark drafts its trained block of {self.ds.block}; max_draft_len={spec.max_draft_len}" + ) + + # ---- parameters: the weight tables, torch.empty only ------------------ + self.w = _weights.MODEL_WEIGHTS.declare(d) + self.w_ds = ( + _weights.DRAFT_WEIGHTS.declare(_weights.DRAFT_WEIGHTS.dims(self)) + if self.ds is not None + else None + ) + self._vision = VisionTower(self) + # llm_args.disable_mm_encoder: an engine that will never see an image + # (e.g. a generation-only disaggregated server) captures no ViT graphs. + self.disable_mm_encoder = bool(getattr(model_config, "disable_mm_encoder", False)) + + self._L: list | None = None + self._rope_tables: dict | None = None + self._rope_positions = 0 + self._eng: dict | None = None + self._const: dict | None = None + self._ar_ws = None + self._pools: dict | None = None + self._pool_mgr = None + self._derive_pending = True + self._spec_checked = False + self._cand_binds_decode = False + self._prefill: Target | None = None + self._decode: Target | None = None + + # Decoder SWA bounded replay: the bounded layers are the ones + # after the last KV-source layer, whose ratio-1 compressor and index K + # every later layer reads; they own no compressed pool, reuse layer + # 20's KV and top-k chain, and hold no Engram. + self.late = max(self.kv_sources) + 1 + assert self.late == 21 and set(ratios[self.late :]) == {1}, (self.late, ratios[self.late :]) + assert max(self.engram_layers) < self.late and self.cand_src < self.late + self.replay_window = int(tc.sliding_window) + self.replay_requested = _env_switch(_REPLAY_ENV) is not False + # Until the first forward takes the cache manager's resolution (_resolve_replay). + self.replay = self.replay_requested + self._replay_mgr = None + self._witness = os.environ.get(_WITNESS_ENV, "").strip() == "1" + # FP8 attention: the sparse-MLA core with E4M3 query and pools (native KV storage only). + self.fp8_attn = _env_switch(_FP8_ATTN_ENV) is True + self._fp8_kv_scale = [None] * self.num_layers + if self.fp8_attn and self.rank == 0: + sys.stdout.write( + f"[modeling_v2 deepseek_v41] {_FP8_ATTN_ENV}=1: sparse attention computes with E4M3 " + f"Q/K/V (per-tensor power-of-two scales) on {self.member}\n" + ) + sys.stdout.flush() + self._tail_records: dict = {} + self._tail_ring: list = [] + self._ar32 = None + + def __pp_init__(self) -> None: + """Pipeline stages are this core's own (``pp_layer_ids``, the stream + handed over in ``forward``); the base class's per-module skip does not + apply to a flat ParameterDict.""" + + def _resolve_replay(self, mgr) -> None: + """Decoder replay as the engine's cache manager resolved it, once per manager. + + With block reuse on it runs only beside encoder replay; the manager resolves + both. It reads the same switch and mapping as the constructor, so it never turns + on what the model turned off (an explicit 0).""" + if self._replay_mgr is not None and self._replay_mgr() is mgr: + return + on = bool(mgr.v41_decoder_replay) + assert self.replay_requested or not on, ( + f"the cache manager resolved decoder replay on where the model resolved {_REPLAY_ENV} off" + ) + self.replay = on + self._replay_mgr = weakref.ref(mgr) + + def _release_tail_records(self, mgr) -> None: + """Drop the between-chunk records, stashes included, of requests the cache manager + no longer holds. A request freed between two of its context chunks (cancelled, or + evicted to be recomputed) never reaches the last chunk that drops its record; one + that comes back starts a new context phase and a new record.""" + held = mgr.kv_cache_map + for rid in [rid for rid in self._tail_records if rid not in held]: + del self._tail_records[rid] + + def _ensure_ar(self, n: int, dev) -> None: + """The int32 0..n-1 table the replay's per-row indices are views of; + regrown only in context steps, which are never captured.""" + if self._ar32 is None or self._ar32.shape[0] < n: + self._ar32 = torch.arange( + max(n, 2 * _WINDOW + 1, 2 * (0 if self._ar32 is None else self._ar32.shape[0])), + dtype=torch.int32, + device=dev, + ) + + # ------------------------------------------------------------------ + def post_load_weights(self) -> None: + """Post-load derivation, then the two targets. + + The per-layer operand table (views of the weights and the folded NVFP4 + call scalars), the rope tables, the Engram hash constants and residency, + the all-reduce workspace and the load-time constant tensors the forward + fills from. Meta is over here. The targets are built last, against the + real weights, and bind every entry they call. + + Named for the shell hook that drives it. Not an override: that hook lives on + DecoderModelForCausalLM, and a core derives from DecoderModel, which has none. + """ + w = self.w + dev = w["final_norm"].device + L = [None] * self.num_layers + if self.pp_prev is not None: + i = self.pp_prev + L[i] = { + "hc_ffn": (w[f"l{i}_hc_ffn_fn"], w[f"l{i}_hc_ffn_scale"], w[f"l{i}_hc_ffn_base"]) + } + for i in self.pp_layer_ids: + d = { + "hc_attn": ( + w[f"l{i}_hc_attn_fn"], + w[f"l{i}_hc_attn_scale"], + w[f"l{i}_hc_attn_base"], + ), + "hc_ffn": (w[f"l{i}_hc_ffn_fn"], w[f"l{i}_hc_ffn_scale"], w[f"l{i}_hc_ffn_base"]), + "attn_norm": w[f"l{i}_attn_norm"], + "ffn_norm": w[f"l{i}_ffn_norm"], + "wqkv_a": (w[f"l{i}_wqkv_a"], w[f"l{i}_wqkv_a_sf"]), + "q_norm": w[f"l{i}_q_norm"], + "kv_norm": w[f"l{i}_kv_norm"], + "wq_b": (w[f"l{i}_wq_b"], w[f"l{i}_wq_b_sf"]), + "sink": w[f"l{i}_sink"], + # [G_loc, 4096, 1024] view: group g's wo_a transposed. + "wo_a_t": torch.transpose( + w[f"l{i}_wo_a"].view(self.g_loc, self.o_lora, self.group_in), 1, 2 + ), + "wo_b": (w[f"l{i}_wo_b"], w[f"l{i}_wo_b_sf"]), + "gate_w_t": w[f"l{i}_gate_w"].t(), + "gate_bias": w[f"l{i}_gate_bias"], + "gate_bias_vl": w[f"l{i}_gate_bias_vl"], + "sh_w13": (w[f"l{i}_sh_w13"], w[f"l{i}_sh_w13_sf"]), + "sh_w2": (w[f"l{i}_sh_w2"], w[f"l{i}_sh_w2_sf"]), + "fc1_w": w[f"l{i}_fc1_w"], + "fc1_s": w[f"l{i}_fc1_s"].view(torch.float8_e4m3fn), + "fc2_w": w[f"l{i}_fc2_w"], + "fc2_s": w[f"l{i}_fc2_s"].view(torch.float8_e4m3fn), + } + if i in self.kv_sources: + d["comp_w_t"] = w[f"l{i}_comp_w"].t() + d["comp_norm"] = w[f"l{i}_comp_norm"] + d["idx_wk_t"] = w[f"l{i}_idx_wk"].t() + d["idx_k_norm"] = w[f"l{i}_idx_k_norm"] + if i in self.index_sources: + d["idx_wq_b"] = (w[f"l{i}_idx_wq_b"], w[f"l{i}_idx_wq_b_sf"]) + d["idx_wproj_t"] = w[f"l{i}_idx_wproj"].t() + if i in self.engram_layers: + if not self.eng_host: + d["eng_table"] = w[f"l{i}_eng_table"] + d["eng_scale"] = w[f"l{i}_eng_scale"] + d["eng_wkv"] = (w[f"l{i}_eng_wkv"], w[f"l{i}_eng_wkv_sf"]) + d["eng_q"] = w[f"l{i}_eng_q"] + d["eng_k"] = w[f"l{i}_eng_k"] + # NVFP4 routed-expert scalars (modelopt stores reciprocals: + # input_scale = amax/(448*6)). FC1's activation global scale is one + # number for the layer: the max over all 384 w1/w3 input scales + # (R-M2), identical on every rank. FC2's is LAYER-level too (C4: + # a per-expert one saturates the NVFP4 block scales). + isc1, isc3 = w[f"l{i}_e_isc1"], w[f"l{i}_e_isc3"] + ws1, ws3 = w[f"l{i}_e_ws1"], w[f"l{i}_e_ws3"] + assert torch.equal(ws1, ws3), ( + f"layer {i}: w1/w3 weight_scale_2 differ; the fused FC1 and its " + "clamp need one dequant scale per expert" + ) + in1 = torch.maximum(isc1.max(), isc3.max()) + in2 = w[f"l{i}_e_isc2"].max() + win = slice(self.e_off, self.e_off + self.e_loc) + a1 = (in1 * ws1[win]).contiguous() + d["g1"] = (1.0 / in1).reshape(1).contiguous() + d["o1g"] = a1 + d["o1"] = (a1 / in2).contiguous() + d["o2"] = (in2 * w[f"l{i}_e_ws2"][win]).contiguous() + # swiglu_limit in the FC1 accumulator domain. + d["clamp"] = (_SWIGLU_LIMIT / a1).contiguous() + L[i] = d + self._L = L + self._rope_positions = 0 + self._ensure_rope(dev, 65536) + tj = os.path.join(str(self._ckpt_dir), "tokenizer.json") + assert os.path.isfile(tj), ( + f"Engram's compressed vocab is built from the checkpoint tokenizer; {tj} not found" + ) + self._eng = _engram_layout(self._tc, tj, dev) + self._engram_residency() + self._const = { + "neg1_i32": torch.full((1, 1), -1, dtype=torch.int32, device=dev), + # The index top-k's per-row offsets (window-local positions: 0). + "zero_i32_rows": torch.zeros((_TOPK_MAX_ROWS,), dtype=torch.int32, device=dev), + "zero_bf16": torch.zeros((1, 1), dtype=torch.bfloat16, device=dev), + "ape": torch.zeros((2, self.head_dim), dtype=torch.float32, device=dev), + "dummy_comp": torch.zeros((1, 1, 64, self.head_dim), dtype=torch.bfloat16, device=dev), + "attn_ws": torch.zeros(_ATTN_WS_BYTES, dtype=torch.uint8, device=dev), + "false_b": torch.zeros((1,), dtype=torch.bool, device=dev), + "true_b": torch.ones((1,), dtype=torch.bool, device=dev), + # The main KV's NVFP4 has no global scale: the dequantizer's is 1. + "one_f32": torch.ones((1,), dtype=torch.float32, device=dev), + } + if self.fp8_attn: + self._derive_fp8_kv_scales() + self._const["dummy_comp8"] = torch.zeros( + (1, 1, 64, self.head_dim), dtype=torch.float8_e4m3fn, device=dev + ) + self._vision.post_load_weights() + if not self.dp and self.ar_one_node: + from tensorrt_llm._torch.distributed.ops import get_allreduce_workspace + from tensorrt_llm.mapping import Mapping + + self._ar_ws = get_allreduce_workspace( + Mapping( + world_size=self.world, + rank=self.rank, + gpus_per_node=self.gpus_per_node, + tp_size=self.tp, + pp_size=self.pp, + ) + ) + # Decoder replay: row j of a tail keeps window columns 0..j (its keys + # r .. r + j); column 0 alone names a row's own window slot. + w128 = torch.arange(_WINDOW, device=dev) + self._const["tri"] = (w128.view(1, -1) <= w128.view(-1, 1)).contiguous() + self._const["col0"] = (w128 == 0).view(1, -1).contiguous() + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + + def _derive_fp8_kv_scales(self) -> None: + """FP8 attention: one power-of-two KV scale per layer, from the norm weights + (``_FP8_KV_BOUND``). One attention call takes one scale for both of its pools, so + every layer that reads a compressed source shares that source's scale, taken over + the source's compressed bound and every reader's window bound; a window-only layer + has its own. With a power-of-two scale the window's MXFP8 values stay exact in + E4M3 wherever they land in E4M3's normal range.""" + w = self.w + # This pipeline stage's layers (a group's source is on the stage with its readers). + bound = { + i: _FP8_KV_BOUND * float(w[f"l{i}_kv_norm"].abs().max()) for i in self.pp_layer_ids + } + groups: dict = {} + for i in self.pp_layer_ids: + groups.setdefault(self.kv_src[i] if self.ratios[i] else -1 - i, []).append(i) + for key, layers in groups.items(): + b = max(bound[i] for i in layers) + if key >= 0: + b = max(b, _FP8_KV_BOUND * float(w[f"l{key}_comp_norm"].abs().max())) + s = 2.0 ** math.ceil(math.log2(b / 448.0)) + for i in layers: + self._fp8_kv_scale[i] = s + if self.rank == 0: + sys.stdout.write( + "[modeling_v2 deepseek_v41] FP8 attention KV scales (log2, per layer): " + + " ".join(str(int(math.log2(s))) for s in self._fp8_kv_scale if s is not None) + + "\n" + ) + sys.stdout.flush() + + def _ensure_rope(self, dev, positions: int) -> None: + """Both rope tables, sized to ``positions`` rows. Every row depends + only on its own position, so a longer table is the same table with + more rows; regrown only outside capture (first forward).""" + if positions <= self._rope_positions: + return + self._rope_positions = positions + self._rope_tables = { + 0: _rope_table(self.theta, positions, None, dev), + 1: _rope_table(self.comp_theta, positions, self.yarn, dev), + } + + def _table(self, ratio: int) -> torch.Tensor: + # Ratio-0 layers: theta 10000, no YaRN. Compressed layers: theta + # 160000 with YaRN (q, window kv, compressed latent, index q/k). + return self._rope_tables[0 if ratio == 0 else 1] + + # ------------------------------------------------------------------ + def _derive_first_step(self, md, position_ids) -> None: + """What the first step fixes, run once and always: the engine-side layout this + assembly is written for (the cache manager's, fixed at engine construction), and + the tables sized by the engine's own limits. + + Unconditional, unlike the targets' step-contract probe: the rope table and the + top-k offsets decide what the model computes, and the layout checks are facts the + routing context cannot see.""" + if not self._derive_pending: + return + assert isinstance(md, DeepseekV4TrtllmAttentionMetadata), type(md) + assert md.is_deepseek_v41, "the cache manager is not in DeepSeek-V4.1 layout mode" + assert position_ids.dtype == torch.int32, position_ids.dtype + assert md.tokens_per_block == _PAGE, md.tokens_per_block + layout = md.kv_cache_manager.v41_layout + assert list(layout.compress_ratios) == self.ratios, "manager layout != checkpoint roles" + for i in range(self.num_layers): + if self.ratios[i]: + assert layout.kv_source(i) == self.kv_src[i], (i, layout.kv_source(i)) + assert layout.index_source(i) == self.idx_src[i], (i, layout.index_source(i)) + assert layout.engram_token_owner == self.engram_layers[0] + assert md.effective_beam_width == 1, "beam search is not implemented" + assert not md.is_cross, "cross attention is not implemented" + self._ensure_rope(self.w["final_norm"].device, int(md.max_seq_len) + 64) + # One zero offset per row of the largest step (outside capture). + max_t = int(md.swa_local_indices_cuda.shape[0]) + if self._const["zero_i32_rows"].shape[0] < max_t: + self._const["zero_i32_rows"] = torch.zeros( + (max_t,), dtype=torch.int32, device=self.w["final_norm"].device + ) + # Candidate masking can bind in a decode step only when a sequence may + # hold more than 2048 x 8 ratio-1 compressed positions: a per-engine + # constant (the paged logits' static width), so the decode graph is + # built with or without it once. + self._cand_binds_decode = md.max_compressed_kv_lens[1] > _CAND_TOPK_BLOCKS * _CAND_BLOCK + self._derive_pending = False + + def _get_pools(self, md) -> dict: + """Pool views of the engine's cache manager, keyed on the manager + object (a KV-size-estimation manager precedes the final one). Views + only: the pools are the manager's and never move. Under native KV + storage this also allocates, once per manager, the two bf16 read + buffers the attention is handed (see ``_native_window``).""" + mgr = md.kv_cache_manager + if self._pools is not None and self._pool_mgr is not None and self._pool_mgr() is mgr: + return self._pools + T = DeepseekV4AttentionType + layout = mgr.v41_layout + native = layout.kv_storage == DEEPSEEK_V41_KV_STORAGE_NATIVE + assert native or layout.kv_storage == DEEPSEEK_V41_KV_STORAGE_BF16, layout.kv_storage + if self.fp8_attn and not native: + raise ValueError( + f"{_FP8_ATTN_ENV}=1 reads the native KV pools; it does not run with KV storage " + f"{layout.kv_storage!r}" + ) + H = self.head_dim + + def view(v, shape): + # A pool is only ever re-viewed, never copied: a copy would take + # the writes away from the pool the attention reads. + out = torch.reshape(v, shape) + assert out.data_ptr() == v.data_ptr() and out.numel() == v.numel(), ( + "pool view is not a view of the manager's pool" + ) + return out + + # bf16 containers: swa3/comp3 are written, swa4/comp4 read by the + # attention. Native: (codes, scales) views of the 528 / 288 B rows are + # written, and read back through the dequantizers. + swa3, swa4, swa_q = {}, {}, {} + for i in self.pp_layer_ids: + if native: + codes, scales = mgr.get_v41_native_kv_views(i, T.SWA) + assert codes.dtype == torch.uint8 and tuple(codes.shape[1:]) == (_PAGE, H), ( + codes.shape + ) + assert tuple(scales.shape[1:]) == (_PAGE, H // _SF32), scales.shape + swa_q[i] = (codes, scales) + else: + v = mgr.get_buffers(i, T.SWA) + assert v.dtype == torch.bfloat16 and tuple(v.shape[1:]) == (_PAGE, H), v.shape + swa3[i] = v + swa4[i] = view(v, [v.shape[0], 1, v.shape[1], v.shape[2]]) + comp3, comp4, comp_q, comp_q4, comp_q1, idxk, ckv, csc = {}, {}, {}, {}, {}, {}, {}, {} + bytes_per_token2 = 0 # twice the persistent KV bytes per token, as allocated + for o in self.kv_sources: + if o not in self.local: + continue + bk = _PAGE // self.ratios[o] + v = mgr.get_buffers(o, T.COMPRESS) + if native: + codes, scales = mgr.get_v41_native_kv_views(o, T.COMPRESS) + assert codes.dtype == torch.uint8 and tuple(codes.shape[1:]) == (bk, H // 2), ( + codes.shape + ) + assert tuple(scales.shape[1:]) == (bk, H // _SF16), scales.shape + comp_q[o] = (codes, scales) + # The dequantizer's [pages, page_size, 1, W] operands. + comp_q4[o] = ( + view(codes, [codes.shape[0], bk, 1, H // 2]), + view(scales, [scales.shape[0], bk, 1, H // _SF16]), + ) + # The same rows as pages of one row: a slot number + # (page * bk + slot) is then its own page index. + comp_q1[o] = ( + view(codes, [codes.shape[0] * bk, 1, 1, H // 2]), + view(scales, [scales.shape[0] * bk, 1, 1, H // _SF16]), + ) + else: + assert v.dtype == torch.bfloat16 and tuple(v.shape[1:]) == (bk, H), v.shape + comp3[o] = v + comp4[o] = view(v, [v.shape[0], 1, v.shape[1], v.shape[2]]) + idxk[o] = mgr.get_indexer_k_cache_buffers(o) + assert idxk[o].dtype == torch.uint8 and tuple(idxk[o].shape[1:]) == (bk, 1, 68), idxk[ + o + ].shape + if self.ratios[o] == 2: + ckv[o] = mgr.get_buffers(o, T.COMPRESSOR_KV) + csc[o] = mgr.get_buffers(o, T.COMPRESSOR_SCORE) + assert ckv[o].dtype == torch.float32 and csc[o].dtype == torch.float32 + # A ratio-r owner stores one main-KV row and one index-K row per r tokens. + row_bytes = v.shape[-1] * v.element_size() + idxk[o].shape[-1] * idxk[o].element_size() + bytes_per_token2 += 2 * row_bytes // self.ratios[o] + tok = None + if self.eng_local: + tok = mgr.get_buffers(layout.engram_token_owner, T.ENGRAM_TOKEN) + assert ( + tok.dtype == torch.int32 + and tok.shape[1] == _PAGE + and tok.numel() == tok.shape[0] * _PAGE + ) + tok = view(tok, [tok.shape[0], tok.shape[1]]) + local = {i: mgr.layer_offsets[i] for i in self.pp_layer_ids} + self._pools = { + "native": native, + "swa3": swa3, + "swa4": swa4, + "swa_q": swa_q, + "comp3": comp3, + "comp4": comp4, + "comp_q": comp_q, + "comp_q4": comp_q4, + "comp_q1": comp_q1, + "idxk": idxk, + "ckv": ckv, + "csc": csc, + "tok": tok, + "local": local, + } + if native: + self._pools.update(self._native_read_buffers(md)) + self._pool_mgr = weakref.ref(mgr) + + # Persistent KV bytes per token as allocated (windowed state excluded). + assert bytes_per_token2 % 2 == 0, bytes_per_token2 + bpt = bytes_per_token2 // 2 + assert bpt == mgr.get_cache_bytes_per_token(), (bpt, mgr.get_cache_bytes_per_token()) + lines = [ + f"MODELING_V2 deepseek_v41 rank {self.rank}: KV storage {layout.kv_storage}, " + f"persistent KV {bpt} B/token\n" + ] + if os.environ.get(_WITNESS_ENV) == "1": + lines.append(f"MODELING_V2_WITNESS kv_bytes_per_token={bpt}\n") + sys.stdout.write("".join(lines)) + sys.stdout.flush() + return self._pools + + def _native_read_buffers(self, md) -> dict: + """The bf16 rows the attention reads under native KV storage, allocated + once per manager (outside capture) and zero-filled, so every row the + kernel may load is finite: + + * ``swa_c`` [max_tokens * 128, 512]: query t's window column c at row + ``t * 128 + c``; ``win_idx`` [max_tokens, 128] is that numbering, the + window half of every renumbered index table; + * ``comp_c``: as many rows as the smaller of its two forms needs at + the engine's capacity, holding either form (``_comp_form`` picks one + per step): + + - per request (``comp_out`` / ``comp_c4_req`` / ``comp_tbl``): + request b's compressed entry e at row ``b * L_r + e``, + ``L_r = cols * 128 / r``, viewed per ratio as the dequantizer's + ``[seqs, L_r, 1, 512]`` output and as the attention's + ``[seqs * cols, 1, 128 / r, 512]`` pool; ``comp_tbl[r]`` [seqs, + cols] (row b = ``b * cols + j``) is the page table that numbering + implies, so the index op maps top-k entry e of request b onto row + b * L_r + e. Each KV source dequantizes every entry of every + request of the step (``_native_compressed``). It holds the first + ``comp_req_fit`` request slots; + - per query (``topk_out`` / ``comp_c4_q``): query t's top-k column k + at row ``t * 512 + k`` (``topk_ren`` is that numbering), + dequantized from the pool slot the index op names + (``_native_topk``) once per top-k set. It holds ``comp_q_fit`` + query rows. + + The attention reads the same bf16 values in the same column order + from either form. + """ + assert not torch.cuda.is_current_stream_capturing(), "read buffers are allocated eagerly" + dev = self.w["final_norm"].device + H = self.head_dim + max_t = md.swa_local_indices_cuda.shape[0] + # FP8 attention reads E4M3 rows (per-tensor scaled); otherwise the exact bf16 values. + rdt = torch.float8_e4m3fn if self.fp8_attn else torch.bfloat16 + swa_c = torch.zeros((max_t * _WINDOW, H), dtype=rdt, device=dev) + win_idx = torch.arange(max_t * _WINDOW, dtype=torch.int32, device=dev).view(max_t, _WINDOW) + # The ratios of this stage's KV sources (a pipeline stage may hold none: + # its buffers are then sized for the smallest ratio of the layout). + ratios = sorted({self.ratios[o] for o in self.kv_sources if o in self.local}) + rmin = min(ratios) if ratios else min(md.compress_block_tables) + seqs, cols = md.compress_block_tables[rmin].shape + for r in ratios: + assert tuple(md.compress_block_tables[r].shape) == (seqs, cols), ( + md.compress_block_tables[r].shape + ) + out = { + "swa_c": swa_c, + "swa_c4": torch.reshape(swa_c, [max_t, 1, _WINDOW, H]), + "win_idx": win_idx, + } + rows_request = seqs * cols * (_PAGE // rmin) + rows_topk = max_t * _TOPK + rows = min(rows_request, rows_topk) + comp_c = torch.zeros((rows, H), dtype=rdt, device=dev) + req_fit = rows // (cols * (_PAGE // rmin)) + q_fit = rows // _TOPK + comp_out, comp_c4_req, comp_c4_q = {}, {}, {} + for r in ratios: + bk = _PAGE // r + n = req_fit * cols * bk + comp_out[r] = torch.reshape(comp_c[:n], [req_fit, cols * bk, 1, H]) + comp_c4_req[r] = torch.reshape(comp_c[:n], [req_fit * cols, 1, bk, H]) + comp_c4_q[r] = torch.reshape(comp_c[: q_fit * _TOPK], [q_fit * _TOPK // bk, 1, bk, H]) + tbl = torch.arange(req_fit * cols, dtype=torch.int32, device=dev).view(req_fit, cols) + out.update( + comp_cols=cols, + comp_req_fit=req_fit, + comp_q_fit=q_fit, + comp_out=comp_out, + comp_c4_req=comp_c4_req, + comp_tbl={r: tbl for r in ratios}, + comp_c4_q=comp_c4_q, + topk_out=torch.reshape(comp_c[: q_fit * _TOPK], [q_fit, _TOPK, 1, H]), + topk_ren=torch.arange(q_fit * _TOPK, dtype=torch.int32, device=dev).view(q_fit, _TOPK), + # The index op's pool slots, contiguous: the dequantizer's page table. + topk_slots=torch.full((q_fit, _TOPK), -1, dtype=torch.int32, device=dev), + topk_cols=torch.full((q_fit,), _TOPK, dtype=torch.int32, device=dev), + ) + sys.stdout.write( + f"MODELING_V2 deepseek_v41 rank {self.rank}: native compressed read buffer {rows} rows: " + f"per request {req_fit} of {seqs} request slots (x cols {cols} x {_PAGE // rmin}), " + f"per query {q_fit} of {max_t} rows (x {_TOPK})\n" + ) + sys.stdout.flush() + return out + + def _comp_form(self, pools, md, st, tail) -> str: + """The native compressed read buffer's form for this step, "request" or + "topk" (``_native_read_buffers``): one per step, as both forms share + the buffer. A form is usable when the step fits it (every request slot + of the step, every query row of each pass); of the usable ones, the one + that dequantizes fewer rows: + + * per request, every compressed entry of every request, once per KV + source: sum over sources and requests of kv_len // ratio; + * per query, 512 rows per query row once per top-k set: the index + sources over the step's rows, plus, in a bounded step, the layers + past L over the decoder rows (whose first layer re-reads layer L's + set). + + A generation-only step is the same at every length in one CUDA graph, + so it prices each request at half ``max_seq_len``: the per-request form + at the short envelopes, the per-query form for long ones. ``st["nc"]`` + is the step's context-request count as the calling target states it.""" + B, T, nc = st["B"], st["T"], st["nc"] + bounded = tail is not None and tail.bound + rows = max(T, min(tail.rows, int(md.swa_local_indices_cuda.shape[0]))) if bounded else T + fits_req = B <= pools["comp_req_fit"] + fits_q = rows <= pools["comp_q_fit"] + assert fits_req or fits_q, (B, rows, pools["comp_req_fit"], pools["comp_q_fit"]) + if not (fits_req and fits_q): + form = "request" if fits_req else "topk" + else: + if nc: + kv = (md.cached_token_lens_cpu[:B] + md.seq_lens_kv[:B]).tolist() + else: + kv = [int(md.max_seq_len) // 2] * B + cost_req = sum( + n // self.ratios[o] for o in self.kv_sources if o in self.local for n in kv + ) + index_local = [i for i in self.index_sources if i in self.local] + early = sum(1 for i in index_local if i < self.late) + if bounded: + cost_q = _TOPK * (early * T + (len(index_local) - early + 1) * tail.rows) + else: + cost_q = _TOPK * len(index_local) * T + form = "request" if cost_req <= cost_q else "topk" + _rv("native_read.form", form, B, T, rows) + return form + + # ------------------------------------------------------------------ + def attach_engram_host_table(self, checkpoint_dir) -> None: + """host_offload: map the node's shared host copy of both Engram tables. + + Collective over the node's ranks (``open_shared_engram_table``). With + the default ``vmm`` backend, local rank 0 imports the table from a + process on this node that already serves it, or creates and fills the + CUDA VMM host-NUMA allocations from the checkpoint; after a node + barrier every rank imports and maps it, and the views are CUDA views of + host memory. With the ``shm`` backend, local rank 0 fills (or reuses) + the ``/dev/shm`` file and every rank maps and registers it, and the + views are pinned CPU tensors whose mapped device pointer is resolved + here, before any CUDA-graph capture. Each table is checked against the + checkpoint geometry. The views must not move or go away while a graph + that captured a gather can replay, so the mapping stays until the + process exits. + """ + assert self.eng_attach and self._eng_table is None + dev = torch.device("cuda", torch.cuda.current_device()) + table = _engram_host.open_shared_engram_table( + str(checkpoint_dir), + local_rank=local_mpi_rank(), + barrier=local_mpi_barrier, + device=dev, + owned=True, + ) + assert tuple(table.layers) == self.engram_layers, (table.layers, self.engram_layers) + for li, i in enumerate(self.engram_layers): + wt, sc = table.weight(i), table.scale(i) + assert tuple(wt.shape) == (self.eng_rows[li], self.eng_dim), (i, tuple(wt.shape)) + assert tuple(sc.shape) == (self.eng_rows[li], self.eng_dim // 32), (i, tuple(sc.shape)) + assert wt.dtype == torch.float8_e4m3fn and sc.element_size() == 1, (wt.dtype, sc.dtype) + table.check_gather_views(dev) + self._eng_table = table + atexit.register(table.close, remove_if_owned=True) + sys.stdout.write( + f"MODELING_V2 deepseek_v41 rank {self.rank}: Engram host table {table.describe()}\n" + ) + sys.stdout.flush() + + def _engram_residency(self) -> None: + """Hand the Engram layers their tables and print this rank's device + memory after the weight load, the residency witness's reading.""" + eng = [i for i in self.engram_layers if i in self.local] + if self.eng_attach and self._eng_table is None: + # A load path that never ran the weight table's load (LoadFormat.DUMMY). + self.attach_engram_host_table(self._ckpt_dir) + if self.eng_host: + for i in eng: + self._L[i]["eng_table"] = self._eng_table.weight(i) + self._L[i]["eng_scale"] = self._eng_table.scale(i).view(torch.uint8) + # Host-resident tables are CUDA views under the vmm backend; none of their bytes are on the device. + dev_bytes = ( + 0 + if self.eng_host + else sum( + self._L[i][k].numel() * self._L[i][k].element_size() + for i in eng + for k in ("eng_table", "eng_scale") + if self._L[i][k].is_cuda + ) + ) + torch.cuda.synchronize() + sys.stdout.write( + f"MODELING_V2 deepseek_v41 rank {self.rank}: Engram tables " + f"{'host (shared pinned)' if self.eng_host else 'device (row shard)'}, " + f"{dev_bytes / 2**20:.1f} MiB of them on the device; after weight load " + f"__DEVICE_MEMORY_MIB__={torch.cuda.memory_allocated() / 2**20:.1f}\n" + ) + sys.stdout.flush() + + def _dp_rows(self, md, num_tokens: int) -> int: + """Uniform row count of the dep4 MoE / Engram round trips: the max of + the engine's per-rank token counts (host ints, identical on every rank).""" + counts = [int(n) for n in md.all_rank_num_tokens] + assert len(counts) == self.tp, counts + assert counts[self.tp_rank] == num_tokens, (self.tp_rank, num_tokens, counts) + return max(counts) + + # ------------------------------------------------------------------ + def _select_target(self, attn_metadata) -> Target: + """Which target runs this step.""" + return self._decode if phase_of(attn_metadata) is Phase.DECODE else self._prefill + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None = None, + position_ids: torch.IntTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + lora_params: dict | None = None, + **kwargs, + ) -> torch.Tensor: + """Derive what the first step fixes, keep the per-manager state current, then + route the step to its target.""" + assert self._prefill is not None, "load_weights must run before forward" + assert lora_params is None, "LoRA is not implemented by this target" + assert inputs_embeds is None, ( + "this target embeds token ids itself; images arrive as multimodal_params" + ) + assert input_ids is not None and position_ids is not None + md = attn_metadata + advance_step_generation() + self._derive_first_step(md, position_ids) + spec_metadata = kwargs.get("spec_metadata") + if spec_metadata is not None and not self._spec_checked: + assert self.ds is not None, ( + "spec_metadata reached a core built without the DSpark drafter" + ) + assert tuple(spec_metadata.layers_to_capture or ()) == self.ds.targets, ( + f"DSpark reads the input stream of layers {self.ds.targets}; the runtime " + f"captures {spec_metadata.layers_to_capture}" + ) + self._spec_checked = True + self._resolve_replay(md.kv_cache_manager) + if self._tail_records: + self._release_tail_records(md.kv_cache_manager) + return self._select_target(md).forward( + md, input_ids=input_ids, position_ids=position_ids, **kwargs + ) + + +# The MXFP8 x MXFP8 linears of a trunk layer, by the operand-table key each binds. +_MX_LINEARS = ("wqkv_a", "wq_b", "wo_b", "idx_wq_b", "sh_w13", "sh_w2", "eng_wkv") + + +class PrefillTarget(Target): + """The general case: context rows, and possibly generation rows beside them. + + Owns everything a context step needs beyond decode: the vision splice, the + encoder replay's compressor rows and Engram history, the context attention, + the context indexer with its candidate-block masks, and the decoder SWA + bounded replay (``_DecoderTail``). + """ + + def __init__(self, core: "DeepseekV41UnifiedModelingV2") -> None: + """Bind every op this target calls, once, against the real weights.""" + super().__init__(core) + L = core._L + eps = core.eps + e = core._eng + + # MXFP8 linears: one activation quantizer, one GEMM per weight role. + self._mxq = Mxfp8Quantize() + self._mx = {name: Mxfp8Mxfp8Gemm() for name in _MX_LINEARS} + self._mx_shape: dict = {} + + # Norms. + self._attn_norm = FlashinferRmsnorm() + self._ffn_norm = FlashinferRmsnorm() + self._q_norm = FlashinferRmsnorm() + self._kv_norm = FlashinferRmsnorm() + self._comp_norm = FlashinferRmsnorm() + self._idx_k_norm = FlashinferRmsnorm() + for op in ( + self._attn_norm, + self._ffn_norm, + self._q_norm, + self._kv_norm, + self._comp_norm, + self._idx_k_norm, + ): + op.bind_const(eps=eps) + self._final_norm = FlashinferRmsnorm() + self._final_norm.bind_const(weight=core.w["final_norm"], eps=eps) + + # Attention: RoPE, the window and compressed KV writes, their native + # read-back, the index table, the sparse MLA core, the grouped output. + self._rope = MlaRopeInplace() + self._swa_write = Mxfp8FakeQuantize() + self._comp_write = Nvfp4FakeQuantize() + self._swa_read = Mxfp8Dequantize() + self._swa_read8 = Mxfp8ToE4m3() + self._comp_read = Nvfp4KvToBf16() + self._comp_read8 = Nvfp4KvToE4m3() + self._l2g = Dsv4LocalToGlobalIndices() + self._q_e4m3 = E4m3QuantizePow2Dynamic() + self._attn = TrtllmSparseMlaDsv4() + self._attn.bind_const( + workspace_buffer=core._const["attn_ws"], + softmax_scale=core.head_dim**-0.5, + output_scale=1.0, + ) + self._attn8 = TrtllmSparseMlaDsv4Fp8() + self._attn8.bind_const(workspace_buffer=core._const["attn_ws"]) + self._wo_a = BmmOut() + + # Compressor and index K (the KV sources). + self._comp_mm = CublasMm() + self._comp_prefill = CompressorPrefillReduction() + self._comp_paged = CompressorPagedKvCompress() + for op in (self._comp_prefill, self._comp_paged): + op.bind_const( + ape=core._const["ape"], page_size=_PAGE, head_dim=core.head_dim, compress_ratio=2 + ) + self._idx_wk = CublasMm() + # MXFP4 (UE8M0 per 32, linear scales): the index K and the index query. + self._mxfp4 = Fp4Quantize() + self._mxfp4.bind_const( + global_scale=None, sf_vec_size=_SF32, sf_use_ue8m0=True, is_sf_swizzled_layout=False + ) + self._idx_scatter = IndexerKCacheScatter() + + # Indexer (the index sources) and the candidate-block masks. + self._idx_wproj = CublasMm() + self._idx_gather = IndexerKCacheGather() + self._idx_gather.bind_const(head_dim=_IDX_DIM // 2) + self._mqa_logits = DeepgemmFp8Fp4MqaLogits() + self._mqa_logits.bind_const(clean_logits=True, max_seqlen_k=0, logits_dtype=torch.float32) + self._paged_logits = DeepgemmFp8Fp4PagedMqaLogits() + self._paged_logits.bind_const(logits_dtype=torch.float32) + self._paged_meta = DeepgemmPagedMqaLogitsMetadata() + self._cand = CandidateBlockSelect() + self._cand.bind_const(topk_blocks=_CAND_TOPK_BLOCKS, block_size=_CAND_BLOCK, next_n=1) + self._topk = FlashinferTopKRaggedDeterministic() + self._topk.bind_const(k=_TOPK) + + # MoE: router (fp32 logits), routing with gate.bias and gate.bias_vl, + # the NVFP4 activation quantizer and runner, the shared expert's SwiGLU. + self._router = CublasMm() + self._router.bind_const(bias=None, out_dtype=torch.float32) + self._gate = GateForward() + self._gate_vl = GateForward() + for op in (self._gate, self._gate_vl): + op.bind_const(topk=core.topk, route_scale=core.route_scale) + self._moe_quant = Fp4Quantize() + self._moe_quant.bind_const( + sf_vec_size=_SF16, sf_use_ue8m0=False, is_sf_swizzled_layout=_SF_LINEAR + ) + self._moe_runner = Fp4BlockScaleMoeRunner() + self._moe_runner.bind_const( + routing_logits=None, + routing_bias=None, + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm2_bias=None, + num_experts=core.num_experts, + top_k=core.topk, + n_group=None, + topk_group=None, + intermediate_size=core.i_loc, + local_expert_offset=core.e_off, + local_num_experts=core.e_loc, + routed_scaling_factor=None, + routing_method_type=_ROUTING_METHOD_INERT, + do_finalize=True, + act_type=_ACT_SWIGLU, + ) + self._silu = SiluAndMul() + self._silu.bind_const(swiglu_limit=_SWIGLU_LIMIT) + + # Collectives: the TP (or attention-DP) group of this stage. + self._allgather = Allgather() + self._allgather.bind_const(sizes=None, group=core.group) + self._reducescatter = Reducescatter() + self._reducescatter.bind_const(sizes=None, group=core.group) + # fp32 sums on AUTO (the IPC workspace) when the group is on one node, + # NCCL across nodes; bf16 sums on NCCL. + self._ar32_strategy = _AR_AUTO if core._ar_ws is not None else _AR_NCCL + self._allreduce32 = Allreduce() + self._allreduce16 = Allreduce() + for op, ws, strategy in ( + (self._allreduce32, core._ar_ws, self._ar32_strategy), + (self._allreduce16, None, _AR_NCCL), + ): + op.bind_const( + residual=None, + norm_weight=None, + scale=None, + bias=None, + workspace=ws, + group=core.group, + strategy=strategy, + op=0, + eps=1e-6, + trigger_completion_at_end=True, + ) + self._pp_send = PpSendTensors() + self._pp_recv = PpRecvTensors() + + # Engram: the hash constants, the table gather, the gate. + self._eng_hash = EngramHash() + self._eng_hash.bind_const( + token_map=e["token_map"], + multipliers=e["multipliers"], + primes=e["primes"], + offsets=e["offsets"], + pad_id=e["pad_id"], + ) + self._eng_gather = EngramGather() + self._eng_gate = EngramGate() + self._eng_gate.bind_const(eps=eps, clamp_value=1e-6) + + # Hyper-connections: per sublayer the mix GEMM + squared sums, the fused + # sinkhorn, and the lagged-pre collapse; the post mapping. + self._hc_attn_gemm = MhcGemmSqrsumFma() + self._hc_ffn_gemm = MhcGemmSqrsumFma() + for op in (self._hc_attn_gemm, self._hc_ffn_gemm): + op.bind_const(N=_MIX, K=core.hc_k) + self._hc_attn_fuse = MhcBigFuse() + self._hc_ffn_fuse = MhcBigFuse() + for op in (self._hc_attn_fuse, self._hc_ffn_fuse): + op.bind_const( + K=core.hc_k, + hidden_size=0, + rms_eps=eps, + hc_pre_eps=core.hc_eps, + hc_sinkhorn_eps=core.hc_eps, + hc_post_mult_value=2.0, + sinkhorn_repeat=20, + num_splits=1, + block_size=128, + ) + self._hc_pre_attn = MhcHcHeadApply() + self._hc_pre_ffn = MhcHcHeadApply() + self._hc_pre_final = MhcHcHeadApply() + for op in (self._hc_pre_attn, self._hc_pre_ffn, self._hc_pre_final): + op.bind_const( + mult=_HC, hidden_size=core.hidden, K=core.hc_k, norm_eps=eps, eps=core.hc_eps + ) + self._hc_post = MhcPostMapping() + + for i in core.pp_layer_ids: + d = L[i] + for name in _MX_LINEARS: + if name in d: + self._bind_mx(name, i, d[name]) + self._attn_norm.bind_layered(i, weight=d["attn_norm"]) + self._ffn_norm.bind_layered(i, weight=d["ffn_norm"]) + self._q_norm.bind_layered(i, weight=d["q_norm"]) + self._kv_norm.bind_layered(i, weight=d["kv_norm"]) + self._attn.bind_layered(i, sinks=d["sink"]) + if core.fp8_attn: + s_kv = core._fp8_kv_scale[i] + self._attn8.bind_layered(i, sinks=d["sink"], bmm2_scale=s_kv) + # The call's QK scale is softmax scale x query scale x KV scale; + # the quantizer writes it with its query scale folded in. + self._q_e4m3.bind_layered(i, scale_mult=core.head_dim**-0.5 * s_kv) + self._wo_a.bind_layered(i, b=d["wo_a_t"]) + self._router.bind_layered(i, mat_b=d["gate_w_t"]) + self._gate.bind_layered(i, bias=d["gate_bias"]) + self._gate_vl.bind_layered(i, bias=d["gate_bias_vl"]) + self._moe_quant.bind_layered(i, global_scale=d["g1"]) + self._moe_runner.bind_layered( + i, + gemm1_weights=d["fc1_w"], + gemm1_weights_scale=d["fc1_s"], + gemm1_clamp_limit=d["clamp"], + gemm2_weights=d["fc2_w"], + gemm2_weights_scale=d["fc2_s"], + output1_scale_scalar=d["o1"], + output1_scale_gate_scalar=d["o1g"], + output2_scale_scalar=d["o2"], + ) + fn, sc, bs = d["hc_attn"] + self._hc_attn_gemm.bind_layered(i, w=fn) + self._hc_attn_fuse.bind_layered(i, hc_scale=sc, hc_base=bs) + # The FFN input collapses with this attention's mixes. + self._hc_pre_ffn.bind_layered(i, scale=sc, base=bs) + fn, sc, bs = d["hc_ffn"] + self._hc_ffn_gemm.bind_layered(i, w=fn) + self._hc_ffn_fuse.bind_layered(i, hc_scale=sc, hc_base=bs) + if i > 0: + # Lagged pre: the attention input collapses with the previous + # layer's FFN mixes (the previous stage's, at a stage's first layer). + _, psc, pbs = L[i - 1]["hc_ffn"] + self._hc_pre_attn.bind_layered(i, scale=psc, base=pbs) + if i in core.kv_sources: + self._comp_mm.bind_layered(i, mat_b=d["comp_w_t"]) + self._comp_norm.bind_layered(i, weight=d["comp_norm"]) + self._idx_wk.bind_layered(i, mat_b=d["idx_wk_t"]) + self._idx_k_norm.bind_layered(i, weight=d["idx_k_norm"]) + if i in core.index_sources: + self._idx_wproj.bind_layered(i, mat_b=d["idx_wproj_t"]) + if i in core.engram_layers: + li = core.engram_layers.index(i) + self._eng_gather.bind_layered( + i, + weight=d["eng_table"], + scale=d["eng_scale"], + row_start=core.eng_start[li], + row_end=core.eng_end[li], + ) + self._eng_gate.bind_layered(i, q_weight=d["eng_q"], k_weight=d["eng_k"]) + if core.last_pp: + # The final norm's input collapses with the last layer's FFN mixes. + _, fsc, fbs = L[core.num_layers - 1]["hc_ffn"] + self._hc_pre_final.bind_const(scale=fsc, base=fbs) + + def _bind_mx(self, name: str, layer: int, wsf) -> None: + weight, weight_sf = wsf + self._mx[name].bind_layered(layer, weight=weight, weight_sf=weight_sf) + self._mx_shape[name, layer] = tuple(weight.shape) + + def step_args(self, md) -> dict: + """The metadata surface this target reads, projected once per step. + + Also the step-contract probe (``Target._check_step_contract``): reading every + field here is what checks that they exist. The context-only fields -- the + context index-K gather, the encoder replay's compute rows and Engram history, + the context attention's window counts -- are this target's alone.""" + return dict( + kv_cache_manager=md.kv_cache_manager, + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + num_seqs=md.num_seqs, + all_rank_num_tokens=md.all_rank_num_tokens, + req_idx_per_token=md.req_idx_per_token, + cu_seq_lens_cuda=md.cu_seq_lens_cuda, + cached_token_lens_cuda=md.cached_token_lens_cuda, + kv_lens_cuda_runtime=md.kv_lens_cuda_runtime, + swa_local_indices_cuda=md.swa_local_indices_cuda, + sparse_mla_topk_lens=md.sparse_mla_topk_lens, + compressed_visible_lens_cuda=md.compressed_visible_lens_cuda, + compressed_kv_lens_cuda=md.compressed_kv_lens_cuda, + max_compressed_kv_lens=md.max_compressed_kv_lens, + compress_block_tables=md.compress_block_tables, + indexer_compress_block_tables=md.indexer_compress_block_tables, + sliding_block_tables=md.sliding_block_tables, + cu_new_comp_kv_cuda=md.cu_new_comp_kv_cuda, + compressed_position_ids_cuda=md.compressed_position_ids_cuda, + num_total_compressed_tokens=md.num_total_compressed_tokens, + max_ctx_compressed_tokens=md.max_ctx_compressed_tokens, + engram_token_block_table=md.engram_token_block_table, + v41_swa_slot_mapping=md.v41_swa_slot_mapping, + v41_compress_slot_mapping=md.v41_compress_slot_mapping, + v41_index_k_slot_data=md.v41_index_k_slot_data, + v41_index_k_slot_scale=md.v41_index_k_slot_scale, + v41_ctx_index_k_len=md.v41_ctx_index_k_len, + v41_ctx_index_k_gather_data=md.v41_ctx_index_k_gather_data, + v41_ctx_index_k_gather_scale=md.v41_ctx_index_k_gather_scale, + v41_ctx_index_ks=md.v41_ctx_index_ks, + v41_ctx_index_ke=md.v41_ctx_index_ke, + v41_max_ctx_q_len=md.v41_max_ctx_q_len, + v41_swa_seq_lens_cuda=md.v41_swa_seq_lens_cuda, + # Encoder replay (always valid; the plain values when a step replays nothing). + v41_num_replay_tokens_total=md.v41_num_replay_tokens_total, + v41_num_replay_tokens_cpu=md.v41_num_replay_tokens_cpu, + v41_num_compute_tokens=md.v41_num_compute_tokens, + v41_compute_token_index_cuda=md.v41_compute_token_index_cuda, + v41_compute_start_pos_cuda=md.v41_compute_start_pos_cuda, + v41_compute_cu_seq_lens_cuda=md.v41_compute_cu_seq_lens_cuda, + v41_replay_history_len=md.v41_replay_history_len, + v41_replay_history_token_ids_cuda=md.v41_replay_history_token_ids_cuda, + v41_replay_history_req_idx_cuda=md.v41_replay_history_req_idx_cuda, + v41_replay_history_cu_seq_lens_cuda=md.v41_replay_history_cu_seq_lens_cuda, + v41_replay_history_start_pos_cuda=md.v41_replay_history_start_pos_cuda, + # Host-side per-request state the decoder replay and the read-buffer + # form plan from. + request_ids=md.request_ids, + seq_lens=md.seq_lens, + seq_lens_kv=md.seq_lens_kv, + cached_token_lens_cpu=md.cached_token_lens_cpu, + tokens_per_block=md.tokens_per_block, + max_seq_len=md.max_seq_len, + ) + + # ------------------------------------------------------------------ + def _mx_linear(self, name: str, x, layer: int, out_dtype=torch.bfloat16): + """MXFP8 x MXFP8 linear: 1x32 UE8M0 activation quantization, then the + GEMM against the checkpoint's 32x32 blocks expanded to 1x32.""" + m = x.shape[0] + xq, xsf = self._mxq(x) + be = _gemm_backend(m) + _rv("mxfp8_mxfp8_gemm", m, self._mx_shape[name, layer], be) + return self._mx[name](xq, xsf, out_dtype=out_dtype, backend=be, layer=layer) + + def _allreduce_fp32(self, x): + """TP-member sum of bf16 per-rank partials in fp32 (the reference's + fp32 all-reduce of bf16-rounded partials); returns fp32. AUTO on the IPC + workspace when the group is on one node, NCCL across nodes.""" + x32 = torch.empty(list(x.shape), dtype=torch.float32, device=x.device) + x32.copy_(x) + _rv("allreduce", _sh(x32), "fp32", self._ar32_strategy) + return self._allreduce32(x32)[0] + + def _engram(self, i, h, hashes, md, num_tokens, dev, keep=None): + """Engram at layer i, in place on the stream h [T, 4, 5120]. ``keep`` + [T] bool is False on image-span rows (their gate is 0), or None.""" + core = self.core + li = core.engram_layers.index(i) + ids = hashes[:, li, :] # [T, 24] int64, row stride 48 + if core.eng_host: + # Every row is in the shared pinned copy: this rank's own tokens, + # no zero fill and no collective, on every member. + vals = self._eng_gather(ids, layer=i) + red = torch.reshape(vals, [num_tokens, core.eng_in]) + elif core.dp: + rows = core._dp_rows(md, num_tokens) + ids_c = torch.empty([num_tokens, core.eng_cols], dtype=torch.int64, device=dev) + ids_c.copy_(ids) + ids32 = ids_c.view(torch.int32) + if rows > num_tokens: + ids32 = nn.functional.pad(ids32, [0, 0, 0, rows - num_tokens]) + _rv("allgather", _sh(ids32), "int32") + ids_all = self._allgather(ids32).view(torch.int64) + vals = self._eng_gather(ids_all, layer=i) + flat = torch.reshape(vals, [vals.shape[0], core.eng_in]) + _rv("reducescatter", _sh(flat), "bf16") + red = self._reducescatter(flat) + if rows > num_tokens: + red = torch.split(red, [num_tokens, rows - num_tokens], 0)[0] + else: + vals = self._eng_gather(ids, layer=i) + flat = torch.reshape(vals, [num_tokens, core.eng_in]) + _rv("allreduce", _sh(flat), "bf16", _AR_NCCL) + red = self._allreduce16(flat)[0] + _rv("engram_gather", num_tokens, li) + kv = self._mx_linear("eng_wkv", red, i) # [T, 25600] + key, value = torch.split(kv, [_HC * core.hidden, core.hidden], -1) + key3 = torch.reshape(key, [num_tokens, _HC, core.hidden]) + _rv("engram_gate" if keep is None else "engram_gate.masked", num_tokens) + return self._eng_gate(h, key3, value, token_mask=keep, out=h, layer=i) + + # ------------------------------------------------------------------ + def _attention(self, i, x, st, pools, cache, dev): + """One attention sublayer on the normed input x [T, 5120] -> [T, 5120] bf16.""" + core = self.core + r = core.ratios[i] + T, tc, gen = st["T"], st["tc"], st["gen"] + nc, G, B, next_n = st["nc"], st["G"], st["B"], st["next_n"] + md = st["md"] + pos = st["pos"] + table = core._table(r) + H = core.h_loc + hd = core.head_dim + + # Low-rank query + the shared window latent. + qkv = self._mx_linear("wqkv_a", x, i) # [T, 1792] + q_lr, kv_raw = torch.split(qkv, [core.q_lora, hd], -1) + qr = self._q_norm(q_lr, layer=i) + kv = self._kv_norm(kv_raw, layer=i) + qrq, qrsf = self._mxq(qr) + q = self._mx["wq_b"](qrq, qrsf, out_dtype=torch.bfloat16, backend=_gemm_backend(T), layer=i) + _rv("mla_rope_inplace", T, H, core.nope, r) + q3 = self._rope(torch.reshape(q, [T, H, hd]), pos, table, H, core.nope, _ROPE, False, False) + kv = torch.reshape( + self._rope( + torch.reshape(kv, [T, 1, hd]), pos, table, 1, core.nope, _ROPE, False, False + ), + [T, hd], + ) + # Window KV: FP8 E4M3 x UE8M0/32 value set into this layer's SWA slots + # (native: the codes and UE8M0 bytes of the 528 B rows). A pool write. + native = pools["native"] + swa_slots = md.v41_swa_slot_mapping[pools["local"][i]][:T] + if native: + codes, scales = pools["swa_q"][i] + _rv("mxfp8_fake_quantize.native", T, _sh(codes)) + self._swa_write(kv, codes, swa_slots, scales) + else: + _rv("mxfp8_fake_quantize", T, _sh(pools["swa3"][i])) + self._swa_write(kv, pools["swa3"][i], swa_slots) + + # Compressed KV + index K (the four KV sources). + per_query = native and st["comp_form"] == "topk" + if i in core.kv_sources: + self._compress(i, x, st, pools, dev) + if native and not per_query: + self._native_compressed(i, st, pools) + + # Compressed positions per query (index sources); readers reuse. + if i in core.index_sources: + cache["topk"][r] = self._index(i, x, qrq, qrsf, st, pools, cache, dev) + topk = cache["topk"].get(r) if r else None + + # Pool offsets -> the attention's index table. Native: the compressed + # columns land on the step's dequantized rows (per request through + # ``comp_tbl``; per query renumbered by ``_native_topk``), the window + # columns are renumbered by ``_native_window``. + req = md.req_idx_per_token[:T] + swa_bt = md.sliding_block_tables[pools["local"][i], DeepseekV4AttentionType.SWA.value] + swa_local = md.swa_local_indices_cuda[:T] + if r: + _rv("dsv4_local_to_global_indices", T, 128, _TOPK, _PAGE // r) + gidx = self._l2g( + req, + swa_bt, + swa_local, + _PAGE, + 0, + pools["comp_tbl"][r] if native and not per_query else md.compress_block_tables[r], + topk, + _PAGE // r, + 0, + ) + if per_query: + gidx = self._native_topk(i, gidx, topk, T, pools, cache) + if native: + comp_pool = pools["comp_c4_q"][r] if per_query else pools["comp_c4_req"][r] + else: + comp_pool = pools["comp4"][core.kv_src[i]] + else: + _rv("dsv4_local_to_global_indices", T, 128) + gidx = self._l2g(req, swa_bt, swa_local, _PAGE, 0) + comp_pool = core._const["dummy_comp8" if core.fp8_attn else "dummy_comp"] + if native: + gidx, swa_pool = self._native_window(i, gidx, T, pools) + else: + swa_pool = pools["swa4"][i] + lens = md.sparse_mla_topk_lens[r][:T] + kv_lens = md.kv_lens_cuda_runtime + o = torch.empty([T, H, hd], dtype=torch.bfloat16, device=dev) + if core.fp8_attn: + # E4M3 query, one power-of-two scale s_q from this step's max|q| (on the device); + # the call's QK scale is scale * s_q * s_kv (written by the quantizer) and its + # PV scale s_kv. + qa = torch.empty([T, H, hd], dtype=torch.float8_e4m3fn, device=dev) + bmm1 = torch.empty([1], dtype=torch.float32, device=dev) + _rv("e4m3_quantize_pow2_dynamic", T, H) + qa, bmm1 = self._q_e4m3(q3, qa, bmm1, layer=i) + attn = self._attn8 + fp8 = dict(bmm1_scale=bmm1) + else: + qa = q3 + attn = self._attn + fp8 = {} + kind = "trtllm_sparse_mla_dsv4_fp8" if core.fp8_attn else "trtllm_sparse_mla_dsv4" + # Both calls write their rows of `o`, which the inverse RoPE reads. + if tc: + _rv(kind + ".varlen", tc, H, gidx.shape[1], nc, st["max_ctx_q"]) + attn( + torch.split(qa, [tc, gen], 0)[0], + swa_pool, + comp_pool, + sparse_indices=torch.split(gidx, [tc, gen], 0)[0], + sparse_topk_lens=torch.split(lens, [tc, gen], 0)[0], + # Window counts from kv length minus the SWA floor: a replay + # row counts only the window keys at or above its replay start. + seq_lens=md.v41_swa_seq_lens_cuda[:nc], + cum_seq_lens_q=md.cu_seq_lens_cuda[: nc + 1], + max_q_len=st["max_ctx_q"], + out=torch.split(o, [tc, gen], 0)[0], + layer=i, + **fp8, + ) + if gen: + _rv(kind + ".decode", G, next_n, H, gidx.shape[1]) + attn( + torch.reshape(torch.split(qa, [tc, gen], 0)[1], [G, next_n, H, hd]), + swa_pool, + comp_pool, + sparse_indices=torch.split(gidx, [tc, gen], 0)[1], + sparse_topk_lens=torch.split(lens, [tc, gen], 0)[1], + seq_lens=kv_lens[nc:B], + out=torch.reshape(torch.split(o, [tc, gen], 0)[1], [G, next_n, H, hd]), + layer=i, + **fp8, + ) + # Remove the query rotation from the output's rope tail. + o = self._rope(o, pos, table, H, core.nope, _ROPE, True, False) + # Grouped low-rank output: per group [T, 4096] x [4096, 1024], into `oa`. + oa = torch.empty([T, core.g_loc * core.o_lora], dtype=torch.bfloat16, device=dev) + _rv("bmm_out", core.g_loc, T, core.group_in, core.o_lora) + self._wo_a( + torch.transpose(torch.reshape(o, [T, core.g_loc, core.group_in]), 0, 1), + out=torch.transpose(torch.reshape(oa, [T, core.g_loc, core.o_lora]), 0, 1), + layer=i, + ) + y = self._mx_linear("wo_b", oa, i) # [T, 5120] bf16 (per-rank partial under TP) + if core.dp: + return y + y32 = self._allreduce_fp32(y) + yb = torch.empty([T, core.hidden], dtype=torch.bfloat16, device=dev) + yb.copy_(y32) + return yb + + def _compress(self, i, x, st, pools, dev): + """KV source i: this step's compressed entries -> index K (native + MXFP4 pool) and main KV (E2M1 x E4M3/16 value set, bf16 pool). + + Encoder replay: the replay rows at the head of a prefix hit's first + chunk rebuild the windows only. The compressor runs over the other rows + (``v41_compute_token_index_cuda``), packed as a step whose requests + start at their reuse boundaries (``v41_compute_start_pos_cuda`` / + ``v41_compute_cu_seq_lens_cuda``); the compressed stream, its positions + and its write slots already start there, so no cached entry is + recomputed or overwritten.""" + core = self.core + r = core.ratios[i] + md = st["md"] + T, nc, G, B, next_n = st["T"], st["nc"], st["G"], st["B"], st["next_n"] + hd = core.head_dim + rows, start, cu = T, md.cached_token_lens_cuda, md.cu_seq_lens_cuda + if st["replay"]: + rows = md.v41_num_compute_tokens + assert rows == T - st["replay"], (rows, T, st["replay"]) + start, cu = md.v41_compute_start_pos_cuda, md.v41_compute_cu_seq_lens_cuda + st["compress_rows"] = rows + if rows == 0: + return # every row of the step replays: nothing to compress + _rv("embedding.compute_rows", rows, T) + x = nn.functional.embedding(md.v41_compute_token_index_cuda[:rows], x) + if r == 2: + n = md.num_total_compressed_tokens[2] + x32 = torch.empty([rows, core.hidden], dtype=torch.float32, device=dev) + x32.copy_(x) + kv_score = self._comp_mm(x32, layer=i) # fp32 [T, 1024] = [kv || score] + pooled = torch.empty([n, hd], dtype=torch.bfloat16, device=dev) + # Reserved decode rows the op leaves unwritten stay finite. + pooled.copy_(core._const["zero_bf16"].expand([n, hd])) + local = pools["local"][i] + bt_kv = md.sliding_block_tables[local, DeepseekV4AttentionType.COMPRESSOR_KV.value] + bt_sc = md.sliding_block_tables[local, DeepseekV4AttentionType.COMPRESSOR_SCORE.value] + cu_comp = md.cu_new_comp_kv_cuda[2] + # Both fill `pooled` and update the compressor state pools. + if nc: + _rv("compressor_prefill_reduction", T, nc, md.max_ctx_compressed_tokens[2]) + self._comp_prefill( + kv_score, + paged_kv=pools["ckv"][i], + paged_score=pools["csc"][i], + block_table_kv=bt_kv[:nc], + block_table_score=bt_sc[:nc], + output=pooled, + kv_lens=md.kv_lens_cuda_runtime[:nc], + start_pos=start[:nc], + cu_seq_lens=cu[: nc + 1], + cu_kv_comp=cu_comp[: nc + 1], + batch_size=nc, + max_outputs=md.max_ctx_compressed_tokens[2], + ) + if G: + _rv("compressor_paged_kv_compress", G, next_n, n) + self._comp_paged( + kv_score, + paged_kv=pools["ckv"][i], + paged_score=pools["csc"][i], + block_table_kv=bt_kv[nc:B], + block_table_score=bt_sc[nc:B], + output=pooled, + kv_lens=md.kv_lens_cuda_runtime[nc:B], + cu_seq_lens=cu[nc : B + 1], + cu_kv_comp=cu_comp[nc : B + 1], + batch_size=G, + next_n=next_n, + ) + # The ops above ran even with no completed group: they store this + # step's pending odd token in the compressor state. + if n == 0: + return + latent = self._comp_norm(pooled, layer=i) + else: + n = rows + assert md.num_total_compressed_tokens[1] == rows, ( + md.num_total_compressed_tokens[1], + rows, + T, + ) + latent = self._comp_norm(self._comp_mm(x, layer=i), layer=i) + cpos = md.compressed_position_ids_cuda[r][:n] + tbl = core._rope_tables[1] + # Index K from the RoPE-free latent (before the rotation below). + k = self._idx_k_norm(self._idx_wk(latent, layer=i), layer=i) + k = torch.reshape( + self._rope( + torch.reshape(k, [n, 1, _IDX_DIM]), + cpos, + tbl, + 1, + _IDX_DIM - _ROPE, + _ROPE, + False, + False, + ), + [n, _IDX_DIM], + ) + kd, ksf = self._mxfp4(k) + _rv("indexer_k_cache_scatter", n, r, _sh(pools["idxk"][i])) + self._idx_scatter( + kd, + torch.reshape(ksf.view(torch.int32), [n, 1]), + pools["idxk"][i], + md.v41_index_k_slot_data[r][:n], + md.v41_index_k_slot_scale[r][:n], + n, + ) + # Main KV: rope at the group-first position j*r, then the E2M1 x E4M3/16 + # value set into the owner's COMPRESS slots (native: the packed codes + # and E4M3 bytes of the 288 B rows). A pool write. + latent = torch.reshape( + self._rope( + torch.reshape(latent, [n, 1, hd]), cpos, tbl, 1, core.nope, _ROPE, False, False + ), + [n, hd], + ) + slots = md.v41_compress_slot_mapping[r][:n] + if pools["native"]: + codes, scales = pools["comp_q"][i] + _rv("nvfp4_fake_quantize.native", n, r, _sh(codes)) + self._comp_write(latent, codes, slots, scales) + else: + _rv("nvfp4_fake_quantize", n, r, _sh(pools["comp3"][i])) + self._comp_write(latent, pools["comp3"][i], slots) + + def _native_compressed(self, i, st, pools): + """Native storage, KV source i (after its write): every compressed + entry of every request of the step, dequantized into the compact + buffer at row ``b * L_r + e`` -- all the rows the layers reading this + source can name (``compressed_kv_lens`` = kv_len // r per request). + The buffer is read through ``comp_c4_req``.""" + core = self.core + r = core.ratios[i] + md = st["md"] + # Rows are the engine's request indices (block-table rows); in the + # decoder replay pass st["B"] counts its compacted request list. + B = st.get("B_engine", st["B"]) + codes, scales = pools["comp_q4"][i] + out = pools["comp_out"][r][:B] + if core.fp8_attn: + _rv("nvfp4_kv_to_e4m3", B, out.shape[1], r, _sh(codes)) + self._comp_read8( + codes, + scales, + md.compress_block_tables[r][:B], + md.compressed_kv_lens_cuda[r][:B], + out, + 1.0 / core._fp8_kv_scale[i], + ) + return + _rv("nvfp4_kv_to_bf16", B, out.shape[1], r, _sh(codes)) + self._comp_read( + codes, + scales, + md.compress_block_tables[r][:B], + md.compressed_kv_lens_cuda[r][:B], + core._const["one_f32"], + out, + ) + + def _native_topk(self, i, gidx, topk, T, pools, cache): + """Native storage, per-query read buffer, layer i: the index table + with its compressed half renumbered onto ``topk_out`` (query t's + column k at row ``t * 512 + k``). The rows are dequantized from the + pool slots the index op named (-1: left unwritten) when the top-k set + or its KV source changes; the layers that share both reuse them. + + Renumbering every column is exact: the index op packs a row's valid + entries first and ``sparse_mla_topk_lens`` counts exactly those, so + every -1 column is masked, and a masked column carries no weight + whatever finite row it names (every row of the buffer is finite: zero + at allocation, dequantized values after).""" + core = self.core + win, comp = torch.split(gidx, [_WINDOW, _TOPK], 1) + src = core.kv_src[i] + last = cache.get("comp_rows") + if last is None or last[0] is not topk or last[1] != src: + slots = pools["topk_slots"][:T] + slots.copy_(comp) + codes, scales = pools["comp_q1"][src] + if core.fp8_attn: + # Every layer reading this source shares its KV scale (``_derive_fp8_kv_scales``). + _rv("nvfp4_kv_to_e4m3.topk", T, _TOPK, core.ratios[i], _sh(codes)) + self._comp_read8( + codes, + scales, + slots, + pools["topk_cols"][:T], + pools["topk_out"][:T], + 1.0 / core._fp8_kv_scale[i], + ) + else: + _rv("nvfp4_kv_to_bf16.topk", T, _TOPK, core.ratios[i], _sh(codes)) + self._comp_read( + codes, + scales, + slots, + pools["topk_cols"][:T], + core._const["one_f32"], + pools["topk_out"][:T], + ) + # Holding the top-k tensor keeps its identity unique for the step. + cache["comp_rows"] = (topk, src) + return torch.cat([win, pools["topk_ren"][:T]], 1) + + def _native_window(self, i, gidx, T, pools): + """Native storage, layer i: dequantize the window rows each query's + index row names into the compact buffer -- query t's column c to row + ``t * 128 + c``, a -1 column to a zero row -- and return the index + table with its window half renumbered onto that buffer (the compressed + half is unchanged), and the buffer as the attention's window pool. + + A zero row stands in for -1 exactly: the kernel reads an attended -1 + as an all-zero row, and a masked column (past the query's window count) + carries no weight whatever finite row it names.""" + core = self.core + codes, scales = pools["swa_q"][i] + if gidx.shape[1] == _WINDOW: + win, rest = gidx, None + else: + win, rest = torch.split(gidx, [_WINDOW, gidx.shape[1] - _WINDOW], 1) + if core.fp8_attn: + _rv("mxfp8_to_e4m3", T * _WINDOW, _sh(pools["swa_c"]), _sh(codes)) + self._swa_read8( + codes, + scales, + pools["swa_c"], + 1.0 / core._fp8_kv_scale[i], + torch.reshape(win, [T * _WINDOW]), + ) + else: + _rv("mxfp8_dequantize", T * _WINDOW, _sh(pools["swa_c"]), _sh(codes)) + self._swa_read(codes, scales, pools["swa_c"], torch.reshape(win, [T * _WINDOW])) + ren = pools["win_idx"][:T] + return (ren if rest is None else torch.cat([ren, rest], 1)), pools["swa_c4"] + + def _gather_index_k(self, pool, slots_data, slots_scale, s_len, r): + """The step's context index-K rows ``[0, s_len)`` out of the paged pool, + as the logits op reads them: (int8 codes [s_len, 64], int32 UE8M0 + words [s_len]). Gathered in calls of at most ``_GATHER_MAX_ROWS`` rows: + the op offsets its output by row index x 64 in 32 bits on builds before + the int64 fix, which wraps from row 2^25 on (a multi-request step at 1M + context reaches it). One call when the rows fit.""" + k8s, k32s = [], [] + for a in range(0, s_len, _GATHER_MAX_ROWS): + n = min(_GATHER_MAX_ROWS, s_len - a) + _rv("indexer_k_cache_gather", n, r) + k8, k32 = self._idx_gather(pool, slots_data, slots_scale, a, n) + k8s.append(k8.view(torch.int8)) + k32s.append(torch.reshape(k32.view(torch.int32), [-1])) + if len(k8s) == 1: + return k8s[0], k32s[0] + return torch.cat(k8s, 0), torch.cat(k32s, 0) + + def _cand_masks(self, cand_q, tiles, ks, ke, kvc, s_len, tag) -> list: + """The candidate-block masks of the context rows of each row tile, + scored as the candidate source scores them -- its index query rows + (``cand_q``: MXFP4 q, UE8M0 words, head weights) against this step's + gathered index K, each row over its own ``[ks, ke)``. A row's logits + and mask do not depend on the other rows of the call or on where its + request's keys sit in the gather, so every row gets the mask the source + layer gives it: a decoder-replay tail row whichever chunk forwarded it + (``tag`` "tail"), and a row whose mask the step did not keep + (``tag`` "rescore", see ``_index_prefill``).""" + q, qs, w = cand_q + masks = [] + for a, b in tiles: + _rv("deepgemm_fp8_fp4_mqa_logits." + tag + "_cand", b - a, s_len, 1) + logits = self._mqa_logits(q[a:b], qs[a:b], kvc[0], kvc[1], w[a:b], ks[a:b], ke[a:b]) + _rv("candidate_block_select." + tag, b - a, s_len) + masks.append(self._cand(logits, ks[a:b], ke[a:b])) + return masks + + def _index_prefill(self, i, q, qs, w, kvc, ks, ke, vis, s_len, st, cache, out) -> None: + """Index source i, the step's context rows: their top-512 into ``out`` + ([tc, 512] int32), scored against the ``s_len`` gathered index keys + ``kvc``, row t over its own ``[ks[t], ke[t])``. + + The rows run in tiles (``_index_row_tiles``), so a long context never + holds more than one tile's logits. Candidate blocks bind only above + 2048 x 8 visible positions (below that the mask is all-visible, exact + to skip). While a step's masks (tc x s_len bytes) fit + ``_CAND_MASK_BYTES``, they are kept, one per tile, for the re-index + layers, which tile their rows the same way: the candidate source keeps + its own; in a bounded step (whose engine rows stop at layer L) the + decoder pass scores its tail rows' masks at its first re-index layer. + Past the budget the source keeps its query rows instead and every + re-index layer re-scores each tile's masks from them (``_cand_masks``), + so the step holds one tile's masks at a time.""" + core = self.core + tc = q.shape[0] + reindex = core.cand_src < i + tiles = _index_row_tiles(tc, s_len) + binds = s_len > _CAND_TOPK_BLOCKS * _CAND_BLOCK + keep = tc * s_len <= _CAND_MASK_BYTES + src_masks = None + if i == core.cand_src and binds and not st["tail_bound"]: + if keep: + src_masks = cache["cand_ctx"] = [] + else: + cache["cand_q"] = (q, qs, w) + masks = None + if reindex and binds: + if keep and "cand_ctx" not in cache: + cache["cand_ctx"] = self._cand_masks( + cache["cand_q"], tiles, ks, ke, kvc, s_len, "tail" + ) + masks = cache.get("cand_ctx") + rescore = reindex and binds and masks is None + r = core.ratios[i] + for n, (a, b) in enumerate(tiles): + _rv("deepgemm_fp8_fp4_mqa_logits", b - a, s_len, r) + logits = self._mqa_logits(q[a:b], qs[a:b], kvc[0], kvc[1], w[a:b], ks[a:b], ke[a:b]) + if src_masks is not None: + _rv("candidate_block_select.prefill", b - a, s_len) + src_masks.append(self._cand(logits, ks[a:b], ke[a:b])) + if masks is not None: + logits = torch.where(masks[n], logits, float("-inf")) + elif rescore: + mask = self._cand_masks(cache["cand_q"], [(a, b)], ks, ke, kvc, s_len, "rescore")[0] + logits = torch.where(mask, logits, float("-inf")) + _rv("flashinfer_top_k_ragged_deterministic.prefill", b - a, s_len) + out[a:b].copy_( + self._topk(logits, core._const["zero_i32_rows"][: b - a], vis[a:b], ks[a:b]) + ) + + def _index(self, i, x, qrq, qrsf, st, pools, cache, dev): + """Index source i: top-512 compressed positions per query, [T, 512] + int32, valid entries first in ascending order, -1 after.""" + core = self.core + r = core.ratios[i] + md = st["md"] + T, tc, gen, nc, G, B, next_n = ( + st["T"], + st["tc"], + st["gen"], + st["nc"], + st["G"], + st["B"], + st["next_n"], + ) + src = core.kv_src[i] + iq = self._mx["idx_wq_b"]( + qrq, qrsf, out_dtype=torch.bfloat16, backend=_gemm_backend(T), layer=i + ) + iq3 = self._rope( + torch.reshape(iq, [T, _IDX_HEADS, _IDX_DIM]), + st["pos"], + core._rope_tables[1], + _IDX_HEADS, + _IDX_DIM - _ROPE, + _ROPE, + False, + False, + ) + qd, qsf = self._mxfp4(iq3) # [T,32,64] u8, [T*128] u8 + q8 = qd.view(torch.int8) + qs32 = torch.reshape(qsf.view(torch.int32), [T, _IDX_HEADS]) + wb = self._idx_wproj(x, layer=i) # bf16 [T, 32], 2^-6 folded in + w32 = torch.empty([T, _IDX_HEADS], dtype=torch.float32, device=dev) + w32.copy_(wb) + idx = torch.empty([T, _TOPK], dtype=torch.int32, device=dev) + idx_ctx, idx_gen = torch.split(idx, [tc, gen], 0) + reindex = core.cand_src < i + if i == core.cand_src and st["tail_bound"]: + # Every row of the step: the decoder pass takes its tail rows (and + # stashes those of a request's earlier chunks). + cache["cand_q"] = (q8, qs32, w32) + if tc: + s_len = md.v41_ctx_index_k_len[r] + if s_len == 0: + idx_ctx.copy_(core._const["neg1_i32"].expand([tc, _TOPK])) + else: + ks, ke = md.v41_ctx_index_ks[r], md.v41_ctx_index_ke[r] + kvc = cache["ctx_k"].get(src) + if kvc is None: + kvc = self._gather_index_k( + pools["idxk"][src], + md.v41_ctx_index_k_gather_data[r], + md.v41_ctx_index_k_gather_scale[r], + s_len, + r, + ) + cache["ctx_k"][src] = kvc + self._index_prefill( + i, + torch.split(q8, [tc, gen], 0)[0], + torch.split(qs32, [tc, gen], 0)[0], + torch.split(w32, [tc, gen], 0)[0], + kvc, + ks, + ke, + md.compressed_visible_lens_cuda[r][:tc], + s_len, + st, + cache, + idx_ctx, + ) + if gen: + vis = md.compressed_visible_lens_cuda[r][tc:T] + ctx_lens = torch.reshape(vis, [G, next_n]) + sched = cache["sched"].get(r) + if sched is None: + sched = self._paged_meta(ctx_lens, _PAGE // r) + cache["sched"][r] = sched + max_ctx = md.max_compressed_kv_lens[r] + _rv("deepgemm_fp8_fp4_paged_mqa_logits", G, next_n, r, max_ctx) + logits = self._paged_logits( + torch.reshape( + torch.split(q8, [tc, gen], 0)[1], [G, next_n, _IDX_HEADS, _IDX_DIM // 2] + ), + torch.reshape(torch.split(qs32, [tc, gen], 0)[1], [G, next_n, _IDX_HEADS]), + pools["idxk"][src], + torch.split(w32, [tc, gen], 0)[1], + ctx_lens, + md.indexer_compress_block_tables[r][nc:B], + sched, + max_ctx, + ) + if core._cand_binds_decode: + if i == core.cand_src: + _rv("candidate_block_select.decode", gen, max_ctx) + cache["cand_gen"] = self._cand(logits, None, vis) + elif reindex: + logits = torch.where(cache["cand_gen"], logits, float("-inf")) + _rv("flashinfer_top_k_ragged_deterministic.decode", gen, max_ctx) + idx_gen.copy_(self._topk(logits, core._const["zero_i32_rows"][:gen], vis, None)) + return idx + + def _route(self, i, chunk, img, dev): + """Router of one token chunk [n, 5120] bf16: the gate GEMM (bf16 x and + the checkpoint's bf16 weight on tensor cores, fp32 accumulation, fp32 + logits; every product is exact in fp32, the summation order is the + tensor cores'), then top-6 by ``sqrt(softplus(logit)) + bias``. With + ``img`` ([n] bool, True on image-span rows) those rows select with + ``gate.bias_vl`` instead: a second routing call over the same logits + and a per-row select of both results. Returns (weights [n, 6] fp32, + ids [n, 6] int32).""" + core = self.core + n = chunk.shape[0] + logits = self._router(chunk, layer=i) # fp32 [n, 384] + _rv("gate_forward", n, core.num_experts, core.topk) + tw, tid = self._gate(logits, layer=i) + if img is None: + return tw, tid + _rv("gate_forward.bias_vl", n, core.num_experts, core.topk) + tw_vl, tid_vl = self._gate_vl(logits, layer=i) + sel = torch.reshape(img, [n, 1]) + return torch.where(sel, tw_vl, tw), torch.where(sel, tid_vl, tid) + + def _moe(self, i, x, rows, dev, img=None): + """Routed NVFP4 experts + FP8 shared expert on the normed input x. + ``img`` is the image-span mask of the rows the router sees (this rank's + rows, or the gathered rows under attention DP), or None. ``rows`` is + the dep4 per-rank row count every rank pads to (None under attention + TP). ``x`` None is a dep4 rank with no row in this pass: it sends zero + rows and returns None.""" + core = self.core + T = 0 if x is None else x.shape[0] + assert T or core.dp, "only a dep4 rank runs the MoE without rows" + if T: + # Shared expert (replicated, this rank's own tokens): [gate ; up]. + h13 = self._mx_linear("sh_w13", x, i) + _rv("silu_and_mul", T, 2 * core.moe_inter) + act = self._silu(h13) + shared = self._mx_linear("sh_w2", act, i) + # Routed: dep4 gathers the token set ahead of the router, so every rank + # routes the identical set and the four expert windows tile it once. + pad_rows = 0 + xr = x + if core.dp: + assert rows is not None and T <= rows, (T, rows) + pad_rows = rows - T + if not T: + xp = torch.empty([rows, core.hidden], dtype=torch.bfloat16, device=dev) + xp.copy_(core._const["zero_bf16"].expand([rows, core.hidden])) + else: + xp = x if not pad_rows else nn.functional.pad(x, [0, 0, 0, pad_rows]) + _rv("allgather", _sh(xp), "bf16") + xr = self._allgather(xp) + parts = [] + sizes = _moe_chunks(xr.shape[0]) + if img is not None: + assert img.shape[0] == xr.shape[0], (tuple(img.shape), tuple(xr.shape)) + img_chunks = torch.split(img, sizes, 0) if img is not None else [None] * len(sizes) + for chunk, img_c in zip(torch.split(xr, sizes, 0), img_chunks): + n = chunk.shape[0] + tw, tid = self._route(i, chunk, img_c, dev) + twb = torch.empty([n, core.topk], dtype=torch.bfloat16, device=dev) + twb.copy_(tw) + hs, hsf = self._moe_quant(chunk, layer=i) + out = torch.empty([n, core.hidden], dtype=torch.bfloat16, device=dev) + _rv("fp4_block_scale_moe_runner", n, core.e_loc, core.e_off, core.i_loc) + # Given `output`, the runner finalizes into it and returns an empty tensor. + self._moe_runner( + hidden_states=hs, + hidden_states_scale=hsf.view(torch.float8_e4m3fn), + topk_weights=twb, + topk_ids=tid, + output=out, + layer=i, + ) + parts.append(out) + routed = parts[0] if len(parts) == 1 else torch.cat(parts, 0) + if core.dp: + r32 = torch.empty(list(routed.shape), dtype=torch.float32, device=dev) + r32.copy_(routed) + _rv("reducescatter", _sh(r32), "fp32") + r32 = self._reducescatter(r32) + if not T: + return None + if pad_rows: + r32 = torch.split(r32, [T, pad_rows], 0)[0] + else: + r32 = self._allreduce_fp32(routed) + m32 = torch.add(r32, shared) # fp32: the shared expert joins after the reduce + m = torch.empty([T, core.hidden], dtype=torch.bfloat16, device=dev) + m.copy_(m32) + return m + + # ------------------------------------------------------------------ + def _mm_step(self, mm_params, T, tc, dev, emb): + """Splice this step's image-span rows into the token embedding ``emb`` + [T, 5120] in place; return the step's image masks, or None when no + context chunk of the step carries images. A later pipeline stage passes + ``emb`` None and gets the masks only (the first stage spliced the rows + into the stream it receives). + + A request with images arrives as one ``MultimodalParams``: its + ``multimodal_data`` is the input processor's (``image``: the tower's + inputs for all its images; ``token_layout``: masks per prompt + position), ``input_ids_start_offset`` is its first row in this step and + ``multimodal_runtime`` the chunk ``[begin, end)`` of the prompt those + rows hold (the whole prompt when absent). Only the images whose spans + meet the chunk run the tower (``VisionTower.image_rows``), one image at + a time, and each image's rows are copied into ``emb`` before the next + image runs: a step's vision memory is one image's working set, however + many images the prompt or the step holds. Returns ``img`` [T] bool (True on + image-span rows: the MoE selects with ``bias_vl``) and ``keep`` [T] + bool (False on image-span rows: Engram DEAD, gate 0).""" + if not mm_params: + return None + core = self.core + reqs = [] + for p in mm_params: + data = p.multimodal_data or {} + image, layout = data.get("image"), data.get("token_layout") + if image is None and layout is None: + continue + assert image is not None and layout is not None, sorted(data) + L = layout["image_mask"].shape[0] + rt = p.multimodal_runtime + begin, end = ( + (0, L) if rt is None else (int(rt.past_seen_token_num), int(rt.chunk_end_pos)) + ) + row0 = int(p.input_ids_start_offset) + assert 0 <= begin <= end <= L and row0 + (end - begin) <= tc, (begin, end, L, row0, tc) + reqs.append((image, layout, begin, end, row0)) + if not reqs: + return None + img = torch.empty([T], dtype=torch.bool, device=dev) + img.copy_(core._const["false_b"].expand([T])) + dead = torch.empty([T], dtype=torch.bool, device=dev) + dead.copy_(core._const["false_b"].expand([T])) + spliced = 0 + for image, layout, begin, end, row0 in reqs: + n = end - begin + img[row0 : row0 + n].copy_(layout["image_mask"][begin:end]) + dead[row0 : row0 + n].copy_(layout["engram_dead_mask"][begin:end]) + # (image, first, end) prompt positions of every span that meets the chunk. + hits = [] + for j, (s, ln) in enumerate(zip(image["span_starts"], image["span_lengths"])): + s, ln = int(s), int(ln) + a, b = max(s, begin), min(s + ln, end) + if a < b: + hits.append((j, a, b, s)) + if emb is None: + continue + for j, a, b, s in hits: + rows = core._vision.image_rows(image, j) + emb[row0 + a - begin : row0 + b - begin].copy_(rows[a - s : b - s]) + spliced += b - a + keep = torch.where(dead, core._const["false_b"], core._const["true_b"]) + _rv("mm.step", T, len(reqs), spliced) + return {"img": img, "keep": keep} + + def _dp_gather_mask(self, loc, rows, T, dev): + """dep4: the image-span mask of the all-gathered token set the router + sees, each rank padded to ``rows`` as the MoE input is. ``loc`` is this + rank's [T] mask, or None for T text rows (T may be 0: a pass that holds + none of this rank's rows).""" + assert T <= rows, (T, rows) + if loc is None or not T: + loc = torch.empty([rows if not T else T], dtype=torch.bool, device=dev) + loc.copy_(self.core._const["false_b"].expand([loc.shape[0]])) + u8 = loc.view(torch.uint8) + if rows > loc.shape[0]: + u8 = nn.functional.pad(u8, [0, rows - loc.shape[0]]) + _rv("allgather", _sh(u8), "uint8") + return self._allgather(u8).view(torch.bool) + + def _dp_late_rows(self, mine: list, dev) -> list: + """dep4 decoder replay: the passes every rank runs over the bounded + layers, as each pass's row count padded to its max over ranks. + + ``mine`` is this rank's pass row counts (a bounded rank's slices, or its + step rows as one pass). The bounded layers' MoE is a collective, so the + ranks must run the same passes in lockstep; a rank's decoder rows can + also outnumber its step rows (a small last chunk still runs its whole + tail, the rest from the stash). Every rank calls this in every eager + step under decoder replay, prefill and decode alike: ``[count, rows_0, + rows_1, ...]`` rides in the int32 [1, 48] form the Engram ids gather + in, and the result has one entry per pass of the rank with the most + (0-row entries never occur: some rank holds the rows of every pass).""" + core = self.core + width = 2 * core.eng_cols # 48 + vals = [len(mine), *mine] + assert len(vals) <= width, (len(mine), width) + core._ensure_ar(max(vals) + 1, dev) + cols = [core._ar32[v : v + 1] for v in vals] + [core._ar32[0:1]] * (width - len(vals)) + v = torch.reshape(torch.cat(cols, 0), [1, width]) + _rv("allgather", _sh(v), "int32") + got = self._allgather(v).tolist() # [world, 48] + assert len(got) == core.tp and got[core.tp_rank][: len(vals)] == vals, ( + core.tp_rank, + vals, + got, + ) + return _dp_pass_rows(got) + + def _dp_idle_pass(self, rows: int, dev) -> None: + """dep4 decoder replay: this rank's place in a pass over the bounded + layers that holds none of its rows -- the mask gather and every + bounded layer's routed round trip, with zero rows.""" + img = self._dp_gather_mask(None, rows, 0, dev) + for i in range(self.core.late, self.core.num_layers): + self._moe(i, None, rows, dev, img) + + # ------------------------------------------------------------------ + def _scratch(self, s, T: int, dev) -> None: + """The per-row scratch of one pass over the T rows of stream ``s.h``.""" + core = self.core + f32 = torch.float32 + s.T = T + s.h_flat = torch.reshape(s.h, [T, core.hc_k]) + s.tile = _tile_n(T) + s.yA = torch.empty([T, _MIX], dtype=f32, device=dev) + s.rA = torch.empty([T], dtype=f32, device=dev) + s.mix4 = torch.empty([T, _HC], dtype=f32, device=dev) + s.post = torch.empty([T, _HC], dtype=f32, device=dev) + s.comb = torch.empty([T, _HC, _HC], dtype=f32, device=dev) + s.no_res = torch.empty([T, _HC, 0], dtype=torch.bfloat16, device=dev) + s.no_li = torch.empty([T, 0], dtype=torch.bfloat16, device=dev) + s.coll = torch.empty([T, core.hidden], dtype=torch.bfloat16, device=dev) + _rv("mhc", T, s.tile) + + def _layer(self, i: int, s, dev) -> None: + """Layer i over the pass ``s`` (its stream, scratch, step dict, cache).""" + core = self.core + T, h, h_flat = s.T, s.h, s.h_flat + # DSpark reads the INPUT stream of its target layers (37-39), every + # token; the metadata averages the 4 hc copies itself. + if s.spec is not None and s.spec.is_layer_capture(i): + s.spec.maybe_capture_hidden_states(i, h_flat) + if i in core.engram_layers: + h = self._engram(i, h, s.hashes, s.md, T, dev, s.keep) + # --- attention sublayer ------------------------------------ + s.yA, s.rA = self._hc_attn_gemm(h_flat, y=s.yA, r=s.rA, M=T, tile_n=s.tile, layer=i) + s.post, s.comb, s.no_li = self._hc_attn_fuse( + s.yA, + s.rA, + s.no_res, + post_mix=s.post, + comb_mix=s.comb, + layer_input=s.no_li, + M=T, + layer=i, + ) + if i == 0: + a_in = s.emb # one-hot pre on copy 0: the embedding itself + else: + # Lagged pre: the previous FFN's mixes collapse this input. + s.mix4.copy_(torch.split(s.yF, [_HC, _MIX - _HC], 1)[0]) + s.coll = self._hc_pre_attn(s.mix4, s.rF, h, s.coll, M=T, layer=i) + a_in = s.coll + x = self._attn_norm(a_in, layer=i) + if s.tail is not None and i == core.pp_layer_ids[0]: + s.tail.count_forwarded(x) + a = self._attention(i, x, s.st, s.pools, s.cache, dev) + h = self._hc_post(h, a, s.post, s.comb, h) + # --- FFN sublayer -------------------------------------------- + s.yF, s.rF = self._hc_ffn_gemm(h_flat, y=s.yF, r=s.rF, M=T, tile_n=s.tile, layer=i) + s.post, s.comb, s.no_li = self._hc_ffn_fuse( + s.yF, + s.rF, + s.no_res, + post_mix=s.post, + comb_mix=s.comb, + layer_input=s.no_li, + M=T, + layer=i, + ) + s.mix4.copy_(torch.split(s.yA, [_HC, _MIX - _HC], 1)[0]) + s.coll = self._hc_pre_ffn(s.mix4, s.rA, h, s.coll, M=T, layer=i) + x = self._ffn_norm(s.coll, layer=i) + m = self._moe(i, x, s.dp_rows, dev, s.img) + self._hc_post(h, m, s.post, s.comb, h) + + def _final(self, s): + """The final collapse with layer 39's FFN mixes, then the final norm.""" + s.mix4.copy_(torch.split(s.yF, [_HC, _MIX - _HC], 1)[0]) + s.coll = self._hc_pre_final(s.mix4, s.rF, s.h, s.coll, M=s.T) + return self._final_norm(s.coll) + + # ------------------------------------------------------------------ + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None = None, + position_ids: torch.IntTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + lora_params: dict | None = None, + **kwargs, + ) -> torch.Tensor: + core = self.core + self._check_step_contract(attn_metadata) + md = attn_metadata + step = self.step_args(md) + spec_metadata = kwargs.get("spec_metadata") + pools = core._get_pools(md) + dev = input_ids.device + T = input_ids.shape[0] + nc = step["num_contexts"] + tc = step["num_ctx_tokens"] + B = step["num_seqs"] + G = B - nc + gen = T - tc + next_n = gen // G if G else 1 + assert next_n * G == gen, (gen, G) + # A DSpark generation row verifies its 5 drafts: 1 + 5 tokens. + assert next_n == 1 or (spec_metadata is not None and next_n == 1 + core.ds.block), ( + f"generation rows of {next_n} tokens (spec_metadata {'set' if spec_metadata is not None else 'None'})" + ) + st = { + "md": md, + "T": T, + "tc": tc, + "gen": gen, + "nc": nc, + "G": G, + "B": B, + "next_n": next_n, + "pos": torch.reshape(position_ids, [-1]), + "max_ctx_q": int(step["v41_max_ctx_q_len"]) if nc else 0, + # Encoder replay rows of this step (host int; 0 without a prefix hit). + "replay": int(step["v41_num_replay_tokens_total"]) if nc else 0, + } + cache = {"topk": {}, "ctx_k": {}, "sched": {}} + if st["replay"]: + _rv("encoder_replay.rows", T, st["replay"], nc, gen) + + # The token embedding (first pipeline stage); a context chunk that holds + # image-span rows gets them spliced in over the placeholder rows (mm + # None: a text step). A later stage starts from the previous stage's + # stream and needs the image masks only. + emb = nn.functional.embedding(input_ids, core.w["embed"]) if core.first_pp else None + mm = self._mm_step(kwargs.get("multimodal_params"), T, tc, dev, emb) + # This rank's per-row image-span mask; under dep4 the router reads the + # gathered one (img_moe). + img_loc = mm["img"] if mm is not None else None + img_moe = img_loc + eager = not torch.cuda.is_current_stream_capturing() + # dep4: every rank pads its MoE rows to dp_rows (past layer L under + # decoder replay, to its pass's count instead: _dp_late_rows). + dp_rows = core._dp_rows(md, T) if core.dp else None + if core.dp and eager: + # The router runs on the all-gathered token set, so it needs every + # rank's mask. Every rank gathers in every eager step (a peer may + # hold image rows); captured steps are generation-only on every + # rank (the graph decision is all-rank), where no row is an image. + img_moe = self._dp_gather_mask(img_loc, dp_rows, T, dev) + + # Encoder replay: the history pool holds nothing before a replay start + # (the manager rebuilds windowed pages on a hit), so the compressed ids + # of the three prompt positions before it are stored first. Their hash + # ids are not used. + hist = step["v41_replay_history_len"] if nc and core.eng_local else 0 + if hist: + _rv("engram_hash.replay_history", hist, nc) + self._eng_hash( + step["v41_replay_history_token_ids_cuda"][:hist], + step["v41_replay_history_req_idx_cuda"][:hist], + step["v41_replay_history_cu_seq_lens_cuda"][: nc + 1], + step["v41_replay_history_start_pos_cuda"][:nc], + pools["tok"], + step["engram_token_block_table"], + ) + # Engram hash ids of every token (int64 [T, 2, 24]); also writes this + # step's compressed ids into the paged history pool (-1 = DEAD on + # image-span rows). + hashes = None + if core.eng_local: + # The history is read and written at the rows' own positions: each + # request's start is the position of its first row in this step. The + # metadata's cached_token_lens_cuda is not a substitute in generation + # steps: under the overlap scheduler with speculation it is the host's + # count before the previous step's acceptance is applied on the device + # (on_update_kv_lens corrects the kv lengths and positions, not it), so + # it runs ahead of the rows by the rejected drafts. + eng_start = torch.reshape( + nn.functional.embedding( + step["cu_seq_lens_cuda"][:B], torch.reshape(st["pos"], [T, 1]) + ), + [B], + ) + _rv("engram_hash" if mm is None else "engram_hash.masked", T, B) + hashes = self._eng_hash( + input_ids, + step["req_idx_per_token"][:T], + step["cu_seq_lens_cuda"][: B + 1], + eng_start, + pools["tok"], + step["engram_token_block_table"], + token_mask=mm["keep"] if mm is not None else None, + ) + + h = torch.empty([T, _HC, core.hidden], dtype=torch.bfloat16, device=dev) + yF = torch.empty([T, _MIX], dtype=torch.float32, device=dev) + rF = torch.empty([T], dtype=torch.float32, device=dev) + if core.first_pp: + h.copy_(torch.reshape(emb, [T, 1, core.hidden]).expand([T, _HC, core.hidden])) + else: + # A later pipeline stage starts from the previous stage's stream and + # its last FFN mixes (the lagged pre of this stage's first attention). + _rv("pp_recv_tensors", T) + self._pp_recv([h, yF, rF]) + + s = SimpleNamespace( + h=h, + yF=yF, + rF=rF, + st=st, + cache=cache, + img=img_moe, + spec=spec_metadata, + emb=emb, + hashes=hashes, + keep=mm["keep"] if mm is not None else None, + md=md, + pools=pools, + tail=None, + dp_rows=dp_rows, + ) + self._scratch(s, T, dev) + # Decoder SWA bounded replay: this context step's decoder rows (and the + # witness counters), planned from host metadata. It lives on the stage + # that holds the bounded layers (the last one; layer 20, which every + # bounded layer reads, sits on it too). + tail = ( + _DecoderTail(core, md, st, dev, self._l2g) + if nc and (core.replay or core._witness) and core.late in core.local + else None + ) + s.tail = tail + bounded = tail is not None and tail.bound + # A bounded step's context rows stop at layer L: the candidate source + # then scores no mask over them, and the decoder pass scores its own. + st["tail_bound"] = bounded + # Native KV storage: the compressed read buffer's form for the whole step. + st["comp_form"] = core._comp_form(pools, md, st, tail) if pools["native"] else None + max_rows = int(md.swa_local_indices_cuda.shape[0]) # the step capacity + # dep4 decoder replay: the passes every rank runs over the bounded + # layers (a bounded rank's slices, or the step rows as one pass). + dp_late = None + if core.dp and core.replay and eager: + dp_late = self._dp_late_rows(tail.pass_rows(max_rows) if bounded else [T], dev) + + stop = core.late if bounded or dp_late is not None else core.num_layers + for i in core.pp_layer_ids: + if i >= stop: + break + if tail is not None and i == core.late: + tail.count_full(s.h) + self._layer(i, s, dev) + if not core.last_pp: + # Hand the stream and the last FFN mixes to the next stage; this + # stage's output rows are never read (its lm_head is skipped). + _rv("pp_send_tensors", T) + self._pp_send([s.h, s.yF, s.rF]) + return s.coll + if bounded: + # Every per-row tensor a bounded layer reads is mapped to the + # decoder rows. The image-span mask the router reads follows the + # decoder rows and the stash too; a text step's is all False (the + # router's select then returns the text routing exactly), so a + # request whose chunks differ in images keeps one. + if img_loc is None: + img_loc = torch.empty([T], dtype=torch.bool, device=dev) + img_loc.copy_(core._const["false_b"].expand([T])) + passes = tail.enter( + {"h": s.h, "yF": s.yF, "rF": s.rF, "img": img_loc}, st, cache, max_rows + ) + for k, (rows, st_d, cache_d) in enumerate(passes): + p = SimpleNamespace( + h=rows["h"], + yF=rows["yF"], + rF=rows["rF"], + st=st_d, + cache=cache_d, + img=rows["img"], + spec=_TailCapture(spec_metadata, tail) if spec_metadata is not None else None, + emb=None, + hashes=None, + keep=None, + md=md, + pools=pools, + tail=None, + dp_rows=None, + ) + if dp_late is not None: + p.dp_rows = dp_late[k] + p.img = self._dp_gather_mask(rows["img"], dp_late[k], st_d["T"], dev) + self._scratch(p, st_d["T"], dev) + tail.activate(k) + for i in range(core.late, core.num_layers): + self._layer(i, p, dev) + tail.scatter(self._final(p)) + for k in range(len(passes), len(dp_late or ())): + self._dp_idle_pass(dp_late[k], dev) + out = None + else: + if dp_late is not None: + # Nothing is bounded on this rank: its step rows are its one + # pass, and it joins its peers' further passes with zero rows. + if tail is not None: + tail.count_full(s.h) + s.dp_rows = dp_late[0] + s.img = self._dp_gather_mask(img_loc, dp_late[0], T, dev) + for i in range(core.late, core.num_layers): + self._layer(i, s, dev) + for k in range(1, len(dp_late)): + self._dp_idle_pass(dp_late[k], dev) + out = self._final(s) + if tail is not None: + if st["replay"]: + # The replay rows the counters report are the rows the step's + # compressors did not run over. + skipped = st["T"] - st["compress_rows"] + assert skipped == sum(r.R for r in tail.reqs), (skipped, [r.R for r in tail.reqs]) + out = tail.close(out, spec_metadata) + return out + + +class DecodeTarget(Target): + """The specialization: routed to only when the step has no context request. + + Every row is a generation row: ``G`` requests of ``next_n`` rows each (1, or + 1 + 5 under DSpark). Binds three fewer ops than the prefill target -- the + context compressor reduction, the context index-K gather and the context + index logits serve rows this target's routing guarantees are absent -- and + holds no vision splice, encoder replay or decoder tail. Under attention DP + an eager step still joins its peers' collectives (the image-mask gather, the + decoder-replay pass plan and any idle pass), since a peer rank may be running + a prefill step at the same time. + """ + + def __init__(self, core: "DeepseekV41UnifiedModelingV2") -> None: + """Bind every op this target calls, once, against the real weights.""" + super().__init__(core) + L = core._L + eps = core.eps + e = core._eng + + # MXFP8 linears: one activation quantizer, one GEMM per weight role. + self._mxq = Mxfp8Quantize() + self._mx = {name: Mxfp8Mxfp8Gemm() for name in _MX_LINEARS} + self._mx_shape: dict = {} + + # Norms. + self._attn_norm = FlashinferRmsnorm() + self._ffn_norm = FlashinferRmsnorm() + self._q_norm = FlashinferRmsnorm() + self._kv_norm = FlashinferRmsnorm() + self._comp_norm = FlashinferRmsnorm() + self._idx_k_norm = FlashinferRmsnorm() + for op in ( + self._attn_norm, + self._ffn_norm, + self._q_norm, + self._kv_norm, + self._comp_norm, + self._idx_k_norm, + ): + op.bind_const(eps=eps) + self._final_norm = FlashinferRmsnorm() + self._final_norm.bind_const(weight=core.w["final_norm"], eps=eps) + + # Attention: RoPE, the window and compressed KV writes, their native + # read-back, the index table, the sparse MLA core, the grouped output. + self._rope = MlaRopeInplace() + self._swa_write = Mxfp8FakeQuantize() + self._comp_write = Nvfp4FakeQuantize() + self._swa_read = Mxfp8Dequantize() + self._swa_read8 = Mxfp8ToE4m3() + self._comp_read = Nvfp4KvToBf16() + self._comp_read8 = Nvfp4KvToE4m3() + self._l2g = Dsv4LocalToGlobalIndices() + self._q_e4m3 = E4m3QuantizePow2Dynamic() + self._attn = TrtllmSparseMlaDsv4() + self._attn.bind_const( + workspace_buffer=core._const["attn_ws"], + softmax_scale=core.head_dim**-0.5, + output_scale=1.0, + ) + self._attn8 = TrtllmSparseMlaDsv4Fp8() + self._attn8.bind_const(workspace_buffer=core._const["attn_ws"]) + self._wo_a = BmmOut() + + # Compressor and index K (the KV sources). + self._comp_mm = CublasMm() + self._comp_paged = CompressorPagedKvCompress() + self._comp_paged.bind_const( + ape=core._const["ape"], page_size=_PAGE, head_dim=core.head_dim, compress_ratio=2 + ) + self._idx_wk = CublasMm() + # MXFP4 (UE8M0 per 32, linear scales): the index K and the index query. + self._mxfp4 = Fp4Quantize() + self._mxfp4.bind_const( + global_scale=None, sf_vec_size=_SF32, sf_use_ue8m0=True, is_sf_swizzled_layout=False + ) + self._idx_scatter = IndexerKCacheScatter() + + # Indexer (the index sources) over the paged index K, and the candidate-block masks. + self._idx_wproj = CublasMm() + self._paged_logits = DeepgemmFp8Fp4PagedMqaLogits() + self._paged_logits.bind_const(logits_dtype=torch.float32) + self._paged_meta = DeepgemmPagedMqaLogitsMetadata() + self._cand = CandidateBlockSelect() + self._cand.bind_const(topk_blocks=_CAND_TOPK_BLOCKS, block_size=_CAND_BLOCK, next_n=1) + self._topk = FlashinferTopKRaggedDeterministic() + self._topk.bind_const(k=_TOPK) + + # MoE: router (fp32 logits), routing with gate.bias and gate.bias_vl, + # the NVFP4 activation quantizer and runner, the shared expert's SwiGLU. + self._router = CublasMm() + self._router.bind_const(bias=None, out_dtype=torch.float32) + self._gate = GateForward() + self._gate_vl = GateForward() + for op in (self._gate, self._gate_vl): + op.bind_const(topk=core.topk, route_scale=core.route_scale) + self._moe_quant = Fp4Quantize() + self._moe_quant.bind_const( + sf_vec_size=_SF16, sf_use_ue8m0=False, is_sf_swizzled_layout=_SF_LINEAR + ) + self._moe_runner = Fp4BlockScaleMoeRunner() + self._moe_runner.bind_const( + routing_logits=None, + routing_bias=None, + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm2_bias=None, + num_experts=core.num_experts, + top_k=core.topk, + n_group=None, + topk_group=None, + intermediate_size=core.i_loc, + local_expert_offset=core.e_off, + local_num_experts=core.e_loc, + routed_scaling_factor=None, + routing_method_type=_ROUTING_METHOD_INERT, + do_finalize=True, + act_type=_ACT_SWIGLU, + ) + self._silu = SiluAndMul() + self._silu.bind_const(swiglu_limit=_SWIGLU_LIMIT) + + # Collectives: the TP (or attention-DP) group of this stage. + self._allgather = Allgather() + self._allgather.bind_const(sizes=None, group=core.group) + self._reducescatter = Reducescatter() + self._reducescatter.bind_const(sizes=None, group=core.group) + # fp32 sums on AUTO (the IPC workspace) when the group is on one node, + # NCCL across nodes; bf16 sums on NCCL. + self._ar32_strategy = _AR_AUTO if core._ar_ws is not None else _AR_NCCL + self._allreduce32 = Allreduce() + self._allreduce16 = Allreduce() + for op, ws, strategy in ( + (self._allreduce32, core._ar_ws, self._ar32_strategy), + (self._allreduce16, None, _AR_NCCL), + ): + op.bind_const( + residual=None, + norm_weight=None, + scale=None, + bias=None, + workspace=ws, + group=core.group, + strategy=strategy, + op=0, + eps=1e-6, + trigger_completion_at_end=True, + ) + self._pp_send = PpSendTensors() + self._pp_recv = PpRecvTensors() + + # Engram: the hash constants, the table gather, the gate. + self._eng_hash = EngramHash() + self._eng_hash.bind_const( + token_map=e["token_map"], + multipliers=e["multipliers"], + primes=e["primes"], + offsets=e["offsets"], + pad_id=e["pad_id"], + ) + self._eng_gather = EngramGather() + self._eng_gate = EngramGate() + self._eng_gate.bind_const(eps=eps, clamp_value=1e-6) + + # Hyper-connections: per sublayer the mix GEMM + squared sums, the fused + # sinkhorn, and the lagged-pre collapse; the post mapping. + self._hc_attn_gemm = MhcGemmSqrsumFma() + self._hc_ffn_gemm = MhcGemmSqrsumFma() + for op in (self._hc_attn_gemm, self._hc_ffn_gemm): + op.bind_const(N=_MIX, K=core.hc_k) + self._hc_attn_fuse = MhcBigFuse() + self._hc_ffn_fuse = MhcBigFuse() + for op in (self._hc_attn_fuse, self._hc_ffn_fuse): + op.bind_const( + K=core.hc_k, + hidden_size=0, + rms_eps=eps, + hc_pre_eps=core.hc_eps, + hc_sinkhorn_eps=core.hc_eps, + hc_post_mult_value=2.0, + sinkhorn_repeat=20, + num_splits=1, + block_size=128, + ) + self._hc_pre_attn = MhcHcHeadApply() + self._hc_pre_ffn = MhcHcHeadApply() + self._hc_pre_final = MhcHcHeadApply() + for op in (self._hc_pre_attn, self._hc_pre_ffn, self._hc_pre_final): + op.bind_const( + mult=_HC, hidden_size=core.hidden, K=core.hc_k, norm_eps=eps, eps=core.hc_eps + ) + self._hc_post = MhcPostMapping() + + for i in core.pp_layer_ids: + d = L[i] + for name in _MX_LINEARS: + if name in d: + self._bind_mx(name, i, d[name]) + self._attn_norm.bind_layered(i, weight=d["attn_norm"]) + self._ffn_norm.bind_layered(i, weight=d["ffn_norm"]) + self._q_norm.bind_layered(i, weight=d["q_norm"]) + self._kv_norm.bind_layered(i, weight=d["kv_norm"]) + self._attn.bind_layered(i, sinks=d["sink"]) + if core.fp8_attn: + s_kv = core._fp8_kv_scale[i] + self._attn8.bind_layered(i, sinks=d["sink"], bmm2_scale=s_kv) + # The call's QK scale is softmax scale x query scale x KV scale; + # the quantizer writes it with its query scale folded in. + self._q_e4m3.bind_layered(i, scale_mult=core.head_dim**-0.5 * s_kv) + self._wo_a.bind_layered(i, b=d["wo_a_t"]) + self._router.bind_layered(i, mat_b=d["gate_w_t"]) + self._gate.bind_layered(i, bias=d["gate_bias"]) + self._gate_vl.bind_layered(i, bias=d["gate_bias_vl"]) + self._moe_quant.bind_layered(i, global_scale=d["g1"]) + self._moe_runner.bind_layered( + i, + gemm1_weights=d["fc1_w"], + gemm1_weights_scale=d["fc1_s"], + gemm1_clamp_limit=d["clamp"], + gemm2_weights=d["fc2_w"], + gemm2_weights_scale=d["fc2_s"], + output1_scale_scalar=d["o1"], + output1_scale_gate_scalar=d["o1g"], + output2_scale_scalar=d["o2"], + ) + fn, sc, bs = d["hc_attn"] + self._hc_attn_gemm.bind_layered(i, w=fn) + self._hc_attn_fuse.bind_layered(i, hc_scale=sc, hc_base=bs) + # The FFN input collapses with this attention's mixes. + self._hc_pre_ffn.bind_layered(i, scale=sc, base=bs) + fn, sc, bs = d["hc_ffn"] + self._hc_ffn_gemm.bind_layered(i, w=fn) + self._hc_ffn_fuse.bind_layered(i, hc_scale=sc, hc_base=bs) + if i > 0: + # Lagged pre: the attention input collapses with the previous + # layer's FFN mixes (the previous stage's, at a stage's first layer). + _, psc, pbs = L[i - 1]["hc_ffn"] + self._hc_pre_attn.bind_layered(i, scale=psc, base=pbs) + if i in core.kv_sources: + self._comp_mm.bind_layered(i, mat_b=d["comp_w_t"]) + self._comp_norm.bind_layered(i, weight=d["comp_norm"]) + self._idx_wk.bind_layered(i, mat_b=d["idx_wk_t"]) + self._idx_k_norm.bind_layered(i, weight=d["idx_k_norm"]) + if i in core.index_sources: + self._idx_wproj.bind_layered(i, mat_b=d["idx_wproj_t"]) + if i in core.engram_layers: + li = core.engram_layers.index(i) + self._eng_gather.bind_layered( + i, + weight=d["eng_table"], + scale=d["eng_scale"], + row_start=core.eng_start[li], + row_end=core.eng_end[li], + ) + self._eng_gate.bind_layered(i, q_weight=d["eng_q"], k_weight=d["eng_k"]) + if core.last_pp: + # The final norm's input collapses with the last layer's FFN mixes. + _, fsc, fbs = L[core.num_layers - 1]["hc_ffn"] + self._hc_pre_final.bind_const(scale=fsc, base=fbs) + + def _bind_mx(self, name: str, layer: int, wsf) -> None: + weight, weight_sf = wsf + self._mx[name].bind_layered(layer, weight=weight, weight_sf=weight_sf) + self._mx_shape[name, layer] = tuple(weight.shape) + + def step_args(self, md) -> dict: + """The metadata surface this target reads, projected once per step. + + States the two context counts instead of reading them back: this target + runs only when the step has no context request, so reading them would ask + the metadata to confirm what the routing already decided. Also the + step-contract probe (``Target._check_step_contract``).""" + return dict( + kv_cache_manager=md.kv_cache_manager, + num_contexts=0, + num_ctx_tokens=0, + num_seqs=md.num_seqs, + all_rank_num_tokens=md.all_rank_num_tokens, + req_idx_per_token=md.req_idx_per_token, + cu_seq_lens_cuda=md.cu_seq_lens_cuda, + kv_lens_cuda_runtime=md.kv_lens_cuda_runtime, + swa_local_indices_cuda=md.swa_local_indices_cuda, + sparse_mla_topk_lens=md.sparse_mla_topk_lens, + compressed_visible_lens_cuda=md.compressed_visible_lens_cuda, + compressed_kv_lens_cuda=md.compressed_kv_lens_cuda, + max_compressed_kv_lens=md.max_compressed_kv_lens, + compress_block_tables=md.compress_block_tables, + indexer_compress_block_tables=md.indexer_compress_block_tables, + sliding_block_tables=md.sliding_block_tables, + cu_new_comp_kv_cuda=md.cu_new_comp_kv_cuda, + compressed_position_ids_cuda=md.compressed_position_ids_cuda, + num_total_compressed_tokens=md.num_total_compressed_tokens, + engram_token_block_table=md.engram_token_block_table, + v41_swa_slot_mapping=md.v41_swa_slot_mapping, + v41_compress_slot_mapping=md.v41_compress_slot_mapping, + v41_index_k_slot_data=md.v41_index_k_slot_data, + v41_index_k_slot_scale=md.v41_index_k_slot_scale, + tokens_per_block=md.tokens_per_block, + max_seq_len=md.max_seq_len, + ) + + def _mx_linear(self, name: str, x, layer: int, out_dtype=torch.bfloat16): + """MXFP8 x MXFP8 linear: 1x32 UE8M0 activation quantization, then the + GEMM against the checkpoint's 32x32 blocks expanded to 1x32.""" + m = x.shape[0] + xq, xsf = self._mxq(x) + be = _gemm_backend(m) + _rv("mxfp8_mxfp8_gemm", m, self._mx_shape[name, layer], be) + return self._mx[name](xq, xsf, out_dtype=out_dtype, backend=be, layer=layer) + + def _allreduce_fp32(self, x): + """TP-member sum of bf16 per-rank partials in fp32 (the reference's + fp32 all-reduce of bf16-rounded partials); returns fp32. AUTO on the IPC + workspace when the group is on one node, NCCL across nodes.""" + x32 = torch.empty(list(x.shape), dtype=torch.float32, device=x.device) + x32.copy_(x) + _rv("allreduce", _sh(x32), "fp32", self._ar32_strategy) + return self._allreduce32(x32)[0] + + def _engram(self, i, h, hashes, md, num_tokens, dev, keep=None): + """Engram at layer i, in place on the stream h [T, 4, 5120]. ``keep`` + [T] bool is False on image-span rows (their gate is 0), or None.""" + core = self.core + li = core.engram_layers.index(i) + ids = hashes[:, li, :] # [T, 24] int64, row stride 48 + if core.eng_host: + # Every row is in the shared pinned copy: this rank's own tokens, + # no zero fill and no collective, on every member. + vals = self._eng_gather(ids, layer=i) + red = torch.reshape(vals, [num_tokens, core.eng_in]) + elif core.dp: + rows = core._dp_rows(md, num_tokens) + ids_c = torch.empty([num_tokens, core.eng_cols], dtype=torch.int64, device=dev) + ids_c.copy_(ids) + ids32 = ids_c.view(torch.int32) + if rows > num_tokens: + ids32 = nn.functional.pad(ids32, [0, 0, 0, rows - num_tokens]) + _rv("allgather", _sh(ids32), "int32") + ids_all = self._allgather(ids32).view(torch.int64) + vals = self._eng_gather(ids_all, layer=i) + flat = torch.reshape(vals, [vals.shape[0], core.eng_in]) + _rv("reducescatter", _sh(flat), "bf16") + red = self._reducescatter(flat) + if rows > num_tokens: + red = torch.split(red, [num_tokens, rows - num_tokens], 0)[0] + else: + vals = self._eng_gather(ids, layer=i) + flat = torch.reshape(vals, [num_tokens, core.eng_in]) + _rv("allreduce", _sh(flat), "bf16", _AR_NCCL) + red = self._allreduce16(flat)[0] + _rv("engram_gather", num_tokens, li) + kv = self._mx_linear("eng_wkv", red, i) # [T, 25600] + key, value = torch.split(kv, [_HC * core.hidden, core.hidden], -1) + key3 = torch.reshape(key, [num_tokens, _HC, core.hidden]) + _rv("engram_gate" if keep is None else "engram_gate.masked", num_tokens) + return self._eng_gate(h, key3, value, token_mask=keep, out=h, layer=i) + + def _attention(self, i, x, st, pools, cache, dev): + """One attention sublayer on the normed input x [T, 5120] -> [T, 5120] bf16, + every row a generation row.""" + core = self.core + r = core.ratios[i] + T, G, B, next_n = st["T"], st["G"], st["B"], st["next_n"] + md = st["md"] + pos = st["pos"] + table = core._table(r) + H = core.h_loc + hd = core.head_dim + + # Low-rank query + the shared window latent. + qkv = self._mx_linear("wqkv_a", x, i) # [T, 1792] + q_lr, kv_raw = torch.split(qkv, [core.q_lora, hd], -1) + qr = self._q_norm(q_lr, layer=i) + kv = self._kv_norm(kv_raw, layer=i) + qrq, qrsf = self._mxq(qr) + q = self._mx["wq_b"](qrq, qrsf, out_dtype=torch.bfloat16, backend=_gemm_backend(T), layer=i) + _rv("mla_rope_inplace", T, H, core.nope, r) + q3 = self._rope(torch.reshape(q, [T, H, hd]), pos, table, H, core.nope, _ROPE, False, False) + kv = torch.reshape( + self._rope( + torch.reshape(kv, [T, 1, hd]), pos, table, 1, core.nope, _ROPE, False, False + ), + [T, hd], + ) + # Window KV: FP8 E4M3 x UE8M0/32 value set into this layer's SWA slots + # (native: the codes and UE8M0 bytes of the 528 B rows). A pool write. + native = pools["native"] + swa_slots = md.v41_swa_slot_mapping[pools["local"][i]][:T] + if native: + codes, scales = pools["swa_q"][i] + _rv("mxfp8_fake_quantize.native", T, _sh(codes)) + self._swa_write(kv, codes, swa_slots, scales) + else: + _rv("mxfp8_fake_quantize", T, _sh(pools["swa3"][i])) + self._swa_write(kv, pools["swa3"][i], swa_slots) + + # Compressed KV + index K (the four KV sources). + per_query = native and st["comp_form"] == "topk" + if i in core.kv_sources: + self._compress(i, x, st, pools, dev) + if native and not per_query: + self._native_compressed(i, st, pools) + + # Compressed positions per query (index sources); readers reuse. + if i in core.index_sources: + cache["topk"][r] = self._index(i, x, qrq, qrsf, st, pools, cache, dev) + topk = cache["topk"].get(r) if r else None + + # Pool offsets -> the attention's index table (see the prefill target). + req = md.req_idx_per_token[:T] + swa_bt = md.sliding_block_tables[pools["local"][i], DeepseekV4AttentionType.SWA.value] + swa_local = md.swa_local_indices_cuda[:T] + if r: + _rv("dsv4_local_to_global_indices", T, 128, _TOPK, _PAGE // r) + gidx = self._l2g( + req, + swa_bt, + swa_local, + _PAGE, + 0, + pools["comp_tbl"][r] if native and not per_query else md.compress_block_tables[r], + topk, + _PAGE // r, + 0, + ) + if per_query: + gidx = self._native_topk(i, gidx, topk, T, pools, cache) + if native: + comp_pool = pools["comp_c4_q"][r] if per_query else pools["comp_c4_req"][r] + else: + comp_pool = pools["comp4"][core.kv_src[i]] + else: + _rv("dsv4_local_to_global_indices", T, 128) + gidx = self._l2g(req, swa_bt, swa_local, _PAGE, 0) + comp_pool = core._const["dummy_comp8" if core.fp8_attn else "dummy_comp"] + if native: + gidx, swa_pool = self._native_window(i, gidx, T, pools) + else: + swa_pool = pools["swa4"][i] + lens = md.sparse_mla_topk_lens[r][:T] + kv_lens = md.kv_lens_cuda_runtime + o = torch.empty([T, H, hd], dtype=torch.bfloat16, device=dev) + if core.fp8_attn: + # E4M3 query, one power-of-two scale s_q from this step's max|q| (on the device); + # the call's QK scale is scale * s_q * s_kv (written by the quantizer) and its + # PV scale s_kv. + qa = torch.empty([T, H, hd], dtype=torch.float8_e4m3fn, device=dev) + bmm1 = torch.empty([1], dtype=torch.float32, device=dev) + _rv("e4m3_quantize_pow2_dynamic", T, H) + qa, bmm1 = self._q_e4m3(q3, qa, bmm1, layer=i) + attn = self._attn8 + fp8 = dict(bmm1_scale=bmm1) + else: + qa = q3 + attn = self._attn + fp8 = {} + kind = "trtllm_sparse_mla_dsv4_fp8" if core.fp8_attn else "trtllm_sparse_mla_dsv4" + if T: + # The call writes `o`, which the inverse RoPE reads. + _rv(kind + ".decode", G, next_n, H, gidx.shape[1]) + attn( + torch.reshape(qa, [G, next_n, H, hd]), + swa_pool, + comp_pool, + sparse_indices=gidx, + sparse_topk_lens=lens, + seq_lens=kv_lens[:B], + out=torch.reshape(o, [G, next_n, H, hd]), + layer=i, + **fp8, + ) + # Remove the query rotation from the output's rope tail. + o = self._rope(o, pos, table, H, core.nope, _ROPE, True, False) + # Grouped low-rank output: per group [T, 4096] x [4096, 1024], into `oa`. + oa = torch.empty([T, core.g_loc * core.o_lora], dtype=torch.bfloat16, device=dev) + _rv("bmm_out", core.g_loc, T, core.group_in, core.o_lora) + self._wo_a( + torch.transpose(torch.reshape(o, [T, core.g_loc, core.group_in]), 0, 1), + out=torch.transpose(torch.reshape(oa, [T, core.g_loc, core.o_lora]), 0, 1), + layer=i, + ) + y = self._mx_linear("wo_b", oa, i) # [T, 5120] bf16 (per-rank partial under TP) + if core.dp: + return y + y32 = self._allreduce_fp32(y) + yb = torch.empty([T, core.hidden], dtype=torch.bfloat16, device=dev) + yb.copy_(y32) + return yb + + def _compress(self, i, x, st, pools, dev): + """KV source i: this step's compressed entries -> index K (native + MXFP4 pool) and main KV (E2M1 x E4M3/16 value set, bf16 pool), over + generation rows only.""" + core = self.core + r = core.ratios[i] + md = st["md"] + T, G, B, next_n = st["T"], st["G"], st["B"], st["next_n"] + hd = core.head_dim + cu = md.cu_seq_lens_cuda + if r == 2: + n = md.num_total_compressed_tokens[2] + x32 = torch.empty([T, core.hidden], dtype=torch.float32, device=dev) + x32.copy_(x) + kv_score = self._comp_mm(x32, layer=i) # fp32 [T, 1024] = [kv || score] + pooled = torch.empty([n, hd], dtype=torch.bfloat16, device=dev) + # Reserved decode rows the op leaves unwritten stay finite. + pooled.copy_(core._const["zero_bf16"].expand([n, hd])) + local = pools["local"][i] + bt_kv = md.sliding_block_tables[local, DeepseekV4AttentionType.COMPRESSOR_KV.value] + bt_sc = md.sliding_block_tables[local, DeepseekV4AttentionType.COMPRESSOR_SCORE.value] + cu_comp = md.cu_new_comp_kv_cuda[2] + if G: + # Fills `pooled` and updates the compressor state pools. + _rv("compressor_paged_kv_compress", G, next_n, n) + self._comp_paged( + kv_score, + paged_kv=pools["ckv"][i], + paged_score=pools["csc"][i], + block_table_kv=bt_kv[:B], + block_table_score=bt_sc[:B], + output=pooled, + kv_lens=md.kv_lens_cuda_runtime[:B], + cu_seq_lens=cu[: B + 1], + cu_kv_comp=cu_comp[: B + 1], + batch_size=G, + next_n=next_n, + ) + # The op above ran even with no completed group: it stores this + # step's pending odd token in the compressor state. + if n == 0: + return + latent = self._comp_norm(pooled, layer=i) + else: + n = T + assert md.num_total_compressed_tokens[1] == T, (md.num_total_compressed_tokens[1], T) + latent = self._comp_norm(self._comp_mm(x, layer=i), layer=i) + cpos = md.compressed_position_ids_cuda[r][:n] + tbl = core._rope_tables[1] + # Index K from the RoPE-free latent (before the rotation below). + k = self._idx_k_norm(self._idx_wk(latent, layer=i), layer=i) + k = torch.reshape( + self._rope( + torch.reshape(k, [n, 1, _IDX_DIM]), + cpos, + tbl, + 1, + _IDX_DIM - _ROPE, + _ROPE, + False, + False, + ), + [n, _IDX_DIM], + ) + kd, ksf = self._mxfp4(k) + _rv("indexer_k_cache_scatter", n, r, _sh(pools["idxk"][i])) + self._idx_scatter( + kd, + torch.reshape(ksf.view(torch.int32), [n, 1]), + pools["idxk"][i], + md.v41_index_k_slot_data[r][:n], + md.v41_index_k_slot_scale[r][:n], + n, + ) + # Main KV: rope at the group-first position j*r, then the E2M1 x E4M3/16 + # value set into the owner's COMPRESS slots. A pool write. + latent = torch.reshape( + self._rope( + torch.reshape(latent, [n, 1, hd]), cpos, tbl, 1, core.nope, _ROPE, False, False + ), + [n, hd], + ) + slots = md.v41_compress_slot_mapping[r][:n] + if pools["native"]: + codes, scales = pools["comp_q"][i] + _rv("nvfp4_fake_quantize.native", n, r, _sh(codes)) + self._comp_write(latent, codes, slots, scales) + else: + _rv("nvfp4_fake_quantize", n, r, _sh(pools["comp3"][i])) + self._comp_write(latent, pools["comp3"][i], slots) + + def _native_compressed(self, i, st, pools): + """Native storage, KV source i (after its write): every compressed + entry of every request of the step, dequantized into the compact + buffer at row ``b * L_r + e`` -- all the rows the layers reading this + source can name (``compressed_kv_lens`` = kv_len // r per request). + The buffer is read through ``comp_c4_req``.""" + core = self.core + r = core.ratios[i] + md = st["md"] + # Rows are the engine's request indices (block-table rows); in the + # decoder replay pass st["B"] counts its compacted request list. + B = st.get("B_engine", st["B"]) + codes, scales = pools["comp_q4"][i] + out = pools["comp_out"][r][:B] + if core.fp8_attn: + _rv("nvfp4_kv_to_e4m3", B, out.shape[1], r, _sh(codes)) + self._comp_read8( + codes, + scales, + md.compress_block_tables[r][:B], + md.compressed_kv_lens_cuda[r][:B], + out, + 1.0 / core._fp8_kv_scale[i], + ) + return + _rv("nvfp4_kv_to_bf16", B, out.shape[1], r, _sh(codes)) + self._comp_read( + codes, + scales, + md.compress_block_tables[r][:B], + md.compressed_kv_lens_cuda[r][:B], + core._const["one_f32"], + out, + ) + + def _native_topk(self, i, gidx, topk, T, pools, cache): + """Native storage, per-query read buffer, layer i: the index table + with its compressed half renumbered onto ``topk_out`` (query t's + column k at row ``t * 512 + k``). The rows are dequantized from the + pool slots the index op named (-1: left unwritten) when the top-k set + or its KV source changes; the layers that share both reuse them. + + Renumbering every column is exact: the index op packs a row's valid + entries first and ``sparse_mla_topk_lens`` counts exactly those, so + every -1 column is masked, and a masked column carries no weight + whatever finite row it names (every row of the buffer is finite: zero + at allocation, dequantized values after).""" + core = self.core + win, comp = torch.split(gidx, [_WINDOW, _TOPK], 1) + src = core.kv_src[i] + last = cache.get("comp_rows") + if last is None or last[0] is not topk or last[1] != src: + slots = pools["topk_slots"][:T] + slots.copy_(comp) + codes, scales = pools["comp_q1"][src] + if core.fp8_attn: + # Every layer reading this source shares its KV scale (``_derive_fp8_kv_scales``). + _rv("nvfp4_kv_to_e4m3.topk", T, _TOPK, core.ratios[i], _sh(codes)) + self._comp_read8( + codes, + scales, + slots, + pools["topk_cols"][:T], + pools["topk_out"][:T], + 1.0 / core._fp8_kv_scale[i], + ) + else: + _rv("nvfp4_kv_to_bf16.topk", T, _TOPK, core.ratios[i], _sh(codes)) + self._comp_read( + codes, + scales, + slots, + pools["topk_cols"][:T], + core._const["one_f32"], + pools["topk_out"][:T], + ) + # Holding the top-k tensor keeps its identity unique for the step. + cache["comp_rows"] = (topk, src) + return torch.cat([win, pools["topk_ren"][:T]], 1) + + def _native_window(self, i, gidx, T, pools): + """Native storage, layer i: dequantize the window rows each query's + index row names into the compact buffer -- query t's column c to row + ``t * 128 + c``, a -1 column to a zero row -- and return the index + table with its window half renumbered onto that buffer (the compressed + half is unchanged), and the buffer as the attention's window pool. + + A zero row stands in for -1 exactly: the kernel reads an attended -1 + as an all-zero row, and a masked column (past the query's window count) + carries no weight whatever finite row it names.""" + core = self.core + codes, scales = pools["swa_q"][i] + if gidx.shape[1] == _WINDOW: + win, rest = gidx, None + else: + win, rest = torch.split(gidx, [_WINDOW, gidx.shape[1] - _WINDOW], 1) + if core.fp8_attn: + _rv("mxfp8_to_e4m3", T * _WINDOW, _sh(pools["swa_c"]), _sh(codes)) + self._swa_read8( + codes, + scales, + pools["swa_c"], + 1.0 / core._fp8_kv_scale[i], + torch.reshape(win, [T * _WINDOW]), + ) + else: + _rv("mxfp8_dequantize", T * _WINDOW, _sh(pools["swa_c"]), _sh(codes)) + self._swa_read(codes, scales, pools["swa_c"], torch.reshape(win, [T * _WINDOW])) + ren = pools["win_idx"][:T] + return (ren if rest is None else torch.cat([ren, rest], 1)), pools["swa_c4"] + + def _index(self, i, x, qrq, qrsf, st, pools, cache, dev): + """Index source i: top-512 compressed positions per generation row, + [T, 512] int32, valid entries first in ascending order, -1 after.""" + core = self.core + r = core.ratios[i] + md = st["md"] + T, G, B, next_n = st["T"], st["G"], st["B"], st["next_n"] + src = core.kv_src[i] + iq = self._mx["idx_wq_b"]( + qrq, qrsf, out_dtype=torch.bfloat16, backend=_gemm_backend(T), layer=i + ) + iq3 = self._rope( + torch.reshape(iq, [T, _IDX_HEADS, _IDX_DIM]), + st["pos"], + core._rope_tables[1], + _IDX_HEADS, + _IDX_DIM - _ROPE, + _ROPE, + False, + False, + ) + qd, qsf = self._mxfp4(iq3) # [T,32,64] u8, [T*128] u8 + q8 = qd.view(torch.int8) + qs32 = torch.reshape(qsf.view(torch.int32), [T, _IDX_HEADS]) + wb = self._idx_wproj(x, layer=i) # bf16 [T, 32], 2^-6 folded in + w32 = torch.empty([T, _IDX_HEADS], dtype=torch.float32, device=dev) + w32.copy_(wb) + idx = torch.empty([T, _TOPK], dtype=torch.int32, device=dev) + reindex = core.cand_src < i + if T: + vis = md.compressed_visible_lens_cuda[r][:T] + ctx_lens = torch.reshape(vis, [G, next_n]) + sched = cache["sched"].get(r) + if sched is None: + sched = self._paged_meta(ctx_lens, _PAGE // r) + cache["sched"][r] = sched + max_ctx = md.max_compressed_kv_lens[r] + _rv("deepgemm_fp8_fp4_paged_mqa_logits", G, next_n, r, max_ctx) + logits = self._paged_logits( + torch.reshape(q8, [G, next_n, _IDX_HEADS, _IDX_DIM // 2]), + torch.reshape(qs32, [G, next_n, _IDX_HEADS]), + pools["idxk"][src], + w32, + ctx_lens, + md.indexer_compress_block_tables[r][:B], + sched, + max_ctx, + ) + if core._cand_binds_decode: + if i == core.cand_src: + _rv("candidate_block_select.decode", T, max_ctx) + cache["cand_gen"] = self._cand(logits, None, vis) + elif reindex: + logits = torch.where(cache["cand_gen"], logits, float("-inf")) + _rv("flashinfer_top_k_ragged_deterministic.decode", T, max_ctx) + idx.copy_(self._topk(logits, core._const["zero_i32_rows"][:T], vis, None)) + return idx + + def _route(self, i, chunk, img, dev): + """Router of one token chunk [n, 5120] bf16: the gate GEMM (bf16 x and + the checkpoint's bf16 weight on tensor cores, fp32 accumulation, fp32 + logits; every product is exact in fp32, the summation order is the + tensor cores'), then top-6 by ``sqrt(softplus(logit)) + bias``. With + ``img`` ([n] bool, True on image-span rows) those rows select with + ``gate.bias_vl`` instead: a second routing call over the same logits + and a per-row select of both results. Returns (weights [n, 6] fp32, + ids [n, 6] int32).""" + core = self.core + n = chunk.shape[0] + logits = self._router(chunk, layer=i) # fp32 [n, 384] + _rv("gate_forward", n, core.num_experts, core.topk) + tw, tid = self._gate(logits, layer=i) + if img is None: + return tw, tid + _rv("gate_forward.bias_vl", n, core.num_experts, core.topk) + tw_vl, tid_vl = self._gate_vl(logits, layer=i) + sel = torch.reshape(img, [n, 1]) + return torch.where(sel, tw_vl, tw), torch.where(sel, tid_vl, tid) + + def _moe(self, i, x, rows, dev, img=None): + """Routed NVFP4 experts + FP8 shared expert on the normed input x. + ``img`` is the image-span mask of the rows the router sees (this rank's + rows, or the gathered rows under attention DP), or None. ``rows`` is + the dep4 per-rank row count every rank pads to (None under attention + TP). ``x`` None is a dep4 rank with no row in this pass: it sends zero + rows and returns None.""" + core = self.core + T = 0 if x is None else x.shape[0] + assert T or core.dp, "only a dep4 rank runs the MoE without rows" + if T: + # Shared expert (replicated, this rank's own tokens): [gate ; up]. + h13 = self._mx_linear("sh_w13", x, i) + _rv("silu_and_mul", T, 2 * core.moe_inter) + act = self._silu(h13) + shared = self._mx_linear("sh_w2", act, i) + # Routed: dep4 gathers the token set ahead of the router, so every rank + # routes the identical set and the four expert windows tile it once. + pad_rows = 0 + xr = x + if core.dp: + assert rows is not None and T <= rows, (T, rows) + pad_rows = rows - T + if not T: + xp = torch.empty([rows, core.hidden], dtype=torch.bfloat16, device=dev) + xp.copy_(core._const["zero_bf16"].expand([rows, core.hidden])) + else: + xp = x if not pad_rows else nn.functional.pad(x, [0, 0, 0, pad_rows]) + _rv("allgather", _sh(xp), "bf16") + xr = self._allgather(xp) + parts = [] + sizes = _moe_chunks(xr.shape[0]) + if img is not None: + assert img.shape[0] == xr.shape[0], (tuple(img.shape), tuple(xr.shape)) + img_chunks = torch.split(img, sizes, 0) if img is not None else [None] * len(sizes) + for chunk, img_c in zip(torch.split(xr, sizes, 0), img_chunks): + n = chunk.shape[0] + tw, tid = self._route(i, chunk, img_c, dev) + twb = torch.empty([n, core.topk], dtype=torch.bfloat16, device=dev) + twb.copy_(tw) + hs, hsf = self._moe_quant(chunk, layer=i) + out = torch.empty([n, core.hidden], dtype=torch.bfloat16, device=dev) + _rv("fp4_block_scale_moe_runner", n, core.e_loc, core.e_off, core.i_loc) + # Given `output`, the runner finalizes into it and returns an empty tensor. + self._moe_runner( + hidden_states=hs, + hidden_states_scale=hsf.view(torch.float8_e4m3fn), + topk_weights=twb, + topk_ids=tid, + output=out, + layer=i, + ) + parts.append(out) + routed = parts[0] if len(parts) == 1 else torch.cat(parts, 0) + if core.dp: + r32 = torch.empty(list(routed.shape), dtype=torch.float32, device=dev) + r32.copy_(routed) + _rv("reducescatter", _sh(r32), "fp32") + r32 = self._reducescatter(r32) + if not T: + return None + if pad_rows: + r32 = torch.split(r32, [T, pad_rows], 0)[0] + else: + r32 = self._allreduce_fp32(routed) + m32 = torch.add(r32, shared) # fp32: the shared expert joins after the reduce + m = torch.empty([T, core.hidden], dtype=torch.bfloat16, device=dev) + m.copy_(m32) + return m + + def _dp_gather_mask(self, loc, rows, T, dev): + """dep4: the image-span mask of the all-gathered token set the router + sees, each rank padded to ``rows`` as the MoE input is. ``loc`` is this + rank's [T] mask, or None for T text rows (T may be 0: a pass that holds + none of this rank's rows).""" + assert T <= rows, (T, rows) + if loc is None or not T: + loc = torch.empty([rows if not T else T], dtype=torch.bool, device=dev) + loc.copy_(self.core._const["false_b"].expand([loc.shape[0]])) + u8 = loc.view(torch.uint8) + if rows > loc.shape[0]: + u8 = nn.functional.pad(u8, [0, rows - loc.shape[0]]) + _rv("allgather", _sh(u8), "uint8") + return self._allgather(u8).view(torch.bool) + + def _dp_late_rows(self, mine: list, dev) -> list: + """dep4 decoder replay: the passes every rank runs over the bounded + layers, as each pass's row count padded to its max over ranks. + + ``mine`` is this rank's pass row counts (a bounded rank's slices, or its + step rows as one pass). The bounded layers' MoE is a collective, so the + ranks must run the same passes in lockstep; a rank's decoder rows can + also outnumber its step rows (a small last chunk still runs its whole + tail, the rest from the stash). Every rank calls this in every eager + step under decoder replay, prefill and decode alike: ``[count, rows_0, + rows_1, ...]`` rides in the int32 [1, 48] form the Engram ids gather + in, and the result has one entry per pass of the rank with the most + (0-row entries never occur: some rank holds the rows of every pass).""" + core = self.core + width = 2 * core.eng_cols # 48 + vals = [len(mine), *mine] + assert len(vals) <= width, (len(mine), width) + core._ensure_ar(max(vals) + 1, dev) + cols = [core._ar32[v : v + 1] for v in vals] + [core._ar32[0:1]] * (width - len(vals)) + v = torch.reshape(torch.cat(cols, 0), [1, width]) + _rv("allgather", _sh(v), "int32") + got = self._allgather(v).tolist() # [world, 48] + assert len(got) == core.tp and got[core.tp_rank][: len(vals)] == vals, ( + core.tp_rank, + vals, + got, + ) + return _dp_pass_rows(got) + + def _dp_idle_pass(self, rows: int, dev) -> None: + """dep4 decoder replay: this rank's place in a pass over the bounded + layers that holds none of its rows -- the mask gather and every + bounded layer's routed round trip, with zero rows.""" + img = self._dp_gather_mask(None, rows, 0, dev) + for i in range(self.core.late, self.core.num_layers): + self._moe(i, None, rows, dev, img) + + def _scratch(self, s, T: int, dev) -> None: + """The per-row scratch of one pass over the T rows of stream ``s.h``.""" + core = self.core + f32 = torch.float32 + s.T = T + s.h_flat = torch.reshape(s.h, [T, core.hc_k]) + s.tile = _tile_n(T) + s.yA = torch.empty([T, _MIX], dtype=f32, device=dev) + s.rA = torch.empty([T], dtype=f32, device=dev) + s.mix4 = torch.empty([T, _HC], dtype=f32, device=dev) + s.post = torch.empty([T, _HC], dtype=f32, device=dev) + s.comb = torch.empty([T, _HC, _HC], dtype=f32, device=dev) + s.no_res = torch.empty([T, _HC, 0], dtype=torch.bfloat16, device=dev) + s.no_li = torch.empty([T, 0], dtype=torch.bfloat16, device=dev) + s.coll = torch.empty([T, core.hidden], dtype=torch.bfloat16, device=dev) + _rv("mhc", T, s.tile) + + def _layer(self, i: int, s, dev) -> None: + """Layer i over the pass ``s`` (its stream, scratch, step dict, cache).""" + core = self.core + T, h, h_flat = s.T, s.h, s.h_flat + # DSpark reads the INPUT stream of its target layers (37-39), every + # token; the metadata averages the 4 hc copies itself. + if s.spec is not None and s.spec.is_layer_capture(i): + s.spec.maybe_capture_hidden_states(i, h_flat) + if i in core.engram_layers: + h = self._engram(i, h, s.hashes, s.md, T, dev, s.keep) + # --- attention sublayer ------------------------------------ + s.yA, s.rA = self._hc_attn_gemm(h_flat, y=s.yA, r=s.rA, M=T, tile_n=s.tile, layer=i) + s.post, s.comb, s.no_li = self._hc_attn_fuse( + s.yA, + s.rA, + s.no_res, + post_mix=s.post, + comb_mix=s.comb, + layer_input=s.no_li, + M=T, + layer=i, + ) + if i == 0: + a_in = s.emb # one-hot pre on copy 0: the embedding itself + else: + # Lagged pre: the previous FFN's mixes collapse this input. + s.mix4.copy_(torch.split(s.yF, [_HC, _MIX - _HC], 1)[0]) + s.coll = self._hc_pre_attn(s.mix4, s.rF, h, s.coll, M=T, layer=i) + a_in = s.coll + x = self._attn_norm(a_in, layer=i) + a = self._attention(i, x, s.st, s.pools, s.cache, dev) + h = self._hc_post(h, a, s.post, s.comb, h) + # --- FFN sublayer -------------------------------------------- + s.yF, s.rF = self._hc_ffn_gemm(h_flat, y=s.yF, r=s.rF, M=T, tile_n=s.tile, layer=i) + s.post, s.comb, s.no_li = self._hc_ffn_fuse( + s.yF, + s.rF, + s.no_res, + post_mix=s.post, + comb_mix=s.comb, + layer_input=s.no_li, + M=T, + layer=i, + ) + s.mix4.copy_(torch.split(s.yA, [_HC, _MIX - _HC], 1)[0]) + s.coll = self._hc_pre_ffn(s.mix4, s.rA, h, s.coll, M=T, layer=i) + x = self._ffn_norm(s.coll, layer=i) + m = self._moe(i, x, s.dp_rows, dev, s.img) + self._hc_post(h, m, s.post, s.comb, h) + + def _final(self, s): + """The final collapse with layer 39's FFN mixes, then the final norm.""" + s.mix4.copy_(torch.split(s.yF, [_HC, _MIX - _HC], 1)[0]) + s.coll = self._hc_pre_final(s.mix4, s.rF, s.h, s.coll, M=s.T) + return self._final_norm(s.coll) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None = None, + position_ids: torch.IntTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + lora_params: dict | None = None, + **kwargs, + ) -> torch.Tensor: + core = self.core + self._check_step_contract(attn_metadata) + md = attn_metadata + step = self.step_args(md) + spec_metadata = kwargs.get("spec_metadata") + pools = core._get_pools(md) + dev = input_ids.device + T = input_ids.shape[0] + B = step["num_seqs"] + next_n = T // B if B else 1 + assert next_n * B == T, (T, B) + # A DSpark generation row verifies its 5 drafts: 1 + 5 tokens. + assert next_n == 1 or (spec_metadata is not None and next_n == 1 + core.ds.block), ( + f"generation rows of {next_n} tokens (spec_metadata {'set' if spec_metadata is not None else 'None'})" + ) + st = { + "md": md, + "T": T, + "tc": step["num_ctx_tokens"], + "gen": T, + "nc": step["num_contexts"], + "G": B, + "B": B, + "next_n": next_n, + "pos": torch.reshape(position_ids, [-1]), + "max_ctx_q": 0, + "replay": 0, + "tail_bound": False, + } + cache = {"topk": {}, "ctx_k": {}, "sched": {}} + + emb = nn.functional.embedding(input_ids, core.w["embed"]) if core.first_pp else None + eager = not torch.cuda.is_current_stream_capturing() + # dep4: every rank pads its MoE rows to dp_rows (past layer L under + # decoder replay, to its pass's count instead: _dp_late_rows). + dp_rows = core._dp_rows(md, T) if core.dp else None + img_moe = None + if core.dp and eager: + # A peer rank may hold image rows in a prefill step: the router's + # gathered mask needs this rank's (all text) rows too. + img_moe = self._dp_gather_mask(None, dp_rows, T, dev) + + # Engram hash ids of every token (int64 [T, 2, 24]); also writes this + # step's compressed ids into the paged history pool. The history is read + # and written at the rows' own positions (see the prefill target for why + # cached_token_lens_cuda is not a substitute). + hashes = None + if core.eng_local: + eng_start = torch.reshape( + nn.functional.embedding( + step["cu_seq_lens_cuda"][:B], torch.reshape(st["pos"], [T, 1]) + ), + [B], + ) + _rv("engram_hash", T, B) + hashes = self._eng_hash( + input_ids, + step["req_idx_per_token"][:T], + step["cu_seq_lens_cuda"][: B + 1], + eng_start, + pools["tok"], + step["engram_token_block_table"], + ) + + h = torch.empty([T, _HC, core.hidden], dtype=torch.bfloat16, device=dev) + yF = torch.empty([T, _MIX], dtype=torch.float32, device=dev) + rF = torch.empty([T], dtype=torch.float32, device=dev) + if core.first_pp: + h.copy_(torch.reshape(emb, [T, 1, core.hidden]).expand([T, _HC, core.hidden])) + else: + # A later pipeline stage starts from the previous stage's stream and + # its last FFN mixes (the lagged pre of this stage's first attention). + _rv("pp_recv_tensors", T) + self._pp_recv([h, yF, rF]) + + s = SimpleNamespace( + h=h, + yF=yF, + rF=rF, + st=st, + cache=cache, + img=img_moe, + spec=spec_metadata, + emb=emb, + hashes=hashes, + keep=None, + md=md, + pools=pools, + dp_rows=dp_rows, + ) + self._scratch(s, T, dev) + # Native KV storage: the compressed read buffer's form for the whole step. + st["comp_form"] = core._comp_form(pools, md, st, None) if pools["native"] else None + # dep4 decoder replay: a decode rank holds one pass over the bounded + # layers (its step rows) and joins any further pass a prefill peer runs. + dp_late = None + if core.dp and core.replay and eager: + dp_late = self._dp_late_rows([T], dev) + + stop = core.late if dp_late is not None else core.num_layers + for i in core.pp_layer_ids: + if i >= stop: + break + self._layer(i, s, dev) + if not core.last_pp: + # Hand the stream and the last FFN mixes to the next stage; this + # stage's output rows are never read (its lm_head is skipped). + _rv("pp_send_tensors", T) + self._pp_send([s.h, s.yF, s.rF]) + return s.coll + if dp_late is not None: + s.dp_rows = dp_late[0] + s.img = self._dp_gather_mask(None, dp_late[0], T, dev) + for i in range(core.late, core.num_layers): + self._layer(i, s, dev) + for k in range(1, len(dp_late)): + self._dp_idle_pass(dp_late[k], dev) + return self._final(s) + + +class VisionTower: + """DeepSeek-ViT + aligner + span assembly (the checkpoint's + ``inference/vision.py`` and ``model.py`` ``merge_image_embeddings``). + + Holds no parameters of its own: its weights are the ``vis*`` rows of the + core's weight table, and it binds its own entries to them in + ``post_load_weights``. Replicated on every rank, no collective. It runs per + context chunk, over the images whose spans meet that chunk, ONE IMAGE PER + CALL: an image's rows depend on that image alone -- not on the chunking, nor + on the request's other images -- and a step's vision memory is one image's + working set. + + Input: one request's ``multimodal_data["image"]`` from the sibling + ``input_processor`` (``P`` patches, ``C`` aligner cells, ``S`` span positions + over all its images, in prompt order): ``pixel_values`` [P, 3, 14, 14] + bf16, ``vit_cu_seqlens`` / ``aligner_cu_cells`` / ``span_lengths`` (host + lists), ``vit_rope_cos`` / ``vit_rope_sin`` [P, 32] fp32, and the + image-local tables ``aligner_index_local`` [C, 9] (value P_i = image i's + zero pad) and ``span_source_local`` [S] (delimiters at C_i, C_i + 1, + C_i + 2). + + Per image (attention stays within it):: + + x = patch_w . pixels + b [P, 1024] + 32 x (pre-norm, eps 1e-6): + qkv = wqkv . rmsnorm(x) + b; 2D rotate-half RoPE on q, k + x += wo . attention(qkv) + b (bidirectional, 16 x 64) + x += w2 . silu_and_mul(w1 . rmsnorm(x)) + v = rmsnorm(x) + cells = v rows gathered 3x3 per cell (zero pad) [C, 9216] + out = w2 . gelu(w1 . cells + b1) + b2 [C, 5120] + span = [out; image_start; image_end; image_newline][span_source_local] + + The packed qkv is roped in place as ``[3P, 16, 64]`` rows: q and k of + patch t read table row t, v reads an identity row (cos 1, sin 0) and is + left unchanged. An image has at most 9189 patches and 1021 cells (the + 1024-position span limit), inside the per-call envelope the entries + certify. Every per-call length / position tensor is a view or a concat + of views of load-time ``arange`` tables. + + **CUDA graphs (switch ``TRTLLM_DSV41_VIT_CUDA_GRAPH``, unset = on; none + under ``llm_args.disable_mm_encoder``).** The + tower call -- patch embed, 32 blocks, final norm -- is captured at load + time once per padded patch count ``B`` in ``_VIT_GRAPH_BUCKETS`` (all + graphs in one memory pool), and an image of P patches replays the + smallest ``B >= P``. Every row-wise op runs over B rows (rows P..B-1 are + padding and never reach a real row); the attention is the same + ``thop_attention`` call over two sequence slots -- the image, then a pad + sequence over rows [P, B) -- with ``max_context_q_len_override = B``, + which keeps the FMHA's launch shape independent of the split, so a replay + follows whatever lengths and offsets the device buffers hold. Before a + replay the image's pixels, RoPE rows and the slot lengths / offsets are + copied into the graphs' static buffers; after it the P output rows are + copied out. The aligner and the span gathers stay eager. With the switch + at 0 every tower call runs eagerly over its P rows. + + Replicated ranks compute the same bits here without any tactic pin: no + call of the tower, the aligner or the span table goes through the + AutoTuner (``cublas_mm`` picks its cuBLASLt algorithm from the shape + alone -- a fixed table, else cuBLASLt's heuristic; the norms, RoPE, + ``thop`` attention, activations and gathers have one kernel per shape), + so nothing is chosen per rank.""" + + def __init__(self, core: "DeepseekV41UnifiedModelingV2"): + self._core = core + self._c: dict | None = None + # Graph mode: {bucket: CUDAGraph} (None = every call eager) and the + # static buffers the graphs read and write. + self._g: dict | None = None + self._g_bufs: dict | None = None + + def post_load_weights(self) -> None: + """Bind every entry the tower, aligner and span table call, then capture + the tower's CUDA graphs (unless disabled).""" + w = self._core.w + dev = w["vis_norm"].device + self._patch = CublasMm() + self._patch.bind_const(mat_b=w["vis_patch_w"].t(), bias=w["vis_patch_b"]) + self._norm1 = FlashinferRmsnorm() + self._norm1.bind_const(weight=w["vis0_norm1"], eps=_VIT_EPS) + self._qkv = CublasMm() + self._wo = CublasMm() + self._norm2 = FlashinferRmsnorm() + self._norm2.bind_const(eps=_VIT_EPS) + self._w1 = CublasMm() + self._w2 = CublasMm() + # Each block ends with the next block's norm1 (the tower's final norm + # after the last block). + self._next_norm = FlashinferRmsnorm() + self._next_norm.bind_const(eps=_VIT_EPS) + for L in range(_VIT_LAYERS): + p = f"vis{L}_" + self._qkv.bind_layered(L, mat_b=w[p + "wqkv_w"].t(), bias=w[p + "wqkv_b"]) + self._wo.bind_layered(L, mat_b=w[p + "wo_w"].t(), bias=w[p + "wo_b"]) + self._norm2.bind_layered(L, weight=w[p + "norm2"]) + self._w1.bind_layered(L, mat_b=w[p + "w1"].t()) + self._w2.bind_layered(L, mat_b=w[p + "w2"].t()) + nxt = w[f"vis{L + 1}_norm1"] if L + 1 < _VIT_LAYERS else w["vis_norm"] + self._next_norm.bind_layered(L, weight=nxt) + self._silu = FlashinferSiluAndMul() + self._rope = MlaRopeInplace() + self._attn = ThopAttention() + # Bidirectional attention within each sequence, no pool, no RoPE: every + # argument but the step's lengths, offsets and buffers is constant. + self._attn.bind_const( + k=None, + v=None, + output_sf=None, + kv_cache_block_offsets=None, + host_kv_cache_pool_pointers=None, + host_kv_cache_pool_mapping=None, + cache_indirection=None, + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + out_scale=None, + rotary_inv_freq=None, + rotary_cos_sin=None, + latent_cache=None, + q_pe=None, + block_ids_per_seq=None, + attention_sinks=None, + is_fused_qkv=True, + update_kv_cache=True, + predicted_tokens_per_seq=1, + local_layer_idx=0, + num_heads=_VIT_HEADS, + num_kv_heads=_VIT_HEADS, + head_size=_VIT_HEAD_DIM, + tokens_per_block=None, + max_context_length=_VIT_ATTN_MAX_LEN, + max_seq_len=_VIT_ATTN_MAX_LEN, + attention_window_size=_VIT_ATTN_MAX_LEN, + beam_width=1, + mask_type=_VIT_MASK_BIDIRECTIONAL, + quant_mode=0, + q_scaling=1.0, + position_embedding_type=0, + rope_dim=0, + rope_base=10000.0, + rope_scale_type=0, + rope_scale=1.0, + rope_short_m_scale=1.0, + rope_long_m_scale=1.0, + rope_max_positions=1024, + rope_original_max_positions=1024, + use_paged_context_fmha=False, + attention_input_type=0, + is_mla_enable=False, + chunked_prefill_buffer_batch_size=1, + q_lora_rank=None, + kv_lora_rank=None, + qk_nope_head_dim=None, + qk_rope_head_dim=None, + v_head_dim=None, + rope_append=None, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + attention_chunk_size=None, + softmax_stats_tensor=None, + is_spec_decoding_enabled=False, + use_spec_decoding=False, + is_spec_dec_tree=False, + spec_decoding_generation_lengths=None, + spec_decoding_position_offsets_for_cpp=None, + spec_decoding_packed_mask=None, + spec_decoding_bl_tree_mask_offset=None, + spec_decoding_bl_tree_mask=None, + spec_bl_tree_first_sparse_mask_offset_kv=None, + sparse_kv_indices=None, + sparse_kv_offsets=None, + sparse_attn_indices=None, + sparse_attn_offsets=None, + sparse_attn_indices_block_size=0, + cu_kv_seqlens=None, + ) + self._al1 = CublasMm() + self._al1.bind_const(mat_b=w["vis_al_w1"].t(), bias=w["vis_al_b1"]) + self._al2 = CublasMm() + self._al2.bind_const(mat_b=w["vis_al_w2"].t(), bias=w["vis_al_b2"]) + self._gelu = MoeGelu() + ident = torch.zeros((1, 2, _VIT_ROPE_HALF), dtype=torch.float32, device=dev) + ident[:, 0] = 1.0 + self._c = { + "span_tail": w["vis_span_tail"], + # Row v of either table holds the value v: per-call lengths, + # offsets and positions are its views. + "ar": torch.arange(_VIT_MAX_PATCHES + 1, dtype=torch.int32, device=dev), + "ar_host": torch.arange(_VIT_MAX_PATCHES + 1, dtype=torch.int32), + "zeros_host": torch.zeros(_VIT_GRAPH_SLOTS, dtype=torch.int32), + "ident": ident, + # thop_attention grows this in place on its first calls (eager + # calls only; the graphs have their own). + "workspace": torch.empty(0, dtype=torch.int8, device=dev), + } + if getattr(self._core, "disable_mm_encoder", False): + # The ViT weights stay loaded, so an image that arrives anyway is + # still answered, by eager tower calls. + sys.stdout.write( + f"MODELING_V2 deepseek_v41 rank {self._core.rank}: disable_mm_encoder: no ViT CUDA " + "graphs captured (an image, if one arrives, runs the tower eagerly)\n" + ) + sys.stdout.flush() + elif _env_switch(_VIT_GRAPH_ENV) is not False: + self._capture_graphs(dev) + + # -- metadata ------------------------------------------------------------- + def _meta(self, image: dict, i: int) -> dict: + """The attention / RoPE metadata of image ``i`` of a request (eager call).""" + c = self._c + ar, arh = c["ar"], c["ar_host"] + cu = image["vit_cu_seqlens"] + p0, P = int(cu[i]), int(cu[i + 1]) - int(cu[i]) + cos = torch.reshape(image["vit_rope_cos"][p0 : p0 + P], [P, 1, _VIT_ROPE_HALF]) + sin = torch.reshape(image["vit_rope_sin"][p0 : p0 + P], [P, 1, _VIT_ROPE_HALF]) + table = torch.cat([torch.cat([cos, sin], 1), c["ident"]], 0) # [P + 1, 2, 32] + t = torch.reshape(ar[:P], [P, 1]) + pos = torch.reshape( + torch.cat([t.expand([P, 2]), torch.reshape(ar[P : P + 1], [1, 1]).expand([P, 1])], 1), + [3 * P], + ) + return { + "P": P, + "N": 1, + "rope_table": table, + "rope_pos": pos, # (t, t, P) per patch t + "seq_lens": ar[P : P + 1], + "cu_seqlens": torch.cat([ar[0:1], ar[P : P + 1]], 0), + "seq_lens_host": arh[P : P + 1], + "total_kv_host": torch.cat([arh[P : P + 1], arh[0:1]], 0), + "types_host": c["zeros_host"][:1], + "q_len_override": None, + "workspace": c["workspace"], + } + + # -- pieces --------------------------------------------------------------- + def _patch_embed(self, px): + """[P, 588] pixels -> [P, 1024]; a P in the split-K band runs on + zero-padded rows and keeps the first P output rows.""" + P = px.shape[0] + rows = _patch_embed_rows(P) + _rv("vision.cublas_mm", rows, _VIT_PATCH_IN, _VIT_DIM) + if rows == P: + return self._patch(px) + x = self._patch(nn.functional.pad(px, [0, 0, 0, rows - P])) + return torch.split(x, [P, rows - P], 0)[0] + + def _attention(self, qkv, m): + """Bidirectional attention within each sequence of the call.""" + P, N = m["P"], m["N"] + out = torch.empty([P, _VIT_DIM], dtype=torch.bfloat16, device=qkv.device) + if m["q_len_override"] is None: + _rv("vision.thop_attention", P, N) + else: + _rv("vision.thop_attention.graph", P, N, m["q_len_override"]) + return self._attn( + q=qkv, + output=out, + workspace_=m["workspace"], + sequence_length=m["seq_lens"], + host_past_key_value_lengths=m["seq_lens_host"], + host_total_kv_lens=m["total_kv_host"], + context_lengths=m["seq_lens"], + host_context_lengths=m["seq_lens_host"], + host_request_types=m["types_host"], + max_context_q_len_override=m["q_len_override"], + max_num_requests=N, + cu_q_seqlens=m["cu_seqlens"], + num_contexts=N, + num_ctx_tokens=P, + ) + + def _block(self, L, x, h, m): + """Block L on the residual stream ``x`` with ``h = norm1(x)``; returns + ``(x_out, next_norm(x_out))``.""" + P = m["P"] + qkv = self._qkv(h, layer=L) # [P, 3072] + qkv = torch.reshape( + self._rope( + torch.reshape(qkv, [3 * P, _VIT_HEADS, _VIT_HEAD_DIM]), + m["rope_pos"], + m["rope_table"], + _VIT_HEADS, + 0, + _VIT_HEAD_DIM, + False, + True, + ), + [P, 3 * _VIT_DIM], + ) + a = self._wo(self._attention(qkv, m), layer=L) + x = torch.add(x, a) + h2 = self._norm2(x, layer=L) + act = self._silu(self._w1(h2, layer=L)) + x = torch.add(x, self._w2(act, layer=L)) + return x, self._next_norm(x, layer=L) + + def _tower(self, px, m): + """One tower call: [P, 588] pixels -> ViT output [P, 1024].""" + x = self._patch_embed(px) + h = self._norm1(x) + for L in range(_VIT_LAYERS): + x, h = self._block(L, x, h, m) + return h + + def _vit(self, image: dict, i: int): + """ViT output [P, 1024] of image ``i`` of a request (one tower call).""" + cu = image["vit_cu_seqlens"] + p0, P = int(cu[i]), int(cu[i + 1]) - int(cu[i]) + assert 0 < P <= _VIT_MAX_IMAGE_PATCHES, (i, P) + pv = image["pixel_values"] + px = torch.reshape(pv, [pv.shape[0], _VIT_PATCH_IN])[p0 : p0 + P] + _rv("vision.vit", P, 1) + if self._g is not None: + return self._vit_graph(image, px, p0) + return self._tower(px, self._meta(image, i)) + + # -- CUDA graphs ---------------------------------------------------------- + def _graph_meta(self, B: int) -> dict: + """The metadata a bucket-B graph is captured with: the static device + slot lengths / offsets (rewritten before each replay), and host values + that hold the B rows as one sequence (read only at capture).""" + g, arh = self._g_bufs, self._c["ar_host"] + return { + "P": B, + "N": _VIT_GRAPH_SLOTS, + "rope_table": g["table"], + "rope_pos": g["pos"][: 3 * B], + "seq_lens": g["seq"], + "cu_seqlens": g["cu"], + "seq_lens_host": torch.cat([arh[B : B + 1], arh[0:1]], 0), + "total_kv_host": torch.cat([arh[B : B + 1], arh[0:1]], 0), + "types_host": self._c["zeros_host"], + "q_len_override": B, + "workspace": g["workspace"], + } + + def _set_graph_slots(self, P: int, B: int) -> None: + """Slot lengths (P, B - P) and offsets (0, P, B) of an image of P + patches in bucket B.""" + g, ar = self._g_bufs, self._c["ar"] + g["seq"].copy_(torch.cat([ar[P : P + 1], ar[B - P : B - P + 1]], 0)) + g["cu"].copy_(torch.cat([ar[0:1], ar[P : P + 1], ar[B : B + 1]], 0)) + + def _graph_body(self, B: int, m: dict) -> None: + """What a bucket-B graph runs: the tower over the static pixel rows, + its output into the static output rows.""" + g = self._g_bufs + g["h"][:B].copy_(self._tower(g["px"][:B], m)) + + def _stage_graph_inputs(self, image: dict, px, p0: int) -> int: + """Copy one image's inputs into the graphs' static buffers; returns its bucket.""" + g = self._g_bufs + P = px.shape[0] + B = next(b for b in _VIT_GRAPH_BUCKETS if b >= P) + g["px"][:P].copy_(px) + cos = torch.reshape(image["vit_rope_cos"][p0 : p0 + P], [P, 1, _VIT_ROPE_HALF]) + sin = torch.reshape(image["vit_rope_sin"][p0 : p0 + P], [P, 1, _VIT_ROPE_HALF]) + g["table"][:P].copy_(torch.cat([cos, sin], 1)) + self._set_graph_slots(P, B) + return B + + def _vit_graph(self, image: dict, px, p0: int): + """One tower call replayed from the smallest bucket that holds it.""" + P = px.shape[0] + B = self._stage_graph_inputs(image, px, p0) + _rv("vision.vit_graph", B, P) + self._g[B].replay() + out = torch.empty([P, _VIT_DIM], dtype=torch.bfloat16, device=px.device) + out.copy_(self._g_bufs["h"][:P]) + return out + + def _capture_graphs(self, dev) -> None: + """Capture one graph per bucket, largest first, into one shared pool. + + Every bucket runs once eagerly before any capture, so the graphs' + attention workspace reaches its final size first (a later resize would + leave an earlier graph pointing at freed memory); the eager calls keep + their own workspace for the same reason.""" + BMAX = _VIT_GRAPH_BUCKETS[-1] + table = torch.zeros((BMAX + 1, 2, _VIT_ROPE_HALF), dtype=torch.float32, device=dev) + table[BMAX:] = self._c["ident"] + t = torch.arange(BMAX, dtype=torch.int32, device=dev).view(BMAX, 1) + pos = torch.cat( + [t.expand(BMAX, 2), torch.full((BMAX, 1), BMAX, dtype=torch.int32, device=dev)], 1 + ) + self._g_bufs = { + "px": torch.zeros((BMAX, _VIT_PATCH_IN), dtype=torch.bfloat16, device=dev), + # Rows 0..BMAX-1: the staged image's RoPE rows; row BMAX: the identity. + "table": table, + "pos": pos.reshape(3 * BMAX).contiguous(), # (t, t, BMAX) per patch t + "seq": torch.zeros(_VIT_GRAPH_SLOTS, dtype=torch.int32, device=dev), + "cu": torch.zeros(_VIT_GRAPH_SLOTS + 1, dtype=torch.int32, device=dev), + "workspace": torch.empty(0, dtype=torch.int8, device=dev), + "h": torch.zeros((BMAX, _VIT_DIM), dtype=torch.bfloat16, device=dev), + } + free0 = torch.cuda.mem_get_info(dev)[0] + t0 = time.perf_counter() + graphs = {} + with torch.inference_mode(): + metas = {B: self._graph_meta(B) for B in _VIT_GRAPH_BUCKETS} + for B in reversed(_VIT_GRAPH_BUCKETS): + self._set_graph_slots(B, B) + self._graph_body(B, metas[B]) + torch.cuda.synchronize(dev) + pool = torch.cuda.graph_pool_handle() + stream = torch.cuda.Stream(dev) + for B in reversed(_VIT_GRAPH_BUCKETS): + self._set_graph_slots(B, B) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph( + graph, pool=pool, stream=stream, capture_error_mode="thread_local" + ): + self._graph_body(B, metas[B]) + graphs[B] = graph + torch.cuda.synchronize(dev) + # The capture-time host tensors are read at capture only; kept with the graphs. + self._g_metas = metas + self._g = graphs + torch.cuda.empty_cache() + held = free0 - torch.cuda.mem_get_info(dev)[0] + sys.stdout.write( + f"MODELING_V2 deepseek_v41 rank {self._core.rank}: ViT CUDA graphs: {len(graphs)} buckets " + f"{_VIT_GRAPH_BUCKETS[0]}..{BMAX} patches captured in {time.perf_counter() - t0:.1f} s; " + f"device memory held {held / 2**20:.1f} MiB (static buffers + graph pool)\n" + ) + sys.stdout.flush() + + # -- aligner and spans ------------------------------------------------------ + def _aligner(self, cells): + """Gathered cells [C, 9, 1024] -> [C, 5120], in chunks of <= 2048 cells.""" + c = self._c + C = cells.shape[0] + outs = [] + for c0 in range(0, C, _AL_MAX_ROWS): + n = min(_AL_MAX_ROWS, C - c0) + flat = torch.reshape(cells[c0 : c0 + n], [n, _VIT_RATIO * _VIT_RATIO * _VIT_DIM]) + hh = self._al1(flat) + _rv("vision.moe_gelu", n) + g = self._gelu(hh, c["ar"][n : n + 1], c["ar"][1:2], n) + outs.append(self._al2(g)) + return outs[0] if len(outs) == 1 else torch.cat(outs, 0) + + def image_rows(self, image: dict, i: int): + """[S_i, 5120] bf16: the whole span of image ``i`` of a request.""" + ccu = image["aligner_cu_cells"] + c0, c1 = int(ccu[i]), int(ccu[i + 1]) + s0 = sum(int(n) for n in image["span_lengths"][:i]) + s1 = s0 + int(image["span_lengths"][i]) + v = self._vit(image, i) + # Row P_i of the padded ViT output is the zero pad. + cells = nn.functional.embedding( + image["aligner_index_local"][c0:c1], nn.functional.pad(v, [0, 0, 0, 1]) + ) + out = self._aligner(cells) + # Rows C_i, C_i + 1, C_i + 2 = start, end, newline. + table = torch.cat([out, self._c["span_tail"]], 0) + return nn.functional.embedding(image["span_source_local"][s0:s1], table) + + def profile_largest_image(self, max_num_tokens: int) -> None: + """The largest vision work of one context step, for the engine's memory + profiling: the tower runs one image per call and each image's rows go + into the step's embedding before the next image runs, so that work is + one image of the largest size -- a 3 x 3063 patch grid (9189 patches, + 1021 aligner cells, a 1024-position span, the ``max_image_tokens`` + limit) -- next to the step's [max_num_tokens, 5120] embedding. Runs it + (graph replay or eager, as the engine will) and keeps nothing.""" + from .input_processor import ( + ImageGrid, + PreprocessedImage, + VisionGeometry, + build_vision_inputs, + ) + + dev = self._c["ar"].device + geometry = VisionGeometry( + patch_size=14, downsample_ratio=_VIT_RATIO, max_image_tokens=1024, min_pixels=544 * 544, + max_wh_ratio=None, hidden_size=_VIT_DIM, num_heads=_VIT_HEADS, num_layers=_VIT_LAYERS, + rope_theta=10000.0, image_token_id=_IMAGE_TOKEN, + ) # fmt: skip + grid = ImageGrid( + n_llm_h=1, n_llm_w=1021, best_height=42, best_width=42882, n_vit_h=3, n_vit_w=3063 + ) + assert grid.num_span_tokens == geometry.max_image_tokens + patches = torch.zeros((grid.n_vit_h * grid.n_vit_w, 3, 14, 14), dtype=torch.bfloat16) + image = build_vision_inputs([PreprocessedImage(patches=patches, grid=grid)], [0], geometry) + image = {k: (v.to(dev) if torch.is_tensor(v) else v) for k, v in image.items()} + S = grid.num_span_tokens + R = max(int(max_num_tokens), 1) + emb = torch.empty([R, self._core.hidden], dtype=torch.bfloat16, device=dev) + rows = self.image_rows(image, 0) + emb[: min(R, S)].copy_(rows[: min(R, S)]) + torch.cuda.synchronize(dev) + sys.stdout.write( + f"MODELING_V2 deepseek_v41 rank {self._core.rank}: memory profiling ran the largest image " + f"({grid.n_vit_h * grid.n_vit_w} patches, {S}-position span, " + f"{'graph' if self._g is not None else 'eager'} tower) beside a {R}-row embedding\n" + ) + sys.stdout.flush() + + +# The drafter's MXFP8 x MXFP8 linears: per stage, then the two shared operands. +_DS_STAGE_LINEARS = ("wqkv_a", "wq_b", "wo_b", "sh_w13", "sh_w2") +_DS_SHARED_LINEARS = ("main_proj", "wkv3") + + +class DSparkDraftModel: + """The object ``DSv4DSparkWorker`` drives as ``draft_model`` + (``tensorrt_llm/_torch/speculative/dspark.py``). + + Holds no parameters and no state of its own: its weights are the core's + ``w_ds`` (``DRAFT_WEIGHTS``: stage ``s`` as ``l{s}_*``) and the trunk's + embedding, its head is the shell's ``lm_head``, its RoPE table is the core's + theta-10000 table 0, and the per-request rings ``[rows, 3, 128, 512]`` are + the worker's. It binds its own entries in ``post_load_weights``, never at + construction (the engine replaces the meta-device parameters when it + materializes them). It takes no attention metadata: the draft step is + generation-only and the prefill seeding writes rings, so it has no phase to + route on. + + **Positions.** The worker hands frames (frame = position + 1). The drafter + ropes at the position (frame - 1) everywhere and writes ring slot + ``(frame - 1) & 127`` -- the reference's ``start_pos % window``. A ring's + ``valid_len`` rows are the most recent contiguous positions ending at the + bonus position ``p``, so a block query attends slots + ``(p - valid_len + 1 + c) & 127`` for ``c < valid_len``: the first + ``valid_len`` columns of its index row, which is all the kernel reads of + the window segment. That is ``0..valid_len-1`` when the request's first + ring row sat at a slot-0 position (a cold prefill), but not after a + prefix-cache hit, whose seeding starts at the first uncached position; a + full ring keeps plain slot order. + + **What one draft step computes** (``forward_batched``, G requests, N = 5G): + ``main_x = main_norm(main_proj(main_hidden))`` once with stage 0's weights; + every stage's ring row for the bonus position, + ``fp8_qdq(rope(kv_norm_s(wkv_s(main_x))))``; block ids ``[bonus, noise x 4]`` + through the trunk embedding into 4 hc copies; three trunk-shaped blocks + (lagged-pre mHC chain, attention, MoE) whose attention reads the stage's + ring plus all 5 block rows, non-causal, with sinks; the head collapse, norm + and the shell's ``lm_head``; then the sequential Markov correction + ``logits[j] += W2 . W1[prev_j]`` in fp32 with ``prev_{j+1} = argmax``. + """ + + block_size = _DS_BLOCK + num_stages = _DS_STAGES + + def __init__(self, core: "DeepseekV41UnifiedModelingV2", lm_head): + assert core.ds is not None, "the core was built without the DSpark drafter" + self.g = core.ds + self._attn_params = { + "window_size": self.g.window, + "head_dim": self.g.head_dim, + "n_heads": self.g.heads, + "eps": self.g.eps, + } + self._core = core + self._lm_head = lm_head + # Written by the worker: the rows the RoPE table must hold. + self._runtime_position_ceiling = 0 + self.dummy_loaded = False + self._c: dict | None = None + self._tabs: dict = {} + + # -- runtime surface ------------------------------------------------------ + def run_moe_lockstep_noop(self, all_rank_draft_tokens, device) -> None: + """The drafter's MoE is replicated on every rank and has no collective: + a rank with no generation rows owes its peers nothing.""" + return None + + def load_weights_from_target_model(self, model) -> None: + """LoadFormat.DUMMY only: the drafter's rows are in the core, which the + dummy loader already filled.""" + self.dummy_loaded = True + + # -- post-load binding ---------------------------------------------------- + def post_load_weights(self) -> None: + """Bind every entry the drafter calls, build its load-time constants, and + run one eager call of every drafter GEMM shape on each ``mxfp8_mxfp8_gemm`` + backend.""" + core, g = self._core, self.g + w, wd = core.w, core.w_ds + dev = wd["norm"].device + eps = g.eps + + self._mxq = Mxfp8Quantize() + self._mx = {name: Mxfp8Mxfp8Gemm() for name in _DS_STAGE_LINEARS + _DS_SHARED_LINEARS} + self._mx_shape: dict = {} + for name in _DS_SHARED_LINEARS: + self._mx[name].bind_const(weight=wd[name], weight_sf=wd[name + "_sf"]) + self._mx_shape[name, None] = tuple(wd[name].shape) + self._q_norm = FlashinferRmsnorm() + self._kv_norm = FlashinferRmsnorm() + self._attn_norm = FlashinferRmsnorm() + self._ffn_norm = FlashinferRmsnorm() + for op in (self._q_norm, self._kv_norm, self._attn_norm, self._ffn_norm): + op.bind_const(eps=eps) + self._main_norm = FlashinferRmsnorm() + self._main_norm.bind_const(weight=wd["main_norm"], eps=eps) + self._norm = FlashinferRmsnorm() + self._norm.bind_const(weight=wd["norm"], eps=eps) + self._rope = MlaRopeInplace() + # The FP8 1x32 value set of a K/V row: ring rows and the block's rows. + self._kv_write = Mxfp8FakeQuantize() + self._attn = TrtllmSparseMlaDsv4() + self._attn.bind_const( + workspace_buffer=core._const["attn_ws"], + softmax_scale=g.head_dim**-0.5, + output_scale=1.0, + ) + self._wo_a = BmmOut() + # The router: fp32 x and the checkpoint's gate weight promoted to fp32. + self._router = CublasMm() + self._gate = GateForward() + self._gate.bind_const(topk=g.topk, route_scale=g.route_scale) + self._silu = SiluAndMul() + self._silu.bind_const(swiglu_limit=g.swiglu) + # The MXFP4 runner reads its MXFP8 activation linear, padded to 512 columns. + self._moe_quant = Mxfp8Quantize() + self._moe_quant.bind_const(swizzled_layout=_SF_LINEAR, alignment=_DS_MOE_ALIGN) + self._moe_runner = Mxe4m3Mxe2m1BlockScaleMoeRunner() + self._hc_attn_gemm = MhcGemmSqrsumFma() + self._hc_ffn_gemm = MhcGemmSqrsumFma() + for op in (self._hc_attn_gemm, self._hc_ffn_gemm): + op.bind_const(N=_MIX, K=_HC * g.hidden) + self._hc_attn_fuse = MhcBigFuse() + self._hc_ffn_fuse = MhcBigFuse() + for op in (self._hc_attn_fuse, self._hc_ffn_fuse): + op.bind_const( + K=_HC * g.hidden, + hidden_size=0, + rms_eps=eps, + hc_pre_eps=core.hc_eps, + hc_sinkhorn_eps=core.hc_eps, + hc_post_mult_value=2.0, + sinkhorn_repeat=20, + num_splits=1, + block_size=128, + ) + self._hc_pre_attn = MhcHcHeadApply() + self._hc_pre_ffn = MhcHcHeadApply() + self._hc_pre_final = MhcHcHeadApply() + for op in (self._hc_pre_attn, self._hc_pre_ffn, self._hc_pre_final): + op.bind_const( + mult=_HC, hidden_size=g.hidden, K=_HC * g.hidden, norm_eps=eps, eps=core.hc_eps + ) + self._hc_post = MhcPostMapping() + self._mk_head = CublasMm() + self._mk_head.bind_const(mat_b=wd["mk_head"].t()) + + for s in range(_DS_STAGES): + p = f"l{s}_" + for name in _DS_STAGE_LINEARS: + self._mx[name].bind_layered(s, weight=wd[p + name], weight_sf=wd[p + name + "_sf"]) + self._mx_shape[name, s] = tuple(wd[p + name].shape) + self._q_norm.bind_layered(s, weight=wd[p + "q_norm"]) + self._kv_norm.bind_layered(s, weight=wd[p + "kv_norm"]) + self._attn_norm.bind_layered(s, weight=wd[p + "attn_norm"]) + self._ffn_norm.bind_layered(s, weight=wd[p + "ffn_norm"]) + self._attn.bind_layered(s, sinks=wd[p + "sink"]) + self._wo_a.bind_layered( + s, + b=torch.transpose( + wd[p + "wo_a"].view(g.o_groups, g.o_lora, g.heads * g.head_dim // g.o_groups), + 1, + 2, + ), + ) + self._router.bind_layered(s, mat_b=wd[p + "gate_w"].t()) + self._gate.bind_layered(s, bias=wd[p + "gate_bias"]) + self._moe_runner.bind_layered( + s, + gemm1_weights=wd[p + "fc1_w"], + gemm1_weights_scale=wd[p + "fc1_s"], + gemm2_weights=wd[p + "fc2_w"], + gemm2_weights_scale=wd[p + "fc2_s"], + ) + self._hc_attn_gemm.bind_layered(s, w=wd[p + "hc_attn_fn"]) + self._hc_attn_fuse.bind_layered( + s, hc_scale=wd[p + "hc_attn_scale"], hc_base=wd[p + "hc_attn_base"] + ) + # The FFN input collapses with this stage's attention mixes. + self._hc_pre_ffn.bind_layered( + s, scale=wd[p + "hc_attn_scale"], base=wd[p + "hc_attn_base"] + ) + self._hc_ffn_gemm.bind_layered(s, w=wd[p + "hc_ffn_fn"]) + self._hc_ffn_fuse.bind_layered( + s, hc_scale=wd[p + "hc_ffn_scale"], hc_base=wd[p + "hc_ffn_base"] + ) + if s > 0: + # Lagged pre: the previous stage's FFN mixes collapse this input. + q = f"l{s - 1}_" + self._hc_pre_attn.bind_layered( + s, scale=wd[q + "hc_ffn_scale"], base=wd[q + "hc_ffn_base"] + ) + # Head collapse with the last stage's FFN mixes. + q = f"l{_DS_STAGES - 1}_" + self._hc_pre_final.bind_const(scale=wd[q + "hc_ffn_scale"], base=wd[q + "hc_ffn_base"]) + + # valid_len -> max(valid_len, 1) as an int32 row lookup: a sparse-MLA + # row with seq_lens 0 attends the WHOLE 128-row window page. + vl_floor = ( + torch.arange(g.window + 1, dtype=torch.int32, device=dev).clamp(min=1).view(-1, 1) + ) + # bonus + 1 -> max(bonus, 0): rejection sampling can hand a -1 bonus, + # which an embedding lookup would fault on. + bonus_shift = ( + (torch.arange(g.vocab + 1, dtype=torch.int64, device=dev) - 1).clamp(min=0).view(-1, 1) + ) + # valid_len -> (128 - max(valid_len, 1), slot mask): the first ring + # slot a block query reads is ((start_pos + shift) & mask), i.e. + # (p - valid + 1) & 127, or slot 0 for a full ring. + v1 = torch.arange(g.window + 1, dtype=torch.int64, device=dev).clamp(min=1) + ring_shift = (g.window - v1).view(-1, 1) + ring_mask = torch.where(v1 < g.window, g.window - 1, 0).view(-1, 1) + self._c = { + "mk_embed": wd["mk_embed"], + "embed": w["embed"], + "noise": torch.full((1, 1), g.noise, dtype=torch.int64, device=dev), + "neg1_i64": torch.full((1,), -1, dtype=torch.int64, device=dev), + "one_i64": torch.full((1,), 1, dtype=torch.int64, device=dev), + "zero_bf16": torch.zeros((1, 1), dtype=torch.bfloat16, device=dev), + "j5": torch.arange(_DS_BLOCK, dtype=torch.int64, device=dev).view(1, _DS_BLOCK), + "stage_base": (torch.arange(_DS_STAGES, dtype=torch.int64, device=dev) * _WINDOW).view( + 1, _DS_STAGES + ), + "vl_floor": vl_floor.contiguous(), + "bonus_shift": bonus_shift.contiguous(), + "ring_shift": ring_shift.contiguous(), + "ring_mask": ring_mask.contiguous(), + "c128": torch.arange(g.window, dtype=torch.int64, device=dev).view(1, g.window), + } + # The routed experts' SwiGLU clamp, one value for every expert. + self._moe_runner.bind_const( + routing_logits=None, + routing_bias=None, + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=torch.full((g.experts,), g.swiglu, dtype=torch.float32, device=dev), + gemm2_bias=None, + num_experts=g.experts, + top_k=g.topk, + n_group=None, + topk_group=None, + intermediate_size=g.inter, + valid_hidden_size=g.hidden, + valid_intermediate_size=g.inter, + local_expert_offset=0, + local_num_experts=g.experts, + routed_scaling_factor=None, + routing_method_type=_ROUTING_METHOD_INERT, + act_type=_ACT_SWIGLU, + ) + self._tabs = {} + self._warm_gemm_backends(dev) + + def _warm_gemm_backends(self, dev) -> None: + """One eager call of every drafter GEMM shape on each backend + (cute-dsl at M <= 32, cudnn above): a backend's first-ever call inside a + CUDA-graph capture breaks the capture, and cudnn's first call moves + FlashInfer's cached workspace once.""" + assert not torch.cuda.is_current_stream_capturing() + linears = [(name, None) for name in _DS_SHARED_LINEARS] + for s in range(_DS_STAGES): + linears += [(name, s) for name in _DS_STAGE_LINEARS] + seen = set() + for name, s in linears: + n, k = self._mx_shape[name, s] + if (n, k) in seen: + continue + seen.add((n, k)) + for m in (1, 33): + x = torch.empty([m, k], dtype=torch.bfloat16, device=dev) + x.copy_(self._c["zero_bf16"].expand([m, k])) + self._mx_linear(name, x, s) + torch.cuda.synchronize() + + def _tables(self, rows: int, dev) -> dict: + """Constant index tables for a ring of ``rows`` slots, built on the + first (eager) call for that ring: stage-page bases, the block-row + columns, the live-column count, and the zero-padded block K/V buffer.""" + t = self._tabs.get(rows) + if t is not None: + return t + assert not torch.cuda.is_current_stream_capturing(), ( + "drafter index tables are built on the first eager call for a ring size; " + "run one eager draft step (the graph runner's warm-up does) before capture" + ) + r = torch.arange(rows, dtype=torch.int64, device=dev) + s3 = torch.arange(_DS_STAGES, dtype=torch.int64, device=dev) + base3 = (r[:, None] * _DS_STAGES + s3[None, :]) * _WINDOW # [rows, 3] + blk = torch.full((rows, _DS_BLK_COLS), -1, dtype=torch.int32, device=dev) + blk[:, :_DS_BLOCK] = ( + r[:, None] * _DS_BLOCK + torch.arange(_DS_BLOCK, device=dev)[None, :] + ).to(torch.int32) + lens = torch.full((rows * _DS_BLOCK,), _DS_IDX_LEN, dtype=torch.int32, device=dev) + blk_kv = torch.zeros( + (rows * _DS_BLOCK + _DS_BLK_PAD_ROWS, self.g.head_dim), dtype=torch.bfloat16, device=dev + ) + t = {"base3": base3.contiguous(), "blk": blk, "lens": lens, "blk_kv": blk_kv} + self._tabs[rows] = t + return t + + def _rope_table(self) -> torch.Tensor: + """The core's theta-10000, no-YaRN table (grown only in its first forward).""" + tab = self._core._rope_tables[0] + need = max(int(self._runtime_position_ceiling), 1) + assert tab.shape[0] >= need, ( + f"RoPE table has {tab.shape[0]} rows, the worker's position ceiling is {need}" + ) + return tab + + # -- compositions ---------------------------------------------------------- + def _mx_linear(self, name: str, x, layer=None, out_dtype=torch.bfloat16): + """MXFP8 x MXFP8 linear (the trunk's recipe and backend choice).""" + m = x.shape[0] + xq, xsf = self._mxq(x) + be = _gemm_backend(m) + _rv("dspark.mxfp8_mxfp8_gemm", m, self._mx_shape[name, layer], be) + return self._mx[name](xq, xsf, out_dtype=out_dtype, backend=be, layer=layer) + + def _main_x(self, hidden): + """main_norm(main_proj(captured rows)), stage 0's weights.""" + return self._main_norm(self._mx_linear("main_proj", hidden)) + + def _ring_write(self, main_x, pos64, idx3, ring3, table) -> None: + """Every stage's ``fp8_qdq(rope(kv_norm_s(wkv_s(main_x)), pos))`` into + ring rows ``idx3[:, s]`` of ``ring3`` (pages of 128 rows); a negative + row is skipped. A write into the worker's rings.""" + M = main_x.shape[0] + hd = self.g.head_dim + kv3 = self._mx_linear("wkv3", main_x) # [M, 3 * 512] + parts = torch.split(kv3, [hd] * _DS_STAGES, -1) + p32 = torch.empty([M], dtype=torch.int32, device=main_x.device) + p32.copy_(pos64) + cols = torch.split(idx3, 1, 1) + _rv("dspark.ring_write", M, _sh(ring3)) + for s in range(_DS_STAGES): + kv = self._kv_norm(parts[s], layer=s) + kv = torch.reshape( + self._rope( + torch.reshape(kv, [M, 1, hd]), p32, table, 1, hd - _ROPE, _ROPE, False, False + ), + [M, hd], + ) + self._kv_write(kv, ring3, torch.reshape(cols[s], [M])) + + def write_context_windows(self, hidden, window_positions, window) -> None: + """Prefill seeding (eager): ``hidden`` [n <= 128, 15360] captured rows, + ``window_positions`` [n] int64 frames, ``window`` the slot's + ``[3, 128, 512]`` ring view.""" + n = hidden.shape[0] + assert n <= _WINDOW and tuple(window.shape) == (_DS_STAGES, _WINDOW, self.g.head_dim), ( + n, + tuple(window.shape), + ) + pos64 = torch.add(window_positions, self._c["neg1_i64"]) + off = torch.bitwise_and(pos64, _WINDOW - 1) + # [n, 3]: s * 128 + (p & 127) + idx3 = torch.add(torch.reshape(off, [n, 1]), self._c["stage_base"]) + self._ring_write(self._main_x(hidden), pos64, idx3, window, self._rope_table()) + + def write_context_windows_batched( + self, interim_hidden, interim_pos, slots, interim_valid, kv_windows + ) -> None: + """Multi-accept back-fill (captured): ``interim_hidden`` [G, 5, 15360], + ``interim_pos`` [G, 5] int64 frames, ``slots`` [G] int64, + ``interim_valid`` [G, 5] bool (masked rows are not written), + ``kv_windows`` [rows, 3, 128, 512].""" + G = interim_hidden.shape[0] + M = G * _DS_BLOCK + dev = interim_hidden.device + rows = kv_windows.shape[0] + tabs = self._tables(rows, dev) + ring3 = torch.reshape(kv_windows, [rows * _DS_STAGES, _WINDOW, self.g.head_dim]) + pos64 = torch.add(interim_pos, self._c["neg1_i64"]) # [G, 5] + off = torch.bitwise_and(pos64, _WINDOW - 1) + bases = nn.functional.embedding(slots, tabs["base3"]) # [G, 3] + # [G, 5, 3] + idx = torch.add( + torch.reshape(bases, [G, 1, _DS_STAGES]), torch.reshape(off, [G, _DS_BLOCK, 1]) + ) + idx = torch.where(torch.reshape(interim_valid, [G, _DS_BLOCK, 1]), idx, self._c["neg1_i64"]) + hid = torch.reshape(interim_hidden, [M, interim_hidden.shape[-1]]) + self._ring_write( + self._main_x(hid), + torch.reshape(pos64, [M]), + torch.reshape(idx, [M, _DS_STAGES]), + ring3, + self._rope_table(), + ) + + def _attention(self, s, x, st): + """One stage's attention over its ring + the transient block.""" + g = self.g + N, dev, table = st["N"], st["dev"], st["table"] + H, hd = g.heads, g.head_dim + qkv = self._mx_linear("wqkv_a", x, s) # [N, 1792] + q_lr, kv_raw = torch.split(qkv, [g.q_lora, hd], -1) + qr = self._q_norm(q_lr, layer=s) + kv = self._kv_norm(kv_raw, layer=s) + qrq, qrsf = self._mxq(qr) + q = self._mx["wq_b"](qrq, qrsf, out_dtype=torch.bfloat16, backend=_gemm_backend(N), layer=s) + q3 = self._rope( + torch.reshape(q, [N, H, hd]), st["bpos"], table, H, hd - _ROPE, _ROPE, False, False + ) + kv = torch.reshape( + self._rope( + torch.reshape(kv, [N, 1, hd]), st["bpos"], table, 1, hd - _ROPE, _ROPE, False, False + ), + [N, hd], + ) + # Block K/V: the FP8 1x32 value set, rows 0..N-1 of the padded buffer. + blk_kv = self._kv_write(kv, st["blk_kv"]) + o = torch.empty([N, 1, H, hd], dtype=torch.bfloat16, device=dev) + _rv("dspark.trtllm_sparse_mla_dsv4", N, H, _DS_IDX_W, _sh(st["ring4"])) + o = self._attn( + torch.reshape(q3, [N, 1, H, hd]), + st["ring4"], + torch.reshape(blk_kv, [N, 1, 1, hd]), + sparse_indices=st["idx"][s], + sparse_topk_lens=st["lens"], + seq_lens=st["seq_lens"], + out=o, + layer=s, + ) + o3 = self._rope( + torch.reshape(o, [N, H, hd]), st["bpos"], table, H, hd - _ROPE, _ROPE, True, False + ) + gin = H * hd // g.o_groups + oa = torch.empty([N, g.o_groups * g.o_lora], dtype=torch.bfloat16, device=dev) + # Grouped low-rank output, into `oa`. + self._wo_a( + torch.transpose(torch.reshape(o3, [N, g.o_groups, gin]), 0, 1), + out=torch.transpose(torch.reshape(oa, [N, g.o_groups, g.o_lora]), 0, 1), + layer=s, + ) + return self._mx_linear("wo_b", oa, s) + + def _moe(self, s, x, dev): + """FP8 shared expert + routed MXFP4 top-3 of 128, replicated (no collective).""" + g = self.g + N = x.shape[0] + h13 = self._mx_linear("sh_w13", x, s) + act = self._silu(h13) + shared = self._mx_linear("sh_w2", act, s) + x32 = torch.empty([N, g.hidden], dtype=torch.float32, device=dev) + x32.copy_(x) + logits = self._router(x32, layer=s) # fp32 [N, 128] + tw, tid = self._gate(logits, layer=s) + twb = torch.empty([N, g.topk], dtype=torch.bfloat16, device=dev) + twb.copy_(tw) + hs, hsf = self._moe_quant(x) + out = torch.empty([N, g.hidden], dtype=torch.bfloat16, device=dev) + _rv("dspark.mxe4m3_mxe2m1_block_scale_moe_runner", N) + # Every rank computes this drafter MoE on the same rows and must produce + # the same bits, but the op's AutoTuner entry tunes per rank + # (DistributedTuningStrategy.INDEPENDENT) inside the engine's tuning + # pass, a cache file loaded at warm-up (TLLM_AUTOTUNER_CACHE_PATH) can + # carry per-rank entries for it, and its tactics differ in output bits. + # In fallback_tactic() the op neither tunes nor reads the cache: it takes + # its fallback tactic on every rank, eager and under capture alike -- + # the cold state its contract certifies. Given `output`, the + # runner writes it and returns an empty tensor. + with fallback_tactic(): + self._moe_runner( + hidden_states=hs, + hidden_states_scale=hsf, + topk_weights=twb, + topk_ids=tid, + output=out, + layer=s, + ) + return torch.add(out, shared) + + def _blocks(self, h, emb, st): + """The three stages on the [N, 4, 5120] stream; returns the normed head + input [N, 5120].""" + g = self.g + N, dev = st["N"], st["dev"] + K = _HC * g.hidden + tile = _tile_n(N) + f32 = torch.float32 + yA, rA = ( + torch.empty([N, _MIX], dtype=f32, device=dev), + torch.empty([N], dtype=f32, device=dev), + ) + yF, rF = ( + torch.empty([N, _MIX], dtype=f32, device=dev), + torch.empty([N], dtype=f32, device=dev), + ) + mix4 = torch.empty([N, _HC], dtype=f32, device=dev) + post = torch.empty([N, _HC], dtype=f32, device=dev) + comb = torch.empty([N, _HC, _HC], dtype=f32, device=dev) + no_res = torch.empty([N, _HC, 0], dtype=torch.bfloat16, device=dev) + no_li = torch.empty([N, 0], dtype=torch.bfloat16, device=dev) + coll = torch.empty([N, g.hidden], dtype=torch.bfloat16, device=dev) + h_flat = torch.reshape(h, [N, K]) + _rv("dspark.mhc", N, tile) + for s in range(_DS_STAGES): + yA, rA = self._hc_attn_gemm(h_flat, y=yA, r=rA, M=N, tile_n=tile, layer=s) + post, comb, no_li = self._hc_attn_fuse( + yA, rA, no_res, post_mix=post, comb_mix=comb, layer_input=no_li, M=N, layer=s + ) + if s == 0: + a_in = emb # one-hot pre on copy 0: the block embedding itself + else: + # Lagged pre: the previous stage's FFN mixes collapse this input. + mix4.copy_(torch.split(yF, [_HC, _MIX - _HC], 1)[0]) + coll = self._hc_pre_attn(mix4, rF, h, coll, M=N, layer=s) + a_in = coll + x = self._attn_norm(a_in, layer=s) + a = self._attention(s, x, st) + h = self._hc_post(h, a, post, comb, h) + yF, rF = self._hc_ffn_gemm(h_flat, y=yF, r=rF, M=N, tile_n=tile, layer=s) + post, comb, no_li = self._hc_ffn_fuse( + yF, rF, no_res, post_mix=post, comb_mix=comb, layer_input=no_li, M=N, layer=s + ) + mix4.copy_(torch.split(yA, [_HC, _MIX - _HC], 1)[0]) + coll = self._hc_pre_ffn(mix4, rA, h, coll, M=N, layer=s) + x = self._ffn_norm(coll, layer=s) + m = self._moe(s, x, dev) + h = self._hc_post(h, m, post, comb, h) + # Head collapse with the last stage's FFN mixes, then its norm. + mix4.copy_(torch.split(yF, [_HC, _MIX - _HC], 1)[0]) + coll = self._hc_pre_final(mix4, rF, h, coll, M=N) + return self._norm(coll) + + def forward_batched( + self, + main_hidden, + bonus, + start_pos, + *, + kv_windows, + slots, + valid_len, + temperature=0.0, + confidence_threshold=0.0, + return_logits=True, + all_rank_num_tokens=None, + ): + """One draft step for G generation requests (captured). + + ``main_hidden`` [G, 15360] (the captured row of the last accepted + position), ``bonus`` [G] int64, ``start_pos`` [G] int64 = the frame of + ``main_hidden``, ``slots`` [G] int64, ``valid_len`` [G] int64 ring rows + valid after this step's bonus write. Writes the bonus row of every + stage's ring first. Returns ``(draft_ids [G, 5] int64, None, + block_logits [G, 5, V] fp32)``; the worker samples from ``[2]``. The + Markov feedback is greedy whatever ``temperature`` (the worker passes + 0.0); confidence truncation is not computed. ``all_rank_num_tokens`` is + unused: the drafter has no collective. + """ + g = self.g + G = main_hidden.shape[0] + N = G * _DS_BLOCK + dev = main_hidden.device + hd = g.head_dim + rows = kv_windows.shape[0] + assert G <= rows, (G, rows) + tabs = self._tables(rows, dev) + table = self._rope_table() + ring3 = torch.reshape(kv_windows, [rows * _DS_STAGES, _WINDOW, hd]) + ring4 = torch.reshape(kv_windows, [rows * _DS_STAGES, 1, _WINDOW, hd]) + bonus = torch.reshape( + nn.functional.embedding(torch.add(bonus, self._c["one_i64"]), self._c["bonus_shift"]), + [G], + ) + + # The bonus position p = start_pos - 1 into every stage's ring. + main_x = self._main_x(main_hidden) + pos64 = torch.add(start_pos, self._c["neg1_i64"]) + off = torch.bitwise_and(pos64, _WINDOW - 1) + bases = nn.functional.embedding(slots, tabs["base3"]) # [G, 3] stage-page bases + self._ring_write(main_x, pos64, torch.add(bases, torch.reshape(off, [G, 1])), ring3, table) + + # Block ids [bonus, noise x 4] -> the trunk embedding -> 4 hc copies. + ids = torch.cat( + [torch.reshape(bonus, [G, 1]), self._c["noise"].expand([G, _DS_BLOCK - 1])], 1 + ) + emb = torch.reshape(nn.functional.embedding(ids, self._c["embed"]), [N, g.hidden]) + h = torch.empty([N, _HC, g.hidden], dtype=torch.bfloat16, device=dev) + h.copy_(torch.reshape(emb, [N, 1, g.hidden]).expand([N, _HC, g.hidden])) + + # Block positions p + 1 + j = start_pos + j; per-row window counts >= 1. + bpos = torch.empty([G, _DS_BLOCK], dtype=torch.int32, device=dev) + bpos.copy_(torch.add(torch.reshape(start_pos, [G, 1]), self._c["j5"])) + vl = torch.empty([G, _DS_BLOCK], dtype=torch.int32, device=dev) + vl.copy_(nn.functional.embedding(valid_len, self._c["vl_floor"]).expand([G, _DS_BLOCK])) + # Index rows [page base + ((first + c) & 127) | g * 5 + 0..4, -1 x 3] per + # stage, first = (p - valid + 1) & 127 (0 for a full ring). + first = torch.bitwise_and( + torch.add( + torch.reshape(start_pos, [G, 1]), + nn.functional.embedding(valid_len, self._c["ring_shift"]), + ), + nn.functional.embedding(valid_len, self._c["ring_mask"]), + ) # [G, 1] + ring_off = torch.bitwise_and(torch.add(first, self._c["c128"]), _WINDOW - 1) # [G, 128] + blk = torch.reshape( + torch.split(tabs["blk"], [G, rows - G], 0)[0], [G, 1, _DS_BLK_COLS] + ).expand([G, _DS_BLOCK, _DS_BLK_COLS]) + stage_bases = torch.split(bases, 1, 1) + idx = [] + for s in range(_DS_STAGES): + ring_cols = torch.empty([G, _WINDOW], dtype=torch.int32, device=dev) + ring_cols.copy_(torch.add(ring_off, stage_bases[s])) + ring_cols = torch.reshape(ring_cols, [G, 1, _WINDOW]).expand([G, _DS_BLOCK, _WINDOW]) + idx.append(torch.reshape(torch.cat([ring_cols, blk], -1), [N, _DS_IDX_W])) + st = { + "N": N, + "dev": dev, + "table": table, + "ring4": ring4, + "bpos": torch.reshape(bpos, [N]), + "seq_lens": torch.reshape(vl, [N]), + "lens": torch.split(tabs["lens"], [N, rows * _DS_BLOCK - N], 0)[0], + "blk_kv": torch.split(tabs["blk_kv"], [N, tabs["blk_kv"].shape[0] - N], 0)[0], + "idx": idx, + } + xn = self._blocks(h, emb, st) + + # The shell's head: fp32 block logits [G, 5, V], gathered under TP. + lg = self._lm_head(xn) + V = lg.shape[-1] + logits = torch.empty([G, _DS_BLOCK, V], dtype=torch.float32, device=dev) + logits.copy_(torch.reshape(lg, [G, _DS_BLOCK, V])) + + # Sequential Markov correction with greedy feedback. + toks = torch.empty([G, _DS_BLOCK], dtype=torch.int64, device=dev) + lcols = torch.split(logits, 1, 1) + tcols = torch.split(toks, 1, 1) + prev = bonus + for j in range(_DS_BLOCK): + lj = torch.reshape(lcols[j], [G, V]) + e = nn.functional.embedding(prev, self._c["mk_embed"]) # [G, 256] bf16 + e32 = torch.empty([G, g.markov_rank], dtype=torch.float32, device=dev) + e32.copy_(e) + lj.copy_(torch.add(lj, self._mk_head(e32))) + prev = torch.argmax(lj, dim=-1) + torch.reshape(tcols[j], [G]).copy_(prev) + return toks, None, logits + + +@register_auto_model("ModelingV2DeepseekV41FlashNvfp4Sm103Unified") +class ModelingV2DeepseekV41FlashNvfp4Sm103Unified( + DecoderModelForCausalLM[DeepseekV41UnifiedModelingV2, PretrainedConfig] +): + # `load_weights` streams the rank-local slices it declares out of the + # mmapped safetensors handles, so the loader hands them over lazily and the + # checkpoint is never materialized in host RAM. + requires_lazy_safetensors = True + + def __init__(self, model_config: ModelConfig): + cfg = model_config.pretrained_config + assert cfg is not None + super().__init__( + DeepseekV41UnifiedModelingV2(model_config), + config=model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + # The DSpark branch, built only when `speculative_config` asked for it + # (the core already refused any mode but the embedded DSpark drafter). + # It is written out rather than inherited from the in-tree one-engine + # shell: that shell builds its drafter through a module-level function + # that would construct TensorRT-LLM's own DSpark model instead of this + # target's. The runtime finds the pieces by attribute: `spec_worker` + # (the engine), `draft_config` (read by the loader without getattr; + # None = the drafter's weights come from the same checkpoint) and + # `load_draft_weights` (the loader's second, lazily mapped pass). + self.spec_config = getattr(model_config, "spec_config", None) + self.draft_config = None + self.draft_model = None + self.spec_worker = None + self._draft_weights_loaded = False + if self.spec_config is not None: + assert self.model.ds is not None, "core built without the DSpark drafter" + self.draft_model = DSparkDraftModel(self.model, self.lm_head) + self.spec_worker = get_spec_worker(self.spec_config, model_config, model_config.mapping) + self.spec_worker.set_draft_model(self.draft_model) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None = None, + position_ids: torch.IntTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + return_context_logits: bool = False, + spec_metadata=None, + lora_params: dict | None = None, + **kwargs, + ): + """One engine step. + + Without a spec worker, a step with no `multimodal_params` delegates to + the inherited base forward with every argument it was handed + (`resource_manager` stays inside `**kwargs`), so the text path is the + base forward exactly. The base forward forwards no `**kwargs` to the + core, so a step that carries images calls the core itself with them and + then the same logits processor. + + With one, the shell owns the logits gather: DSpark sets + `without_logits`, so the engine applies no second gather, and the worker + takes the trunk's hidden states ungathered beside the gathered logits. + `position_ids` goes on in the engine's `[1, T]` shape; the worker + squeezes it itself. `multimodal_params` reaches the core in `**kwargs`.""" + if self.spec_worker is None: + assert spec_metadata is None, ( + "spec_metadata arrived without a spec worker; this engine was " + "built without a speculative_config" + ) + mm_params = kwargs.get("multimodal_params") + if mm_params: + hidden = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + lora_params=lora_params, + multimodal_params=mm_params, + ) + return self.logits_processor.forward( + hidden, self.lm_head, attn_metadata, return_context_logits + ) + return super().forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + return_context_logits, + spec_metadata, + lora_params, + **kwargs, + ) + assert spec_metadata is not None, ( + "the spec worker is built but the engine passed no spec_metadata" + ) + hidden = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + lora_params=lora_params, + **kwargs, + ) + if not self.model.last_pp: + # An earlier pipeline stage: the drafter, the logits and the verify + # run on the last stage, whose sampled and drafted tokens the engine + # hands back to every stage. This stage's output is never read. + return hidden + self._seed_tail_rings() + # `gather_ids`: one row per context request (its last token) and + # 1 + 5 per generation request. + logits = self.logits_processor.forward( + nn.functional.embedding(spec_metadata.gather_ids, hidden), + self.lm_head, + attn_metadata, + True, + ) + return self.spec_worker( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=kwargs.get("resource_manager"), + ) + + def profile_inline_mm_encoder(self, max_num_tokens: int) -> None: + """The engine's memory profiling (``KvCacheCreator``) hook for the + vision tower, which runs inside the context step: run the largest + vision work one step can hold, so the measured peak covers it. The + step keeps nothing of it past the embedding the text dummy also holds. + Nothing to profile under ``disable_mm_encoder``.""" + if self.model.disable_mm_encoder: + return + self.model._vision.profile_largest_image(max_num_tokens) + + def load_weights(self, weights, *args, **kwargs): + _weights.MODEL_WEIGHTS.load(self, weights) + + def load_draft_weights(self, weights, weight_mapper=None): + """The loader's second pass over the same checkpoint (lazily mapped): + the drafter's `mtp.*` rows.""" + assert self.draft_model is not None, "load_draft_weights on an engine built without DSpark" + _weights.DRAFT_WEIGHTS.load(self, weights) + self._draft_weights_loaded = True + + def post_load_weights(self): + super().post_load_weights() + self.model.post_load_weights() + if self.draft_model is not None: + assert self._draft_weights_loaded or self.draft_model.dummy_loaded, ( + "DSpark is on but the loader never ran load_draft_weights: the drafter's " + "parameters would be uninitialized memory" + ) + self.draft_model.post_load_weights() + + def _seed_tail_rings(self) -> None: + """Decoder replay: a tail's rows from earlier chunks exist only in this + step's decoder pass. Their captured rows go into the request's ring + (the worker's slot for it) before the worker seeds this chunk's rows, + so the ring holds the whole tail.""" + pending = self.model._tail_ring + if not pending: + return + self.model._tail_ring = [] + worker = self.spec_worker + slots = getattr(worker, "_req_to_slot", None) or {} + for rid, hidden, frames in pending: + slot = slots.get(rid) + if slot is not None: + self.draft_model.write_context_windows(hidden, frames, worker._kv_windows[slot]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/__init__.py new file mode 100644 index 000000000000..1059445198db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/activation/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/activation/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/activation/moe_gelu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/activation/moe_gelu.py new file mode 100644 index 000000000000..75b69d80dd19 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/activation/moe_gelu.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exact (erf) GELU over the rows of a tile-laid-out 2-D tensor via the MoE activation kernel. + +The op is `torch.ops.trtllm.moe_gelu`, registered in +`cpp/tensorrt_llm/thop/moe/cuteDsl/cuteDslMoeUtilsOp.cpp`. +""" + +from __future__ import annotations + +import math + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The DeepSeek-V4.1 vision aligner applies GELU to w1's output, [C, 5120]. +_WIDTH = 5120 + +#: Row counts run against the reference with one full tile, at N(0, 1) and N(0, 3^2): +#: decode-sized to multi-image prefill-sized. +_DENSE_ROWS = (1, 9, 100, 1024, 4096, 8192) + +#: The aligner's own row counts: the golden cases' C and the MMMU maximum. +_ALIGNER_ROWS = ( + (182, "golden-case row count"), + (546, "golden-case row count"), + (601, "golden-case row count"), + (986, "golden-case row count"), + (1902, "the largest C over the MMMU val documents"), + (2048, "the aligner's row-count ceiling"), +) + + +class MoeGelu(OpWrapper): + """Exact GELU `0.5 * x * (1 + erf(x / sqrt(2)))`, fp32 math, rounded to `x.dtype`. + + `x` is `[num_tiles * tile_tokens_dim, width]`. Row `r` of tile + `t = r // tile_tokens_dim` is computed only when + `t < num_non_exiting_tiles[0]` and `r < tile_idx_to_mn_limit[t]`; every + other output row is left unwritten. For a plain dense activation over all + `rows` rows use one tile: `tile_idx_to_mn_limit = [rows]` (int32), + `num_non_exiting_tiles = [1]` (int32), `tile_tokens_dim = rows`. + + Fusion boundary: only the activation is inside the call. The caller owns the + `w1` projection and its bias that produce `x`, and the `w2` projection that + consumes the output; no gating, bias or quantization here. Returns a newly + allocated contiguous tensor of `x`'s shape and dtype; `x` is not mutated. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"rows={rows} N(0,{sigma:g}^2): [rows, 5120] bf16, one full tile", + spec=dict(kind="randn", rows=rows, limit=rows, sigma=sigma), + ) + for rows in _DENSE_ROWS + for sigma in (1.0, 3.0) + ) + + tuple( + Cell( + why=f"aligner C={rows} N(0,3^2): {reason}", + spec=dict(kind="randn", rows=rows, limit=rows, sigma=3.0), + ) + for rows, reason in _ALIGNER_ROWS + ) + + ( + Cell( + why="every bf16 bit pattern: [13, 5120], +-inf/NaN/+-0 included", + spec=dict(kind="bits", rows=13, limit=13), + ), + Cell( + why="limit 37 < rows 64: rows at or past the limit unwritten", + spec=dict(kind="randn", rows=64, limit=37, sigma=3.0), + ), + ) + ) + + note = """ + Metadata consumed: the tile metadata, read on the device by the kernel. + `num_non_exiting_tiles[0] * tile_tokens_dim` bounds the rows walked and + `tile_idx_to_mn_limit[t]` is the exclusive end row of tile `t`. The vision + aligner builds the single-tile values host-side once per request, `[C]` and + `[1]` int32 on CUDA with `tile_tokens_dim = C`; no runtime or attention + metadata is read. + + Preconditions: `x` 2-D, contiguous, bf16 on CUDA with `width % 8 == 0`; + `x.shape[0] == tile_tokens_dim * tile_idx_to_mn_limit.size(0)`; both metadata + tensors int32 on CUDA, `num_non_exiting_tiles` one element. The kernel + ignores strides (see `is_valid`). Rows outside the metadata are not written + and the output is a fresh `torch::empty`: a stale limit one row short (left + over from another request's C) leaves the last row holding whatever the + allocator block held, with no error, so the limit tensor must be rebuilt or + rewritten for every C. Under CUDA graphs the metadata is read at replay from + the captured addresses: a graph is valid only for the C its metadata holds. + The op raises RuntimeError on a row count that is not + `tile_tokens_dim * num_tiles`, non-int32 limits, `width % 8 != 0`, an fp32 + input and a non-2-D input. + + Numerics: the erf form, not the tanh approximation (they differ at bf16 + only for x < -2.3). Measured on sm_103 over all 65536 bf16 bit patterns the + output is bit-identical to `F.gelu(x, approximate="none")` on the bf16 + tensor, which is what the aligner computes; against fp64 erf GELU the 115 + mismatching patterns all lie in x in [-8.3125, -4.25], where fp32 + `1 + erf(...)` cancels, at most 2.4e-7 absolute. `+inf -> inf`, + `-inf -> NaN`, `NaN -> NaN`, `+-0 -> +-0`, as `F.gelu`. + + One launch of one ahead-of-time compiled kernel at every row count 1..2048, + with the PDL launch attribute; no autotuning or tactic choice. Every row + count 1..2048 in the aligner's call form returns exactly the first rows of + the 2048-row result. + + Not certified: fp16 input, widths other than 5120, multi-tile layouts, + zero rows, rows 2049..4095, 4097..8191 and above 8192, CUDA-graph capture at + row counts other than 182, 546, 601, 986 and 1902, and PDL behaviour under + a producer kernel that triggers early completion. + """ + + def raw_call( + self, + x: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + tile_tokens_dim: int, + ) -> torch.Tensor: + return torch.ops.trtllm.moe_gelu( + x, tile_idx_to_mn_limit, num_non_exiting_tiles, tile_tokens_dim + ) + + def reference( + self, + x: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + tile_tokens_dim: int, + ) -> torch.Tensor: + """Exact GELU in fp64, rounded once to `x.dtype`, on the rows the metadata selects. + + Rows the op leaves unwritten (tile at or past `num_non_exiting_tiles[0]`, + or row at or past its tile's limit) are NaN here: `compare` then matches + them only where the output block held NaN before the call, which is how + the test shows those rows were left alone. + """ + xd = x.double() + ref = (0.5 * xd * (1.0 + torch.erf(xd / math.sqrt(2.0)))).to(x.dtype) + rows = torch.arange(x.shape[0], device=x.device) + tile = rows // tile_tokens_dim + limits = tile_idx_to_mn_limit.to(device=x.device, dtype=torch.int64) + written = (tile < int(num_non_exiting_tiles.reshape(-1)[0].item())) & ( + rows < limits[tile.clamp_max(limits.numel() - 1)] + ) + return ref.masked_fill(~written.unsqueeze(-1), float("nan")) + + def is_valid( + self, + x: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + tile_tokens_dim: int, + ) -> None: + # The kernel addresses row r at `data_ptr + r * width` and ignores the + # tensor's strides, so a non-contiguous view is read as if it were + # contiguous: on sm_103 a row-strided view (row stride width + 256) gave a + # bit-exact fraction of 0.01607971265912056 and a transposed view + # 0.0014892577892169356 against the reference of the view, with no error. + assert x.is_contiguous(), ( + f"x must be contiguous; got strides {x.stride()} for shape {tuple(x.shape)} " + f"-- the kernel ignores strides and would read the wrong elements" + ) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """`assert_close` at its default bf16 tolerance against the fp64 erf reference, NaN equal to NaN. + + Nothing is loosened: measured on sm_103 every case passes the default + bf16 band (bit-exact fraction 0.998 over all bit patterns, 0.938 on + N(0, 3^2)), while the tanh-approximated GELU fails it on 123 of the + 66560 bit-pattern elements and 51261 of 5242880 N(0, 1) elements. + `equal_nan` is what the bit-pattern cell needs (`-inf` and NaN inputs + give NaN) and what the unwritten-rows convention of `reference` uses. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + torch.testing.assert_close(out, ref, equal_nan=True) + + +moe_gelu = MoeGelu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/activation/silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/activation/silu_and_mul.py new file mode 100644 index 000000000000..38e631a79c30 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/activation/silu_and_mul.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SwiGLU with an optional clamp (`swiglu_limit`) via the in-tree Triton silu_and_mul kernel. + +The op is `torch.ops.trtllm.silu_and_mul`, registered in +`tensorrt_llm/_torch/custom_ops/torch_custom_ops.py` around `silu_and_mul_kernel` +in `tensorrt_llm/_torch/modules/swiglu.py`. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: DeepSeek-V4.1's shared-expert clamp. +_LIMIT = 10.0 + +#: d = 2304 is the shared expert's intermediate width with the expert replicated +#: per rank; d = 576 is the same expert TP-split 4 ways. Rows span decode (1, 8, +#: 32) and prefill chunks (4096, 8192). Neither d is a multiple of the kernel's +#: 1024-wide column block, so the masked tail block runs in every case. +_SHAPES = ( + (1, 2304, "decode"), + (8, 2304, "decode"), + (32, 2304, "decode"), + (4096, 2304, "prefill chunk"), + (8192, 2304, "prefill chunk"), + (1, 576, "the 4-way TP-split width"), + (4096, 576, "the 4-way TP-split width"), +) + +#: Row counts the core target produced in its runtime-value dumps (boot, GSM8K +#: and long-context runs), then primes and values either side of 128-row edges. +_RUNTIME_ROWS = (2, 3, 4, 5, 6, 7, 11, 14, 16, 24, 48, 323, 576, 1935, 2048, 3305, 4617, 6523, 8179) +_EDGE_ROWS = (13, 31, 127, 129, 1021, 4093, 8191) + + +def _cell(why: str, rows: int, d: int = 2304, **overrides) -> Cell: + spec = dict( + rows=rows, + d=d, + x_dtype=torch.bfloat16, + dtype=None, + swiglu_limit=_LIMIT, + swiglu_alpha=None, + swiglu_beta=None, + row_pad=0, + ) + spec.update(overrides) + return Cell(why=why, spec=spec) + + +class SiluAndMul(OpWrapper): + """`silu(min(g, L)) * clamp(u, -L, L)` over a packed `[gate; up]` last dim, as a new tensor. + + `x` is `[rows, 2 * d]`; `g = x[:, :d]`, `u = x[:, d:]`, loaded and computed + in fp32. The clamp applies only when `swiglu_limit` is a positive float; the + gate is clamped from above only, the up projection on both sides. + `swiglu_alpha` / `swiglu_beta` (default 1.0 / 0.0) generalise the body to + `g * sigmoid(alpha * g) * (u + beta)`. The output dtype is `dtype` if given, + else `x.dtype`; with `scale` the result is divided by it and saturated to + `dtype`'s finite range before the cast. + + Fusion boundary: the clamp, the activation, the gating multiply and the + optional per-tensor output scale with its saturating cast are inside the one + Triton launch. The caller owns the gate/up projection that produces `x` and + whatever consumes the output (for example an MXFP8 quantize before the down + projection); no block-scale quantization happens here. Returns a newly + allocated contiguous `[rows, d]` tensor; `x` is not mutated. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + _cell(f"bf16->bf16 limit 10 [{rows}, 2x{d}]: {what}", rows, d) + for rows, d, what in _SHAPES + ) + + tuple( + _cell( + f"fp32->bf16 limit 10 [{rows}, 2x{d}]: dtype=bf16, rounded at the store", + rows, + d, + x_dtype=torch.float32, + dtype=torch.bfloat16, + ) + for rows, d, _ in _SHAPES + ) + + tuple( + _cell( + f"fp32->fp32 limit 10 [{rows}, 2x2304]: fp32 output", + rows, + x_dtype=torch.float32, + ) + for rows in (8, 4096) + ) + + tuple( + _cell(f"bf16->bf16 limit 10 [{rows}, 2x2304]: a runtime-dump row count", rows) + for rows in _RUNTIME_ROWS + ) + + tuple( + _cell(f"bf16->bf16 limit 10 [{rows}, 2x2304]: a prime or 128-row edge", rows) + for rows in _EDGE_ROWS + ) + + ( + _cell("bf16->bf16 limit 10 [32768, 2x2304]: the decoder bounded-replay top", 32768), + _cell( + "alpha=1.0 beta=0.0 passed: bit-identical to None, [64, 2x2304]", + 64, + swiglu_alpha=1.0, + swiglu_beta=0.0, + ), + _cell("row stride 2d + 256, unit inner stride: [256, 2x2304]", 256, row_pad=256), + _cell("swiglu_limit=None: no clamp, [64, 2x2304]", 64, swiglu_limit=None), + _cell("swiglu_limit=0.0: no clamp, [64, 2x2304]", 64, swiglu_limit=0.0), + _cell("zero rows: [0, 2x2304] gives an empty [0, 2304] output", 0), + ) + ) + + note = """ + Metadata consumed: none. Stateless. + + With `swiglu_limit=10.0` and `swiglu_alpha` / `swiglu_beta` left `None` this + is DeepSeek-V4.1's expert activation. It is not GPT-OSS's `swigluoai` body + `g * sigmoid(1.702 g) * (u + 1)`, reachable only with `swiglu_alpha=1.702, + swiglu_beta=1.0`. The first half of the last dim is the gate: an FC1 weight + stacked `[up; gate]` computes the wrong function without any error. + + Preconditions: `x` 2-D on CUDA with an even last dim; a unit inner stride + (see `is_valid`), any row stride >= `2 * d` is read correctly. Only a + positive `swiglu_limit` clamps: `None`, `0.0` and negative values all mean no + clamp. For bf16 output from fp32 input pass `dtype=torch.bfloat16`. The op + raises ValueError on a 3-D `x` and a bare AssertionError on an odd last dim. + + `swiglu_limit`, `swiglu_alpha`, `swiglu_beta` and the dtype pair are Triton + constexprs: each distinct value compiles its own kernel on first use. The + row count is not part of the compiled kernel's key (rows are only the size + of grid axis 0, program `(i, j)` reads and writes row `i`): every prefix + `x[:r]`, r = 1..32768, at d = 2304, bf16, limit 10 returns the first `r` rows + of the 32768-row result bitwise, with no new compiled variant. A CUDA-graph + capture of a call whose constants have never run eagerly was not tested; + after one eager call the capture works at row counts never run eagerly. + + Row counts the target hands this op reach 32768 (16384 per rank at the + largest `max_num_tokens`; the decoder bounded-replay pass adds finishing + requests' 128-row tails, up to 128 x 256). The target calls it only at + d = 2304; d = 576 is the 4-way TP split it does not use. CUDA graphs are + certified at d = 2304, bf16, limit 10 at 114 row counts: G, 5G and 6G + (decode, DSpark drafter block, DSpark target verify) over the captured batch + sizes 1..32, 40..128 step 8, 192 and 256, one shared pool, largest first. + + Hazards: with the clamp on, a NaN gate gives `silu(10) * u` and a NaN up is + treated as `+10` where `torch.clamp` propagates NaN (probe, not certified). + Inputs beyond the limit are what separate the clamp variants: a symmetric + gate clamp and the unclamped body are indistinguishable from this op on + inputs within +-9.5. + + Not certified: the `scale` path, fp16 input, `swiglu_alpha` / `swiglu_beta` + other than 1.0 / 0.0, `swiglu_limit` other than 10.0, d other than 2304 and + 576, more than 32768 rows, d = 576 at row counts other than 1 and 4096, and + CUDA-graph capture at d = 576 or at a row count outside the 114. + """ + + def raw_call( + self, + x: torch.Tensor, + scale: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, + swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None, + ) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, + scale=scale, + dtype=dtype, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + ) + + def reference( + self, + x: torch.Tensor, + scale: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, + swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None, + ) -> torch.Tensor: + """fp32: gate = min(g, L), up = clamp(u, -L, L), silu(gate) * up, rounded once to the output dtype. + + With alpha / beta at None or 1.0 / 0.0 the body is `F.silu(gate) * up`; + otherwise `gate * sigmoid(alpha * gate) * (up + beta)`. `scale`, when + given, divides and saturates to the output dtype's finite range before + the cast (not certified). + """ + out_dtype = dtype if dtype is not None else x.dtype + d = x.shape[-1] // 2 + gate = x[:, :d].float() + up = x[:, d:].float() + if swiglu_limit is not None and swiglu_limit > 0: + gate = torch.clamp(gate, max=swiglu_limit) + up = torch.clamp(up, min=-swiglu_limit, max=swiglu_limit) + if swiglu_alpha in (None, 1.0) and swiglu_beta in (None, 0.0): + r = F.silu(gate) * up + else: + alpha = 1.0 if swiglu_alpha is None else swiglu_alpha + beta = 0.0 if swiglu_beta is None else swiglu_beta + r = gate * torch.sigmoid(alpha * gate) * (up + beta) + if scale is not None: + top = torch.finfo(out_dtype).max + r = torch.clamp(r / scale.float().reshape(-1)[0], -top, top) + return r.to(out_dtype) + + def is_valid( + self, + x: torch.Tensor, + scale: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, + swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None, + ) -> None: + # The kernel walks each row with a unit element stride and takes only + # `x.stride(0)` from the tensor, so an inner stride other than 1 is read + # as if it were 1: a transposed view returned a bit-exact fraction of + # 0.00047 (max abs diff 200) and no error on this machine. Row-strided + # views with a unit inner stride are handled correctly and are not + # rejected. + assert x.stride(-1) == 1, ( + f"x must have a unit stride along its last dim; got strides {x.stride()} " + f"-- the kernel ignores the inner stride and would read the wrong elements" + ) + + +silu_and_mul = SiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/candidate_block_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/candidate_block_select.py new file mode 100644 index 000000000000..70a38574b882 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/candidate_block_select.py @@ -0,0 +1,527 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Candidate block pre-filter (level one of DeepSeek-V4.1's two-level index top-k), one Triton launch. + +Per row of index scores: max-pool the row's visible positions in blocks of ``block_size``, pin +the block that holds the newest visible position, keep the ``topk_blocks`` best blocks, and +return a bool mask over the score columns that is True on every visible position of a kept +block. + +The kernel is defined here (``_candidate_block_select_kernel``). No engine kernel implements +this step, so this entry is the implementation, not a wrapper around one. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# Launch configuration: (candidate blocks scored per loop iteration, warps). One program owns +# one row, so few rows (decode) take wide iterations and many rows (prefill) narrow ones. The +# result does not depend on the choice. +_FEW_ROWS = 256 +_FEW_ROWS_CONFIG = (1024, 8) +_MANY_ROWS_CONFIG = (256, 4) +_FILL_COLS = 1024 + + +@triton.jit +def _block_keys( + row_ptr, + col_stride, + start, + length, + num_cols, + bidx, + bvalid, + BLOCK_SIZE: tl.constexpr, + UNIT_COL_STRIDE: tl.constexpr, +): + """(amax, order-preserving uint32 key) of the candidate blocks ``bidx`` of one row. + + A block's positions past ``length`` count as -inf. The key orders like the float, + with -0.0 folded onto +0.0 so the two tie. + """ + pos = bidx[:, None] * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)[None, :] + col = start + pos + live = bvalid[:, None] & (pos < length) & (col >= 0) & (col < num_cols) + if UNIT_COL_STRIDE: + ptrs = row_ptr + col + else: + ptrs = row_ptr + col.to(tl.int64) * col_stride + x = tl.load(ptrs, mask=live, other=float("-inf")).to(tl.float32) + amax = tl.max(x, axis=1) + bits = amax.to(tl.int32, bitcast=True) + bits = tl.where(bits == -2147483648, 0, bits) + u = bits.to(tl.uint32, bitcast=True) + key = tl.where(bits < 0, ~u, u | 0x80000000) + return amax, key + + +@triton.jit +def _candidate_block_select_kernel( + logits_ptr, + logits_row_stride, + logits_col_stride, + row_starts_ptr, + row_ends_ptr, + mask_ptr, + mask_row_stride, + num_cols, + topk_blocks, + NEXT_N: tl.constexpr, + HAS_ROW_STARTS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + UNIT_COL_STRIDE: tl.constexpr, + CHUNK_BLOCKS: tl.constexpr, + FILL_COLS: tl.constexpr, +): + row = tl.program_id(0) + end = tl.load(row_ends_ptr + row // NEXT_N).to(tl.int32) - (NEXT_N - 1 - row % NEXT_N) + if HAS_ROW_STARTS: + start = tl.load(row_starts_ptr + row).to(tl.int32) + else: + start = row * 0 + length = tl.maximum(end - start, 0) + end = start + length + num_blocks = (length + BLOCK_SIZE - 1) // BLOCK_SIZE + newest = num_blocks - 1 # the pinned block; -1 (none) for an empty row + binding = num_blocks > topk_blocks + row_ptr = logits_ptr + row.to(tl.int64) * logits_row_stride + out_ptr = mask_ptr + row.to(tl.int64) * mask_row_stride + + # Level-one selection when the row has more blocks than topk_blocks: the pinned block + # takes one slot, the other topk_blocks - 1 go to the best of blocks [0, newest). + # thr = the (topk_blocks - 1)-th largest key among them (radix select, 4 bits a pass); + # above = how many keys are strictly greater; n_tie = how many keys equal to thr to keep. + need = topk_blocks - 1 + thr = tl.zeros((), dtype=tl.uint32) + n_tie = row * 0 + if binding: + digits = tl.arange(0, 16) + above = row * 0 + prefix = tl.zeros((), dtype=tl.uint32) + for p in tl.static_range(8): + hist = tl.zeros((16,), dtype=tl.int32) + for b0 in range(0, newest, CHUNK_BLOCKS): + bidx = b0 + tl.arange(0, CHUNK_BLOCKS) + bvalid = bidx < newest + _, key = _block_keys( + row_ptr, + logits_col_stride, + start, + length, + num_cols, + bidx, + bvalid, + BLOCK_SIZE, + UNIT_COL_STRIDE, + ) + if p == 0: + sel = bvalid + else: # the digits above this one must equal the prefix found so far + sel = bvalid & ((key >> (32 - 4 * p)) == (prefix >> (32 - 4 * p))) + d = ((key >> (28 - 4 * p)) & 15).to(tl.int32) + onehot = (d[:, None] == digits[None, :]) & sel[:, None] + hist += tl.sum(onehot.to(tl.int32), axis=0) + ge = tl.sum(tl.where(digits[None, :] >= digits[:, None], hist[None, :], 0), axis=1) + ok = (above + ge) >= need + dsel = tl.max(tl.where(ok, digits, -1), axis=0) + above += tl.sum(tl.where(digits > dsel, hist, 0), axis=0) + prefix = prefix | (dsel.to(tl.uint32) << (28 - 4 * p)) + thr = prefix + n_tie = need - above + + # Expand the kept blocks over the row's visible columns. Ties at thr go to the lowest + # block indices; a block whose visible scores are all -inf is dropped unless pinned. + tie_seen = row * 0 + for b0 in range(0, num_blocks, CHUNK_BLOCKS): + bidx = b0 + tl.arange(0, CHUNK_BLOCKS) + bvalid = bidx < num_blocks + amax, key = _block_keys( + row_ptr, + logits_col_stride, + start, + length, + num_cols, + bidx, + bvalid, + BLOCK_SIZE, + UNIT_COL_STRIDE, + ) + older = bvalid & (bidx < newest) + is_tie = (older & (key == thr) & binding).to(tl.int32) + tie_rank = tie_seen + tl.cumsum(is_tie, axis=0) - is_tie + chosen = (key > thr) | ((is_tie != 0) & (tie_rank < n_tie)) | (num_blocks <= topk_blocks) + keep = (bidx == newest) | (older & (amax > float("-inf")) & chosen) + tie_seen += tl.sum(is_tie, axis=0) + pos = bidx[:, None] * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)[None, :] + col = start + pos + live = bvalid[:, None] & (pos < length) & (col >= 0) & (col < num_cols) + val = tl.broadcast_to(keep[:, None], (CHUNK_BLOCKS, BLOCK_SIZE)).to(tl.uint8) + tl.store(out_ptr + col, val, mask=live) + + # Every column outside [start, end) is False. + zeros = tl.zeros((FILL_COLS,), dtype=tl.uint8) + for c0 in range(0, start, FILL_COLS): + col = c0 + tl.arange(0, FILL_COLS) + tl.store(out_ptr + col, zeros, mask=(col < start) & (col < num_cols)) + for c0 in range(end, num_cols, FILL_COLS): + col = c0 + tl.arange(0, FILL_COLS) + tl.store(out_ptr + col, zeros, mask=(col >= end) & (col >= 0) & (col < num_cols)) + + +_TOPK_BLOCKS = 2048 # the checkpoint's candidate_topk_blocks +_BLOCK = 8 # the checkpoint's candidate_block_size +_ENV_KV = 16384 +_ENVELOPE_LM = 1048576 + 14 # the longest request with DSpark: 1,048,590 keys + +#: Decode rows' visible counts: tiny, both sides of the 16384 binding edge, and far past it. +_DECODE_LENS = (1, 5, 8, 9, 16383, 16384, 16385, 16391, 16392, 16393, 16400, 20000, 32768) +_DECODE_LENS += (65536, 131072, 262144) + +#: Packed prefill calls made untiled (max_num_tokens 8192, max_batch_size 32, max_seq_len 16384): +#: (what, query rows per request, keys per request). Request b's last q_b positions are the rows. +_UNTILED_STEPS = ( + ("maximal: 32 requests x 16384 keys, 256 rows each", (256,) * 32, (_ENV_KV,) * 32), + ("widest step, one request's rows only", (8192 - 31,) + (1,) * 31, (_ENV_KV,) * 32), + ("W 16384 (the target skips it: exact)", (4096, 4096), (8192, 8192)), + ("W 16384, one request, 8192 rows", (8192,), (16384,)), + ("W 16385, 2 rows", (1, 1), (8193, 8192)), + ("W 16385, 256 rows", (128, 128), (8193, 8192)), + ("W 16385, 257 rows", (129, 128), (8193, 8192)), + ("W 16385, 8192 rows", (4096, 4096), (8193, 8192)), + ("W 16385, 32 requests", (256,) * 32, (512,) * 31 + (513,)), + ("dumped (7674, 24058)", (7674,), (24058,)), + ("dumped (6523, 24058)", (6523,), (24058,)), + ("dumped (8192, 17535)", (8192,), (17535,)), +) + +#: Packed prefill calls in the target's row tiles (max_num_tokens 16384, max_batch_size 256, +#: max_seq_len up to 1,048,576): one call per tile of at most 2^28 fp32 logits elements. A request +#: with 0 rows only holds keys in the gather (the decoder-replay tail's neighbours). +_TILED_STEPS = ( + ("L 32768 DSpark: one request at Lm 32782, 16384 rows", (16384,), (32782,)), + ("L 32768 no DSpark: one request at 32768, 16384 rows", (16384,), (32768,)), + ("L 65536 DSpark: one request at Lm 65550, 16384 rows", (16384,), (65550,)), + ("L 131072 DSpark: one request at Lm 131086, 16384 rows", (16384,), (131086,)), + ("L 262144 DSpark: one request at Lm 262158, 16384 rows", (16384,), (262158,)), + ("L 1048576 DSpark: one request at Lm 1048590, 1024 rows", (1024,), (_ENVELOPE_LM,)), + ("L 1048576 no DSpark: one request at 1048576, 1024 rows", (1024,), (1048576,)), + ("tail (128, 1048219): keys of another request first", (0, 128), (448219, 600000)), + ("tail (128, 523951)", (0, 128), (23951, 500000)), + ("tail (128, 261860)", (128, 0), (200000, 61860)), + ("tail (128, 130773)", (0, 0, 128), (30000, 773, 100000)), + ( + "tail 3 requests x 128 rows at S_1 1048219 (tiles 256 + 128)", + (128, 128, 0, 128), + (300000, 200000, 48219, 500000), + ), + ("256 requests x 1 row, W 1048576", (1,) * 256, (4096,) * 256), + ( + "two requests at Lm, 256 rows, S_1 2097180 (tiles of 64)", + (128, 128), + (_ENVELOPE_LM, _ENVELOPE_LM), + ), + ("four requests at Lm, 128 rows, S_1 4194360 (tiles of 63)", (32,) * 4, (_ENVELOPE_LM,) * 4), + ( + "sixteen requests at Lm, 64 rows, S_1 16777440 (tiles of 15)", + (4,) * 16, + (_ENVELOPE_LM,) * 16, + ), + ("W 16384 at 16384 rows (the target skips it)", (16384,), (16384,)), + ("W 16385 at 16384 rows", (8192, 8192), (8193, 8192)), + ("W 16385, 256 requests, 16384 rows", (64,) * 256, (65,) + (64,) * 255), + ("dumped (16292, 16625)", (10441, 5851), (10441, 6184)), + ("dumped (16378, 16695)", (10451, 5927), (10451, 6244)), +) + + +class CandidateBlockSelect(OpWrapper): + """Return a fresh bool ``[N, W]`` candidate mask for index scores ``logits`` ``[N, W]``. + + Row ``i`` sees columns ``[s, e)``:: + + s = row_starts[i] (0 when row_starts is None) + e = row_ends[i // next_n] - (next_n - 1 - i % next_n) + L = max(e - s, 0) visible columns s .. s + L - 1 + + Its visible positions are cut into blocks of ``block_size`` starting at ``s``; each block scores + the fp32 max of its visible scores; the block holding the newest position (``e - 1``) is always + kept; the other ``topk_blocks - 1`` places go to the highest-scoring remaining blocks (all of + them when there are at most ``topk_blocks``), ties to the lower block index; a block whose + visible scores are all -inf is dropped unless pinned. ``mask[i, c]`` is True exactly for the + visible columns of kept blocks; every other element is written False. + + Fusion boundary: block max-pooling, the pin, the block top-k, the -inf drop, the expansion to + columns and clearing the columns outside each row's range are inside the call. Producing the + scores, and applying the keep mask to a consumer layer's scores (``torch.where(mask, s, -inf)`` + before its own top-512), are the caller's. Returns the mask (fresh, contiguous). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"prefill one request seq={seq} starts={'explicit' if explicit else 'None'}: " + "every query row of a fresh prompt, visible 1..seq", + spec=dict(kind="prefill_one", seq=seq, explicit_starts=explicit), + ) + for seq in (4096, 20480) + for explicit in (True, False) + ) + + ( + Cell( + why="prefill chunk 16000..20095 row stride W+13: a later chunk of a long context, " + "past the binding edge", + spec=dict(kind="prefill_chunk", first=16000, q=4096, pad=13), + ), + Cell( + why="prefill 3 requests sharing the columns: ragged row_starts, 5 spare columns", + spec=dict( + kind="prefill_ragged", + keys=(18000, 700, 21000), + queries=((16300, 1700), (0, 700), (19000, 2000)), + spare=5, + ), + ), + Cell( + why="decode visible 1..262144, NaN/+inf unread: the binding edge and far past it", + spec=dict(kind="decode", lens=_DECODE_LENS), + ), + Cell( + why="decode next_n=6 DSpark verify rows, 3 empty: row j sees one more than j - 1", + spec=dict(kind="decode_next_n", visible=(16386, 16390, 20000, 50000, 3), next_n=6), + ), + Cell( + why="decode bf16 scores: frequent ties, cut by the lowest block index", + spec=dict(kind="bf16", lens=(16385, 20000, 40000)), + ), + Cell( + why="decode heavy ties {-1,-0,0,1,2}: the cut inside a +-0 tie", + spec=dict(kind="ties", lens=(16392, 20000, 40001, 30000)), + ), + Cell( + why="decode -inf blocks: dropped though < 2048 remain, the pin held", + spec=dict(kind="minus_inf", lens=(20000, 20000, 30000)), + ), + Cell( + why="empty rows L=0 around a full row: all False", + spec=dict(kind="empty_rows", lens=(0, 64, 0), width=64), + ), + Cell( + why="no rows N=0: an empty [0, W] mask, no launch", + spec=dict(kind="no_rows", width=64), + ), + Cell( + why="column-strided logits (stride 8): visible 20000 and 16000", + spec=dict(kind="col_strided", lens=(20000,) * 4 + (16000,) * 4), + ), + ) + + tuple( + Cell( + why=f"untiled {what}", + spec=dict(kind="packed", q_lens=q, kv_lens=kv, tiled=False), + ) + for what, q, kv in _UNTILED_STEPS + ) + + tuple( + Cell( + why=f"tiled {what}", + spec=dict(kind="packed", q_lens=q, kv_lens=kv, tiled=True), + ) + for what, q, kv in _TILED_STEPS + ) + ) + + note = """ + Every cell runs at the checkpoint's `topk_blocks = 2048`, `block_size = 8`: the pre-filter + binds only on rows with more than 16384 visible positions, and at or below that the mask is + the row's range minus its non-pinned all--inf blocks. The packed cells fill every element + outside a row's range with NaN and set 1 in 32 whole blocks to -inf, so each row's addressing + shows in its mask; each fresh mask lands on a block pre-filled with 1s. Call sites: prefill + with `row_starts = ks`, `row_ends = ke` on DeepGEMM's `[rows, S_1]` logits (row stride + `align(S_1 + 256, 256)`), whole or per row tile; decode with `row_starts = None` on the paged + `[G * next_n, max_ctx]` logits (row stride `ceil256(max_ctx)`) and a row-offset view of the + step's visible-length buffer, `next_n` 1 (each verify row carries its own length) or 6. + + Relation to the checkpoint's `select_candidate_blocks`: equal on every visible position on + distinct scores; on tied block scores the kept block-score multiset is equal (its + `torch.topk` keeps an unspecified member of the tie). It also marks the invisible tail of the + pinned block in prefill, which this op leaves False; consumers hold -inf there either way. + + Metadata consumed: none; the row ranges are the caller's. + + Preconditions: `0 <= s <= e <= W` per row (reads and writes are masked to `[0, W)`, so a + violation touches nothing outside the two tensors, but is not certified); `block_size` a power + of two (12 fails to compile); CUDA tensors (a CPU tensor raises). Columns outside a row's range + are never read (NaN and +inf there are certified). A row's mask depends only on its own scores + and range: the same row alone, inside a larger call, or with its keys moved within the gather + gets the same mask. + + The kernel JIT-compiles on the first call of each specialization (input dtype, `next_n`, + `row_starts` given, `block_size`, unit column stride, launch configuration, Triton's + divisibility and pointer-alignment classes); capture in a CUDA graph after an eager call of + the same specialization is certified, a first-ever call inside capture is not. The launch + configuration follows the row count (at most 256 rows: 1024 blocks an iteration and 8 warps, + else 256 and 4); the mask does not depend on it. + + Measured on sm_103: one program per row runs an 8-pass radix select over the block scores, so + decode rows above 16384 visible are slower than the checkpoint's torch composition (0.89x at + 20480, 0.32x at 131072); prefill is faster. + + Not certified: `topk_blocks` other than 2048 or `block_size` other than 8; `next_n` other than + 1 and 6, or `next_n > 1` with `row_starts`; decode calls of more than 1536 rows; CUDA-graph + capture at a batch outside 1..32, 40..128 step 8, 192, 256, or at a width other than 16389, + 16390 and L, L + 14, L + 19 for L in 16384, 32768, 65536, 131072, 262144, 1048576; + `W > 33,329,476`; more than 20480 rows in one call; more than 2^32 logits elements + or mask bytes in one call; NaN or +inf inside a visible range; float16 logits and int64 lengths + (observed working only). + """ + + def raw_call( + self, + logits: torch.Tensor, + row_starts: Optional[torch.Tensor], + row_ends: torch.Tensor, + topk_blocks: int, + block_size: int, + next_n: int, + ) -> torch.Tensor: + num_rows, num_cols = logits.shape + mask = torch.empty((num_rows, num_cols), dtype=torch.bool, device=logits.device) + if num_rows == 0: + return mask + chunk_blocks, num_warps = _FEW_ROWS_CONFIG if num_rows <= _FEW_ROWS else _MANY_ROWS_CONFIG + _candidate_block_select_kernel[(num_rows,)]( + logits, + logits.stride(0), + logits.stride(1), + row_ends if row_starts is None else row_starts, + row_ends, + mask.view(torch.uint8), + mask.stride(0), + num_cols, + topk_blocks, + NEXT_N=next_n, + HAS_ROW_STARTS=row_starts is not None, + BLOCK_SIZE=block_size, + UNIT_COL_STRIDE=logits.stride(1) == 1, + CHUNK_BLOCKS=chunk_blocks, + FILL_COLS=_FILL_COLS, + num_warps=num_warps, + ) + return mask + + def is_valid( + self, + logits: torch.Tensor, + row_starts: Optional[torch.Tensor], + row_ends: torch.Tensor, + topk_blocks: int, + block_size: int, + next_n: int, + ) -> None: + num_rows, num_cols = logits.shape + # The kernel reads row_ends[i // next_n] and row_starts[i] at unit stride with no bound: a + # shorter or strided tensor is read wrongly with no error. Measured with the kernel called + # directly on 8 rows x 20000: a 4-element row_ends got rows 4-7 wrong; an 8-element + # row_ends with next_n = 6 gave rows 0-4, 6 and 7 a shorter range; a 4-element row_starts + # got rows 4 and 6 wrong; a stride-2 view of either got rows 1, 3, 5 and 7 wrong. + assert row_ends.is_contiguous() and row_ends.numel() * next_n == num_rows, ( + "row_ends must be contiguous with row_ends.numel() * next_n == logits.size(0)" + ) + assert row_starts is None or ( + row_starts.is_contiguous() and row_starts.numel() == num_rows + ), "row_starts must be None or contiguous with row_starts.numel() == logits.size(0)" + + def reference( + self, + logits: torch.Tensor, + row_starts: Optional[torch.Tensor], + row_ends: torch.Tensor, + topk_blocks: int, + block_size: int, + next_n: int, + ) -> torch.Tensor: + """The checkpoint's `select_candidate_blocks` with this op's tie rule, from torch ops. + + Per row: gather the visible scores `[s, s + L)` (everything else -inf, so NaN outside the + range is never read), -inf-pad to whole blocks, take the block amax in fp32 (-0.0 folded + onto +0.0), pin the newest block to +inf, keep the first `topk_blocks` of a stable + descending sort (ties to the lower block index) whose value is above -inf, expand to + positions and scatter back to the row's columns. Rows are processed in chunks of at most + 2^27 gathered elements, so the reference runs at the widths the cells reach. On distinct + scores this is the checkpoint's `torch.topk` selection exactly. + """ + num_rows, num_cols = logits.shape + dev = logits.device + mask = torch.zeros((num_rows, num_cols), dtype=torch.bool, device=dev) + if num_rows == 0: + return mask + i = torch.arange(num_rows, device=dev) + ends = row_ends.long().repeat_interleave(next_n) - (next_n - 1 - i % next_n) + starts = torch.zeros_like(ends) if row_starts is None else row_starts.long() + lens = (ends - starts).clamp(min=0) + lens_host = lens.tolist() + r0 = 0 + while r0 < num_rows: + r1, widest = r0, 1 + while r1 < num_rows: + w = max(widest, lens_host[r1]) + padded = -(-w // block_size) * block_size + if r1 > r0 and ( + (r1 + 1 - r0) * padded > 1 << 27 or (r1 + 1 - r0) * (num_cols + 1) > 1 << 30 + ): + break + widest, r1 = w, r1 + 1 + width = -(-widest // block_size) * block_size + pos = torch.arange(width, device=dev) + cols = starts[r0:r1, None] + pos + live = (pos < lens[r0:r1, None]) & (cols >= 0) & (cols < num_cols) + vals = logits[r0:r1].gather(1, cols.clamp(0, num_cols - 1)).float() + vals = vals.masked_fill(~live, float("-inf")) + blk = vals.view(r1 - r0, -1, block_size).amax(dim=-1) + blk = torch.where(blk == 0, torch.zeros_like(blk), blk) + newest = (lens[r0:r1, None] + block_size - 1) // block_size - 1 + blk = blk.masked_fill(torch.arange(blk.size(1), device=dev) == newest, float("inf")) + k = min(topk_blocks, blk.size(1)) + order = torch.sort(blk, dim=1, descending=True, stable=True) + keep_blk = torch.zeros_like(blk, dtype=torch.bool).scatter_( + 1, order.indices[:, :k], order.values[:, :k] > float("-inf") + ) + keep = keep_blk.repeat_interleave(block_size, dim=1) & live + part = torch.zeros((r1 - r0, num_cols + 1), dtype=torch.bool, device=dev) + part.scatter_(1, torch.where(live, cols, num_cols), keep) + mask[r0:r1] = part[:, :num_cols] + r0 = r1 + return mask + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise on the whole `[N, W]` mask, every byte 0 or 1. + + The selection is discrete -- block maxima of the same fp32 values, one order, one tie rule + -- so there is nothing to round differently: any differing element is a wrong block, range + or pin. Controls measured on sm_103 (decode rows above 16384 visible / every row of a + 20480-token prefill): no pin of the newest block changes 8 / 1443 rows, a block size of 16 + 13 / 4096, a top-16384 over positions instead of blocks 13 / 4096; each is identical to + this op on every row with at most 16384 visible positions. + """ + assert out.dtype == ref.dtype == torch.bool, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + if out.numel(): + assert int(out.view(torch.uint8).max()) <= 1, "a mask byte is neither 0 nor 1" + if not torch.equal(out, ref): + bad = (out != ref).any(dim=1).nonzero().flatten() + raise AssertionError( + f"{bad.numel()} of {out.size(0)} rows differ from the reference, first {int(bad[0])}" + ) + + +candidate_block_select = CandidateBlockSelect() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/compressor_paged_kv_compress.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/compressor_paged_kv_compress.py new file mode 100644 index 000000000000..378e2bd9bdd5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/compressor_paged_kv_compress.py @@ -0,0 +1,355 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode-side KV compressor: every generation request appends `next_n` tokens +to its paged compressor state, and each window of `compress_ratio` tokens that +the step completes is pooled (per-channel softmax over the window) into one +output row. + +The op is `torch.ops.trtllm.compressor_paged_kv_compress`, registered in +`cpp/tensorrt_llm/thop/compressorOp.cpp`. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_ELEMENT_TYPES = (torch.float32, torch.bfloat16) + +#: Histories (tokens already stored) that put a step's first new token at both +#: parities and on both sides of a 128- and a 256-token page boundary. +_CACHED = (1, 2, 127, 128, 255, 256, 1001, 8191) + + +def _spec(next_n: int, steps: int, **overrides) -> dict: + spec = dict( + cached=_CACHED, + steps=steps, + next_n=next_n, + page_size=128, + out_dtype=torch.float32, + shared_pool=True, + poison_history=False, + lead_rows=0, + lead_out=0, + ape="zero", + ) + spec.update(overrides) + return spec + + +class CompressorPagedKvCompress(OpWrapper): + """Store each generation request's `next_n` new tokens in the paged + compressor state, then pool every window of `R = compress_ratio` positions + the step completes into one output row. + + For request `b < batch_size`, with `kv_len = kv_lens[b]` its length after + the step and `sp = kv_len - next_n`, the new positions `[sp, kv_len)` come + from rows `cu_seq_lens[b] ..` of `kv_score` (`[kv || score]`, `S` wide each):: + + paged_kv [block_table_kv [b, p // page_size], p % page_size] = kv[p] + paged_score[block_table_score[b, p // page_size], p % page_size] = score[p] + ape[p % R] + + # every window k in [sp // R, kv_len // R), both positions read back from the pools + output[cu_kv_comp[b] + k - sp // R, d] = sum_p paged_kv[p, d] * softmax_p(paged_score[p, d]) + + so a window that started in an earlier step (odd `sp` at `R = 2`) reads the + token that step parked at `sp - 1`. bf16 `output` gets the round-to-nearest + -even of the fp32 pooled value. + + Fusion boundary: the state writes, the per-channel window softmax and the + weighted sum are inside the call. The caller owns the `wkv || wgate` + projection that makes `kv_score`, the RMSNorm, RoPE, quantization and the + compressed-KV cache write, and prefill pooling (`compressor_prefill_reduction`, + which leaves the state this op reads). + + Out-variant: writes `output`, `paged_kv` and `paged_score` in place and + returns `(output, paged_kv, paged_score)` in that order (the same tensor + twice when the two pools are one tensor, as in DeepSeek-V4.1). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"decode next_n={n} page={page}: 5 steps over 8 histories", + spec=_spec(n, 5, page_size=page), + ) + for n in range(1, 9) + for page in (128, 256) + ) + + tuple( + Cell( + why=f"poisoned history next_n={n}: only start-1 is read back", + spec=_spec(n, 4, poison_history=True), + ) + for n in (1, 6) + ) + + ( + Cell( + why="separate kv and score pools: two tensors, next_n 6", + spec=_spec(6, 4, shared_pool=False), + ), + Cell( + why="context lead 1537 rows / 700 outputs: mixed-step slices", + spec=_spec(6, 3, lead_rows=1537, lead_out=700), + ), + ) + + tuple( + Cell( + why=f"bf16 output next_n={n}: RNE of the fp32 pooled value", + spec=_spec(n, 6, out_dtype=torch.bfloat16), + ) + for n in (1, 6) + ) + + tuple( + Cell( + why=f"nonzero ape next_n={n}: baked into the stored score", + spec=_spec(n, 4, ape="random"), + ) + for n in (1, 6) + ) + ) + + note = """ + Cells. DeepSeek-V4.1's ratio-2 compressor of the KV-source layers 2/8/14: + `compress_ratio` 2, `head_dim` 512, fp32 `kv_score`, fp32 state with kv and + score on distinct pages of one tensor, zero `ape`, fp32 or bf16 `output`; + plain decode (`next_n` 1), DSpark verify (6), and 1..8 in all, at page + sizes 128 and 256. Every cell runs several steps, so later steps read the + token the op itself parked. Kv ~ N(0, 2^2), score ~ N(0, 4^2), plus one + exact-tie channel and one channel whose pair gap of 240 underflows `exp`. + + Metadata consumed: none. The target passes, per owner layer, the pools + `get_buffers(owner, COMPRESSOR_KV / COMPRESSOR_SCORE)` (one fp32 + `[pages, 128, 512]` base for both roles), the tables + `sliding_block_tables[layer_offsets[owner], COMPRESSOR_KV / SCORE][nc:nc + B]`, + `kv_lens_cuda_runtime[nc:nc + B]`, `cu_seq_lens_cuda[nc:nc + B + 1]`, + `cu_new_comp_kv_cuda[2][nc:nc + B + 1]`, `num_gen_tokens_per_seq` and + `tokens_per_block`, with the step's whole `kv_score` and the whole + `num_total_compressed_tokens[2]`-row output (context rows ahead are neither + read nor written). + + Preconditions: every request feeds exactly `next_n` tokens (the op never + reads `cu_seq_lens[b + 1]`); `cu_kv_comp[b + 1] - cu_kv_comp[b]` is + `kv_len // R - sp // R`; for an odd `sp` the state at `sp - 1` holds that + token's kv and score + ape (left by the prefill op or an earlier call); the + tables map every page covering `[sp - sp % R, kv_len)`, kv and score to + different pages, and have the same width (`block_table_kv.size(1)` is the + row stride of both); `page_size` even; `ape` fp32 `[R, S]`. No table entry + is validated. + + Hazards (is_valid): an fp16 `kv_score` or `output` is read or written as + bf16 bits, and a strided view is walked with the dense row stride, all + without an error. The pools are not guarded. + + Rejected by the op: `compress_ratio` other than 2, 4 or 128 (V4.1's + ratio-1 layer, which pools nothing, cannot route here), `next_n` 0, a + `head_dim` with no instantiation, bf16 `ape`, int64 `kv_lens`, pools of + different element types. The op's optional `new_tokens_per_seq` (ragged + verify) is not passed by this entry. + + Precision: built with `--use_fast_math`, an online fp32 softmax that + divides once, so not bitwise torch's softmax-then-sum. DSpark: a verify + step emits rows for every pair its tokens complete, drafts included; the + next step starting at `sp + 1 + accepted` rewrites rejected positions and, + when it starts odd, re-emits the pair `(sp' - 1, sp')`, so the caller's + compressed-cache write must overwrite, not append. + + No allocation, no host synchronization: graph-capturable. The test also + certifies, on the engine's own pools and managers: every batch 1..256 at + table widths 132 .. 8196, every context lead up to a 16384-token step, the + capture batch set at every width, state pools up to 247806 pages (element + offsets past 2^33), positions up to 1,049,087, and the DSpark tail at + max_seq_len 16384 (positions up to 16397 on 132-column tables). + + Not certified: bf16 `kv_score` or state, `head_dim` 128, ratios 4 and 128, + other page sizes; `next_n` above 8 (under capture, other than 1 and 6); more + than 256 requests; tables wider than 8196 columns; pools above 247806 + pages; unmapped or negative table entries; block reuse and pages released + behind the window on the real manager; performance. + """ + + def raw_call( + self, + kv_score: torch.Tensor, + ape: torch.Tensor, + paged_kv: torch.Tensor, + paged_score: torch.Tensor, + block_table_kv: torch.Tensor, + block_table_score: torch.Tensor, + output: torch.Tensor, + kv_lens: torch.Tensor, + cu_seq_lens: torch.Tensor, + cu_kv_comp: torch.Tensor, + batch_size: int, + page_size: int, + head_dim: int, + compress_ratio: int, + next_n: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + torch.ops.trtllm.compressor_paged_kv_compress( + kv_score, + ape, + paged_kv, + paged_score, + block_table_kv, + block_table_score, + output, + kv_lens, + cu_seq_lens, + cu_kv_comp, + batch_size, + page_size, + head_dim, + compress_ratio, + next_n, + ) + return output, paged_kv, paged_score + + def is_valid( + self, + kv_score: torch.Tensor, + ape: torch.Tensor, + paged_kv: torch.Tensor, + paged_score: torch.Tensor, + block_table_kv: torch.Tensor, + block_table_score: torch.Tensor, + output: torch.Tensor, + kv_lens: torch.Tensor, + cu_seq_lens: torch.Tensor, + cu_kv_comp: torch.Tensor, + batch_size: int, + page_size: int, + head_dim: int, + compress_ratio: int, + next_n: int, + ) -> None: + # These four are read through raw pointers sized only by element width and + # by an assumed dense row stride: an fp16 tensor is read/written as bf16 bits + # and a strided view is walked as if contiguous, with no error either way. + assert kv_score.dtype in _ELEMENT_TYPES and output.dtype in _ELEMENT_TYPES, ( + f"kv_score/output must be float32 or bfloat16, got {kv_score.dtype}/{output.dtype}" + ) + assert kv_score.is_contiguous() and output.is_contiguous(), ( + "kv_score and output must be contiguous (the op assumes dense rows)" + ) + + def reference( + self, + kv_score: torch.Tensor, + ape: torch.Tensor, + paged_kv: torch.Tensor, + paged_score: torch.Tensor, + block_table_kv: torch.Tensor, + block_table_score: torch.Tensor, + output: torch.Tensor, + kv_lens: torch.Tensor, + cu_seq_lens: torch.Tensor, + cu_kv_comp: torch.Tensor, + batch_size: int, + page_size: int, + head_dim: int, + compress_ratio: int, + next_n: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The whole of `(output, paged_kv, paged_score)` after the call, from + their contents before it. + + A documented exception, like `load_paged_kv_cache_for_mla.reference`'s + `stored`: for an odd start the pooled row is a function of the token an + earlier call parked at `sp - 1`, which only the pools hold. This reads it + from `paged_kv` / `paged_score` as passed, so it must see them as they + were before the call -- call it before the op, which overwrites the + step's slots in place. Nothing is mutated: the expected tensors are + copies (one copy when the two pools are one tensor). + + The checkpoint's decode rule (`Compressor.forward` at `start_pos > 0`) + applied to the step's tokens in order: store kv and `score + ape[p % R]` + in the state's element type, then pool each completed window from the + stored values, `sum(kv * softmax_over_the_window(score))` per channel in + fp32, rounded once to `output`'s dtype. `head_dim` is the state width + (the non-overlap ratios 2 and 128); the ratio-4 overlap mode is not + modeled. + """ + if compress_ratio == 4: + raise NotImplementedError("the ratio-4 overlap mode is not modeled") + R, S = compress_ratio, head_dim + dev = kv_score.device + exp_kv = paged_kv.clone() + shared = paged_score.data_ptr() == paged_kv.data_ptr() + exp_score = exp_kv if shared else paged_score.clone() + exp_out = output.clone() + kv_rows, score_rows = exp_kv.view(-1, S), exp_score.view(-1, S) + ape32 = ape.float().reshape(R, S) + lens = kv_lens[:batch_size].tolist() + first_row = cu_seq_lens[:batch_size].tolist() + first_out = cu_kv_comp[:batch_size].tolist() + + def slots(table: torch.Tensor, b: int, pos: torch.Tensor) -> torch.Tensor: + return table[b].long()[pos // page_size] * page_size + pos % page_size + + for b in range(batch_size): + kv_len = lens[b] + sp = kv_len - next_n + pos = torch.arange(sp, kv_len, device=dev) + new = kv_score[first_row[b] : first_row[b] + next_n].float() + kv_rows[slots(block_table_kv, b, pos)] = new[:, :S].to(kv_rows.dtype) + score_rows[slots(block_table_score, b, pos)] = (new[:, S:] + ape32[pos % R]).to( + score_rows.dtype + ) + w0, w1 = sp // R, kv_len // R + if w1 == w0: + continue + p = torch.arange(w0 * R, w1 * R, device=dev) + k = kv_rows[slots(block_table_kv, b, p)].float().unflatten(0, (w1 - w0, R)) + s = score_rows[slots(block_table_score, b, p)].float().unflatten(0, (w1 - w0, R)) + pooled = (k * s.softmax(dim=1)).sum(dim=1) + exp_out[first_out[b] : first_out[b] + w1 - w0] = pooled.to(output.dtype) + return exp_out, exp_kv, exp_score + + def compare( + self, + out: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ref: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ) -> None: + """Pools bitwise; fp32 rows at torch's default fp32 band; bf16 rows at + >= 0.999 bit-identical. + + The state rows are copies (`kv`, and `score + ape` in fp32), so any + differing pool byte is a wrong slot or a write the op must not make. + + The kernel is built with `--use_fast_math` and runs an online softmax + that divides once at the end, so it is not bit-identical to torch's + softmax-then-sum: over the 168 rows of next_n 1 and 6 with the cells' + inputs the largest abs error measured on sm_103 is 1.431e-06, inside + the default fp32 band (rtol 1.3e-6, atol 1e-5). A bf16 element can then + differ only where the fp32 value straddles a bf16 rounding midpoint: + 1.000000 of 86016 elements measured bit-identical; 0.999 is the gate, + which a bf16-operand control (0.717) misses. Rows the op must leave + alone are compared bit for bit (the test pre-fills them with NaN, so the + bf16 fraction is over the finite, written rows). + """ + (o, kv, sc), (ro, rkv, rsc) = out, ref + for name, got, exp in (("paged_kv", kv, rkv), ("paged_score", sc, rsc)): + assert got.dtype == exp.dtype and got.shape == exp.shape, (name, got.dtype, exp.dtype) + bad = (got.view(torch.uint8) != exp.view(torch.uint8)).sum().item() + assert bad == 0, f"{name}: {bad} of {got.numel() * got.element_size()} bytes differ" + assert o.dtype == ro.dtype, (o.dtype, ro.dtype) + assert o.shape == ro.shape, (o.shape, ro.shape) + if o.dtype == torch.float32: + torch.testing.assert_close(o, ro, equal_nan=True) + return + written = torch.isfinite(ro) + same = o.view(torch.int16) == ro.view(torch.int16) + assert same[~written].all(), "output written outside the rows the op owns" + torch.testing.assert_close(o, ro, equal_nan=True) + if written.any(): + exact = same[written].float().mean().item() + assert exact >= 0.999, f"bf16 output bit-exact fraction {exact}" + + +compressor_paged_kv_compress = CompressorPagedKvCompress() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/compressor_prefill_reduction.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/compressor_prefill_reduction.py new file mode 100644 index 000000000000..aca272de5a72 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/compressor_prefill_reduction.py @@ -0,0 +1,365 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Prefill-side KV compressor: per-channel softmax pooling of each window of +`compress_ratio` tokens, with the per-token compressor state persisted to a +paged pool so a later call (next chunk, or decode) can finish an open window. + +The op is `torch.ops.trtllm.compressor_prefill_reduction`, registered in +`cpp/tensorrt_llm/thop/compressorOp.cpp`. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_ELEMENT_TYPES = (torch.float32, torch.bfloat16) + +#: One-shot batches (every start_pos 0): odd and even lengths, page edges, up to +#: the 16384-token step. +_ONE_SHOT = ( + (1,), + (2,), + (3,), + (7, 8), + (5, 128, 255, 256, 1000), + (8191,), + (8192,), + (16383,), + (16384,), +) + +#: Five calls over five sequences (301, 130, 1027, 5 and 8192 tokens): one row per +#: sequence, its chunk size in each call, so members at different start parities +#: share every call, 1-token chunks only park a pending token (zero outputs for +#: that member), and boundaries fall at odd positions. +_CHUNKED = ( + (1, 2, 97, 1, 200), + (33, 31, 1, 1, 64), + (128, 129, 256, 257, 257), + (1, 1, 1, 1, 1), + (4095, 1, 1, 4094, 1), +) + + +def _spec(plan: tuple, **overrides) -> dict: + spec = dict(plan=plan, page_size=128, out_dtype=torch.float32, ape="zero") + spec.update(overrides) + return spec + + +def _one_shot(lengths: tuple) -> tuple: + return tuple((n,) for n in lengths) + + +class CompressorPrefillReduction(OpWrapper): + """Pool each complete window of `R = compress_ratio` tokens that gains a new + token in this call into one output row, and store every new token's + (kv, score + ape) in the paged compressor state. + + For context sequence `b < batch_size`, with `sp = start_pos[b]` and + `kv_len = kv_lens[b]`, the new positions `[sp, kv_len)` are rows + `cu_seq_lens[b] ..` of `kv_score` (`[kv || score]`, `S` wide each):: + + s'[p] = score[p] + ape[p % R] + # every window k in [sp // R, kv_len // R) + output[cu_kv_comp[b] + k - sp // R, d] = sum_p kv[p, d] * softmax_p(s'[p, d]) + paged_kv [block_table_kv [b, p // page_size], p % page_size] = kv[p] # p in [sp, kv_len) + paged_score[block_table_score[b, p // page_size], p % page_size] = s'[p] + + New positions are pooled straight from `kv_score`; only a position below + `sp` (the first window, when `sp % R != 0`) is read back from the pools, + whose score already holds `+ ape`. A trailing incomplete window is only + stored. bf16 `output` gets the round-to-nearest-even of the fp32 value. A + zero-row request (`kv_len == sp`) gets no row and touches nothing. + + Fusion boundary: the window softmax, the weighted sum and the state writes + are inside the call. The caller owns the `wkv || wgate` projection that + makes `kv_score`, the RMSNorm, RoPE, quantization and the compressed-KV + cache write, and decode-time pooling (`compressor_paged_kv_compress`). + + Out-variant: writes `output`, `paged_kv` and `paged_score` in place and + returns `(output, paged_kv, paged_score)` in that order (the same tensor + twice when the two pools are one tensor, as in DeepSeek-V4.1). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"one-shot page={page} L={'-'.join(map(str, lengths))}: start_pos 0", + spec=_spec(_one_shot(lengths), page_size=page), + ) + for page in (128, 256) + for lengths in _ONE_SHOT + ) + + tuple( + Cell( + why=f"chunked page={page}: 5 calls, mixed parities, pending token carried", + spec=_spec(_CHUNKED, page_size=page), + ) + for page in (128, 256) + ) + + ( + Cell( + why="bf16 output: RNE of the fp32 pooled value, one-shot", + spec=_spec(_one_shot((5, 128, 255, 256, 1000, 8191)), out_dtype=torch.bfloat16), + ), + Cell( + why="nonzero ape, chunked: baked into the stored score", + spec=_spec(_CHUNKED, ape="random"), + ), + ) + ) + + note = """ + Cells. DeepSeek-V4.1's ratio-2 compressor of the KV-source layers 2/8/14 + at prefill: `compress_ratio` 2, `head_dim` 512, fp32 `kv_score`, fp32 + state (kv and score pools on randomly permuted, distinct pages, a trap page + behind every unused table column), zero `ape` (one cell nonzero), fp32 or + bf16 `output`, page sizes 128 and 256. Kv ~ N(0, 2^2), score ~ N(0, 4^2), + plus one exact-tie channel and one whose pair gap of 240 underflows `exp`. + + Metadata consumed: none. The target passes the step's whole fp32 + `[T, 1024]` `kv_score` (context rows first; the generation rows after them + are not read), the pools `get_buffers(owner, COMPRESSOR_KV / SCORE)` (one + fp32 `[pages, 128, 512]` tensor for both roles and owners 2, 8, 14), the + tables `sliding_block_tables[layer_offsets[owner], COMPRESSOR_KV / SCORE][:nc]`, + a zero-filled `[num_total_compressed_tokens[2], 512]` output (bf16; rows from + `cu_new_comp_kv_cuda[2][nc]` on are the decode op's and are not written), + `kv_lens_cuda_runtime[:nc]`, `cached_token_lens_cuda[:nc]`, + `cu_seq_lens_cuda[:nc + 1]`, `cu_new_comp_kv_cuda[2][:nc + 1]`, + `num_contexts`, `max_ctx_compressed_tokens[2]` and `tokens_per_block`. + Under encoder SWA bounded replay `kv_score` is the projection of the compute + rows, `start_pos` is `v41_compute_start_pos_cuda[:nc]` and `cu_seq_lens` + `v41_compute_cu_seq_lens_cuda[:nc + 1]`; a first chunk that holds only its + replay rows is then a zero-row request. + + Preconditions: `kv_len - sp` is the sequence's row count in `kv_score` + (the op never reads `cu_seq_lens[b + 1]`); `cu_kv_comp` is the prefix sum + of `kv_len // R - sp // R`; `max_outputs >= max_b(kv_len // R - sp // R)` + -- it sizes the grid, and an undersized value silently drops the rows past + it and the trailing pending token's state (not guarded); the tables map + every page covering `[sp, kv_len)` and, for `sp % R != 0`, the page of + `sp - 1`, which holds that token's kv and score + ape; both tables have + the same width (`block_table_kv.size(1)` is the row stride); `ape` fp32 + `[R, S]`. Pages behind the window may alias: only `sp - 1` is read back. + No table entry is validated. The launcher sends any `head_dim` other than + 512 to the 128 instantiation unchecked. + + Hazards (is_valid): an fp16 `kv_score` or `output` is read or written as + bf16 bits, and a strided view is walked with the dense row stride, all + without an error. The pools are not guarded. + + Rejected by the op: `compress_ratio` other than 2, 4 or 128 (V4.1's + ratio-1 layer cannot route here), bf16 `ape`, int64 `kv_lens`, pools of + different element types. + + Precision: built with `--use_fast_math`, an online fp32 softmax that + divides once, so not bitwise torch's softmax-then-sum. + + The test also certifies, on the engine's own pools and managers: every + row count 1..16384, the target's call (dumped steps, and steps on + managers of the envelope below; shared pool), offset hazards and pools up + to 247806 pages, the multi-request envelope (1..256 contexts, T 1..16384, + `max_outputs` 0..8192, tables 132 .. 8196 columns, kv lengths to + 1,048,590), every dumped prefill step, and zero-row requests anywhere in + the batch. + + Not certified: CUDA-graph capture (prefill runs eagerly); bf16 `kv_score` + or state, `head_dim` 128, ratios 4 and 128, other page sizes; more than + 16384 tokens or 256 contexts per call; kv lengths past 1,048,590 or tables + wider than 8196 columns; pools above 247806 pages; page sharing under block + reuse; a zero-row request at `sp` 0 or at an odd `sp` other than 7; a call + whose `kv_score` has no row at all; performance. + """ + + def raw_call( + self, + kv_score: torch.Tensor, + ape: torch.Tensor, + paged_kv: torch.Tensor, + paged_score: torch.Tensor, + block_table_kv: torch.Tensor, + block_table_score: torch.Tensor, + output: torch.Tensor, + kv_lens: torch.Tensor, + start_pos: torch.Tensor, + cu_seq_lens: torch.Tensor, + cu_kv_comp: torch.Tensor, + batch_size: int, + page_size: int, + head_dim: int, + compress_ratio: int, + max_outputs: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + torch.ops.trtllm.compressor_prefill_reduction( + kv_score, + ape, + paged_kv, + paged_score, + block_table_kv, + block_table_score, + output, + kv_lens, + start_pos, + cu_seq_lens, + cu_kv_comp, + batch_size, + page_size, + head_dim, + compress_ratio, + max_outputs, + ) + return output, paged_kv, paged_score + + def is_valid( + self, + kv_score: torch.Tensor, + ape: torch.Tensor, + paged_kv: torch.Tensor, + paged_score: torch.Tensor, + block_table_kv: torch.Tensor, + block_table_score: torch.Tensor, + output: torch.Tensor, + kv_lens: torch.Tensor, + start_pos: torch.Tensor, + cu_seq_lens: torch.Tensor, + cu_kv_comp: torch.Tensor, + batch_size: int, + page_size: int, + head_dim: int, + compress_ratio: int, + max_outputs: int, + ) -> None: + # These four are read through raw pointers sized only by element width and + # by an assumed dense row stride: an fp16 tensor is read/written as bf16 bits + # and a strided view is walked as if contiguous, with no error either way. + assert kv_score.dtype in _ELEMENT_TYPES and output.dtype in _ELEMENT_TYPES, ( + f"kv_score/output must be float32 or bfloat16, got {kv_score.dtype}/{output.dtype}" + ) + assert kv_score.is_contiguous() and output.is_contiguous(), ( + "kv_score and output must be contiguous (the op assumes dense rows)" + ) + + def reference( + self, + kv_score: torch.Tensor, + ape: torch.Tensor, + paged_kv: torch.Tensor, + paged_score: torch.Tensor, + block_table_kv: torch.Tensor, + block_table_score: torch.Tensor, + output: torch.Tensor, + kv_lens: torch.Tensor, + start_pos: torch.Tensor, + cu_seq_lens: torch.Tensor, + cu_kv_comp: torch.Tensor, + batch_size: int, + page_size: int, + head_dim: int, + compress_ratio: int, + max_outputs: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The whole of `(output, paged_kv, paged_score)` after the call, from + their contents before it. + + A documented exception, like `load_paged_kv_cache_for_mla.reference`'s + `stored`: when `sp % R != 0` the first pooled row is a function of the + token(s) an earlier call parked below `sp`, which only the pools hold. + This reads them from `paged_kv` / `paged_score` as passed, so it must + see them as they were before the call -- call it before the op, which + overwrites the chunk's slots in place. Nothing is mutated: the expected + tensors are copies (one copy when the two pools are one tensor). + + The checkpoint's `Compressor.forward` pooling, `sum(kv * softmax_over_the_window(score + ape))` + per channel in fp32, rounded once to `output`'s dtype; new positions + from `kv_score`, parked ones from the pools. `head_dim` is the state + width (the non-overlap ratios 2 and 128); the ratio-4 overlap mode is not + modeled. `max_outputs` only sizes the launch grid and is ignored: this + is the result under its precondition. + """ + if compress_ratio == 4: + raise NotImplementedError("the ratio-4 overlap mode is not modeled") + R, S = compress_ratio, head_dim + dev = kv_score.device + old_kv, old_score = paged_kv.view(-1, S), paged_score.view(-1, S) + exp_kv = paged_kv.clone() + shared = paged_score.data_ptr() == paged_kv.data_ptr() + exp_score = exp_kv if shared else paged_score.clone() + exp_out = output.clone() + kv_rows, score_rows = exp_kv.view(-1, S), exp_score.view(-1, S) + ape32 = ape.float().reshape(R, S) + lens = kv_lens[:batch_size].tolist() + starts = start_pos[:batch_size].tolist() + first_row = cu_seq_lens[:batch_size].tolist() + first_out = cu_kv_comp[:batch_size].tolist() + + def slots(table: torch.Tensor, b: int, pos: torch.Tensor) -> torch.Tensor: + return table[b].long()[pos // page_size] * page_size + pos % page_size + + for b in range(batch_size): + sp, kv_len = starts[b], lens[b] + if kv_len <= sp: + continue + pos = torch.arange(sp, kv_len, device=dev) + rows = kv_score[first_row[b] : first_row[b] + kv_len - sp].float() + kv_new, score_new = rows[:, :S], rows[:, S:] + ape32[pos % R] + w0, w1 = sp // R, kv_len // R + if w1 > w0: + parked = torch.arange(w0 * R, sp, device=dev) + k = torch.cat([old_kv[slots(block_table_kv, b, parked)].float(), kv_new]) + s = torch.cat([old_score[slots(block_table_score, b, parked)].float(), score_new]) + n = (w1 - w0) * R + k, s = k[:n].unflatten(0, (w1 - w0, R)), s[:n].unflatten(0, (w1 - w0, R)) + pooled = (k * s.softmax(dim=1)).sum(dim=1) + exp_out[first_out[b] : first_out[b] + w1 - w0] = pooled.to(output.dtype) + kv_rows[slots(block_table_kv, b, pos)] = kv_new.to(kv_rows.dtype) + score_rows[slots(block_table_score, b, pos)] = score_new.to(score_rows.dtype) + return exp_out, exp_kv, exp_score + + def compare( + self, + out: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ref: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ) -> None: + """Pools bitwise; fp32 rows at torch's default fp32 band; bf16 rows at + >= 0.999 bit-identical. + + The state rows are copies (`kv`, and `score + ape` in fp32), so any + differing pool byte is a wrong slot or a write the op must not make. + + The kernel is built with `--use_fast_math` and runs an online softmax + that divides once at the end, so it is not bit-identical to torch's + softmax-then-sum: over the five-sequence chunked batch with the cells' + inputs the largest abs error measured on sm_103 is 1.907e-06, inside the + default fp32 band (rtol 1.3e-6, atol 1e-5). A bf16 element can then + differ only where the fp32 value straddles a bf16 rounding midpoint: + 0.999985 of 2516992 elements measured bit-identical; 0.999 is the gate, + which a bf16-operand control (0.719) misses. Rows the op must leave + alone are compared bit for bit (the test pre-fills them with NaN, so the + bf16 fraction is over the finite, written rows). + """ + (o, kv, sc), (ro, rkv, rsc) = out, ref + for name, got, exp in (("paged_kv", kv, rkv), ("paged_score", sc, rsc)): + assert got.dtype == exp.dtype and got.shape == exp.shape, (name, got.dtype, exp.dtype) + bad = (got.view(torch.uint8) != exp.view(torch.uint8)).sum().item() + assert bad == 0, f"{name}: {bad} of {got.numel() * got.element_size()} bytes differ" + assert o.dtype == ro.dtype, (o.dtype, ro.dtype) + assert o.shape == ro.shape, (o.shape, ro.shape) + if o.dtype == torch.float32: + torch.testing.assert_close(o, ro, equal_nan=True) + return + written = torch.isfinite(ro) + same = o.view(torch.int16) == ro.view(torch.int16) + assert same[~written].all(), "output written outside the rows the op owns" + torch.testing.assert_close(o, ro, equal_nan=True) + if written.any(): + exact = same[written].float().mean().item() + assert exact >= 0.999, f"bf16 output bit-exact fraction {exact}" + + +compressor_prefill_reduction = CompressorPrefillReduction() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/deepgemm_fp8_fp4_mqa_logits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/deepgemm_fp8_fp4_mqa_logits.py new file mode 100644 index 000000000000..d39919ea4189 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/deepgemm_fp8_fp4_mqa_logits.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Indexer logits over packed MXFP4 q/k with contiguous (non-paged) keys: the prefill form. + +The op is the vendored DeepGEMM Python API `tensorrt_llm.deep_gemm.fp8_fp4_mqa_logits` +(built from the `deepgemm` entry of `3rdparty/fetch_content.json`), a pybind-backed callable +rather than a `torch.ops.trtllm` op. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm import deep_gemm +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + +#: (T, S, H, ratio, what else it covers) of one request's causal prefill from empty: S = T // ratio. +_CAUSAL = ( + (1, 1, 32, 1, ""), + (5, 2, 32, 2, ", row 0 sees no key"), + (64, 64, 32, 1, ""), + (333, 166, 32, 2, ", T not a multiple of the 4-row query block"), + (4096, 2048, 32, 2, ""), + (4096, 4096, 32, 1, ""), + (8192, 4096, 32, 2, ""), + (8192, 8192, 32, 1, ""), + (64, 64, 8, 1, ", a TP4 rank's quarter of the heads"), + (4096, 4096, 8, 1, ", a TP4 rank's quarter of the heads"), +) + +#: (cached, new, ratio) of one request's context chunk over its cached keys. +_CHUNKED = ((8000, 6, 1), (8190, 1, 2), (3000, 1000, 2), (12288, 4096, 1)) + +#: Long-context prefill shapes, as per-request (cached, new) lists and the ratio: the 24K passkey +#: prompt's last chunk (T 7674 x S 24058, a runtime dump), S 24058 at ratio 2, S 32768 at 4096 and +#: 8192 rows, and three packed requests with their own key ranges (S 32576 / 16288). +_LONG = ( + ([(16384, 7674)], 1), + ([(44020, 4096)], 2), + ([(28672, 4096)], 1), + ([(24576, 8192)], 1), + ([(61440, 4096)], 2), + ([(16384, 4096), (0, 3000), (8000, 1096)], 1), + ([(16384, 4096), (0, 3000), (8000, 1096)], 2), +) + +#: (T, S, ratio) of the recipe-regime causal chunks whose last row sees key S - 1. +_RECIPE = ( + (1, 4096, 1), + (512, 1024, 1), + (4096, 2048, 2), + (4096, 4096, 1), + (7674, 24058, 1), + (4096, 32768, 1), +) + + +def _extent(reqs, ratio: int) -> tuple[int, int]: + """(query rows, packed key rows) of a list of (cached, new) requests at a ratio.""" + return sum(n for _, n in reqs), sum((c + n) // ratio for c, n in reqs) + + +def _decode(packed: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """int8 `[..., 64]` packed E2M1 + int32 `[...]` packed UE8M0 -> float64 `[..., 128]`, exactly. + + Element `2j` is the low nibble of byte `j`; byte `b` of the little-endian + word is the exponent of elements `32b .. 32b+31`. + """ + lut = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + p = packed.view(torch.uint8).long() + v = torch.stack([lut[p & 0xF], lut[p >> 4]], dim=-1).flatten(-2) + e = sf.contiguous().view(torch.uint8).view(*sf.shape, 4).long() - 127 + scale = ((e + 1023) << 52).view(torch.float64) # 2^e from the exponent bits + return (v.unflatten(-1, (4, 32)) * scale.unsqueeze(-1)).flatten(-2) + + +class DeepgemmFp8Fp4MqaLogits(OpWrapper): + """Weighted-ReLU MQA logits of every query row against its range of contiguous key rows. + + The callable is `tensorrt_llm.deep_gemm.fp8_fp4_mqa_logits` (vendored DeepGEMM, pybind), + one call. With `ks = cu_seq_len_k_start`, `ke = cu_seq_len_k_end`: + + logits[t, s] = sum_h relu(dec(q[t, h]) . dec(kv[s])) * weights[t, h] if ks[t] <= s < ke[t] + = -inf otherwise, when clean_logits=True + = not specified otherwise, when clean_logits=False + + `q` is int8 `[T, H, 64]` packed E2M1 with int32 `q_sf [T, H]` (four UE8M0 + bytes, one per 32 elements), `kv` int8 `[S, 64]` with int32 `kv_sf [S]`, + `weights` fp32 `[T, H]`. The relu is per head, before the weight; the head + sum is inside the call. Products come from the block-scaled FP4 MMA with + fp32 accumulation; relu, weighting and the head sum are fp32. + + `max_seqlen_k > 0` selects the compressed form (requires + `clean_logits=False`): column `j` of row `t` is key `ks[t] + j` for + `j < ke[t] - ks[t]`, the rest not specified. + + Fusion boundary: decoding both operands, the per-head dot products, relu, + weighting, the head sum and the `-inf` fill. The caller owns quantizing q + and k, RoPE, the weights projection and its scale fold, gathering k out of + the paged pool, candidate masking, the top-k and any cross-rank reduction. + + Returns a fresh `[T, S]` (or `[T, max_seqlen_k]`) tensor of + `logits_dtype`: a strided view whose row stride is padded (see `note`). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"exact T={t} S={s} H={h} r={r}: causal prefill{extra}", + spec=dict(regime="exact", H=h, r=r, reqs=[(0, t)]), + ) + for t, s, h, r, extra in _CAUSAL + ) + + tuple( + Cell( + why=f"exact T={n} S={(c + n) // r} r={r}: chunk after {c} cached", + spec=dict(regime="exact", H=32, r=r, reqs=[(c, n)]), + ) + for c, n, r in _CHUNKED + ) + + tuple( + Cell( + why=f"exact packed r={r}: 4 requests at own key offsets", + spec=dict(regime="exact", H=32, r=r, reqs=[(0, 700), (0, 1), (512, 2500), (0, 37)]), + ) + for r in (1, 2) + ) + + tuple( + Cell( + why=( + f"exact long T={_extent(reqs, r)[0]} S={_extent(reqs, r)[1]} r={r}" + f" x{len(reqs)}: allocation = the logits buffer" + ), + spec=dict(regime="exact", H=32, r=r, reqs=reqs, alloc_check=True), + ) + for reqs, r in _LONG + ) + + tuple( + Cell( + why=f"recipe T={t} S={s} r={r}: bf16 randn via per_token_cast_to_fp4", + spec=dict(regime="recipe", H=32, r=r, reqs=[(s * r - t, t)]), + ) + for t, s, r in _RECIPE + ) + + ( + Cell( + why="wide UE8M0 bytes 90..160, per token and per key", + spec=dict(regime="wide_scale", H=32, r=1, reqs=[(256, 256)]), + ), + Cell( + why="clean_logits=False: in-range values, shifted ranges", + spec=dict(regime="exact", H=32, T=300, S=600, ranges="shifted"), + ), + Cell( + why="compressed form: column j is key ks + j, empty rows", + spec=dict(regime="exact", H=32, T=200, S=1000, ranges="compressed"), + ), + ) + ) + + note = """ + Cell regimes: "exact" packs random E2M1 bytes with UE8M0 bytes 126..128 and + weights `n * 2^-10` (|n| <= 8) times the 2^-6 fold, so every head term and + partial sum is exact in fp32 (the test asserts that premise) and the op is + bit-identical to the float64 reference; "recipe" packs bf16 randn with + DeepGEMM's `per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, + use_packed_ue8m0=True)` and uses fp32 randn weights times 2^-6. + + Metadata consumed: none. The launch grid uses DeepGEMM's process-global + `deep_gemm.get_num_sms()`. For DeepSeek-V4.1: `weights = weights_proj(x) * + 128^-0.5 * 32^-0.5` (exactly 2^-6; 32 is the global index head count even + on a rank holding fewer heads), `ks[t]` the request's first key row and + `ke[t] = ks[t] + (p + 1) // ratio` for row `t`'s absolute position `p`. The + checkpoint's reference scores in bf16; these fp32 logits are more precise, + so top-k near-ties can resolve differently. + + Output: row stride `align(S + 256, 256)` (compressed: `align(max_seqlen_k, + 256)`), so the result is not contiguous and its consumer must accept the + padded stride. One allocation of `align(T, 128 / H) * stride * 4` bytes and + no other workspace; the first call in a process also keeps 32 MiB. + + JIT: the one kernel (`sm100_mqa_logits`, which with `clean_logits=True` + also writes the -inf columns; there is no separate clean pass) is compiled + with NVCC on the first call of each configuration (heads, dtypes, + compressed or not) and cached on disk under + `$HOME/.deep_gemm` (`DG_JIT_CACHE_DIR`). Make that call outside CUDA-graph + capture; after it, capture and replay on new inputs are bit-exact. + + Preconditions: `T >= 1` and `S >= 1` (otherwise a non-sticky + CUDA_ERROR_INVALID_VALUE; skip the call when no key exists yet); + `0 <= ks[t]` and `ke[t] <= S` (`ke > S` is out of contract); q and k in the + same byte layout (a layout swapped in both is invisible, in one is wrong + everywhere); `kv_sf` 16-byte aligned and a `weights` row stride a multiple + of 4 floats (otherwise CUDA_ERROR_INVALID_VALUE); the index-K gather op's + float8 / float32 outputs viewed as int8 / int32 first. Every mismatch probed + raised, so there is no `is_valid`. H must be a multiple of 4 up to 128 (the + op raises otherwise); H 24 and 128 are measured against the reference by a + test, without being certified. + + Certified at H 32, `clean_logits=True`, full width (the target's call): + one call per context step of 1..8192 rows from 1..32 requests of up to + 16384 tokens (S up to 524288 at ratio 1, 262144 at ratio 2; every T in + 1..8192 and every S in 1..32768 for one request), and the row-tiled + long-context form -- tiles of a step of up to 16384 rows from up to 256 + requests of up to 1,048,590 tokens over up to 2^25 keys, each tile within + 2^28 logits elements, rows passed as row-offset views. A row's logits do + not depend on the call's extents or on where its keys sit (bitwise, recipe + regime included), and repeated calls are bitwise identical. One compiled + variant serves every T and S. The target never captures + the call (context steps are not graph-replayed). + + Not certified: other architectures (DeepGEMM dispatches on the arch + major); `q_sf=None` (the FP8-q path); bf16 or fp16 weights; bf16 logits; + H other than 32, and 8 beyond its cells; `clean_logits=False` and the + compressed form beyond their cells; anything outside the two envelopes. + """ + + def raw_call( + self, + q: torch.Tensor, + q_sf: Optional[torch.Tensor], + kv: torch.Tensor, + kv_sf: torch.Tensor, + weights: torch.Tensor, + cu_seq_len_k_start: torch.Tensor, + cu_seq_len_k_end: torch.Tensor, + clean_logits: bool, + max_seqlen_k: int, + logits_dtype: torch.dtype, + ) -> torch.Tensor: + return deep_gemm.fp8_fp4_mqa_logits( + (q, q_sf), + (kv, kv_sf), + weights, + cu_seq_len_k_start, + cu_seq_len_k_end, + clean_logits=clean_logits, + max_seqlen_k=max_seqlen_k, + logits_dtype=logits_dtype, + ) + + def reference( + self, + q: torch.Tensor, + q_sf: Optional[torch.Tensor], + kv: torch.Tensor, + kv_sf: torch.Tensor, + weights: torch.Tensor, + cu_seq_len_k_start: torch.Tensor, + cu_seq_len_k_end: torch.Tensor, + clean_logits: bool, + max_seqlen_k: int, + logits_dtype: torch.dtype, + ) -> torch.Tensor: + """`sum_h relu(q_h . k) * w_h` over exactly decoded operands, in float64, rounded once. + + NaN marks every element the op does not specify: outside `[ks, ke)` + when `clean_logits=False`, and the compressed form's columns at or past + a row's length. `compare` skips those. `q_sf` must be given (the MXFP4 + path); nothing is mutated. + """ + qd = _decode(q, q_sf) + kd = _decode(kv, kv_sf) + num_q, num_k = qd.shape[0], kd.shape[0] + wd = weights.double() + full = torch.empty(num_q, num_k, dtype=torch.float64, device=q.device) + for t0 in range(0, num_q, 512): + dots = torch.einsum("thd,kd->thk", qd[t0 : t0 + 512], kd) + full[t0 : t0 + 512] = (dots.clamp_min(0) * wd[t0 : t0 + 512].unsqueeze(-1)).sum(1) + ks = cu_seq_len_k_start.long() + ke = cu_seq_len_k_end.long() + if max_seqlen_k > 0: + j = torch.arange(max_seqlen_k, device=q.device) + inside = j[None, :] < (ke - ks)[:, None] + src = (ks[:, None] + j[None, :]).clamp(0, max(num_k - 1, 0)) + out = full.gather(1, src).masked_fill(~inside, float("nan")) + else: + col = torch.arange(num_k, device=q.device) + vis = (col[None, :] >= ks[:, None]) & (col[None, :] < ke[:, None]) + out = full.masked_fill(~vis, float("-inf") if clean_logits else float("nan")) + return out.to(logits_dtype) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """`-inf` exactly where the reference has it, default fp32 tolerance on visible elements. + + Elements the reference marks NaN (not specified by the op) are skipped. + The band is torch's default for fp32 (rtol 1.3e-6, atol 1e-5): on bf16 + randn packed by the recipe with fp32 randn weights (T 512, S 1024, H + 32) the op's largest difference from float64 was 5.9e-7, against 4.9e-7 + for an fp32 torch einsum of the same decoded values, and 43% of visible + elements were bit-exact. In the exact regime the op is bit-identical, + which the test asserts on top of this gate. Controls (relu after the + head sum, an 8-head weight fold, swapped k nibbles, reversed scale bytes) + put nearly every visible element outside the band. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + specified = ~torch.isnan(ref) + masked = torch.isneginf(ref) + wrong_inf = int(((torch.isneginf(out) & specified) != masked).sum()) + assert wrong_inf == 0, f"{wrong_inf} elements disagree on -inf (outside [ks, ke))" + visible = specified & ~masked + torch.testing.assert_close(out[visible], ref[visible]) + + +deepgemm_fp8_fp4_mqa_logits = DeepgemmFp8Fp4MqaLogits() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/deepgemm_fp8_fp4_paged_mqa_logits.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/deepgemm_fp8_fp4_paged_mqa_logits.py new file mode 100644 index 000000000000..adeeeaad2ce8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/deepgemm_fp8_fp4_paged_mqa_logits.py @@ -0,0 +1,353 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode indexer logits over a paged MXFP4 key cache, and the work schedule they require. + +Two entries over two callables of the vendored DeepGEMM (`tensorrt_llm.deep_gemm`, a pybind +module built from the `deepgemm` entry of `3rdparty/fetch_content.json`), one call each: + +- `deepgemm_fp8_fp4_paged_mqa_logits` -> `tensorrt_llm.deep_gemm.fp8_fp4_paged_mqa_logits` +- `deepgemm_paged_mqa_logits_metadata` -> `tensorrt_llm.deep_gemm.get_paged_mqa_logits_metadata` + +The schedule is a precondition of the logits call and has no other consumer, so both live here. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm import deep_gemm +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) +_CODE_BYTES = 64 # packed E2M1 bytes per key +_KEY_BYTES = 68 # 64 code bytes + one 4-byte UE8M0 word per key + +#: (ratio, page entries): the INDEXER_COMPRESS page holds tokens_per_block / ratio entries; +#: tokens_per_block 128 gives (2, 64) and (1, 128), tokens_per_block 256 gives (2, 128). +_PAGES = ((2, 64), (1, 128), (2, 128)) + +#: First-row positions of a decode step's requests, by regime. +_REGIMES = ( + ("boundary", "lengths on page and 256-column split edges"), + ("gate", "positions 1024..16383, the accuracy gates"), + ("long", "positions 20480..40959, long context"), + ("64k", "about 64K compressed positions"), +) + +#: The logits cells; the schedule is certified at exactly these, as their input. +_LOGITS_CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"{regime} B={batch} next_n={next_n} r={ratio} page {bk}: {what}", + spec=dict( + ratio=ratio, block_kv=bk, next_n=next_n, heads=32, batch=batch, regime=regime + ), + ) + for regime, what in _REGIMES + for batch in (1, 8, 32) + for next_n in (1, 6) + for ratio, bk in _PAGES + ) + + tuple( + Cell( + why=f"H={heads} next_n={next_n} r=2 page 64: a head split's share", + spec=dict( + ratio=2, block_kv=64, next_n=next_n, heads=heads, positions=[3, 130, 999, 4100] + ), + ) + for heads in (8, 16, 64) + for next_n in (1, 6) + ) + + ( + Cell( + why="bf16 logits with fp32 weights: next_n=6 B=5 r=2 page 64", + spec=dict( + ratio=2, + block_kv=64, + next_n=6, + heads=32, + positions=[3, 130, 999, 4100, 20001], + logits_dtype=torch.bfloat16, + ), + ), + ) +) + + +def _decode(packed: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """int8 `[..., 64]` packed E2M1 + int32 `[...]` packed UE8M0 -> float64 `[..., 128]`, exactly. + + Element `2j` is the low nibble of byte `j`; byte `b` of the little-endian + word is the exponent of elements `32b .. 32b+31`. + """ + lut = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + p = packed.view(torch.uint8).long() + v = torch.stack([lut[p & 0xF], lut[p >> 4]], dim=-1).flatten(-2) + e = sf.contiguous().view(torch.uint8).view(*sf.shape, 4).long() - 127 + scale = ((e + 1023) << 52).view(torch.float64) # 2^e from the exponent bits + return (v.unflatten(-1, (4, 32)) * scale.unsqueeze(-1)).flatten(-2) + + +class DeepgemmFp8Fp4PagedMqaLogits(OpWrapper): + """Weighted-ReLU MQA logits of `B * next_n` query rows against each request's paged keys. + + The callable is `tensorrt_llm.deep_gemm.fp8_fp4_paged_mqa_logits` (vendored DeepGEMM, pybind), + one call. For query row `i = b * next_n + j` of request `b`: + + logits[i, s] = sum_h relu(dec(q[b, j, h]) . dec(k[s])) * weights[i, h] for s < context_lens[b, j] + + where `k[s]` is slot `s % block_kv` of pool page `block_table[b, s // block_kv]` + and every other column of the row is not specified (nothing is masked to + `-inf`). `q` is int8 `[B, next_n, H, 64]` packed E2M1 with int32 + `q_sf [B, next_n, H]` (four UE8M0 bytes), `kv_cache` uint8 + `[pages, block_kv, 1, 68]` holding per page `block_kv x 64` code bytes then + `block_kv x 4` scale bytes, `weights` fp32 `[B * next_n, H]`. No global + scale is applied: the `128^-0.5 * 32^-0.5` fold belongs in `weights`. + + `clean_logits` is fixed to False (the only value 2-D `context_lens` + accepts) and the varlen `indices` mode is not exposed. + + Fusion boundary: the paged key gather, the MXFP4 MMA, relu, head weighting + and the head reduction. The caller owns q / k quantization and packing, + writing k into the pool, the weights projection and its scale fold, the + schedule (`deepgemm_paged_mqa_logits_metadata`), the candidate-block mask + and the top-k. + + Returns a fresh `[B * next_n, max_context_len]` tensor of `logits_dtype`, + a row-strided view (see `note`). Inputs are not modified. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = _LOGITS_CELLS + + note = """ + Cells: queries `randn * 2` and keys `randn` packed by a native-torch MXFP4 + recipe (the checkpoint's `fp4_act_quant`: UE8M0 `2^ceil(log2(amax / 6))` + per 32-block, E2M1 round-to-nearest-even), weights `randn * 2^-4` rounded + to bf16 and passed as fp32; per-row lengths `(p + j + 1) // ratio`; a padding + request (length 0, an all `-1` table row) in the middle of every batch of 8 + or more; pools of 3x the pages the batch needs plus 8 under a random + permutation, with the slots past each request's length NaN-scaled. + + Metadata consumed: none as an object. For DeepSeek-V4.1, `context_lens[b, j] + = (p_b + j + 1) // ratio` for request `b`'s first query position `p_b` + (padding requests 0); `block_table` is the index-source layer's per-ratio + compressed-page table in pool pages, the pool its INDEXER_COMPRESS buffer; + `schedule_meta` comes from `deepgemm_paged_mqa_logits_metadata` over the + same `context_lens` in the same step and can be shared by that ratio's + layers. `weights = weights_proj(x) * 2^-6`. + + Output: row stride `max_context_len` rounded up to 256 (fp32) or 512 + (bf16). Columns at or past a row's own length can hold NaN, +inf or finite + scores of keys it must not see; bound every row by its own length. + Numerics: per-head dots are exact; the heads are reduced in fp32 in a fixed + order (four FMA chains of `2 * relu` over heads `h = l mod 4`, then + `((c0 + c2) + (c1 + c3)) / 2`), bit-identical to that emulation. bf16 logits + are that value rounded once. + + Preconditions no metadata check can see (each measured silently wrong, so + there is no `is_valid`): the pool in the kernel's page layout, not + key-interleaved; `block_kv` the page's true entry count (32, 64 or 128); + a valid page in every table entry below `ceil(max_j context_lens[b, j] / + block_kv)` (`-1` past it is fine; a narrower table is read past); + `max_context_len >= max(context_lens)` -- unchecked, a smaller width + writes out of bounds, so pass the static bound in captured steps; and a + schedule from the same lengths in the same step (a stale one gave wrong + logits or a sticky illegal-address error). Loud rejections: 1-D lengths, + H not a multiple of 4 in 4..128, bf16 weights with fp32 logits, fp16 logits, + `weights.stride(1) != 1`, int64 tables or lengths, uint8 q, an int8 pool, + a 256-entry page, a schedule built for another SM count. H 24 is taken and + measured against the reference by a test, without being certified. + + JIT: both kernels compile on the first call per (`next_n`, page); one + variant then serves every batch and width. Make that call outside capture; + the schedule and the logits capture together and replay with new lengths, + tables, keys, queries and weights. + + Certified beyond the cells: every B in 1..256 eagerly; every width up to + 1,048,832 (ratio 1) / 524,544 (ratio 2) with lengths up to 1,048,595 / + 524,297; capture at G in 1..32, 40..128 step 8, 192, 256 at the envelope's + graph-copy widths; mixed steps as row-offset views (q, q_sf, weights after + `tc` context rows; the lengths a view of the metadata's buffer; table rows + from `nc`); pools up to 3,245,959 blocks at ratio 1 and 2,615,990 at ratio + 2 in the cache manager's layout (each ratio-2 owner's view), past 2^34 bytes. + + Not certified: other architectures; `block_kv` 32; `next_n` other than 1 + and 6; H other than 32 beyond the cells; B above 256; capture at other + batches or widths, of a first-ever call, or of a mixed step; bf16 weights + (a coarser bf16 reduction); the MXFP8, FP8 and varlen `indices` modes; a + page stride larger than `block_kv * 68`; pools above the bounds. + """ + + def raw_call( + self, + q: torch.Tensor, + q_sf: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + schedule_meta: torch.Tensor, + max_context_len: int, + logits_dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + return deep_gemm.fp8_fp4_paged_mqa_logits( + (q, q_sf), + kv_cache, + weights, + context_lens, + block_table, + schedule_meta, + max_context_len, + False, + logits_dtype, + None, + ) + + def reference( + self, + q: torch.Tensor, + q_sf: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + schedule_meta: torch.Tensor, + max_context_len: int, + logits_dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + """`sum_h relu(q_h . k) * w_h` over exactly decoded operands, in float64, rounded once. + + Each request's keys are read out of `kv_cache` through its table row in + the kernel's page layout. NaN marks every column the op does not + specify (at or past the row's own length); `compare` skips those. + `schedule_meta` is ignored: it only splits the work. Nothing is mutated. + """ + batch, next_n = q.shape[0], q.shape[1] + block_kv = kv_cache.shape[1] + out = torch.full( + (batch * next_n, max_context_len), float("nan"), dtype=torch.float64, device=q.device + ) + wd = weights.double() + lens = context_lens.long() + for b in range(batch): + longest = int(lens[b].max()) + if longest == 0: + continue + npg = -(-longest // block_kv) + pages = kv_cache[block_table[b, :npg].long()].reshape(npg, block_kv * _KEY_BYTES) + codes = pages[:, : block_kv * _CODE_BYTES].reshape(npg * block_kv, _CODE_BYTES) + sf = pages[:, block_kv * _CODE_BYTES :].reshape(npg * block_kv, 4) + sf = sf[:longest].contiguous().view(torch.int32).squeeze(-1) + kd = _decode(codes[:longest].view(torch.int8), sf) + dots = torch.einsum("jhd,sd->jhs", _decode(q[b], q_sf[b]), kd) + rows = slice(b * next_n, (b + 1) * next_n) + score = (dots.clamp_min(0) * wd[rows].unsqueeze(-1)).sum(1) + for j in range(next_n): + length = int(lens[b, j]) + out[b * next_n + j, :length] = score[j, :length] + return out.to(logits_dtype) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """torch's default band for the logits dtype, on every element the op specifies. + + The per-head dots are exact and the fp32 head reduction is the only + rounding: the largest max abs difference from float64 the probes printed + was 1.9e-6, inside the default fp32 band (rtol 1.3e-6, atol 1e-5). The + op is also bit-identical to an fp32 emulation of its documented + reduction order, which the test asserts on top of this gate. bf16 + logits are that fp32 value rounded once. Controls (rolled table rows, + 64-entry pages read as 128, a key-interleaved pool, no relu, reversed + head weights) put 0.99 to 1.0 of the visible elements outside the band. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + specified = ~torch.isnan(ref) + torch.testing.assert_close(out[specified], ref[specified]) + + +class DeepgemmPagedMqaLogitsMetadata(OpWrapper): + """The per-SM work schedule `deepgemm_fp8_fp4_paged_mqa_logits` walks over these lengths. + + The callable is `tensorrt_llm.deep_gemm.get_paged_mqa_logits_metadata` (vendored DeepGEMM, + pybind), one call, one small kernel. `context_lens` is int32 `[B, next_n]`, the same tensor + holding the same values the logits call receives. `num_sms=None` resolves to + `deep_gemm.get_num_sms()` on the host, the only value the logits call accepts. + + Fusion boundary: the split of the lengths' work over the SMs and nothing else. + Returns a fresh int32 `[num_sms + 1, 2]` tensor. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = _LOGITS_CELLS + + note = """ + The schedule's values are DeepGEMM's own split of the work and have no + plain-torch reference: they are certified only through their consumer. + `reference` therefore states the contract the logits call checks (shape, + dtype, device), and every cell gates the values end to end -- the schedule + built from that cell's lengths drives `deepgemm_fp8_fp4_paged_mqa_logits` + to its reference. + + Metadata consumed: `context_lens`, as the logits call gets it. Build one + schedule per `context_lens` tensor per step; a ratio's schedule can be + shared by all index-source layers of that ratio in the step. + + Preconditions: rebuild it whenever the lengths change -- a schedule from + other lengths was measured to give wrong logits (all-zero lengths) or a + sticky illegal-address error (2x lengths); that some other mismatches + happened to give right logits is not a contract. Use the device's SM count. + On sm_103 `block_kv` is only range-checked (32, 64 or 128; 256 raises); the + kernel takes its page size from the pool. 1-D lengths raise. + + Captured together with the logits call at every capture the logits entry + certifies. + + Not certified: other architectures; a schedule consumed by anything else. + """ + + def raw_call( + self, + context_lens: torch.Tensor, + block_kv: int, + num_sms: Optional[int] = None, + ) -> torch.Tensor: + if num_sms is None: + num_sms = deep_gemm.get_num_sms() + return deep_gemm.get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms) + + def reference( + self, + context_lens: torch.Tensor, + block_kv: int, + num_sms: Optional[int] = None, + ) -> torch.Tensor: + """The schedule's contract, not its values: an int32 `[num_sms + 1, 2]` tensor on the lengths' device. + + Returned as an uninitialized tensor of that shape, dtype and device; + `compare` reads only those. The values are gated by the consumer: see + `note`. `block_kv` is ignored. Nothing is mutated. + """ + if num_sms is None: + num_sms = deep_gemm.get_num_sms() + return torch.empty((num_sms + 1, 2), dtype=torch.int32, device=context_lens.device) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Shape, dtype, device and contiguity: what the logits call asserts of a schedule. + + A schedule built for another SM count is rejected by the logits call + with `_schedule_meta_size == num_sms / num_ctas_per_cluster + 1`; its values + are gated through the logits (see `note`). + """ + assert out.dtype == ref.dtype == torch.int32, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + assert out.device == ref.device, (out.device, ref.device) + assert out.is_contiguous() + + +deepgemm_fp8_fp4_paged_mqa_logits = DeepgemmFp8Fp4PagedMqaLogits() +deepgemm_paged_mqa_logits_metadata = DeepgemmPagedMqaLogitsMetadata() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/dsv4_local_to_global_indices.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/dsv4_local_to_global_indices.py new file mode 100644 index 000000000000..9069174a6cce --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/dsv4_local_to_global_indices.py @@ -0,0 +1,261 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sparse-attention index table: per-token local SWA / compressed indices -> pool rows. + +The op is `torch.ops.trtllm.dsv4_local_to_global_indices`, a thin torch custom op registered in +`tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py` over that file's Triton +kernel `_deepseek_v4_local_to_global_kernel` (the kernel the in-tree Python launcher +`deepseek_v4_local_to_global_indices` also runs). One Triton launch per call. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.kernels # noqa: F401 — registers torch.ops.trtllm.dsv4_local_to_global_indices +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_SWA_TPB = 128 +#: Pages per pool and rows per block table in the cells: plainly larger than any batch needs, +#: with the batch at scattered table rows so the table row is never the batch index. +_POOL_PAGES = 4096 +#: (name, [(cached, new) per request], max_seq_len, replay_tail, what the step is) +_STEPS = ( + ( + "prefill_mixed", + ((0, 1), (0, 127), (0, 128), (0, 129), (0, 1000), (0, 3000)), + 16384, + None, + "6 fresh prompts of 1..3000 tokens, 4385 rows", + ), + ("prefill_8192", ((0, 8192),), 16384, None, "one 8192-token prefill"), + ( + "chunked_prefill", + ((8192, 2048), (130, 5), (1, 1)), + 16384, + None, + "2048 rows on an 8192-token cached prefix, plus 5- and 1-token chunks", + ), + ( + "decode", + tuple((c, 1) for c in (0, 1, 2, 126, 127, 128, 129, 255, 256, 4095, 8191, 16383)) + + tuple((c, 1) for c in range(300, 16000, 797)), + 16384, + None, + "32 decode rows at positions 0..16383", + ), + ( + "dspark_verify_6", + tuple((c, 6) for c in range(0, 32 * 509, 509)), + 16384, + None, + "32 DSpark verify runs of 6 rows, 192 rows", + ), + ( + "long_context", + tuple((c, 1) for c in (20479, 20480, 32767, 40000, 65534)) + ((65529, 6),), + 65536, + None, + "positions 20479..65534 on 512-page tables", + ), + ( + "replay_tail", + ((872, 128),), + 16384, + 872, + "the last 128 rows of a 1000-token prompt, window cut at the replay start 872", + ), +) +#: (name, SWA offset in pages, compressed offset in pages): offsets 0, and owner-style offsets. +_OFFSETS = (("offset 0", 0, 0), ("owner offsets", 3 * _POOL_PAGES, 2 * _POOL_PAGES)) + + +class Dsv4LocalToGlobalIndices(OpWrapper): + """Translate each query token's local SWA and compressed indices into pool rows. + + For token `t`, table row `r = req_id[t]`, and a segment `(table, local, tpb, offset)`:: + + map(l) = offset + table[r, l // tpb] * tpb + l % tpb if l >= 0 and l // tpb < table width + = -1 otherwise + out[t, :W_swa] = map_swa(swa_local_indices[t]) + out[t, W_swa:] = map_compressed(compressed_local_indices[t]) + + The compressed segment is present exactly when `compressed_block_table` and + `compressed_local_indices` are both given. Fusion boundary: address translation only; the + window construction, the index top-k, the pool writes and the attention call are the + caller's. Returns a fresh contiguous int32 `[T, W_swa + W_compressed]` tensor. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"{name} r{ratio} {offset_name}: {what}", + spec=dict( + ratio=ratio, + seqs=seqs, + max_seq_len=max_seq_len, + replay_tail=replay_tail, + pool_pages=_POOL_PAGES, + table_rows=64, + swa_tokens_per_block=_SWA_TPB, + swa_row_offset=swa_pages * _SWA_TPB, + compressed_tokens_per_block=_SWA_TPB // ratio if ratio else 0, + compressed_row_offset=comp_pages * _SWA_TPB // ratio if ratio else 0, + ), + ) + for name, seqs, max_seq_len, replay_tail, what in _STEPS + for ratio in (0, 1, 2) + for offset_name, swa_pages, comp_pages in _OFFSETS + ) + + note = """ + Column order is preserved: the op does not sort, compact or de-duplicate, + so the valid-first order the attention kernel needs is whatever the local + indices already have. Only a negative local index or one whose block is + past the table's width becomes -1. The op knows no sequence lengths: the + caller masks invisible positions, and a page id of -1 in a table is used + as page -1. + + DeepSeek-V4.1 calls it once per layer per step with offsets 0: ratio 0 + window-only (`[T, 128]`, 128-row SWA pages); ratios 2 and 1 add the + layer's `[T, 512]` index top-k through the compress block table at + `128 // ratio` entries per page (`[T, 640]`). Under decoder bounded replay + the target also makes a window-only slot call whose local table is + `[row position, -1, ..., -1]`. Under native KV storage the target + may pass a synthetic compressed table `arange(M * W).view(M, W)` so entry + `e` of request `b` names row `b * W * tpb + e` of its compact read buffer; + `req_id` must then be each token's request position. The output is + `sparse_indices` for trtllm_sparse_mla_dsv4 when each offset is the row + where page 0 of this layer's buffer starts in the pool tensor handed to + that call; it reads an attended -1 as a zero row, so the -1s must lie + outside the lengths it is given. + + Metadata consumed: none; every input is an argument. + + Preconditions: the op raises ValueError before launching on non-int32 + tensors, a `W_swa` or `W_compressed` that is not a power of two, a + compressed table without its indices or the reverse, a `req_id` that is + not a contiguous 1-D tensor, a non-positive compressed page size, and an + offset outside `[0, 2**31)`. Not checked: `req_id[t]` below both tables' + row counts (else it reads past the table), and every reachable page id a + real page. Every result must fit int32, `offset + page * tpb + tpb - 1 < + 2**31`: one page past the edge comes out wrapped (`-2**31`) with no error. + Block tables and index tensors are read through their strides; `req_id` + needs stride 1. `T = 0` returns an empty `[0, W]` tensor without a launch. + + Widths, page sizes, the compressed segment's presence, `enable_pdl` and + each tensor's 16-byte alignment are Triton specializations (an unaligned + table view compiles one more variant); `T` is only the grid. Make one + eager call of each form, on the tensors a capture will use, before + CUDA-graph capture. `enable_pdl=True` waits on the preceding grid before + reading and releases dependents after its stores; results are identical + either way, and the consumer must honour PDL ordering. + + The in-tree V4 launcher runs the same kernel but compiles the compressed + segment only for `compress_ratio > 1`: a V4.1 ratio-1 layer through it + gets compressed columns that are never written. This op takes the page + size directly and passes strided inputs without a copy. + + Beyond the cells, the test certifies on engine-size tables: every prefix + `T` = 1..32768, capture at `T = G` and `6 G` for G in 1..32, 40..128 step + 8, 192 and 256, random mixed steps up to 256 requests and 16384 rows, the + dumped `(T, ratio)` cells, page ids up to the int32 edge (2**24 pages of + 128 rows, 2**25 of 64), generation positions up to `Lm + 5` for + `max_seq_len` 16384..1048576 with and without DSpark (Lm = L + 14 with + DSpark), unaligned table views, the slot call, and the native-storage synthetic + table up to 256 requests at every width; plus wrong-variant controls. + + Not certified: widths other than 128 / 512, page sizes other than 128 / + 64; results at or above 2**31; more than 32768 rows or 256 requests per + call; block-table widths other than 129, 132 and the twelve + `ceil(Lm / 128) + 3` widths 131..8196 (plus 2, 4, 128, 512); a context + position at or past L or a generation position past Lm + 5; CUDA-graph + capture at any other `T`, or for the native-storage table at an `(M, W)` pair + not run; the native-storage table above 256 rows or with any other `req_id`; + tracing through torch.compile. + """ + + def raw_call( + self, + req_id: torch.Tensor, + swa_block_table: torch.Tensor, + swa_local_indices: torch.Tensor, + swa_tokens_per_block: int, + swa_row_offset: int, + compressed_block_table: Optional[torch.Tensor] = None, + compressed_local_indices: Optional[torch.Tensor] = None, + compressed_tokens_per_block: int = 0, + compressed_row_offset: int = 0, + enable_pdl: bool = False, + ) -> torch.Tensor: + return torch.ops.trtllm.dsv4_local_to_global_indices( + req_id, + swa_block_table, + swa_local_indices, + swa_tokens_per_block, + swa_row_offset, + compressed_block_table, + compressed_local_indices, + compressed_tokens_per_block, + compressed_row_offset, + enable_pdl, + ) + + def reference( + self, + req_id: torch.Tensor, + swa_block_table: torch.Tensor, + swa_local_indices: torch.Tensor, + swa_tokens_per_block: int, + swa_row_offset: int, + compressed_block_table: Optional[torch.Tensor] = None, + compressed_local_indices: Optional[torch.Tensor] = None, + compressed_tokens_per_block: int = 0, + compressed_row_offset: int = 0, + enable_pdl: bool = False, + ) -> torch.Tensor: + """Pool rows by torch table lookup, in int64 and never cast. + + A row at or past 2**31, which the kernel's int32 arithmetic wraps, stays visible as + its true value rather than being matched by an equally wrapped cast. `enable_pdl` is + ignored (it changes the launch, not the table). + """ + + def segment(table, local, tpb, offset): + width = table.shape[1] + blk = torch.div(local, tpb, rounding_mode="floor").long() + valid = (local >= 0) & (blk < width) + page = table[req_id.long()[:, None], blk.clamp(0, width - 1)] + rows = offset + page.long() * tpb + torch.remainder(local, tpb).long() + return torch.where(valid, rows, -1) + + out = segment(swa_block_table, swa_local_indices, swa_tokens_per_block, swa_row_offset) + if compressed_block_table is not None and compressed_local_indices is not None: + comp = segment( + compressed_block_table, + compressed_local_indices, + compressed_tokens_per_block, + compressed_row_offset, + ) + out = torch.cat([out, comp], dim=1) + return out + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Exact: every entry equals the int64 reference. + + Address translation has nothing to round, so any differing entry is a wrong row. The + controls (window off by one, the other ratio's page size, the neighbouring table row, + the two offsets swapped, page ids truncated to 18 bits, one page past the int32 edge) + each differ on hundreds to hundreds of thousands of entries. + """ + assert out.dtype == torch.int32, out.dtype + assert ref.dtype == torch.int64, ref.dtype + assert out.shape == ref.shape, (out.shape, ref.shape) + bad = (out.long() != ref).sum().item() + assert bad == 0, f"{bad} of {out.numel()} entries differ" + + +dsv4_local_to_global_indices = Dsv4LocalToGlobalIndices() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/flashinfer_top_k_ragged_deterministic.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/flashinfer_top_k_ragged_deterministic.py new file mode 100644 index 000000000000..eb7cc9fa53b4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/flashinfer_top_k_ragged_deterministic.py @@ -0,0 +1,353 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Deterministic per-row top-k over each row's own column range (FlashInfer ragged top-k). + +The callable is the FlashInfer Python API ``flashinfer.top_k_ragged_transform``, with +``deterministic=True``, ``tie_break=TopKTieBreak.SMALL`` and ``dsa_graph_safe=True`` pinned: on +that path one call launches FlashInfer's one-block-per-row ``FilteredTopKUnifiedKernel`` (scalar +loads) and then its ``FinalizeTopKIndicesKernel``, which sorts each row's selection into +ascending order. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +from flashinfer import top_k_ragged_transform +from flashinfer.topk import TopKTieBreak + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_K = 512 # the checkpoint's index_topk + + +def _engine_widths() -> tuple[tuple[int, int], ...]: + """(W, compress ratio) of every static generation width: the 8192-token envelope's eleven, then + for each long max_seq_len L the DSpark eager `(L + 14) // r` and graph `(L + 19) // r` widths and + the DSpark-off `L // r`.""" + widths = [ + (8192, 2), + (16384, 1), + (32768, 1), + (8194, 2), + (16386, 2), + (16389, 1), + (32773, 1), + (8199, 2), + (8201, 2), + (16398, 1), + (16403, 1), + ] + for length in _MAX_SEQ_LENS[1:]: + for w, r in ( + (length + 14, 1), + (length + 19, 1), + ((length + 14) // 2, 2), + ((length + 19) // 2, 2), + (length, 1), + (length // 2, 2), + ): + if (w, r) not in widths: + widths.append((w, r)) + return tuple(widths) + + +#: max_seq_len values of the envelope; with DSpark the cache manager's length is L + 14. +_MAX_SEQ_LENS = (16384, 32768, 65536, 131072, 262144, 1048576) +_ENGINE_WIDTHS = _engine_widths() +#: The generation rows' scores: the paged-logits view at every width, and the re-index layers' +#: candidate-masked contiguous copy at the ratio-1 widths where candidate blocks can bind. +_GEN_FORMS = tuple((w, r, "paged_view") for w, r in _ENGINE_WIDTHS) + tuple( + (w, r, "candidate_masked") for w, r in _ENGINE_WIDTHS if r == 1 and w > 16384 +) +#: The last 16384-token chunk of one request at each manager length Lm, both ratios, in the +#: prefill-logits layout, and at ratio 1 above 16384 positions as the candidate-masked copy. +_FAR_END = tuple((n + d, r) for n in _MAX_SEQ_LENS for d in (14, 0) for r in (1, 2)) +_FAR_END_FORMS = tuple((lm, r, "mqa_view") for lm, r in _FAR_END) + tuple( + (lm, 1, "candidate_masked") for lm, r in _FAR_END if r == 1 and lm > 16384 +) +#: Long-context steps the target runs in row tiles, as (cached, new) requests in batch order. +_TILE_STEPS = ( + ("far_Lm1048590", ((1048590 - 16384, 16384),)), + ("far_Lm1048576", ((1048576 - 16384, 16384),)), + ("s1048219", ((1048219 - 8192, 8192),)), + ("s16384", ((0, 16384),)), + ("mixed_long", ((262158 - 4000, 4000), (131086 - 8000, 8000), (0, 4000), (65550 - 384, 384))), + ("two_1M", ((1048590 - 8192, 8192),) * 2), + ("b256_1M", ((1048590 - 64, 64),) * 256), +) +_TILE_MASKED = ("far_Lm1048590", "s1048219", "mixed_long", "two_1M") +#: Where the K-th value is tied: (rows, W, ratio, layout). +_TIE_FORMS = ( + ("decode", 16384, 1, "paged_view"), + ("decode", 16403, 1, "candidate_masked"), + ("decode", 32768, 1, "paged_view"), + ("decode", 32773, 1, "candidate_masked"), + ("decode", 8201, 2, "paged_view"), + ("prefill", 8192, 1, "mqa_view"), + ("prefill", 4096, 2, "mqa_view"), + ("decode", 1048595, 1, "paged_view"), + ("decode", 1048595, 1, "candidate_masked"), + ("decode", 524297, 2, "paged_view"), + ("prefill", 131072, 1, "mqa_view"), +) +_EVERY_LEN_LM = 262158 +_EVERY_LEN_CHUNKS = -(-_EVERY_LEN_LM // 16384) + + +class FlashinferTopKRaggedDeterministic(OpWrapper): + """Return a fresh int32 ``[N, k]`` of each row's top-k columns, in ascending column order. + + For row ``i`` with ``s = row_starts[i]`` (0 when None), ``L = lengths[i]``, ``o = offsets[i]``: + + - ``L <= 0``: all -1; + - ``0 < L <= k``: ``o, o + 1, ..., o + L - 1`` then -1 (the scores are not read); + - ``L > k``: the ``k`` columns ``j < L`` with the largest ``scores[i, s + j]`` (fp32 bit-pattern + order: -0.0 below +0.0; -inf entries tie), ties at the k-th value to the smaller ``j``, + sorted ascending, plus ``o``. No -1. + + Fusion boundary: selection, the ascending sort and adding ``offsets`` are inside the call. + Producing the scores (masking applied as -inf beforehand), the lengths and column starts, and + turning window-local indices into pool slots are the caller's. Returns the new tensor; the + caller copies it into its rows of the index buffer when a step has both kinds of rows. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + ( + Cell( + why="row semantics: lengths -3..4096, offsets and row_starts, 96 rows", + spec=dict(kind="semantics"), + ), + ) + + tuple( + Cell( + why=f"gen W={w} r={r} {layout}: 1536 verify rows, every prefix, next_n 1", + spec=dict(kind="gen", width=w, ratio=r, layout=layout), + ) + for w, r, layout in _GEN_FORMS + ) + + tuple( + Cell( + why=f"ctx every length r={r} chunk={c}: 16384 rows of a 262158-token request", + spec=dict(kind="every_len", ratio=r, chunk=c), + ) + for r in (1, 2) + for c in range(_EVERY_LEN_CHUNKS) + ) + + tuple( + Cell( + why=f"ctx widest step r={r}: 32 requests x 256 rows, the row past 2^{33 - r} elements", + spec=dict(kind="widest", ratio=r), + ) + for r in (1, 2) + ) + + tuple( + Cell( + why=f"ctx 24K passkey step {form}: 7674 rows x 24058, two requests", + spec=dict(kind="long_ctx", layout=form), + ) + for form in ("mqa_view", "candidate_masked") + ) + + tuple( + Cell( + why=f"ctx far end Lm={lm} r={r} {form}: the last 16384-token chunk", + spec=dict(kind="far_end", lm=lm, ratio=r, layout=form), + ) + for lm, r, form in _FAR_END_FORMS + ) + + tuple( + Cell( + why=f"ctx tiles {name} r={r} {form}: every tile as the target calls it", + spec=dict(kind="tiles", name=name, reqs=reqs, ratio=r, layout=form), + ) + for form, ratios, steps in ( + ("mqa_view", (1, 2), _TILE_STEPS), + ("candidate_masked", (1,), tuple(s for s in _TILE_STEPS if s[0] in _TILE_MASKED)), + ) + for name, reqs in steps + for r in ratios + ) + + tuple( + Cell( + why=f"ties {kind} {rows} W={w} r={r} {layout}: the K-th value straddles the cut", + spec=dict(kind="ties", rows=rows, width=w, ratio=r, layout=layout, scores=kind), + ) + for kind in ("bf16", "levels", "zero_plateau", "tie_at_kth") + for rows, w, r, layout in _TIE_FORMS + ) + + ( + Cell( + why="signed zero: -0.0 ranks below +0.0, not a tie", + spec=dict(kind="signed_zero"), + ), + Cell( + why="fewer than 512 finite scores in a selecting row, the rest -inf", + spec=dict(kind="few_finite"), + ), + ) + ) + + note = """ + Call sites. Context rows: `row_starts = ks`, `lengths = ke - ks` (the metadata's + `compressed_visible_lens_cuda[r]`, `(p + 1) // r`), `offsets` a `[:rows]` view of a zero + buffer, `scores` the prefill-logits view `[rows, S]` (row stride `align(S + 256, 256)`), or at + re-index layers when `S > 16384` the contiguous `torch.where(candidate_mask, logits, -inf)` + copy, whole or per row tile (a tile of at most 2^28 padded logits, or one row) with `[a:b]` + views. Generation rows: `row_starts = None`, one length per row (`next_n` 1 and 6 alike), the + paged-logits view `[G * next_n, W]` (row stride `ceil256(W)`) or the masked copy. Cells cover + the 40 static generation widths (8192-token envelope and L = 32768 .. 1048576, DSpark on and + off, ratio 1 and 2) in both layouts at every row count 1..1536, and context steps of up to + 16384 rows and 256 requests with rows seeing up to 1,048,590 positions, ties straddling the cut + in every tie cell (up to 357,542 columns tied with a row's K-th score). + + Deterministic by construction: the result is a function of the in-range scores, `lengths`, + `row_starts`, `offsets` and `k` alone -- not of the call, the process, the GPU, the row count, + the row stride or a CUDA-graph replay. Certified bitwise across 5 calls and 3 fresh processes on + each of 2 GPUs, captured replays included. The knobs are pinned for this: a tie-break other than + NONE sends every shape to the filtered kernel, `deterministic=True` adds the ascending finalize, + and `dsa_graph_safe=True` pins scalar loads (with it off the kernel picks a vector width from + the row stride alone and faulted with a misaligned address on a 4-byte-aligned view). One call + launches `FilteredTopKUnifiedKernel` and `FinalizeTopKIndicesKernel`, once each, at every shape. + + Metadata consumed: none; the target passes metadata views as arguments. + + Preconditions: `scores` fp32 (fp64 raises; bf16 is accepted but not certified), contiguous or a + row-padded view (unit column stride, row stride > W) whose last row's padding lies inside the + storage (otherwise the view raises); a transposed view raises. `row_starts[i] + lengths[i] <= W` + is not checked: a longer range reads the next row's columns, or the padding. `k <= 2048` + (4096 raises). A short `row_starts` and a strided `lengths` raise. Columns outside each row's + range are never read (NaN there is certified). The first call in a process loads FlashInfer's + JIT-compiled top-k module; a first call inside a CUDA-graph capture is certified. + + Measured on sm_103: deterministic costs time. Generation rows take 0.99-2.55x the CuTe DSL decode + top-k's time, context rows 1.68-4.33x the C++ prefill top-k's; neither alternative fixes the + order or the tie choice. + + Not certified: `k` other than 512; scores other than fp32; more than 16384 context rows or 1536 + generation rows in one call; more than 256 requests or a request longer than 1,048,590 tokens; + an untiled context call above 16384 x 1,049,088 fp32 logits (random steps hold at most 2^32); + generation widths other than the 40; CUDA-graph capture at row counts other than G and 6G for + G in 1..32, 40..128 step 8, 192, 256, or of context rows; NaN inside a row's range; `S = 0` or + zero rows. + """ + + def raw_call( + self, + scores: torch.Tensor, + offsets: torch.Tensor, + lengths: torch.Tensor, + row_starts: Optional[torch.Tensor], + k: int, + ) -> torch.Tensor: + x = scores + if ( + not scores.is_contiguous() + and scores.stride(1) == 1 + and scores.stride(0) > scores.size(1) + ): + # A row-padded view (the logits kernels' outputs). The op takes only contiguous input, + # so hand it the padded rows as they are in memory (a view: no kernel, no copy); + # row_starts and lengths keep every read inside the view's own columns. + x = scores.as_strided((scores.size(0), scores.stride(0)), (scores.stride(0), 1)) + return top_k_ragged_transform( + x, + offsets, + lengths, + k, + deterministic=True, + tie_break=TopKTieBreak.SMALL, + dsa_graph_safe=True, + row_starts=row_starts, + ) + + def is_valid( + self, + scores: torch.Tensor, + offsets: torch.Tensor, + lengths: torch.Tensor, + row_starts: Optional[torch.Tensor], + k: int, + ) -> None: + # The op reads these three through an int32 pointer and never compares the lengths of + # offsets / lengths with the row count: an int64 tensor or a short one gives wrong rows + # without an error. Measured on 8 rows: int64 lengths, offsets or row_starts, and a + # 4-entry view of lengths or of offsets (the op read the next 4 entries), each gave 4 of 8 + # rows wrong with no error. + assert offsets.dtype == torch.int32 and lengths.dtype == torch.int32, ( + "offsets and lengths must be int32" + ) + assert row_starts is None or row_starts.dtype == torch.int32, "row_starts must be int32" + assert offsets.numel() == scores.size(0) and lengths.numel() == scores.size(0), ( + "offsets and lengths must hold one entry per row of scores" + ) + + def reference( + self, + scores: torch.Tensor, + offsets: torch.Tensor, + lengths: torch.Tensor, + row_starts: Optional[torch.Tensor], + k: int, + ) -> torch.Tensor: + """Each row's top-k from torch, in the op's order, as int32 ``[N, k]``. + + Rows with ``L <= k`` are ``0 .. L - 1`` then -1. A selecting row gathers its own range + ``[s, s + L)`` in fp64 with -0.0 moved to -1e-300 (so it ranks below +0.0 and above every + negative fp32 value, the bit-pattern order the op ranks by), takes the first ``k`` of a + stable descending sort (equal scores keep column order: ties to the smaller column) and + sorts them ascending; ``offsets`` is then added to every valid entry. Selecting rows are + ranked a chunk at a time (at most 2^26 gathered scores), so the reference runs at the + widths the cells reach; columns outside a row's range are never read. + """ + n, w = scores.shape + dev = scores.device + lens = lengths.long().clamp(min=0) + pos = torch.arange(k, device=dev) + idx = torch.where(pos[None, :] < lens[:, None], pos[None, :], -1) + starts = torch.zeros_like(lens) if row_starts is None else row_starts.long() + rows = (lens > k).nonzero().flatten().tolist() + lens_host = lens.tolist() + a = 0 + while a < len(rows): + b, widest = a, 1 + while b < len(rows): + cand = max(widest, lens_host[rows[b]]) + if b > a and (b + 1 - a) * cand > 1 << 26: + break + widest, b = cand, b + 1 + rr = torch.tensor(rows[a:b], device=dev) + col = torch.arange(widest, device=dev)[None, :] + cols = starts[rr][:, None] + col + live = (col < lens[rr][:, None]) & (cols >= 0) & (cols < w) + x = scores[rr[:, None], cols.clamp(0, w - 1)].double() + x = torch.where((x == 0) & torch.signbit(x), torch.full_like(x, -1e-300), x) + x = x.masked_fill(~live, float("-inf")) + order = torch.sort(x, dim=1, descending=True, stable=True).indices[:, :k] + idx[rr] = order.sort(dim=1).values + del x, order + a = b + return torch.where(idx >= 0, idx + offsets.long()[:, None], -1).to(torch.int32) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Exact, element for element: the selected set, its ascending order and the -1 padding. + + The op is deterministic by construction and the reference ranks by the same order with the + same tie rule, so there is nothing to round: any differing entry is a wrong selection. + Controls measured on sm_103: lengths one past each row's own read one out-of-range column + and change 1534 of 1536 rows; ties to the larger column change exactly the rows whose K-th + value straddles the cut; the non-deterministic mode changes hundreds to thousands of rows + from call to call; one request's `ks` shifted by one column is seen in every step. + """ + assert out.dtype == ref.dtype == torch.int32, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + if not torch.equal(out, ref): + bad = (out != ref).any(dim=1).nonzero().flatten() + raise AssertionError( + f"{bad.numel()} of {out.size(0)} rows differ from the reference, first {int(bad[0])}" + ) + + +flashinfer_top_k_ragged_deterministic = FlashinferTopKRaggedDeterministic() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/indexer_k_cache_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/indexer_k_cache_gather.py new file mode 100644 index 000000000000..87152cda6342 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/indexer_k_cache_gather.py @@ -0,0 +1,243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gather index-K rows out of a paged indexer-K pool into contiguous tensors, via the trtllm op. + +The op is `torch.ops.trtllm.indexer_k_cache_gather_op`, registered from C++ +(`cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp`, kernel in +`cpp/tensorrt_llm/kernels/indexerKCacheGather.cu`). +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_SCALE_BYTES = 4 # one 4-byte word of four UE8M0 exponents per key + +#: One request's whole index K, gathered at once: (block_kv, S). +_ROUNDTRIP = ( + (64, 1), + (64, 5), + (64, 4096), + (128, 1), + (128, 333), + (128, 8192), + (128, 32768), + (256, 3000), +) + +#: A packed prefill batch and the chunk loop's runs over its slot mappings, as +#: (k_token_start, num_tokens): each request on its own, the whole batch, a +#: mid-request run, single rows at a request's end and at the batch's end, and a +#: 4000-row run at an offset. +_PACKED_LENS = [700, 1, 2500, 37, 4096] +_PACKED_CHUNKS = [ + (0, 700), + (700, 1), + (701, 2500), + (3201, 37), + (3238, 4096), + (0, 7334), + (350, 1000), + (3200, 1), + (7333, 1), + (3255, 4000), +] + + +class IndexerKCacheGather(OpWrapper): + """Copy rows `[k_token_start, k_token_start + num_tokens)` of a paged index-K pool into two + fresh contiguous tensors. + + One call, one kernel launch, no arithmetic. For output row `i`, with + `r = k_token_start + i`: + + k[i, :] = k_cache.flat[slot_mapping_fp8[r] : + head_dim] + k_scale[i, 0] = k_cache.flat[slot_mapping_scale[r] : + 4] (as one 4-byte word) + + where `.flat` is `k_cache`'s logical `[pages, block_kv, 1, 68]` shape + unravelled innermost-first and addressed through its strides (a strided + owner view is read correctly). A row whose data or scale slot is negative is + skipped in both outputs: its bytes are whatever the fresh allocation held. + The pool is only read. + + Returns `(k, k_scale)`: `k` is `float8_e4m3fn [num_tokens, head_dim]` and + `k_scale` is `float32 [num_tokens, 1]`, both fresh and contiguous and both + only carrying raw bytes. On the MXFP4 pool reinterpret them (no copy) before + DeepGEMM: `k.view(torch.int8)` (packed E2M1, element `2j` in the low nibble + of byte `j`) and `k_scale.view(torch.int32).reshape(-1)` (byte `b` of the + little-endian word scales elements `32b .. 32b+31`). + + Fusion boundary: the gather only. The caller owns quantizing K and writing + the pool, the slot mappings, the int8 / int32 reinterpretation, the logits + call, and masking of skipped rows. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"one request S={s}, block_kv {bk}: its whole index K", + spec=dict(block_kv=bk, lens=[s], chunks=[(0, s)]), + ) + for bk, s in _ROUNDTRIP + ) + + tuple( + Cell( + why=f"prefill chunks at k_token_start, block_kv {bk}", + spec=dict(block_kv=bk, lens=_PACKED_LENS, chunks=_PACKED_CHUNKS), + ) + for bk in (64, 128) + ) + + tuple( + Cell( + why=f"strided owner view, block_kv {bk}: page stride 4 pages", + spec=dict(block_kv=bk, lens=[1500, 3], chunks=[(0, 1503)], layers=4), + ) + for bk in (64, 128) + ) + + ( + Cell( + why="num_tokens 0: empty outputs", + spec=dict(block_kv=64, lens=[10], chunks=[(0, 0)]), + ), + ) + ) + + note = """ + Cells: index K is bf16 [S, 128] with per-row magnitudes 2^-12..2^12, + quantized by `torch.ops.trtllm.fp4_quantize(x, None, 32, True, False)` + (MXFP4, UE8M0, linear scale layout) and written with torch indexing into a + uint8 `[pages, block_kv, 1, 68]` pool holding 3x the pages the batch needs + plus 8 over random bytes, pages drawn by a random permutation. The gathered + bytes are the quantizer's, and their decode is the checkpoint's + `fp4_act_quant(k, 32, True)` value set exactly. + + Metadata consumed: none. The caller builds the mappings; for the page-blocked + layout (`block_kv x 64` code bytes, then `block_kv x 4` scale bytes per page), + position `p` of a request with page-table row `table` is + page = table[p // block_kv]; slot = p % block_kv + slot_mapping_fp8 = page * block_kv * 68 + slot * 64 + slot_mapping_scale = page * block_kv * 68 + block_kv * 64 + slot * 4 + with the logical page stride `block_kv * 68` even when the view's `stride(0)` + is larger. The op cannot tell a right mapping from a wrong one: another + table, page + 1, slot + 1, key-interleaved addressing, a page stride without + the scale region, another owner's view and 32-bit-wrapped offsets are all + accepted and measured wrong on every row. + + Preconditions: the innermost dim has stride 1 and every index is a multiple + of 4 inside the pool -- neither is checked (read, not probed). Mappings are + 1-D contiguous int64 CUDA tensors of at least `k_token_start + num_tokens` + entries. Mask or drop skipped rows before use. The op rejects loudly a + head_dim other than 64 or 128, a non-4-D, non-uint8 or `size(2) != 1` pool, + and int32, CPU, 2-D, non-contiguous or too-short mappings. + + Row counts: every count up to 268,439,040 (ratio 1) / 134,219,520 (ratio 2), + the per-step maximum of up to 256 context requests at the longest request, + the rows spread over the whole pool; exhaustively 1..1,048,590 / 1..524,295. + One precompiled kernel serves every count. Counts above 2^25 need the + kernel's output-row offset formed in 64 bits + (`static_cast(token_idx)` in `indexerKCacheGather.cu`): with a + 32-bit product every row from 2^25 on is silently misplaced and can land + outside the output. Pools: any size up to 6,539,977 blocks (the largest any + DeepSeek-V4.1 cache manager can have on a GB300), the far end past 2^35 + bytes included. + + CUDA graphs: not captured. The op allocates its outputs, and the target + calls it only in steps with context requests, which are never captured. + + Not certified: other architectures; the FP8 mode (head_dim 128 over a + 132-byte entry); `block_kv` other than 64, 128 and 256; more rows than the + per-step maxima above; `k_token_start + num_tokens` near 2^31 (rows are + indexed in int32); pools above 6,539,977 blocks; out-of-range, misaligned or + overlapping indices; CUDA-graph capture. + """ + + def raw_call( + self, + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + k_token_start: int, + num_tokens: int, + head_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.indexer_k_cache_gather_op( + k_cache, slot_mapping_fp8, slot_mapping_scale, k_token_start, num_tokens, head_dim + ) + + def is_valid( + self, + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + k_token_start: int, + num_tokens: int, + head_dim: int, + ) -> None: + # Pure-metadata guard: a head_dim that is not the pool's per-token payload width is + # accepted by the op and silently reads the neighbouring key's bytes -- observed with + # head_dim=128 on the 68-byte MXFP4 pool (bytes 64..127 of each row were the next key's + # codes), never an error. + assert k_cache.dim() != 4 or k_cache.size(3) == head_dim + 4, ( + f"head_dim ({head_dim}) must equal k_cache.size(3) - 4 ({k_cache.size(3) - 4}); " + "the op would silently read past each row's payload" + ) + + def reference( + self, + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + k_token_start: int, + num_tokens: int, + head_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """The two outputs, gathered with native torch indexing from the pool's logical bytes. + + A skipped row (either slot negative) is left zero here; the op leaves it + unwritten, so only cells without skipped rows compare whole outputs. + Nothing is mutated. + """ + flat = k_cache.contiguous().view(-1) + rows = slice(k_token_start, k_token_start + num_tokens) + d_off = slot_mapping_fp8[rows].long() + s_off = slot_mapping_scale[rows].long() + keep = (d_off >= 0) & (s_off >= 0) + dev = flat.device + k = torch.zeros(num_tokens, head_dim, dtype=torch.uint8, device=dev) + k_scale = torch.zeros(num_tokens, _SCALE_BYTES, dtype=torch.uint8, device=dev) + k[keep] = flat[d_off[keep].unsqueeze(1) + torch.arange(head_dim, device=dev)] + k_scale[keep] = flat[s_off[keep].unsqueeze(1) + torch.arange(_SCALE_BYTES, device=dev)] + return k.view(torch.float8_e4m3fn), k_scale.view(torch.float32) + + def compare( + self, + out: tuple[torch.Tensor, torch.Tensor], + ref: tuple[torch.Tensor, torch.Tensor], + ) -> None: + """Bitwise on both outputs' bytes. + + The op is a byte copy, so there is nothing to round: any byte that + differs is a row read from the wrong place. + """ + (k, k_scale), (ref_k, ref_scale) = out, ref + assert k.dtype == ref_k.dtype == torch.float8_e4m3fn, (k.dtype, ref_k.dtype) + assert k_scale.dtype == ref_scale.dtype == torch.float32, (k_scale.dtype, ref_scale.dtype) + assert k.shape == ref_k.shape and k_scale.shape == ref_scale.shape, ( + k.shape, + ref_k.shape, + k_scale.shape, + ref_scale.shape, + ) + bad_k = (k.view(torch.uint8) != ref_k.view(torch.uint8)).any(-1).sum().item() + bad_s = (k_scale.view(torch.uint8) != ref_scale.view(torch.uint8)).any(-1).sum().item() + assert bad_k == 0 and bad_s == 0, ( + f"{bad_k} code rows and {bad_s} scale words of {k.size(0)} rows differ" + ) + + +indexer_k_cache_gather = IndexerKCacheGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/indexer_k_cache_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/indexer_k_cache_scatter.py new file mode 100644 index 000000000000..259904398bdb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/indexer_k_cache_scatter.py @@ -0,0 +1,227 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Paged index-K cache write, via the trtllm op. + +The op is `torch.ops.trtllm.indexer_k_cache_scatter_op`, registered from C++ +(`cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp`, kernel in +`cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu`). +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Page sizes in compressed entries at tokens_per_block 128: 64 is the ratio-2 +#: page, 128 the ratio-1 page. +_BLOCK_KV = (64, 128) + + +def _steps(block_kv: int) -> tuple[tuple[str, dict], ...]: + """(why, spec) of the DeepSeek-V4.1 steps the target writes, at one page size. + + `cached` / `new` are compressed entries per request; `cached_below` / + `new_upto` with `requests` draw them at random (seeded by the test). + """ + r = 128 // block_kv + return ( + ("decode B=1, one entry mid-page", dict(cached=[3 * block_kv + 5], new=[1])), + ( + "decode B=32, one entry each, cached < 8 pages", + dict(requests=32, cached_below=8 * block_kv, new=[1] * 32), + ), + ( + f"DSpark verify B=32, 0..{6 // r} entries each", + dict(requests=32, cached_below=8 * block_kv, new_upto=6 // r), + ), + (f"prefill of 8192 tokens, {8192 // r} entries", dict(cached=[0], new=[8192 // r])), + (f"prefill of 16384 tokens, {16384 // r} entries", dict(cached=[0], new=[16384 // r])), + ( + "chunked mixed: mid-page continuation, page edge", + dict(cached=[1003, 0, 64 * 3 - 1], new=[3000, 77, 2]), + ), + ) + + +class IndexerKCacheScatter(OpWrapper): + """Copy rows `0 .. num_tokens-1` of packed index keys and their scale words into a paged pool. + + One call, one kernel launch, no arithmetic. For each token `t`: + + k_cache.flat[slot_mapping_data[t] : + D] = k_data[t, :] # D = k_data.size(1) bytes + k_cache.flat[slot_mapping_scale[t] : + 4] = bytes of k_scale[t, 0] + + where `.flat` is `k_cache`'s logical `[pages, block_kv, 1, 68]` shape + unravelled innermost-first and addressed through its strides (a page-strided + view is written correctly). A token whose data or scale offset is negative is + skipped entirely; every other byte of the pool is left as it was. + + MXFP4 mode: `k_data` uint8 `[T, 64]` (two E2M1 codes per byte, element `2i` + in the low nibble of byte `i`), `k_scale` int32 `[T, 1]` (four UE8M0 bytes, + byte `g` scaling elements `32g .. 32g+31`), 68-byte pool entries. + + Fusion boundary: the copy only. The caller owns the index-K projection, + norm and RoPE, the quantization, the page table, the offsets and the pool's + allocation. Writes `k_cache` in place and returns it. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell(why=f"block_kv {bk}: {why}", spec=dict(block_kv=bk, **spec)) + for bk in _BLOCK_KV + for why, spec in _steps(bk) + ) + + tuple( + Cell( + why=f"skip {which}: -1 offsets on every 5th row skip the whole token", + spec=dict(block_kv=64, cached=[5, 130, 0], new=[40, 17, 64], skip=which), + ) + for which in ("both", "data_only", "scale_only") + ) + + ( + Cell( + why="num_tokens 150 of 256 rows: the buffer tail is never read", + spec=dict(block_kv=128, cached=[7, 300], new=[200, 56], num_tokens=150), + ), + Cell( + why="num_tokens 0: a no-op on a populated pool", + spec=dict(block_kv=128, cached=[7, 300], new=[200, 56], num_tokens=0), + ), + Cell( + why="page-strided pool view, stride(0) = 64 * 68 + 128", + spec=dict(block_kv=64, cached=[10, 200], new=[100, 33], page_pad=128), + ), + ) + ) + + note = """ + Cells: keys are bf16 [T, 128] with a different power-of-two magnitude per + 32-block, quantized by `torch.ops.trtllm.fp4_quantize(k, None, 32, True, + False)`; its data passes as is and its linear 1-D uint8 scales as + `sf.view(torch.int32).view(T, 1)`, a reinterpretation. Pools hold 3x the + pages a batch needs plus 8, pages drawn by a random permutation. + + Metadata consumed: none. The caller builds the offsets; for the page-blocked + layout the DeepGEMM paged indexer-logits kernel reads (`block_kv x 64` code + bytes, then `block_kv x 4` scale bytes per page) they are + slot_mapping_data = page * block_kv * 68 + slot * 64 + slot_mapping_scale = page * block_kv * 68 + block_kv * 64 + slot * 4 + The op cannot tell a right layout from a wrong one: a rolled page table, + key-interleaved offsets, 64-entry pages addressed as 128, byte-reversed or + group-major scale bytes, and another owner's view are all accepted and + measured wrong on every row (byte-reversed scales: 99%). + + Preconditions: pass the KV-source owner's own pool view -- the ratio-2 + owners (layers 2, 8, 14) interleave their pages in one allocation, so a page + id means something only through its own owner's view. Offsets must be in + range and multiples of 4: the kernel neither bounds-checks nor aligns them + (read, not probed). Two tokens given overlapping offsets race. A `-1` + page-table column yields negative offsets and is skipped, but mask such rows + to `-1` explicitly. The op rejects loudly a non-uint8, non-4-D or + `size(2) != 1` pool, a `k_scale` that is not 2-D with 4-byte elements, + non-int64 slot mappings, a non-contiguous `k_data`, CPU tensors, and data + widths other than 64 and 128. + + CUDA graphs: captures and replays correctly; `num_tokens` is baked into the + graph, so pad with `-1` offsets rather than changing it. Every `num_tokens` + in 1..16384 is certified eagerly at both page sizes; captured at n = G and, + for the DSpark verify step, 6G (ratio 1) and 3G (ratio 2), for G in 1..32, + 40..128 step 8, 192, 256. One precompiled kernel serves every row count. + + Pools: any size up to 6,539,977 blocks (the largest any DeepSeek-V4.1 cache + manager can have on a GB300), the far end past 2^35 bytes included; the + kernel forms 64-bit byte offsets. A caller that wraps offsets to 32 bits + gets silently wrong placement. + + Not certified: other architectures; the FP8 mode (128-byte rows, float32 + scale, 132-byte entries); `block_kv` other than 64 and 128; more than 16384 + rows in one call; pools above 6,539,977 blocks; out-of-range, misaligned or + overlapping offsets; capture at row counts outside the sets above; graph + replay on a live cache manager's own metadata buffers. + """ + + def raw_call( + self, + k_data: torch.Tensor, + k_scale: torch.Tensor, + k_cache: torch.Tensor, + slot_mapping_data: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_tokens: int, + ) -> torch.Tensor: + torch.ops.trtllm.indexer_k_cache_scatter_op( + k_data, k_scale, k_cache, slot_mapping_data, slot_mapping_scale, num_tokens + ) + return k_cache + + def is_valid( + self, + k_data: torch.Tensor, + k_scale: torch.Tensor, + k_cache: torch.Tensor, + slot_mapping_data: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_tokens: int, + ) -> None: + # Both guards are pure metadata and cover inputs the op was observed to accept silently: + # 128-byte FP8-width rows written into the 68-byte-per-entry FP4 pool overwrite the + # neighbouring entries, and a num_tokens larger than the rows or slot entries supplied + # reads past the end of those tensors. + assert k_data.size(-1) + 4 == k_cache.size(-1), ( + "k_data row bytes + 4 scale bytes must equal the pool's per-entry bytes " + "(64 + 4 = 68 for MXFP4, 128 + 4 = 132 for FP8)" + ) + assert num_tokens <= min( + k_data.size(0), k_scale.size(0), slot_mapping_data.numel(), slot_mapping_scale.numel() + ), "num_tokens must not exceed the rows of k_data / k_scale or the slot-mapping lengths" + + def reference( + self, + k_data: torch.Tensor, + k_scale: torch.Tensor, + k_cache: torch.Tensor, + slot_mapping_data: torch.Tensor, + slot_mapping_scale: torch.Tensor, + num_tokens: int, + ) -> torch.Tensor: + """The pool after the call, placed by native torch indexing on flat byte offsets. + + `k_cache` is the pool as it was before the call: the op writes it in + place, so the test passes a snapshot taken before calling it. Nothing is + mutated; the result is a new contiguous tensor of `k_cache`'s logical + shape. Each scale element's bytes are placed in memory order, as the op + copies them. + """ + expected = k_cache.clone(memory_format=torch.contiguous_format) + flat = expected.view(-1) + n = num_tokens + d_off = slot_mapping_data[:n].long() + s_off = slot_mapping_scale[:n].long() + keep = (d_off >= 0) & (s_off >= 0) + data_bytes = k_data.size(1) * k_data.element_size() + scale_bytes = k_scale.size(1) * k_scale.element_size() + data = k_data[:n].contiguous().view(torch.uint8).reshape(n, data_bytes)[keep] + scale = k_scale[:n].contiguous().view(torch.uint8).reshape(n, scale_bytes)[keep] + dev = flat.device + flat[d_off[keep].unsqueeze(1) + torch.arange(data.size(1), device=dev)] = data + flat[s_off[keep].unsqueeze(1) + torch.arange(scale.size(1), device=dev)] = scale + return expected + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise over the whole pool. + + The op is a byte copy, so there is nothing to round: any byte that + differs is a row written to the wrong place, a skipped row written, or a + byte outside the written ranges touched. + """ + assert out.dtype == ref.dtype == torch.uint8, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + bad = (out != ref).sum().item() + assert bad == 0, f"{bad} of {out.numel()} pool bytes differ" + + +indexer_k_cache_scatter = IndexerKCacheScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/mla_rope_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/mla_rope_inplace.py new file mode 100644 index 000000000000..e6e4bf88679d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/mla_rope_inplace.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""In-place RoPE on the trailing rope_dim channels of every head, from a caller-built cos/sin table. + +The op is `torch.ops.trtllm.mla_rope_inplace`, registered by `tensorrt_llm._torch.custom_ops`. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Rows of the largest table the target builds: `max_seq_len + 64` at the envelope's largest +#: per-request length 1048576 with DSpark on (+14). Every position below it is certified. +_TABLE_ROWS = 1048654 +_RATIO0, _COMPRESSED = "ratio0_theta1e4", "compressed_theta160k_yarn" +#: (name, H, head_dim, nope_dim, inverse) of every tensor the checkpoint rotates with a recipe +#: table: q, window kv, index q and index k forward, the attention output inversely. +_TABLE_FORMS = ( + ("q_tp4", 16, 512, 448, False), + ("q_dep4", 64, 512, 448, False), + ("window_kv", 1, 512, 448, False), + ("index_q_tp4", 8, 128, 64, False), + ("index_q_32", 32, 128, 64, False), + ("index_k", 1, 128, 64, False), + ("attn_out_tp4", 16, 512, 448, True), + ("attn_out_dep4", 64, 512, 448, True), +) +#: Row counts and how the test places their positions on the table. +_TABLE_T = { + 1: "the table's last row", + 6: "a consecutive run, as in a DSpark verify step", + 8192: "random over the whole table, 0 and the last row included", +} +#: Patch counts of the DeepSeek-V4.1 vision tower's per-request call: the prototype's +#: whole-tower graph sizes, the MMMU maximum 17037, and the two smallest. +_VIT_P = (1, 2, 1517, 4680, 5325, 8772, 17037) + + +class MlaRopeInplace(OpWrapper): + """Rotate `data[:, :, nope_dim:]` in place by the table rows at `position_ids`. + + With `c, s = cos_sin_cache[position_ids[t], 0], cos_sin_cache[position_ids[t], 1]` + (`s` negated when `inverse`), every head's trailing `rope_dim` channels are rotated pair by + pair: adjacent channels `(2i, 2i+1)` when `is_neox=False` (GPT-J), the two halves + `(i, i + rope_dim/2)` when `is_neox=True`. Computed in fp32, rounded once to `data`'s + dtype. The leading `nope_dim` channels are not touched. + + Fusion boundary: one rotation of one tensor. The projection that produced `data`, the + table (frequencies, YaRN, any mscale), the positions and everything after are the + caller's. Writes `data` in place and returns `data`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=( + f"{name} {table.split('_')[0]} T={t}: [T,{h},{d}] nope {nope}" + f"{', inverse' if inverse else ''}, {where}" + ), + spec=dict( + form="table", + table=table, + rows=_TABLE_ROWS, + T=t, + H=h, + D=d, + nope=nope, + inverse=inverse, + is_neox=False, + ), + ) + for name, h, d, nope, inverse in _TABLE_FORMS + for table in (_RATIO0, _COMPRESSED) + for t, where in _TABLE_T.items() + ) + + tuple( + Cell( + why=( + f"group-first {name} r={ratio}: [{_TABLE_ROWS // ratio},1,{d}] at positions" + " j*r, every group whose first position is in the compressed table" + ), + spec=dict( + form="group_first", + table=_COMPRESSED, + rows=_TABLE_ROWS, + ratio=ratio, + H=1, + D=d, + nope=d - 64, + inverse=False, + is_neox=False, + ), + ) + for ratio in (1, 2) + for name, d in (("compressed_kv", 512), ("index_k", 128)) + ) + + tuple( + Cell( + why=( + f"neox full head T={t} {'inverse' if inverse else 'forward'}: [T,16,64]" + " nope 0, half-split pairs, random angles per row" + ), + spec=dict( + form="random_angles", + T=t, + H=16, + D=64, + nope=0, + inverse=inverse, + is_neox=True, + ), + ) + for t in (1296, 4096) + for inverse in (False, True) + ) + + tuple( + Cell( + why=( + f"vit packed qkv P={p}: [3P,16,64] NeoX, own [P+1] table, positions" + " (t, t, P) so v reads the identity row and comes back unchanged" + ), + spec=dict(form="vit_packed", P=p, H=16, D=64, nope=0, inverse=False, is_neox=True), + ) + for p in _VIT_P + ) + ) + + note = """ + Measured on sm_103 against an fp32 torch rotation: at least 0.99997 of + elements bit-exact, every difference inside bf16's default band, largest + 3.1e-2 on unit-normal inputs (consistent with FMA contraction changing the + final rounding). `inverse=True` undoes a forward rotation to within two + bf16 roundings. + + The table is the caller's: RoPE base, YaRN interpolation and any mscale + live in it; the op neither normalizes nor derives frequencies, and a table + scaled by an mscale scales the rotated channels. The certified tables are + DeepSeek-V4.1's two `[rows, 2, 32]` fp32 recipes, no mscale: ratio-0 (base + 10000, no YaRN; decoder layers 0-1 and the DSpark stages) and compressed + (base 160000, YaRN factor 16, original 65536, beta_fast 32, beta_slow 1; + layers 2-39). The target builds them once per engine with + `max_seq_len + 64` rows; a shorter table is bitwise the first rows of a + longer one. The compressed latent and index k of a ratio-r layer rotate at + their group's first position `j * r`. The ViT forms pass their own table. + + Metadata consumed: none. `position_ids[t]` is only a row index into the + table; the op reads no batch layout or sequence boundary. + + Preconditions: every position in `[0, rows)`. There is no bounds check, so + an out-of-range position (a negative padding position too) is an + out-of-bounds read; the row offset is 32-bit, exact below position 2^25. + Alignment is not checked: `nope_dim` and the head width must be multiples + of 8 elements and `data` must start 16-byte aligned (a contiguous row + slice of a contiguous buffer does), otherwise the kernel faults with a + sticky misaligned-address error that leaves the CUDA context unusable. + The op raises RuntimeError, leaving `data` untouched, on: non-contiguous + `data` (a `q[..., 448:]` tail view), non-int32 positions, data that is not + bf16/fp16, a non-fp32 table, a table not shaped `[P, 2, rope_dim / 2]`, + non-3D `data`, `data.size(1) != num_heads`, `data.size(2) != nope_dim + + rope_dim`, and a `rope_dim` that is odd or not a multiple of 4. `T = 0` is + accepted. + + In place: a caller that needs the pre-rope tensor keeps its own copy (the + DeepSeek-V4.1 indexer reads the compressed latent before its rotation). + For a q + kv pair at the same positions, one FlashInfer + `apply_rope_with_cos_sin_cache_inplace` call measured 1.4-2.4 us faster + than two calls of this op. + + Beyond the cells, the test certifies: every row count 1..32768 for the 21 + core form/table calls (q and the output at H 8/16/32/64, window kv, index + q at 32 heads, index k, compressed latent), the boot-gate row counts, one + `mlaRoPEInplaceKernel` launch per call at every probed T, CUDA-graph + capture of all 21 calls at the 126 row counts `m * G` (m in {1, 3, 5, 6}, + G in 1..32, 40..128 step 8, 192, 256), every envelope engine's own table + with positions at its far end, and the packed ViT form at every row count + 1..51111 and under graph capture at T = 3P for its graph sizes. + + Not certified: fp16 numerics (accepted, not compared); `rope_dim` other + than 64; `nope_dim` other than 0, 64 or 448; heads other than 1, 8, 16, 32 + or 64; positions at or past 1048654; T above 32768 at H > 1, except the + packed ViT form (to 51111, tables up to 17038 rows); half-split pairs with + `nope_dim > 0`; graph capture of the core forms at a T outside the + 126-value union. + """ + + def raw_call( + self, + data: torch.Tensor, + position_ids: torch.Tensor, + cos_sin_cache: torch.Tensor, + num_heads: int, + nope_dim: int, + rope_dim: int, + inverse: bool, + is_neox: bool, + ) -> torch.Tensor: + torch.ops.trtllm.mla_rope_inplace( + data, + position_ids, + cos_sin_cache, + num_heads, + nope_dim, + rope_dim, + inverse, + is_neox, + ) + return data + + def reference( + self, + data: torch.Tensor, + position_ids: torch.Tensor, + cos_sin_cache: torch.Tensor, + num_heads: int, + nope_dim: int, + rope_dim: int, + inverse: bool, + is_neox: bool, + ) -> torch.Tensor: + """`data` with its rope channels rotated, in fp32, rounded once to `data`'s dtype. + + `data` is not mutated: the result is a new tensor. `num_heads` is ignored (it is + `data.size(1)`). + """ + y = data.float().clone() + cos = cos_sin_cache[position_ids.long(), 0][:, None, :] + sin = cos_sin_cache[position_ids.long(), 1][:, None, :] + if inverse: + sin = -sin + tail = y[..., nope_dim : nope_dim + rope_dim] + half = rope_dim // 2 + if is_neox: + x1, x2 = tail[..., :half].clone(), tail[..., half:].clone() + tail[..., :half] = x1 * cos - x2 * sin + tail[..., half:] = x1 * sin + x2 * cos + else: + x1, x2 = tail[..., 0::2].clone(), tail[..., 1::2].clone() + tail[..., 0::2] = x1 * cos - x2 * sin + tail[..., 1::2] = x1 * sin + x2 * cos + return y.to(data.dtype) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """torch's default band, plus at least 0.999 of elements bit-exact. + + Kernel and reference both rotate in fp32 and round once; they differ only where FMA + contraction flips the final rounding, measured on sm_103 at 0.99997 or more of the + elements bit-exact in every probed shape. The fraction gate is on top of the band, not + a loosening of it. The wrong-variant controls (pairing, rotated channels, YaRN, mscale, + table, direction, position off by one) put 0.04-0.24 of the elements outside the band. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + torch.testing.assert_close(out, ref) + bitexact = (out == ref).float().mean().item() + assert bitexact >= 0.999, f"bit-exact fraction {bitexact}" + + +mla_rope_inplace = MlaRopeInplace() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/trtllm_sparse_mla_dsv4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/trtllm_sparse_mla_dsv4.py new file mode 100644 index 000000000000..219f3bb678bf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/trtllm_sparse_mla_dsv4.py @@ -0,0 +1,318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DeepSeek-V4-family sparse MQA over [sliding window | compressed] KV rows with attention sinks. + +The kernel is the FlashInfer Python callable +`flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4`, not a `torch.ops.trtllm` +op. On sm_103 one call launches one trtllm-gen FMHA kernel of the H512 +`DynamicTokenSparse` family (at 64 heads on the split-KV path, followed by +FlashInfer's `fmhaReductionKernel`); the entry pins `backend="trtllm-gen"` and +the `HND` pool layout, which is the path its cells certify. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +from flashinfer.mla import trtllm_batch_decode_sparse_mla_dsv4 + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_HEAD_DIM = 512 +_WINDOW = 128 +#: The gate's relative band: torch's bf16 default. +_RTOL = 1.6e-2 +#: The gate's absolute band per unit of max|V|: one bf16 ulp at the largest V element. +_ATOL_PER_VMAX = 2.0**-8 +#: Query rows per chunk of the reference (bounds its [rows, width, 512] fp32 gather). +_REF_CHUNK = 256 + +#: Per-rank query heads of the target's members: 64 / attention TP -- 16 at tep4 / tp4, +#: 64 under attention DP and on the DSpark drafter, 8 at tep8, 32 at pp2 x tep2 / tp2. +_ALL_HEADS = (8, 16, 32, 64) +#: Decode positions: window counts 1..128, compressed counts 0..512 at ratio 2. +_DECODE_POSITIONS = (0, 3, 60, 126, 127, 128, 500, 1500, 5000) +_DECODE_REQUESTS = tuple((p, 1) for p in _DECODE_POSITIONS) +#: DSpark verify: the last accepted token plus 5 drafts per request. +_VERIFY_REQUESTS = ((0, 6), (70, 6), (125, 6), (2000, 6)) +#: Varlen prefill steps, (requests as (cached, q_len), index width, what the step covers). +_PREFILL_STEPS = ( + (((0, 4096),), 640, "1 x 4096 fresh, a prefill-sized prompt"), + (((0, 8192),), 640, "1 x 8192 fresh, the largest prefill chunk"), + (((0, 1500), (3000, 700), (90, 1)), 640, "1500 fresh + 700 at 3000 + 1 row: chunked"), + (((0, 3000), (200, 5)), 128, "3000 fresh + 5 at 200, window-only layer"), +) + + +class TrtllmSparseMlaDsv4(OpWrapper): + """Sparse MQA (K == V, head_dim 512) of each query row over the pool rows its index row names. + + For query row `r` (request `b`, query `j` of that request's `q_len_b`), head `h`:: + + window_count[r] = min(128, seq_lens[b] - q_len_b + j + 1) + attended[r] = {c : c < sparse_topk_lens[r] and (c >= 128 or c < window_count[r])} + kv(c) = swa rows[idx[r, c]] for c < 128, compressed rows[idx[r, c]] for c >= 128, + an all-zero row where idx[r, c] == -1 + s_c = softmax_scale * (fp32, natural-log units) + out[r, h] = output_scale * sum_c exp(s_c - m) kv(c) / (sum_c exp(s_c - m) + exp(sinks[h] - m)) + + Pool rows are in each pool's own flattened row numbering. The sink joins + the denominator only; duplicates are attended as many times as they are + named. Probabilities enter the PV product as bf16; scores, max and sum are + fp32. + + Fusion boundary: one call does the gather, QK, the softmax with sink and + PV. The caller owns RoPE of q and of the KV rows, the writes into both + pools, the index rows and lengths, and the inverse RoPE / output projection + afterwards. Returns the bf16 output, shaped like `query` (`out` when given). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"decode H={h} topk={topk}: q_len 1, window counts 1..128", + spec=dict(kind="decode", H=h, topk=topk, requests=_DECODE_REQUESTS), + ) + for h in _ALL_HEADS + for topk in (128, 640) + ) + + tuple( + Cell( + why=f"verify H={h}: dense q_len 6 (DSpark verify), 4 requests, topk 640", + spec=dict(kind="verify", H=h, topk=640, requests=_VERIFY_REQUESTS), + ) + for h in _ALL_HEADS + ) + + tuple( + Cell( + why=f"prefill H={h} topk={topk}: {label}", + spec=dict(kind="prefill", H=h, topk=topk, requests=requests), + ) + for h in (8, 16, 32) + for requests, topk, label in _PREFILL_STEPS + ) + + ( + Cell( + why="prefill H=64 topk=640: 1100 fresh + 300 at 2500", + spec=dict(kind="prefill", H=64, topk=640, requests=((0, 1100), (2500, 300))), + ), + Cell( + why="window-only H=16: NaN compressed placeholder never read", + spec=dict(kind="window_only", H=16, topk=128, requests=_DECODE_REQUESTS), + ), + Cell( + why="draft block H=16: 5 non-causal queries, virtual seq_lens", + spec=dict(kind="draft_block", H=16, ring_rows=(3, 77, 128), block=5, width=136), + ), + ) + ) + + note = """ + Metadata consumed: none. Everything the kernel reads is an argument; the + caller builds the index rows (window positions and compressed top-k mapped + to pool rows) and the lengths. + + Preconditions, none of them checked by the call: + - Compact each segment and bound it with the lengths: attended rows occupy + the first `window_count[r]` window columns and the first + `sparse_topk_lens[r] - 128` compressed columns; `-1` is the safe filler + after them. A `-1` inside the counts is an all-zero row (score 0, joins + the denominator), not a mask. + - `window_count` comes from `seq_lens` alone: a caller that needs another + count per row passes virtual `seq_lens` (each query its own q_len-1 + request). `seq_lens = 0` is not an empty window: it attends all 128 + window columns, silently. A floored window passes `kv_len - floor`. + - `128 <= sparse_topk_lens[r] <= sparse_indices.shape[1]`; past the width + the call ends in an illegal memory access that poisons the CUDA context. + Lengths below 128 are outside FlashInfer's contract. + - Every attended index is a real row of its pool (an index past the end + returns garbage). + - Masked columns can still be loaded: the whole 128-column window and the + rest of the last group of 4 that `sparse_topk_lens` cuts into. A NaN or + Inf row named there poisons the output. A loaded `-1` reads one row in + 96..127 of its pool (then discards it), so every pool view a `-1` is + loaded from must have at least 128 rows of mapped memory from its first + row (measured with guard pages; a shorter allocation can fault). The + DSpark block buffer sits at the head of an allocation of >= 128 rows, or + repeats column 132's row in its masked tail instead of `-1`. + - Window-only layers still pass a compressed pool; at index width 128 it + is never read. + - Index width a multiple of 4 and >= 128, int32; bf16 query and pools; + fp32 sinks of exactly [H], natural-log units; `softmax_scale` the plain + scale (no log2(e) folded in). FlashInfer raises ValueError on these. + - `workspace_buffer`: the target's one 128 MiB buffer, zeroed once and + shared by every call. The call checks no size and uses at most + SMs x 64 x 1032 bytes of it; its content does not matter. + + Measured on sm_103: deterministic, and a CUDA-graph replay equals the eager + call bitwise; no 32-bit addressing wrap up to pools of 275,251,200 rows. + First use fetches the selected cubins from FlashInfer's artifactory unless + FLASHINFER_CUBIN_DIR points at a prefilled cache. + + The cells are the small-pool calls. The test also drives, outside the + cells, the target's engine-size pools, the engine envelope (context + calls up to 16384 rows and 256 requests, decoder replay passes up to 32768 + rows, kv lengths up to 1,048,595), the native-storage compact read forms + (`t * 128 + c`, `b * L_r + e`, `t * 512 + k`), the DSpark drafter block at + every G 1..256, CUDA-graph capture at the 46-batch capture set, and the + workspace high-water mark, at H 8, 16, 32 and 64. + + Not certified: `output_scale != 1`; H outside {8, 16, 32, 64}; + `sparse_topk_lens < 128`; explicit `enable_pdl`; context calls beyond the + envelope above; decode / verify batches above 256 or `next_n` other than 1 + and 6; capture at a batch outside the capture set; the draft block at H + other than 64 (beyond the one H=16 cell) or width other than 136; a + workspace other than the 128 MiB buffer; E4M3 operands (that is + `trtllm_sparse_mla_dsv4_fp8`). + """ + + def raw_call( + self, + query: torch.Tensor, + swa_kv_cache: torch.Tensor, + compressed_kv_cache: torch.Tensor, + workspace_buffer: torch.Tensor, + sparse_indices: torch.Tensor, + sparse_topk_lens: torch.Tensor, + seq_lens: torch.Tensor, + sinks: Optional[torch.Tensor], + softmax_scale: float, + output_scale: float = 1.0, + cum_seq_lens_q: Optional[torch.Tensor] = None, + max_q_len: Optional[int] = None, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> torch.Tensor: + return trtllm_batch_decode_sparse_mla_dsv4( + query=query, + swa_kv_cache=swa_kv_cache, + workspace_buffer=workspace_buffer, + sparse_indices=sparse_indices, + compressed_kv_cache=compressed_kv_cache, + sparse_topk_lens=sparse_topk_lens, + seq_lens=seq_lens, + out=out, + bmm1_scale=softmax_scale, + bmm2_scale=output_scale, + sinks=sinks, + kv_layout="HND", + cum_seq_lens_q=cum_seq_lens_q, + max_q_len=max_q_len, + enable_pdl=enable_pdl, + backend="trtllm-gen", + ) + + def reference( + self, + query: torch.Tensor, + swa_kv_cache: torch.Tensor, + compressed_kv_cache: torch.Tensor, + workspace_buffer: torch.Tensor, + sparse_indices: torch.Tensor, + sparse_topk_lens: torch.Tensor, + seq_lens: torch.Tensor, + sinks: Optional[torch.Tensor], + softmax_scale: float, + output_scale: float = 1.0, + cum_seq_lens_q: Optional[torch.Tensor] = None, + max_q_len: Optional[int] = None, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """The attention read through the index rows, and the V scale the gate needs. + + A function of the arguments: the pools and the index rows are both + passed. Gathers the rows each index row names, fp32 scores scaled by + `softmax_scale`, fp32 max and sum, probabilities cast to bf16 for the + PV product (the precision of DeepSeek-V4.1's reference `sparse_attn`), + the per-head sink added to the denominator only, in natural-log units. + `workspace_buffer`, `max_q_len`, `out` and `enable_pdl` are ignored. + + Returns `(expected, v_max)`: `expected` fp32, shaped like `query` and + not rounded to bf16 (the gate was measured against the unrounded + value); `v_max` a 0-d fp32 tensor, `|output_scale|` times the largest + |element| of the K/V rows any query row attends. + """ + dev = query.device + h = query.shape[-2] + q = query.reshape(-1, h, _HEAD_DIM) + n = q.shape[0] + window_count = _window_counts(query, seq_lens, cum_seq_lens_q) + swa_f = swa_kv_cache.reshape(-1, _HEAD_DIM) + comp_f = compressed_kv_cache.reshape(-1, _HEAD_DIM) + width = sparse_indices.shape[1] + col = torch.arange(width, device=dev) + if sinks is None: # no sink: nothing joins the denominator + sk = torch.full((1, h), float("-inf"), dtype=torch.float32, device=dev) + else: + sk = sinks[None, :].float() + expected = torch.empty(n, h, _HEAD_DIM, dtype=torch.float32, device=dev) + v_max = torch.zeros((), dtype=torch.float32, device=dev) + for s0 in range(0, n, _REF_CHUNK): + s1 = min(n, s0 + _REF_CHUNK) + ib = sparse_indices[s0:s1].long() + kv = torch.cat( + [swa_f[ib[:, :_WINDOW].clamp_min(0)], comp_f[ib[:, _WINDOW:].clamp_min(0)]], 1 + ).float() + attended = (col[None] < sparse_topk_lens[s0:s1, None]) & ( + (col[None] >= _WINDOW) | (col[None] < window_count[s0:s1, None]) + ) + kv = kv * (ib >= 0)[..., None] + lo, hi = torch.aminmax(kv, dim=-1) # no |kv| temporary: the gather is the big tensor + row_max = torch.where(attended, torch.maximum(hi, -lo), torch.zeros((), device=dev)) + v_max = torch.maximum(v_max, row_max.amax()) + s = torch.einsum("nhd,ntd->nht", q[s0:s1].float(), kv) * softmax_scale + s = s.masked_fill(~attended[:, None, :], float("-inf")) + m = s.amax(-1) + m = torch.where(torch.isfinite(m), m, torch.zeros_like(m)) + p = torch.exp(s - m[..., None]) + denom = p.sum(-1) + torch.exp(sk - m) + num = torch.einsum("nht,ntd->nhd", p.to(torch.bfloat16).float(), kv) + o = num / denom[..., None] + expected[s0:s1] = torch.where(attended.any(-1)[:, None, None], o, torch.zeros_like(o)) + expected = expected * output_scale + return expected.reshape(query.shape), v_max * abs(output_scale) + + def compare(self, out: torch.Tensor, ref: tuple[torch.Tensor, torch.Tensor]) -> None: + """`rtol` 1.6e-2 plus `atol` 2^-8 * v_max, against the fp32 reference. + + `rtol` is torch's bf16 default. The kernel casts its unnormalized + probabilities to bf16 for PV against its own running max, which differs + from the reference's, so each probability carries up to 2^-9 relative + rounding the reference cannot reproduce; summed into a convex + combination of the attended V rows that is at most 2^-8 * max|V| + absolute -- one bf16 ulp at the largest V element. Measured on sm_103: + the correct reference sits at 0.11-0.39 of this gate across the decode, + verify and drafter cases, and the wrong readings at 19-61 times it + (sink ignored, sink also in the numerator, sink times log2(e) or + softmax_scale, duplicates de-duplicated, an attended `-1` read as a + mask). + """ + expected, v_max = ref + assert out.dtype == torch.bfloat16, out.dtype + assert out.shape == expected.shape, (out.shape, expected.shape) + torch.testing.assert_close( + out.float(), expected, rtol=_RTOL, atol=_ATOL_PER_VMAX * float(v_max) + ) + + +def _window_counts( + query: torch.Tensor, seq_lens: torch.Tensor, cum_seq_lens_q: Optional[torch.Tensor] +) -> torch.Tensor: + """Per query row: min(128, seq_lens[b] - q_len_b + j + 1); dense when `cum_seq_lens_q` is None.""" + dev = query.device + if cum_seq_lens_q is None: + q_len = query.shape[1] if query.dim() == 4 else 1 + cu = torch.arange(seq_lens.numel() + 1, device=dev) * q_len + else: + cu = cum_seq_lens_q.long() + q_lens = cu[1:] - cu[:-1] + req = torch.repeat_interleave(torch.arange(seq_lens.numel(), device=dev), q_lens) + j = torch.arange(req.numel(), device=dev) - cu[:-1][req] + return (seq_lens.long()[req] - q_lens[req] + j + 1).clamp(max=_WINDOW) + + +trtllm_sparse_mla_dsv4 = TrtllmSparseMlaDsv4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/trtllm_sparse_mla_dsv4_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/trtllm_sparse_mla_dsv4_fp8.py new file mode 100644 index 000000000000..0bcc98d6ca5e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/attention/trtllm_sparse_mla_dsv4_fp8.py @@ -0,0 +1,339 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DeepSeek-V4-family sparse MQA over [sliding window | compressed] KV rows with attention sinks, FP8 +compute: E4M3 query and E4M3 pools, each with one per-tensor scale, bf16 output. + +The kernel is the FlashInfer Python callable +`flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4`, not a `torch.ops.trtllm` +op. On sm_103 one call launches one trtllm-gen FMHA kernel of the H512 +`DynamicTokenSparse` family with `QkvE4m3OBfloat16` operands (FP8 tensor-core +MMA for QK and PV); the entry pins `backend="trtllm-gen"` and the `HND` pool +layout. +""" + +from __future__ import annotations + +import math +from typing import Optional, Union + +import torch +from flashinfer.mla import trtllm_batch_decode_sparse_mla_dsv4 + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_HEAD_DIM = 512 +_WINDOW = 128 +#: The gate, per query row and head: ||o - ref|| <= _REL * ||ref|| + _ABS * sqrt(512) * v_max. +_REL = 2.0**-4 +_ABS = 2.0**-7 +#: Query rows per chunk of the reference (bounds its [rows, width, 512] fp32 gather). +_REF_CHUNK = 256 + +#: The target's (q_scale, kv_scale) in every cell but the scale cells. +_SCALES = (2.0**-3, 2.0**-6) +_GS = (1, 2, 3, 8, 16, 24, 32) +_HEADS = (16, 64) +#: Layer kinds: 0 window-only (index width 128), 2 and 1 compressed at that ratio (width 640). +_RATIOS = (0, 2, 1) +#: Varlen prefill steps, (label, requests as (cached, q_len)). +_PREFILL_STEPS = ( + ("1", ((0, 1),)), + ("37", ((0, 37),)), + ("2048", ((0, 2048),)), + ("8192", ((0, 8192),)), + ("cached8192+4096", ((8192, 4096),)), + ("mixed4", ((0, 300), (5000, 240), (0, 1), (16000, 64))), + ("32x256", ((0, 256),) * 32), + ("mixed5", ((4000, 100), (0, 3000), (12000, 700), (0, 5), (60, 129))), +) +_SCALE_PAIRS = ((2.0**-3, 2.0**-6), (2.0**-1, 2.0**-2), (1.0, 1.0), (2.0**-5, 2.0**-8)) +#: The per-query compressed read buffer, (label, max_num_tokens, requests, dense decode call). +_PER_QUERY_STEPS = ( + ("d32-2048", 2048, tuple((9000 + 37 * b, 1) for b in range(32)), True), + ("v6x32-2048", 2048, tuple((20000 + 91 * b, 6) for b in range(32)), True), + ("v6x32-8192", 8192, tuple((130000 + b, 6) for b in range(32)), True), + ("p2048", 2048, ((0, 2048),), False), + ("mixed-8192", 8192, ((1 << 20, 300), (0, 4000), (50000, 3000), (7, 1)), False), +) + + +class TrtllmSparseMlaDsv4Fp8(OpWrapper): + """The sparse attention of `trtllm_sparse_mla_dsv4`, on E4M3 query and pools. + + Every query row attends the first `window_count[r] = min(128, seq_lens[b] + - q_len_b + j + 1)` columns of its index row in the SWA pool and the + columns `[128, sparse_topk_lens[r])` in the compressed pool; K == V + (head_dim 512); an attended `-1` reads as an all-zero row; one fp32 sink + per head (natural-log units) joins the softmax denominator only. With + `q8`, `kv8` the E4M3 codes and `q_scale` / `kv_scale` their per-tensor + dequantization scales:: + + bmm1_scale = softmax_scale * q_scale * kv_scale (float, or an fp32 [1] device tensor) + bmm2_scale = kv_scale + s_c = bmm1_scale * (fp32) + out[r, h] = bmm2_scale * sum_c e4m3(p_c) kv8(c) / (sum_c p_c + exp(sink_h - m)), + p_c = exp(s_c - m) + + Both MMAs are E4M3: PV takes the probabilities rounded to E4M3 (scaled by + 448 first, the factor folded back). One `bmm1`/`bmm2` pair per call, so + both pools carry the same `kv_scale`; no block scale is read. + + Fusion boundary: as the bf16 entry -- the gather, QK, the softmax with + sink and PV are inside; the caller owns RoPE, the pool writes and their + E4M3 quantization, the index rows and lengths, and what follows the + output. Returns the bf16 output, shaped like `query` (`out` when given). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=( + f"{'decode' if nn == 1 else 'verify q_len 6'} G={g} H={h} ratio={r}: " + "dense, positions up to 16397" + ), + spec=dict( + kind="decode", + G=g, + H=h, + ratio=r, + next_n=nn, + q_scale=_SCALES[0], + kv_scale=_SCALES[1], + ), + ) + for g in _GS + for h in _HEADS + for r in _RATIOS + for nn in (1, 6) + ) + + tuple( + Cell( + why=f"prefill {label} H={h} ratio={r}: varlen", + spec=dict( + kind="prefill", + requests=requests, + H=h, + ratio=r, + q_scale=_SCALES[0], + kv_scale=_SCALES[1], + ), + ) + for label, requests in _PREFILL_STEPS + for h in _HEADS + for r in _RATIOS + ) + + tuple( + Cell( + why=f"scales q_scale={sq!r} kv_scale={skv!r}: decode H=16 ratio=1", + spec=dict( + kind="scales", + requests=((700, 1), (3, 1), (16000, 1)), + H=16, + ratio=1, + q_scale=sq, + kv_scale=skv, + ), + ) + for sq, skv in _SCALE_PAIRS + ) + + tuple( + Cell( + why=f"per-query {label} H={h} ratio={r}: compressed row t*512+k, decoys", + spec=dict( + kind="per_query", + max_t=max_t, + requests=requests, + decode=decode, + H=h, + ratio=r, + q_scale=_SCALES[0], + kv_scale=_SCALES[1], + ), + ) + for label, max_t, requests, decode in _PER_QUERY_STEPS + for h in _HEADS + for r in (2, 1) + ) + ) + + note = """ + Metadata consumed: none. Everything is an argument. + + Preconditions: every precondition of `trtllm_sparse_mla_dsv4` holds here + with bf16 read as E4M3 -- compact segments bounded by the lengths, the + window count from `seq_lens` (0 attends the whole window), lengths within + [128, index width], attended indices in range, NaN/Inf rows never named + where a column is loaded, and at least 128 rows of mapped memory from the + first row of any pool a `-1` is loaded from. In addition: query and both + pools float8_e4m3fn (FlashInfer raises otherwise); both pools quantized + with the same per-tensor `kv_scale`, with no value the caller needs above + `448 * kv_scale`. An MXFP8 or NVFP4 pool is re-quantized to plain E4M3 + first: no block scale is read. + + `bmm1_scale` may be a float or an fp32 [1] device tensor (FlashInfer + multiplies a tensor by log2(e) in a small torch op inside the call); the + two forms give bitwise-identical output, and the tensor form is what a + CUDA graph captures. `bmm2_scale` is a float. + + Each probability carries up to 2^-4 relative rounding (its E4M3 cast) that + the reference does not reproduce; `compare` states the band. + + Beyond the cells the test certifies: a device `bmm1_scale` bitwise equal + to the float, and decode / verify CUDA-graph capture at G 1, 2, 4, 8, 16, + 24, 32 x q_len 1 / 6 x H 16 / 64, each replay (new q, index rows, lengths + and a changed `bmm1_scale`) bitwise equal to a fresh eager call. + + Not certified: the bf16 entry's pool-size, far-end, guard-page and + kernel-band sweeps (the FP8 cubins are a different variant set); H other + than 16 / 64; `bmm2_scale` as a tensor; an output other than bf16. + """ + + def raw_call( + self, + query: torch.Tensor, + swa_kv_cache: torch.Tensor, + compressed_kv_cache: torch.Tensor, + workspace_buffer: torch.Tensor, + sparse_indices: torch.Tensor, + sparse_topk_lens: torch.Tensor, + seq_lens: torch.Tensor, + sinks: Optional[torch.Tensor], + bmm1_scale: Union[float, torch.Tensor], + bmm2_scale: float, + cum_seq_lens_q: Optional[torch.Tensor] = None, + max_q_len: Optional[int] = None, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> torch.Tensor: + return trtllm_batch_decode_sparse_mla_dsv4( + query=query, + swa_kv_cache=swa_kv_cache, + workspace_buffer=workspace_buffer, + sparse_indices=sparse_indices, + compressed_kv_cache=compressed_kv_cache, + sparse_topk_lens=sparse_topk_lens, + seq_lens=seq_lens, + out=out, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + sinks=sinks, + kv_layout="HND", + cum_seq_lens_q=cum_seq_lens_q, + max_q_len=max_q_len, + enable_pdl=enable_pdl, + backend="trtllm-gen", + ) + + def reference( + self, + query: torch.Tensor, + swa_kv_cache: torch.Tensor, + compressed_kv_cache: torch.Tensor, + workspace_buffer: torch.Tensor, + sparse_indices: torch.Tensor, + sparse_topk_lens: torch.Tensor, + seq_lens: torch.Tensor, + sinks: Optional[torch.Tensor], + bmm1_scale: Union[float, torch.Tensor], + bmm2_scale: float, + cum_seq_lens_q: Optional[torch.Tensor] = None, + max_q_len: Optional[int] = None, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """fp32 attention on the same E4M3 values, and the V scale the gate needs. + + A function of the arguments: the pools and the index rows are both + passed. The E4M3 codes are exact in fp32; scores are `bmm1_scale` + times the code dot product, fp32 max and sum, fp32 probabilities (no + E4M3 rounding: that is the kernel's own arithmetic the gate bounds), + the per-head sink in the denominator only, and the output + `bmm2_scale` times the probability-weighted code rows. A `None` sink + adds nothing. `workspace_buffer`, `max_q_len`, `out` and `enable_pdl` + are ignored. + + Returns `(expected, v_max)`: `expected` fp32, shaped like `query`; + `v_max` a 0-d fp32 tensor, `|bmm2_scale|` times the largest |code| of + the K/V rows any query row attends (max|V| in dequantized units). + """ + dev = query.device + h = query.shape[-2] + q = query.reshape(-1, h, _HEAD_DIM).float() + n = q.shape[0] + bmm1 = float(bmm1_scale) + window_count = _window_counts(query, seq_lens, cum_seq_lens_q) + swa_f = swa_kv_cache.reshape(-1, _HEAD_DIM) + comp_f = compressed_kv_cache.reshape(-1, _HEAD_DIM) + col = torch.arange(sparse_indices.shape[1], device=dev) + if sinks is None: # no sink: nothing joins the denominator + sk = torch.full((1, h, 1), float("-inf"), dtype=torch.float32, device=dev) + else: + sk = sinks.float()[None, :, None] + expected = torch.empty(n, h, _HEAD_DIM, dtype=torch.float32, device=dev) + v_max = torch.zeros((), dtype=torch.float32, device=dev) + for r0 in range(0, n, _REF_CHUNK): + r1 = min(n, r0 + _REF_CHUNK) + idx = sparse_indices[r0:r1].long() + att = (col[None] < sparse_topk_lens[r0:r1, None]) & ( + (col[None] >= _WINDOW) | (col[None] < window_count[r0:r1, None]) + ) + rows = torch.cat( + [swa_f[idx[:, :_WINDOW].clamp_min(0)], comp_f[idx[:, _WINDOW:].clamp_min(0)]], 1 + ).float() + rows = torch.where((idx >= 0)[..., None], rows, 0.0) + lo, hi = torch.aminmax(rows, dim=-1) + row_max = torch.where(att, torch.maximum(hi, -lo), torch.zeros((), device=dev)) + v_max = torch.maximum(v_max, row_max.amax()) + s = torch.einsum("rhd,rcd->rhc", q[r0:r1], rows) * bmm1 + s = torch.where(att[:, None, :], s, -torch.inf) + m = torch.maximum(s.amax(-1, keepdim=True), sk) + m = torch.where(torch.isfinite(m), m, torch.zeros_like(m)) + p = torch.exp(s - m) + den = p.sum(-1, keepdim=True) + torch.exp(sk - m) + o = torch.einsum("rhc,rcd->rhd", p, rows) / den + o = torch.where(att.any(-1)[:, None, None], o, torch.zeros_like(o)) + expected[r0:r1] = o * bmm2_scale + return expected.reshape(query.shape), v_max * abs(bmm2_scale) + + def compare(self, out: torch.Tensor, ref: tuple[torch.Tensor, torch.Tensor]) -> None: + """Per query row and head: ||o - ref||_2 <= 2^-4 ||ref||_2 + 2^-7 sqrt(512) v_max. + + 2^-4 is the largest relative rounding of one E4M3 probability, which is + the whole error of a row whose softmax sits on one key; the absolute + term covers rows whose output nearly cancels. Measured on sm_103: the + worst correct cell sits at 0.59 of this gate (0.44 on the mixed control + step), and the wrong readings at 13.6-113 times it (sink ignored, + `q_scale` dropped, every window count 128, sink times log2(e), + `kv_scale` applied twice, compressed segment dropped). + """ + expected, v_max = ref + assert out.dtype == torch.bfloat16, out.dtype + assert out.shape == expected.shape, (out.shape, expected.shape) + err = (out.float() - expected).norm(dim=-1) + tol = _REL * expected.norm(dim=-1) + _ABS * math.sqrt(_HEAD_DIM) * float(v_max) + within = err <= tol + worst = (err / tol.clamp_min(1e-30)).max().item() + assert bool(within.all()), ( + f"{int((~within).sum())} (row, head) outputs outside the gate, worst {worst:.3f}" + ) + + +def _window_counts( + query: torch.Tensor, seq_lens: torch.Tensor, cum_seq_lens_q: Optional[torch.Tensor] +) -> torch.Tensor: + """Per query row: min(128, seq_lens[b] - q_len_b + j + 1); dense when `cum_seq_lens_q` is None.""" + dev = query.device + if cum_seq_lens_q is None: + q_len = query.shape[1] if query.dim() == 4 else 1 + cu = torch.arange(seq_lens.numel() + 1, device=dev) * q_len + else: + cu = cum_seq_lens_q.long() + q_lens = cu[1:] - cu[:-1] + req = torch.repeat_interleave(torch.arange(seq_lens.numel(), device=dev), q_lens) + j = torch.arange(req.numel(), device=dev) - cu[:-1][req] + return (seq_lens.long()[req] - q_lens[req] + j + 1).clamp(max=_WINDOW) + + +trtllm_sparse_mla_dsv4_fp8 = TrtllmSparseMlaDsv4Fp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/comm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/comm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/comm/allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/comm/allreduce.py new file mode 100644 index 000000000000..eedacef058b1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/comm/allreduce.py @@ -0,0 +1,340 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sum a tensor across a group of ranks; every rank gets the whole sum back. + +The op is `torch.ops.trtllm.allreduce` (MPI-session communicator) or +`torch.ops.trtllm.allreduce_pg` (a c10d ProcessGroup), both registered from +`cpp/tensorrt_llm/thop/allreduceOp.cpp`. +""" + +from __future__ import annotations + +from typing import List, Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_HIDDEN = 5120 +_ENGRAM_ROW = (24, 256) +_ENGRAM_HIDDEN = 6144 + +#: Strategies the op implements for a plain sum, by the integer it takes. +_STRATEGY_NAMES = { + 0: "NCCL", + 8: "NCCL_SYMMETRIC", + 3: "AUTO", + 4: "ONESHOT", + 5: "TWOSHOT", + 1: "MIN_LATENCY", +} +#: The strategies that read the IPC workspace; 0 and 8 take `workspace=None`. +_WORKSPACE_STRATEGIES = (3, 4, 5, 1) +_AUTO, _NCCL = 3, 0 + +#: (layout, world size, group size, nodes). Groups are the contiguous blocks a +#: pipeline-parallel Mapping makes (`Mapping.tp_group`), all running at once. +#: ws4 is tep4/tp4; ws2 one pp2 x tep2/tp2 stage alone; ws4_groups2 the pp2 x +#: tep2/tp2 layout; ws8 tep8 over two 4-GPU NVL72 trays; ws8_groups4 pp2 x +#: tep4/tp4, one group per tray. +_LAYOUTS = ( + ("ws4", 4, 4, 1), + ("ws2", 2, 2, 1), + ("ws4_groups2", 4, 2, 1), + ("ws8", 8, 8, 2), + ("ws8_groups4", 8, 4, 2), +) + +#: Row counts: decode graph sizes, prefill chunks, TWOSHOT's fewer-rows-than- +#: ranks corner (1..3), and both sides of the 64 MiB edge at which a workspace +#: strategy hands an fp32 [T, 5120] call to NCCL (3276 fits, 3277 does not). +_FP32_ROWS = (1, 2, 3, 4, 5, 8, 16, 32, 64, 128, 256, 1024, 2048, 3276, 3277, 4096, 8192) +_BF16_ROWS = (1, 3, 8, 32, 256, 4096, 6553, 6554, 8192) +_ENGRAM_ROWS = (1, 8, 32, 4096) + +#: The core target's two call forms, at every row count: fp32 [T, 5120] up to +#: the decoder pass's 128 x 256 rows, bf16 [T, 6144] engram rows up to +#: max_num_tokens. +_CORE_FP32_ROWS = range(1, 32768 + 1) +_CORE_ENGRAM_ROWS = range(1, 16384 + 1) + +#: (suite, dtype, trailing dims, row counts, label). The suite is the process +#: a form runs in: fp32 and bf16 never share a workspace's one-shot path. +_FORMS = ( + ("fp32", torch.float32, (_HIDDEN,), _FP32_ROWS, "fp32 [T, 5120]"), + ("bf16", torch.bfloat16, (_HIDDEN,), _BF16_ROWS, "bf16 [T, 5120]"), + ("bf16", torch.bfloat16, _ENGRAM_ROW, _ENGRAM_ROWS, "bf16 [T, 24, 256]"), +) + + +def _group_on_one_node(world_size: int, group_size: int, nodes: int) -> bool: + """A group with an IPC workspace: all of its ranks on one node.""" + return group_size <= world_size // nodes + + +def _layout_cells() -> tuple[Cell, ...]: + cells = [] + for layout, world_size, group_size, nodes in _LAYOUTS: + one_node = _group_on_one_node(world_size, group_size, nodes) + base = dict(layout=layout, world_size=world_size, group_size=group_size, nodes=nodes) + strategies = [s for s in _STRATEGY_NAMES if one_node or s not in _WORKSPACE_STRATEGIES] + for strategy in strategies: + for suite, dtype, trailing, rows, label in _FORMS: + cells.append( + Cell( + why=f"{layout} {_STRATEGY_NAMES[strategy]} {label}: rows {rows}", + spec=dict( + base, + suite=suite, + dtype=dtype, + trailing=trailing, + rows=rows, + strategy=strategy, + ), + ) + ) + core_fp32 = _AUTO if one_node else _NCCL + cells.append( + Cell( + why=( + f"{layout} core fp32 [T, 5120] on {_STRATEGY_NAMES[core_fp32]}: every T " + "1..32768, the target's strategy for this group" + ), + spec=dict( + base, + suite="core", + dtype=torch.float32, + trailing=(_HIDDEN,), + rows=_CORE_FP32_ROWS, + strategy=core_fp32, + ), + ) + ) + cells.append( + Cell( + why=f"{layout} core bf16 [T, 6144] on NCCL: every T 1..16384, the engram rows", + spec=dict( + base, + suite="core", + dtype=torch.bfloat16, + trailing=(_ENGRAM_HIDDEN,), + rows=_CORE_ENGRAM_ROWS, + strategy=_NCCL, + ), + ) + ) + return tuple(cells) + + +class Allreduce(OpWrapper): + """All-reduce (sum) `input` over `group` with the `op` fusion epilogue. + + Every rank in `group` passes a tensor of the same shape and dtype and gets + back `[reduced]` (for `op = 0`, the certified value): a freshly allocated + tensor of `input`'s shape and dtype holding the elementwise sum over the + group, the same bytes on every rank. `input` is untouched. + + Fusion boundary: the collective and the output allocation are inside the + call. Everything that produced `input` (a row-parallel GEMM's or the routed + experts' partial sums), any cast into or out of fp32 around the sum, any + bias added after it, residual add and normalization are the caller's. + `strategy` picks the transport; the workspace strategies need the IPC + `workspace` (see `note`). Returns the op's list unchanged. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Layout x strategy x call form. A cell's `rows` are all run; the core + # cells are the core target's two forms at every row count, on the strategy + # it picks for the group (AUTO with the workspace on one node, NCCL across + # the trays). The ws8 / ws8_groups4 cells span two nodes: see `note`. + CELLS: tuple[Cell, ...] = _layout_cells() + + note = """ + Metadata consumed: none. Process state: the communicator -- trtllm's MPI + session communicator when `process_group` is None (`group` names its + ranks), else the boxed c10d group, whose global ranks `group` lists -- and, + for strategies 1, 3, 4 and 5, the `[3 * len(group) + 2]` int64 IPC + `workspace` from `get_allreduce_workspace(Mapping(...))`, built collectively + on every rank of the group, cached per Mapping, 64 MiB base. A group larger + than gpus_per_node gets no IPC buffers, so a group spanning the two trays + runs strategies 0 and 8 only, with `workspace=None`. On the `process_group` + path strategies 0 and 8 both run the group's own allreduce: the NCCL + window-buffer form of 8 exists only on the MPI path. + + Preconditions: + - Every rank of the group calls in the same order with the same shape, + dtype, `strategy` and `op`; nothing cross-checks ranks. + - One element size per workspace's one-shot path. The one-shot kernel sizes + the reset of its previous call's buffer in the current call's element + size: fp32, then bf16, then fp32 on one workspace returned wrong values on + every on-time rank, and bf16 [4096, 24, 256] then fp32 never completed. + AUTO (fp32 up to 1023 rows in a group of 4), TWOSHOT below len(group) + rows and MIN_LATENCY up to 128 rows all reach that kernel. A forward that + sums both dtypes keeps one off the workspace strategies; DeepSeek-V4.1 + sums fp32 on AUTO and its bf16 engram rows, as a 2-D [T, 6144] view, on + NCCL. + - A 3-D input on a workspace strategy stays within 64 MiB of real size: the + size check is `shape[0] * shape[-1] * element_size`, and bf16 + [8192, 24, 256] on ONESHOT never completed. Pass a 2-D view. + - A group's first call is made eagerly, before any CUDA-graph capture. + - `workspace=None` on 1, 3, 4 or 5 raises `bad optional access`; strategies + 2, 6, 7 and 9 raise at `op = 0`. + + Resolution, measured on sm_103 from the op's debug log: AUTO on fp32 + [T, 5120] is ONESHOT for T 1..1023, TWOSHOT 1024..3276, NCCL from 3277 in a + group of 4, and ONESHOT 1..3276 then NCCL in a group of 2. Every workspace + strategy hands a message over 64 MiB to NCCL; TWOSHOT falls back to ONESHOT + below len(group) rows. + + Numerics: exact whenever the true sum is representable in the dtype (the + matrix's payloads are built so). Unrestricted fp32 partials are summed in + fp32 but not correctly rounded and not bit-identical across strategies + (within assert_close's fp32 defaults, the same bits on every rank); bf16 + sums with cancellation differ by strategy, by up to 0.0625 at 4096 rows. + Rows owned by one rank (zeros elsewhere) come back bitwise on every + strategy. + + CUDA graphs are matrix checks, not cells: fp32 on every strategy of the + layout with bf16 sites on NCCL at 1..128 rows, and the core forward's 82 + sites at the 84 decode row counts T = G, 5G, 6G of the capture batch set. + Strategy 8 under capture only on payloads whose sums are exact. + + Not certified: `op` other than 0; layouts other than the cells' five; + groups that are not contiguous blocks of the world; a workspace base other + than 64 MiB; rows above the cells' ranges; capture at other row counts; a + group's first call captured. + + The ws8 and ws8_groups4 cells span two 4-GPU nodes of one NVLink domain and + cannot be collected on one node. They were certified on the MPI-session + symbol; the `process_group` path at world 8 has not been run. Run each + suite by hand from this target's comm test directory under + tests/unittest/_torch/modeling_v2/models/deepseek_v41/: + TLLM_DISABLE_MPI=1 PYTHONPATH= torchrun --nnodes 2 + --nproc-per-node 4 --node-rank <0|1> --master-addr + --master-port _allreduce_op_matrix.py --rank-worker + [--group-size 4] + with each of fp32, bf16 and core, plus mixed and mixed_core with + --group-size 4. + """ + + def raw_call( + self, + input: torch.Tensor, + residual: Optional[torch.Tensor], + norm_weight: Optional[torch.Tensor], + scale: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + workspace: Optional[torch.Tensor], + group: List[int], + strategy: int, + op: int, + eps: float, + trigger_completion_at_end: bool, + process_group: Optional[object] = None, + ) -> List[torch.Tensor]: + """`process_group` selects the communicator flavor, mirroring the + product router in `_torch/distributed/ops.py`: None runs the + MPI-session symbol, and a boxed c10d ProcessGroup runs the `_pg` + symbol the non-MPI backends use. That symbol also takes this process's + global rank, which must sit at the group's own rank in sorted `group`. + """ + if process_group is None: + return torch.ops.trtllm.allreduce( + input, + residual, + norm_weight, + scale, + bias, + workspace, + group, + strategy, + op, + eps, + trigger_completion_at_end, + ) + return torch.ops.trtllm.allreduce_pg( + input, + residual, + norm_weight, + scale, + bias, + workspace, + group, + torch.distributed.get_rank(), + process_group, + strategy, + op, + eps, + trigger_completion_at_end, + ) + + def reference( + self, + per_rank_inputs: Sequence[torch.Tensor], + group: Optional[List[int]] = None, + ) -> List[torch.Tensor]: + """The sum the collective is supposed to make, as the op's `[reduced]`. + + The documented exception to mirroring `raw_call`: a collective's output + is not a function of what the calling rank holds, so this takes the + group's inputs in ascending rank order instead of `input`. Summed in + fp64 and rounded once to the inputs' dtype. Covers `op = 0` only; + `residual`, `norm_weight`, `scale`, `bias`, `eps` are unread there, and + `workspace`, `strategy`, `trigger_completion_at_end` and + `process_group` choose the transport, not the value, so none is taken. + """ + parts = list(per_rank_inputs) + total = parts[0].double() + for part in parts[1:]: + total = total + part.double() + return [total.to(parts[0].dtype)] + + def compare(self, out: List[torch.Tensor], ref: List[torch.Tensor]) -> None: + """A gate of exactly zero, per returned tensor. + + The certified payloads are built so that every sum is exact in the + input dtype (fp32 multiples of 1/256 below 4096 in magnitude, bf16 + multiples of 1/8 of at most 4), and on those every strategy returned + the correctly rounded fp64 sum bit for bit. On such a payload one + rank missing or counted twice, the mean instead of the sum, and fp32 + partials rounded to bf16 first each differ on more than 90% of + elements. Unrestricted payloads are not gated here: their sums are + not correctly rounded (see `note`). + """ + assert len(out) == len(ref), (len(out), len(ref)) + for got, want in zip(out, ref): + assert got.dtype == want.dtype, (got.dtype, want.dtype) + assert got.shape == want.shape, (got.shape, want.shape) + torch.testing.assert_close(got, want, rtol=0, atol=0) + + def is_valid( + self, + input: torch.Tensor, + residual: Optional[torch.Tensor], + norm_weight: Optional[torch.Tensor], + scale: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + workspace: Optional[torch.Tensor], + group: List[int], + strategy: int, + op: int, + eps: float, + trigger_completion_at_end: bool, + process_group: Optional[object] = None, + ) -> None: + # Pure-metadata guards, each for a violation measured on this machine to + # return a wrong answer without raising. + assert input.is_contiguous(), ( + "input must be contiguous; the op reads it as packed memory and a " + "strided view is silently reduced over the wrong elements" + ) + assert input.dtype is not torch.float8_e4m3fn, ( + "float8_e4m3fn is accepted by the op but summed as raw unsigned bytes, " + "not as floats, so the result is meaningless; reduce in fp32/bf16 and " + "quantize afterwards" + ) + + +allreduce = Allreduce() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/comm/pp_recv_tensors.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/comm/pp_recv_tensors.py new file mode 100644 index 000000000000..0f3966c1655c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/comm/pp_recv_tensors.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Receive tensors from the previous pipeline-parallel stage, in place. + +The op is `torch.ops.trtllm.pp_recv_tensors`, registered in +`tensorrt_llm/_torch/distributed/communicator.py`. +""" + +from __future__ import annotations + +from typing import List, Sequence + +import torch + +import tensorrt_llm._torch.distributed.communicator # noqa: F401 — registers torch.ops.trtllm.pp_* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The DeepSeek-V4.1 target's stage hand-over, in order, as (name, trailing +#: dims, dtype) of a [T, ...] tensor: the hyper-connection stream h, the last +#: FFN mixes yF and their RMS statistic rF. +_HANDOVER = ( + ("h", (4, 5120), torch.bfloat16), + ("yF", (24,), torch.float32), + ("rF", (), torch.float32), +) +#: Eager row counts: every T up to max_num_tokens per rank. +_EAGER_ROWS = range(1, 16384 + 1) +#: The 46 CUDA-graph batch sizes (every G in 1..32 plus the padded default list +#: up to 256); a captured step hands over T = G (DSpark off) or 6 G (DSpark +#: verify), 82 values in 1..1536. +_GRAPH_BATCHES = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_GRAPH_ROWS = tuple(sorted(set(_GRAPH_BATCHES) | {6 * g for g in _GRAPH_BATCHES})) +#: (form, world size, TP size per stage, nodes). Stage 1 rank r receives from +#: rank r - TP. On two trays stage 0 is one 4-GPU node and stage 1 the other. +_FORMS = (("one node", 4, 2, 1), ("two trays", 8, 4, 2)) + + +def _handover_cells() -> tuple[Cell, ...]: + cells = [] + for form, world_size, tp, nodes in _FORMS: + base = dict(world_size=world_size, tp_size=tp, pp_size=2, nodes=nodes, tensors=_HANDOVER) + cells += [ + Cell( + why=f"{form} pp2 x tp{tp} eager: every T 1..16384, all pairs at the same T", + spec=dict(base, attention_dp=False, mode="eager", rows=_EAGER_ROWS), + ), + Cell( + why=f"{form} pp2 x dep{tp} eager: every T 1..16384, the pairs at different T", + spec=dict(base, attention_dp=True, mode="eager", rows=_EAGER_ROWS), + ), + Cell( + why=f"{form} pp2 x tp{tp} graph: the 82 captured T, replayed in any order", + spec=dict(base, attention_dp=False, mode="graph", rows=_GRAPH_ROWS), + ), + ] + return tuple(cells) + + +class PpRecvTensors(OpWrapper): + """Fill each tensor in `tensors`, in order, from `mapping.prev_pp_rank()`. + + The counterpart of `pp_send_tensors`: the sender's tensors have the same + shapes and dtypes in the same order, and each receiving tensor is filled in + place with the sender's bytes. Nothing is computed. Fusion boundary: the + transfer only; allocating `tensors` is the caller's. + + Returns `tensors`, the list it filled (the same list object, the same + tensors). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The same cells as `pp_send_tensors`: one hand-over certifies both halves. + CELLS: tuple[Cell, ...] = _handover_cells() + + note = """ + Metadata consumed: none. The source is `mapping.prev_pp_rank()` of the + mapping `init_pp_comm` was last called with; the engine's PP communicator + (`PPCommNCCL`, the MPI path) must exist. + + The receive is enqueued on the current stream, eagerly and under CUDA-graph + capture, so work after it on that stream reads the received bytes. A + receive into a strided view writes the packed bytes into the first + `numel x itemsize` bytes of the view's storage and leaves the view's own + elements unwritten, silently. + + The rest of the hand-over contract -- what the eager and captured sends + guarantee, the preconditions, the measured cost, what is not certified and + the hand-run two-tray command -- is `pp_send_tensors`' note: neither half + runs without the other, and one matrix certifies both. + """ + + def raw_call(self, tensors: List[torch.Tensor]) -> List[torch.Tensor]: + torch.ops.trtllm.pp_recv_tensors(tensors) + return tensors + + def reference(self, sent: Sequence[torch.Tensor]) -> List[torch.Tensor]: + """What `tensors` must hold after the call: the sender's bytes. + + The documented exception for a collective: what a receive returns is + not a function of the receiving buffers, so this takes the tensors the + previous stage's `pp_send_tensors` sent, in order, and returns copies. + """ + return [t.clone() for t in sent] + + def compare(self, out: Sequence[torch.Tensor], ref: Sequence[torch.Tensor]) -> None: + """Bitwise per tensor. + + The op moves bytes and computes nothing, so any differing byte is a + wrong delivery. The matrix's payloads differ in every 32-bit word from + a missed receive (a 0xFF-filled buffer), the previous step's bytes and + the other sender's bytes, and the gate fails each of them. + """ + assert len(out) == len(ref), (len(out), len(ref)) + for k, (got, want) in enumerate(zip(out, ref)): + assert got.dtype == want.dtype, (k, got.dtype, want.dtype) + assert got.shape == want.shape, (k, got.shape, want.shape) + bad = (got.contiguous().view(torch.uint8) != want.contiguous().view(torch.uint8)).sum() + assert bad.item() == 0, f"tensor {k}: {bad.item()} of {want.nbytes} bytes differ" + + def is_valid(self, tensors: List[torch.Tensor]) -> None: + for t in tensors: + assert t.is_contiguous(), ( + "every tensor must be contiguous (it is filled as packed bytes)" + ) + assert t.is_cuda, "the PP communicator moves device tensors only" + + +pp_recv_tensors = PpRecvTensors() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/comm/pp_send_tensors.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/comm/pp_send_tensors.py new file mode 100644 index 000000000000..0b95009f893b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/comm/pp_send_tensors.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Send tensors to the next pipeline-parallel stage (the matching rank of the next stage). + +The op is `torch.ops.trtllm.pp_send_tensors`, registered in +`tensorrt_llm/_torch/distributed/communicator.py`. +""" + +from __future__ import annotations + +from typing import List, Sequence + +import torch + +import tensorrt_llm._torch.distributed.communicator # noqa: F401 — registers torch.ops.trtllm.pp_* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The DeepSeek-V4.1 target's stage hand-over, in order, as (name, trailing +#: dims, dtype) of a [T, ...] tensor: the hyper-connection stream h, the last +#: FFN mixes yF and their RMS statistic rF. +_HANDOVER = ( + ("h", (4, 5120), torch.bfloat16), + ("yF", (24,), torch.float32), + ("rF", (), torch.float32), +) +#: Eager row counts: every T up to max_num_tokens per rank. +_EAGER_ROWS = range(1, 16384 + 1) +#: The 46 CUDA-graph batch sizes (every G in 1..32 plus the padded default list +#: up to 256); a captured step hands over T = G (DSpark off) or 6 G (DSpark +#: verify), 82 values in 1..1536. +_GRAPH_BATCHES = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_GRAPH_ROWS = tuple(sorted(set(_GRAPH_BATCHES) | {6 * g for g in _GRAPH_BATCHES})) +#: (form, world size, TP size per stage, nodes). Stage 0 rank r sends to rank +#: r + TP. On two trays stage 0 is one 4-GPU node and stage 1 the other. +_FORMS = (("one node", 4, 2, 1), ("two trays", 8, 4, 2)) + + +def _handover_cells() -> tuple[Cell, ...]: + cells = [] + for form, world_size, tp, nodes in _FORMS: + base = dict(world_size=world_size, tp_size=tp, pp_size=2, nodes=nodes, tensors=_HANDOVER) + cells += [ + Cell( + why=f"{form} pp2 x tp{tp} eager: every T 1..16384, all pairs at the same T", + spec=dict(base, attention_dp=False, mode="eager", rows=_EAGER_ROWS), + ), + Cell( + why=f"{form} pp2 x dep{tp} eager: every T 1..16384, the pairs at different T", + spec=dict(base, attention_dp=True, mode="eager", rows=_EAGER_ROWS), + ), + Cell( + why=f"{form} pp2 x tp{tp} graph: the 82 captured T, replayed in any order", + spec=dict(base, attention_dp=False, mode="graph", rows=_GRAPH_ROWS), + ), + ] + return tuple(cells) + + +class PpSendTensors(OpWrapper): + """Send each tensor in `tensors`, in order, to `mapping.next_pp_rank()`. + + The receiving rank posts `pp_recv_tensors` with tensors of the same shapes + and dtypes in the same order, and each is filled with the sender's bytes. + Nothing is computed. Fusion boundary: the transfer only; allocating the + receive buffers and choosing what to send are the caller's. + + Returns None: the call writes nothing on this rank. Its result is the + receiver's tensors, which `reference` and `compare` describe. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Form x layout x mode. The tensors are always the target's three. The + # attention-DP layout shares the attention-TP layout's communicator through + # `init_pp_comm`'s reuse of a live one of the same world size. + CELLS: tuple[Cell, ...] = _handover_cells() + + note = """ + Metadata consumed: none. The destination is `mapping.next_pp_rank()` of the + mapping `init_pp_comm` was last called with; the engine's PP communicator + must exist (`init_pp_comm(mapping)`, which the model engine calls when the + mapping has more than one stage). The certified communicator is + `PPCommNCCL`, one world-size NCCL communicator (the MPI path). + + Eagerly each tensor is copied and the copy sent on the communicator's own + stream after the current stream's pending work, so the caller may + overwrite or free its tensors right after the call: with the receiver held + back ~20 ms while the sender overwrote, dropped and reallocated its + tensors, every delivery was the bytes at the time of the call. Under + CUDA-graph capture the send is recorded on the capturing stream with no + copy: a replay sends the static tensors' bytes as they are at replay time, + and later writes to them on that stream do not reach the receiver. An + eager send followed at once by a replay that sends arrives in call order. + + Preconditions: the receiver posts the same count, shapes and dtypes in the + same order (a mismatch was not probed); every tensor contiguous and on a + CUDA device -- a captured send of a strided view sends the first + `numel x itemsize` bytes of its storage, silently; under CUDA graphs both + stages capture and replay the same step. + + Measured on sm_103: about 0.25 ms per hand-over at decode sizes and 1.3 ms + at T = 16384, with or without a tray boundary; T x 41,060 bytes per step. + + Not certified: T = 0 or T > 16384; any other tensor list; more than two + stages, other TP widths per stage, context parallelism; the `PPCommTorch` + path (`init_pp_comm` with MPI disabled); two nodes not in one NVLink + domain; a sender and receiver that disagree; a captured send replayed + without its captured receive; other collectives on the same GPUs during a + hand-over; CPU tensors; a receiver held back longer than ~20 ms; more than + 8 eager sends in flight. + + The two-tray cells span two 4-GPU nodes of one NVL72 rack and cannot be + collected on one node. Run them by hand from this target's comm test + directory under tests/unittest/_torch/modeling_v2/models/deepseek_v41/: + MODELING_V2_TWO_NODE_MPIRUN="" python _pp_send_recv_op_matrix.py --two-node + where starts 8 ranks of the command after it, ranks 0-3 on one + node and 4-7 on the other, with PYTHONPATH at the repo root (for example + `mpirun -np 8 --host :4,:4 -x PYTHONPATH`). It must print + PP_SEND_RECV_OK. The same run certifies `pp_recv_tensors`. + """ + + def raw_call(self, tensors: List[torch.Tensor]) -> None: + torch.ops.trtllm.pp_send_tensors(tensors) + + def reference(self, tensors: Sequence[torch.Tensor]) -> List[torch.Tensor]: + """What the receiving rank's tensors must hold: the sender's bytes. + + The documented exception for a collective: a send's result is on + another rank, so this states what the matching `pp_recv_tensors` must + leave there, given the tensors this rank sends (copies, in order). + """ + return [t.clone() for t in tensors] + + def compare(self, out: Sequence[torch.Tensor], ref: Sequence[torch.Tensor]) -> None: + """Bitwise per tensor, on the receiver's tensors. + + The op moves bytes and computes nothing, so any differing byte is a + wrong delivery. The matrix's payloads differ in every 32-bit word from + a missed receive (a 0xFF-filled buffer), the previous step's bytes and + the other sender's bytes, and the gate fails each of them. + """ + assert len(out) == len(ref), (len(out), len(ref)) + for k, (got, want) in enumerate(zip(out, ref)): + assert got.dtype == want.dtype, (k, got.dtype, want.dtype) + assert got.shape == want.shape, (k, got.shape, want.shape) + bad = (got.contiguous().view(torch.uint8) != want.contiguous().view(torch.uint8)).sum() + assert bad.item() == 0, f"tensor {k}: {bad.item()} of {want.nbytes} bytes differ" + + def is_valid(self, tensors: List[torch.Tensor]) -> None: + for t in tensors: + # A strided view is sent as its packed storage bytes, not as the view. + assert t.is_contiguous(), "every tensor must be contiguous" + assert t.is_cuda, "the PP communicator moves device tensors only" + + +pp_send_tensors = PpSendTensors() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/embedding/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/embedding/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/embedding/engram_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/embedding/engram_gather.py new file mode 100644 index 000000000000..502d8848b780 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/embedding/engram_gather.py @@ -0,0 +1,380 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Row gather + dequantization of an FP8 E4M3 table with one E8M0 scale per 32 channels. + +The DeepSeek-V4.1 Engram lookup: each hash id names one row of a +``[rows, 256]`` float8_e4m3fn table with a ``[rows, 8]`` float8_e8m0fnu scale +table; the row is dequantized in fp32 and rounded once to bf16. Ids outside +this caller's row range ``[row_start, row_end)`` produce all-zero rows, which +is what a row-sharded table's summing all-reduce expects. + +The kernel is a Triton kernel defined in this file (a port of SGLang's +``engram_gather``), not a ``torch.ops.trtllm`` op: no engine op reads FP8 rows +with per-32 E8M0 scales. It takes the two tables as raw addresses, so one +kernel serves a table in device memory and a table in pinned host memory; for +a pinned host tensor the wrapper passes the address returned by +``cudaHostGetDevicePointer`` and the kernel reads host memory directly (UVA +zero-copy) -- no staging copy. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# Launch configuration: fixed, no autotune. ROWS table rows per program, +# NUM_WARPS warps per program. +_ROWS_PER_PROGRAM = 4 +_NUM_WARPS = 4 + + +@triton.jit +def _engram_gather_kernel( + w_addr, + s_addr, + ids_ptr, + out_ptr, + n_rows, + lookups, + ids_row_stride, + row_start, + row_end, + DIM: tl.constexpr, + BLOCK: tl.constexpr, + ROWS: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + r = pid * ROWS + tl.arange(0, ROWS).to(tl.int64) # flat output row = t * lookups + k + live = r < n_rows + t = r // lookups + k = r - t * lookups + gid = tl.load(ids_ptr + t * ids_row_stride + k, mask=live, other=0).to(tl.int64) + owned = live & (gid >= row_start) & (gid < row_end) + local = tl.where(owned, gid - row_start, 0) + + w = w_addr.to(tl.int64).to(tl.pointer_type(tl.float8e4nv)) + s = s_addr.to(tl.int64).to(tl.pointer_type(tl.uint8)) + c = tl.arange(0, DIM) + vals = tl.load(w + local[:, None] * DIM + c[None, :], mask=owned[:, None], other=0.0).to( + tl.float32 + ) + e = tl.load( + s + local[:, None] * (DIM // BLOCK) + (c // BLOCK)[None, :], mask=owned[:, None], other=0 + ).to(tl.int32) + # E8M0 byte e -> 2^(e-127) built from the fp32 exponent bits (exact, no exp2): + # e == 0 is 2^-127, an fp32 subnormal (bits 0x00400000); e == 255 is NaN (0x7FC00000). + bits = tl.where(e == 0, 0x00400000, e << 23) + bits = bits | ((e == 255).to(tl.int32) << 22) + scale = bits.to(tl.float32, bitcast=True) + res = tl.where(owned[:, None], vals * scale, 0.0) + tl.store(out_ptr + r[:, None] * DIM + c[None, :], res.to(tl.bfloat16), mask=live[:, None]) + + +def _kernel_address(t: torch.Tensor, device: torch.device) -> int: + """Address the kernel dereferences: the device pointer of a CUDA tensor, or the + mapped device pointer (``cudaHostGetDevicePointer``) of a pinned host tensor.""" + if t.is_cuda: + return t.data_ptr() + from cuda.bindings import runtime as cudart + + with torch.cuda.device(device): + err, dptr = cudart.cudaHostGetDevicePointer(t.data_ptr(), 0) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError( + f"cudaHostGetDevicePointer failed ({err}): the host table must be pinned (page-locked) memory" + ) + return int(dptr) + + +#: DeepSeek-V4.1's table geometry: D channels per row, one E8M0 byte per B, K lookups per token. +_V41_DBK = dict(D=256, B=32, K=24) + +#: The layer-14 table's rank-3 row_start under four-way row sharding, ceil(384016682 / 4) * 3. +_L14_RANK3_START = 288_012_513 + + +class EngramGather(OpWrapper): + """Gather and dequantize rows ``ids`` of an FP8 E4M3 table with per-32 E8M0 scales. + + The DeepSeek-V4.1 Engram lookup (``ParallelEngramEmbedding.forward`` of the + checkpoint's ``inference/model.py``), one launch. For ``g = ids[t, k]``:: + + out[t, k, c] = bf16(fp32(weight[g - row_start, c]) * 2^(scale[g - row_start, c // B] - 127)) + if row_start <= g < row_end, else +0.0 + + with ``D = weight.shape[1]`` and ``B = D // scale.shape[1]``, the product in + fp32 rounded once to bf16. Inside the call: the row gather, the + dequantization, the zero fill of ids outside ``[row_start, row_end)``. The + caller owns the hash ids (per-layer offsets applied), the reduction across + ranks that turns zero-filled row shards into the full lookup, and the + ``wkv`` projection that consumes the rows. + + ``ids`` is ``[T, K]`` int64/int32 on the GPU (unit stride along ``K``, any + row stride); ``weight`` / ``scale`` hold global rows ``[row_start, row_start + + weight.shape[0])`` as ``[R, D]`` float8_e4m3fn and ``[R, D // 32]`` + float8_e8m0fnu (or uint8) bytes, each a CUDA tensor (device memory, or a + CUDA view of host memory) or a pinned CPU tensor, whose mapped address + ``cudaHostGetDevicePointer`` gives on every call (zero-copy, no staging). + Returns ``[T, K, D]`` bf16: ``out`` when given (the bf16 values stored into + its dtype), else a new tensor. ``T = 0`` returns without a launch. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"every E4M3 x E8M0 byte pair, {res} table: NaN, inf, subnormals", + spec=dict(kind="byte_pairs", residency=res, T=11, rows=256, **_V41_DBK), + ) + for res in ("device", "host") + ) + + tuple( + Cell( + why=f"device row shards T={t}: 4 ranks of 100003 rows, zero fill, sum", + spec=dict(kind="device_shards", T=t, rows=100_003, world=4, **_V41_DBK), + ) + for t in (1, 6, 32, 4096, 8192, 16384) + ) + + tuple( + Cell( + why=f"pinned host shared table T={t}: whole 200000 rows, no device copy", + spec=dict(kind="host_shared", T=t, rows=200_000, **_V41_DBK), + ) + for t in (1, 32, 4096) + ) + + tuple( + Cell( + why=f"pinned host row shards T={t}: offset views of one buffer, sum", + spec=dict(kind="host_shards", T=t, rows=100_003, world=4, **_V41_DBK), + ) + for t in (1, 4096) + ) + + tuple( + Cell( + why=f"byte offsets past 2 GiB, {res}: 9e6-row shard at rank-3 start", + spec=dict( + kind="offsets_2gib", + residency=res, + T=32, + rows=9_000_000, + row_start=_L14_RANK3_START, + **_V41_DBK, + ), + ) + for res in ("device", "host") + ) + + ( + Cell( + why="layouts: T=0, [:, layer, :] view, int32 ids, uint8 scale, out view", + spec=dict(kind="layouts", T=37, rows=30_000, **_V41_DBK), + ), + ) + ) + + note = """ + Metadata consumed: none; the call is stateless. + + DeepSeek-V4.1: the layer-1 and layer-14 tables have 384006168 and 384016682 + rows (98 GB of FP8 each), D = 256, B = 32, K = 24 lookups per token. Device, + row-sharded (tep4 / tp4): each rank passes its own [ceil(rows / 4), 256] + shard with row_start = rank * part, row_end = min(row_start + part, rows), + the same ids on every rank, then a bf16 summing all-reduce; dep4 does the + same over the all-gathered ids (T up to 4 x 16384) and reduce-scatters. Host + (host_offload, the default): every member passes the whole node-shared + table, row_start = 0, row_end = rows, over its own tokens (the [:, layer, :] + view of the hash output, T = 1..16384); no reduction. The rank sum equals the + unsharded gather bit for bit except that -0.0 becomes +0.0. + + Residency: the kernel takes the tables as raw addresses. A CUDA tensor + (device memory, or the vmm backend's CUDA views of host-NUMA memory) passes + data_ptr(); a pinned CPU tensor (the shm backend's views) passes the mapped + pointer cudaHostGetDevicePointer returns, resolved on every call, legal + during capture. A graph bakes the address in: a host table must stay alive + and in place for as long as a graph that captured a call on it can replay. + A CPU view of the vmm table's memory passes the pointer query and gathers + wrong rows. Pageable host memory is accepted (nothing detects it) and not + certified. + + Preconditions beyond is_valid: D a power of two; a device table on ids' + device; any id value (out-of-range ids read nothing and give zeros). + Arithmetic: an E4M3 value times a power of two is exact in fp32, so the + result is bit-identical to the torch reference, including E8M0 byte 0 + (2^-127, subnormal products kept) and 255 (NaN), the E4M3 NaN codes, inf + on overflow, and -0.0 from 0x80. + + Launch: fixed, 4 table rows per program, 4 warps, grid ceil(T * K / 4), no + host synchronization. Triton specializes n_rows = 24 T on parity, and + row_start, row_end and the ids' row stride on their own classes: warm up + each configuration before capturing it. + + Measured on sm_103 beyond the cells, by the entry's test: every T in + 1..16384 (the tp [:, 1, :] view) and 1..65536 (the dep4 gathered view) as + prefixes of one call; any device shard of up to 96004171 rows (the largest + the target allocates) and the far end of every rank's shard; capture at the + 82 graph row counts (T = G and 6 G, 1..1536; dep4 at 4 x those) on four + rank shards; the checkpoint's own two tables (202.76 GB) in the node-shared + host copy under both backends, read whole, at every T to 16384 and every + captured count, by one process, by four processes on four GPUs, and by a + one-GPU process on each socket; real checkpoint rows reproducing the golden + engram_value. + + Not certified: D other than 256, B other than 32, K other than 24; a device + table on another GPU; pageable host tables; a host table above 202.76 GB; + other host backings at the real size or per-rank host shards of the real + tables; out dtypes other than bf16 (fp16 stores infs where bf16 exceeds + its range, with no error); T above 65536 (above 16384 for the tp view and + host_offload); capture at an unlisted row count; concurrent use from + several streams. + """ + + def raw_call( + self, + ids: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + row_start: int, + row_end: int, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + dim = weight.shape[1] + t_len, lookups = ids.shape + if out is None: + out = torch.empty((t_len, lookups, dim), dtype=torch.bfloat16, device=ids.device) + n_rows = t_len * lookups + if n_rows == 0: + return out + _engram_gather_kernel[(triton.cdiv(n_rows, _ROWS_PER_PROGRAM),)]( + _kernel_address(weight, ids.device), + _kernel_address(scale, ids.device), + ids, + out, + n_rows, + lookups, + ids.stride(0), + row_start, + row_end, + DIM=dim, + BLOCK=dim // scale.shape[1], + ROWS=_ROWS_PER_PROGRAM, + num_warps=_NUM_WARPS, + ) + return out + + def is_valid( + self, + ids: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + row_start: int, + row_end: int, + out: Optional[torch.Tensor] = None, + ) -> None: + """The layouts and dtypes the kernel reads or writes as if they were the expected ones. + + Each was run through the launch without these checks on sm_103 and + returned wrong values, or wrote outside ``out``, with no error: a + transposed ``ids`` view (bit-exact 0.043), float32 ids (0.020), an E5M2 + table (0.015), fp32 scales (0.015), a column-slice weight (0.004) or + scale (0.048), a scale with half the weight's rows (the rows past it took + the parent buffer's bytes), a ``[R, 7]`` scale (0.57), a 100-row view + under a 200-row range (row 150 of the parent buffer returned), a + last-dim-slice ``out`` (21334 elements written into its other half) and + an ``out`` of 3 of 7 token rows (24576 elements written past its end). + CPU ``ids`` are not guarded: Triton's launcher refuses them loudly. + """ + # Every guard below is a layout the kernel would otherwise read or write + # as if it were the expected one, without an error (probed on sm_103). + assert ids.dim() == 2 and ids.stride(1) == 1, ( + f"ids must be [T, K] with unit inner stride, got {ids.shape} {ids.stride()}" + ) + assert ids.dtype in (torch.int64, torch.int32), ( + f"ids must be int64 or int32, got {ids.dtype}" + ) + assert weight.dtype == torch.float8_e4m3fn, ( + f"weight must be float8_e4m3fn, got {weight.dtype}" + ) + assert scale.dtype in (torch.float8_e8m0fnu, torch.uint8), ( + f"scale must be float8_e8m0fnu or uint8, got {scale.dtype}" + ) + assert weight.dim() == 2 and weight.is_contiguous(), ( + f"weight must be a contiguous [R, D] tensor, got {weight.shape} {weight.stride()}" + ) + assert scale.dim() == 2 and scale.is_contiguous(), ( + f"scale must be a contiguous [R, D // block] tensor, got {scale.shape} {scale.stride()}" + ) + rows, dim = weight.shape + assert scale.shape[0] == rows and dim % scale.shape[1] == 0, ( + f"scale {tuple(scale.shape)} does not match weight {tuple(weight.shape)}" + ) + assert 0 <= row_end - row_start <= rows, ( + f"row range [{row_start}, {row_end}) exceeds the {rows} rows held" + ) + t_len, lookups = ids.shape + if out is not None: + assert out.is_contiguous() and tuple(out.shape) == (t_len, lookups, dim), ( + f"out must be a contiguous [{t_len}, {lookups}, {dim}] tensor, got {tuple(out.shape)} {out.stride()}" + ) + + def reference( + self, + ids: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + row_start: int, + row_end: int, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``ParallelEngramEmbedding.forward`` on one shard, in torch ops in its order. + + Mask, local index, lookup, fp32 multiply by the per-``B`` scale, one bf16 + cast, zero the masked rows. The named rows are gathered where the table + lives -- on the CPU for a CPU tensor, through the process's own mapping + rather than the GPU's -- and only they are moved to ``ids``' device, so a + host table of any size can be its own reference. ``out`` is ignored. + """ + dev = ids.device + ids64 = ids.long() + mask = (ids64 < row_start) | (ids64 >= row_end) + local = (ids64 - row_start).masked_fill(mask, 0) + block = weight.shape[1] // scale.shape[1] + w8 = weight.view(torch.uint8)[local.to(weight.device)].to(dev) + s8 = scale.view(torch.uint8)[local.to(scale.device)].to(dev) + values = w8.view(torch.float8_e4m3fn).float().unflatten(-1, (-1, block)) + values = values * s8.view(torch.float8_e8m0fnu).float().unsqueeze(-1) + return values.flatten(-2).to(torch.bfloat16).masked_fill(mask.unsqueeze(-1), 0) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Bit equality of the bf16 patterns, NaN positions equal whatever the payload. + + Every product of an E4M3 value and a power of two is exact in fp32 + (subnormals included, as the all-byte-pairs cell shows), and the single + fp32 -> bf16 rounding is round-to-nearest-even on both sides, so any + difference is a defect. The wrong variants the test runs through it -- + scale per 16 channels, out-of-shard rows not zeroed, E4M3 bytes read as + E5M2, 32-bit byte offsets -- are each rejected. + """ + assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16, ( + out.shape, + ref.shape, + out.dtype, + ) + nan_a, nan_e = torch.isnan(out), torch.isnan(ref) + assert torch.equal(nan_a, nan_e), ( + f"NaN positions differ at {(nan_a != nan_e).sum().item()} elements" + ) + a, e = out.view(torch.int16)[~nan_a], ref.view(torch.int16)[~nan_e] + bad = (a != e).sum().item() + same = (out.view(torch.int16) == ref.view(torch.int16)) | (nan_a & nan_e) + assert bad == 0, ( + f"{bad} of {a.numel()} elements differ in their bf16 bits " + f"(mismatch fraction {1.0 - same.float().mean().item():.6f})" + ) + + +engram_gather = EngramGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/embedding/engram_hash.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/embedding/engram_hash.py new file mode 100644 index 000000000000..2a29cb0f2e56 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/embedding/engram_hash.py @@ -0,0 +1,531 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Engram n-gram hash ids on the GPU, with the per-request token history in a paged pool. + +The Triton kernel ``_engram_hash_kernel`` below is this entry's implementation +(no ``torch.ops`` op computes it). One launch per call: it maps every token of +the step through the compressed-vocab table, stores the compressed id into a +position-indexed paged pool, reads the up-to-``N - 1`` predecessors that lie +before the step from that pool, and emits the int64 hash ids of every n-gram +(2 .. N) ending at each token, for every Engram layer. + +The arithmetic is the checkpoint's ``inference/engram.py`` ``NgramHashState``: +cumulative DEAD / sequence-start blocking to ``pad_id``, int64 products with +per-(layer, shift) multipliers, a rolling XOR, and one prime-sized bucket range +per (n-gram size, head) addressed by ``offsets``. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_BLOCK_T = 64 +_NUM_WARPS = 4 +_INT_DTYPES = (torch.int32, torch.int64) + + +@triton.jit +def _compressed_id(ids_ptr, mask_ptr, map_ptr, vocab_size, idx, ok, HAS_MASK: tl.constexpr): + """Compressed id of token row ``idx``, or -1 (DEAD) for a masked / out-of-range id.""" + tok = tl.load(ids_ptr + idx, mask=ok, other=-1).to(tl.int64) + alive = ok & (tok >= 0) & (tok < vocab_size) + if HAS_MASK: + alive = alive & (tl.load(mask_ptr + idx, mask=ok, other=0) != 0) + return tl.load(map_ptr + tok, mask=alive, other=-1).to(tl.int64) + + +@triton.jit +def _engram_hash_kernel( + ids_ptr, + mask_ptr, + req_ptr, + cu_ptr, + start_ptr, + pool_ptr, + pool_row_stride, + bt_ptr, + bt_row_stride, + max_blocks, + page_size, + map_ptr, + vocab_size, + mult_ptr, + primes_ptr, + offsets_ptr, + out_ptr, + num_tokens, + pad_id, + HAS_MASK: tl.constexpr, + L: tl.constexpr, + N: tl.constexpr, + H: tl.constexpr, + BLOCK_T: tl.constexpr, +): + COLS: tl.constexpr = (N - 1) * H + t = tl.program_id(0).to(tl.int64) * BLOCK_T + tl.arange(0, BLOCK_T).to(tl.int64) + live = t < num_tokens + r = tl.load(req_ptr + t, mask=live, other=0).to(tl.int64) + first_row = tl.load(cu_ptr + r, mask=live, other=0).to(tl.int64) + start = tl.load(start_ptr + r, mask=live, other=0).to(tl.int64) + off = t - first_row # offset of the token inside its request's run + pos = start + off # absolute position in the request's sequence + bt_row = bt_ptr + r * bt_row_stride + heads = tl.arange(0, H) + + # Shift 0: the token itself. Its compressed id (-1 = DEAD) is stored at its + # position so that later calls can read it as history. + cur = _compressed_id(ids_ptr, mask_ptr, map_ptr, vocab_size, t, live, HAS_MASK) + blk = pos // page_size + page = tl.load(bt_row + blk, mask=live & (blk < max_blocks), other=-1).to(tl.int64) + tl.store( + pool_ptr + page * pool_row_stride + pos % page_size, + cur.to(tl.int32), + mask=live & (page >= 0), + ) + + for layer in tl.static_range(L): + blocked = cur < 0 + rolling = tl.where(blocked, pad_id, cur) * tl.load(mult_ptr + layer * N) + for s in tl.static_range(1, N): + q = pos - s + in_run = off >= s + # predecessor inside this call's run: read from input_ids + from_run = _compressed_id( + ids_ptr, mask_ptr, map_ptr, vocab_size, t - s, live & in_run, HAS_MASK + ) + # predecessor before the run: read from the pool (written by an earlier call) + hist_ok = live & (off < s) & (q >= 0) + qblk = q // page_size + qpage = tl.load(bt_row + qblk, mask=hist_ok & (qblk < max_blocks), other=-1).to( + tl.int64 + ) + from_hist = tl.load( + pool_ptr + qpage * pool_row_stride + q % page_size, + mask=hist_ok & (qpage >= 0), + other=-1, + ).to(tl.int64) + src = tl.where(in_run, from_run, from_hist) + # once a shift is blocked (sequence start or DEAD) every older shift is too + blocked = blocked | (q < 0) | (src < 0) + rolling = rolling ^ (tl.where(blocked, pad_id, src) * tl.load(mult_ptr + layer * N + s)) + prime = tl.load(primes_ptr + layer * COLS + (s - 1) * H + heads) + offset = tl.load(offsets_ptr + layer * COLS + (s - 1) * H + heads) + val = rolling[:, None] % prime[None, :] + offset[None, :] + tl.store( + out_ptr + t[:, None] * (L * COLS) + layer * COLS + (s - 1) * H + heads[None, :], + val, + mask=live[:, None], + ) + + +def _check(cond: bool, msg: str) -> None: + if not cond: + raise ValueError(f"engram_hash: {msg}") + + +def _check_1d(name: str, x: torch.Tensor, n: Optional[int], dtypes) -> None: + _check( + x.dim() == 1 and x.stride(0) == 1, + f"{name} must be a 1-D tensor with unit stride, got shape {tuple(x.shape)} stride {x.stride()}", + ) + _check(x.dtype in dtypes, f"{name} dtype must be one of {dtypes}, got {x.dtype}") + if n is not None: + _check(x.shape[0] == n, f"{name} must have {n} elements, got {x.shape[0]}") + + +def _launch_engram_hash( + input_ids: torch.Tensor, + req_idx_per_token: torch.Tensor, + cu_seq_lens: torch.Tensor, + start_pos: torch.Tensor, + token_pool: torch.Tensor, + block_table: torch.Tensor, + token_map: torch.Tensor, + multipliers: torch.Tensor, + primes: torch.Tensor, + offsets: torch.Tensor, + pad_id: int, + token_mask: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Validate (``ValueError``, prefixed ``engram_hash:``), allocate ``out`` when absent, launch once.""" + device = input_ids.device + _check(input_ids.is_cuda, "input_ids must be a CUDA tensor") + _check_1d("input_ids", input_ids, None, _INT_DTYPES) + T = input_ids.shape[0] + _check_1d("req_idx_per_token", req_idx_per_token, T, _INT_DTYPES) + _check_1d("start_pos", start_pos, None, _INT_DTYPES) + B = start_pos.shape[0] + _check_1d("cu_seq_lens", cu_seq_lens, B + 1, _INT_DTYPES) + _check_1d("token_map", token_map, None, _INT_DTYPES) + _check( + token_pool.dim() == 2 and token_pool.stride(1) == 1, + "token_pool must be a [pages, page_size] tensor with unit column stride", + ) + _check( + token_pool.dtype == torch.int32, + f"token_pool dtype must be torch.int32, got {token_pool.dtype}", + ) + _check( + block_table.dim() == 2 and block_table.stride(1) == 1, + "block_table must be a [>=B, max_blocks] tensor with unit column stride", + ) + _check( + block_table.shape[0] >= B, + f"block_table must have at least B={B} rows, got {block_table.shape[0]}", + ) + _check( + block_table.dtype in _INT_DTYPES, + f"block_table dtype must be int32 or int64, got {block_table.dtype}", + ) + _check( + multipliers.dim() == 2 and multipliers.is_contiguous() and multipliers.dtype == torch.int64, + "multipliers must be a contiguous int64 [L, N] tensor", + ) + L, N = multipliers.shape + _check(N >= 2, f"multipliers must have N >= 2 columns (max n-gram size), got {N}") + _check( + primes.dim() == 3 and primes.is_contiguous() and primes.dtype == torch.int64, + "primes must be a contiguous int64 [L, N - 1, H] tensor", + ) + _check( + primes.shape[0] == L and primes.shape[1] == N - 1, + f"primes must be [L={L}, N-1={N - 1}, H], got {tuple(primes.shape)}", + ) + H = primes.shape[2] + _check(H >= 1 and (H & (H - 1)) == 0, f"the head count H must be a power of two, got {H}") + _check( + offsets.is_contiguous() + and offsets.dtype == torch.int64 + and tuple(offsets.shape) == (L, (N - 1) * H), + f"offsets must be a contiguous int64 [{L}, {(N - 1) * H}] tensor", + ) + if token_mask is not None: + _check_1d("token_mask", token_mask, T, (torch.bool,)) + if out is None: + out = torch.empty((T, L, (N - 1) * H), dtype=torch.int64, device=device) + else: + _check( + out.is_contiguous() + and out.dtype == torch.int64 + and tuple(out.shape) == (T, L, (N - 1) * H), + f"out must be a contiguous int64 [{T}, {L}, {(N - 1) * H}] tensor", + ) + for name, x in ( + ("req_idx_per_token", req_idx_per_token), + ("cu_seq_lens", cu_seq_lens), + ("start_pos", start_pos), + ("token_pool", token_pool), + ("block_table", block_table), + ("token_map", token_map), + ("multipliers", multipliers), + ("primes", primes), + ("offsets", offsets), + ("out", out), + ): + _check(x.device == device, f"{name} is on {x.device}, input_ids on {device}") + if token_mask is not None: + _check( + token_mask.device == device, + f"token_mask is on {token_mask.device}, input_ids on {device}", + ) + if T == 0: + return out + _engram_hash_kernel[(triton.cdiv(T, _BLOCK_T),)]( + input_ids, + token_mask if token_mask is not None else input_ids, + req_idx_per_token, + cu_seq_lens, + start_pos, + token_pool, + token_pool.stride(0), + block_table, + block_table.stride(0), + block_table.shape[1], + token_pool.shape[1], + token_map, + token_map.shape[0], + multipliers, + primes, + offsets, + out, + T, + pad_id, + HAS_MASK=token_mask is not None, + L=L, + N=N, + H=H, + BLOCK_T=_BLOCK_T, + num_warps=_NUM_WARPS, + ) + return out + + +#: DeepSeek-V4.1's (layers, max n-gram size, heads): the shapes of the constant tables. +_V41_LNH = dict(L=2, N=4, H=8) + +#: Prompt lengths driven as one-shot prefill runs in one call. +_PREFILL_LENS = ((1,), (2,), (3,), (4,), (5, 17, 64, 129, 1000), (8192,)) + + +class EngramHash(OpWrapper): + """Engram n-gram hash ids of one step's tokens, with each request's history in a paged pool. + + One launch. A call carries ``B`` runs, one per request, concatenated in batch + order: token row ``t`` belongs to run ``b = req_idx_per_token[t]`` at absolute + position ``p = start_pos[b] + t - cu_seq_lens[b]``. Inside the call: the + token-map lookup (DEAD = -1 where ``token_mask`` is False or the id is outside + ``[0, token_map.numel())``), the store of every token's compressed id at + ``token_pool[block_table[b, p // P], p % P]`` (``P = token_pool.shape[1]``; + skipped where the page entry is negative), the read of each shift-``s`` + predecessor from ``input_ids`` inside the run or from the pool before it, + cumulative blocking to ``pad_id`` at the sequence start and at DEAD tokens, + and the hashes for every Engram layer:: + + r_0 = tok_0 * multipliers[l, 0]; r_s = r_{s-1} XOR tok_s * multipliers[l, s] + out[t, l, (n - 2) * H + h] = r_{n-1} % primes[l, n - 2, h] + offsets[l, (n - 2) * H + h] + + all in int64. The caller owns the constant tables, the pool and its block + table, ``token_mask``, and everything after the ids (table gather, the + reduction across row shards, the projection and the gate). Returns the int64 + ``[T, L, (N - 1) * H]`` hash ids (``out`` when given); ``token_pool`` is + updated in place. ``T = 0`` returns without a launch. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"prefill {'/'.join(map(str, lens))} P={p}: one-shot runs from 0", + spec=dict( + kind="prefill", lens=lens, page_size=p, index_dtype=torch.int32, **_V41_LNH + ), + ) + for lens in _PREFILL_LENS + for p in (128, 256) + ) + + tuple( + Cell( + why=f"chunked+decode P={p} {str(dt).removeprefix('torch.')}: six requests", + spec=dict(kind="chunked", page_size=p, index_dtype=dt, **_V41_LNH), + ) + for p, dt in ( + (128, torch.int32), + (256, torch.int32), + (3, torch.int32), + (128, torch.int64), + ) + ) + + tuple( + Cell( + why=f"verify D={d}: runs of {d + 1}, random acceptance, rewind", + spec=dict( + kind="verify", drafts=d, page_size=128, index_dtype=torch.int32, **_V41_LNH + ), + ) + for d in (1, 4, 5) + ) + + ( + Cell( + why="DEAD via token_mask spans and out-of-range ids", + spec=dict(kind="dead", page_size=128, index_dtype=torch.int32, **_V41_LNH), + ), + Cell( + why="windowed table: unmapped / shared scratch pages", + spec=dict(kind="windowed", page_size=128, index_dtype=torch.int32, **_V41_LNH), + ), + Cell( + why="T=0: empty call, no launch, pool untouched", + spec=dict(kind="empty", page_size=128, index_dtype=torch.int32, **_V41_LNH), + ), + ) + ) + + note = """ + Metadata consumed: none; every input is an argument. The DeepSeek-V4 + attention metadata carries req_idx_per_token, cu_seq_lens_cuda and + cached_token_lens_cuda (as start_pos) in this convention; the pool and its + block table are the cache manager's ENGRAM_TOKEN role + (metadata.engram_token_block_table, tokens_per_block = P). + + Constants, built once at load from the checkpoint recipe: token_map is + engram.py's build_compressed_token_map over the tokenizer (129280 ids -> 99092 + classes); pad_id is the compressed id token_map[2] (equal to 2 for this + tokenizer, so passing the raw id is invisible on real data and wrong on any + other map); multipliers[l] are odd, numpy PCG64 seeded 10007 * layer_id with + the layer ids (1, 14), not their index; primes are 48 distinct, in layer -> + n-gram -> head order, each the next unused prime above 15999999; offsets[l] + is the exclusive cumulative sum of the layer's 24 primes. (L, N, H) = + (2, 4, 8). int64 end to end: the largest product, 99091 * 82619226485591, + lies between 2^62 and 2^63. + + History: the pool holds one int32 compressed id per position (-1 = DEAD). + The op writes every position its runs cover and reads only + [start_pos - N + 1, start_pos) of each run, so chunked prefill, decode and + speculative verify are the same launch: a verify run writes all its + positions, the next run starts after the accepted ones, and the rejected + tail is overwritten before it is read (no commit call). The op cannot create + history it never saw: a run past position 0 whose N - 1 predecessors never + went through the op in this pool (a prefix-cache hit, a disaggregated + hand-off, an encoder bounded replay) reads whatever those slots hold and is + wrong without an error. The caller stores them first with a call whose runs + are only those positions; empty runs are allowed and their start_pos and + block-table row are not read. + + Preconditions the op does not check: req_idx_per_token, cu_seq_lens and + start_pos agree (each run's rows contiguous, positions consecutive); row b + of block_table maps every page covering [start_pos[b] - N + 1, start_pos[b]) + and the positions later calls will read (other pages may be unmapped or a + shared scratch page); no run's read pages are written by the same call; no + request appears twice in one call. It raises ValueError, prefixed + "engram_hash:", on a wrong dtype, device, rank, stride or length of any + tensor and on H not a power of two, before writing anything. + + Launch: fixed, 64 tokens per program, 4 warps, grid ceil(T / 64), no host + synchronization (CUDA-graph capturable). B is not a launch argument. Triton + specializes num_tokens (1, divisible by 16, other), the block-table width + and row stride, the mask and the index dtypes. + + Measured on sm_103 beyond the cells, by the entry's test: every prefix of a + mixed 16384-token step; the engine envelope (B <= 256 runs, T <= 16384 rows, + positions below max_seq_len + 14 with DSpark for max_seq_len up to + 1048576, at the engine's block-table widths 132 .. 8196), unmasked and with + the multimodal target's keep mask on every context run; CUDA-graph capture + of decode (T = G) and verify (T = 6 G) steps at the 46 graph batches at + every max_seq_len; ENGRAM_TOKEN pools up to [57999360, 128] (element + offsets past 2^32); block tables 128 .. 8196 columns, row stride = width or + width + 5; the encoder-replay history call with empty runs; bit identity + with the checkpoint's NgramHashState and the golden hash ids. + + Not certified: page sizes other than 128, 256 and 3; T > 16384, B > 256 or + positions at or past 1048590; table widths or capture batches other than + those run; pools larger than [57999360, 128]; other (L, N, H); a masked call + under CUDA-graph capture; a False mask entry on a generation row; empty runs + in a masked, captured or int64 call; a block table narrower than a run's + last logical page + 1. + """ + + def raw_call( + self, + input_ids: torch.Tensor, + req_idx_per_token: torch.Tensor, + cu_seq_lens: torch.Tensor, + start_pos: torch.Tensor, + token_pool: torch.Tensor, + block_table: torch.Tensor, + token_map: torch.Tensor, + multipliers: torch.Tensor, + primes: torch.Tensor, + offsets: torch.Tensor, + pad_id: int, + token_mask: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return _launch_engram_hash( + input_ids, + req_idx_per_token, + cu_seq_lens, + start_pos, + token_pool, + block_table, + token_map, + multipliers, + primes, + offsets, + pad_id, + token_mask, + out, + ) + + def reference( + self, + input_ids: torch.Tensor, + req_idx_per_token: torch.Tensor, + cu_seq_lens: torch.Tensor, + start_pos: torch.Tensor, + token_pool: torch.Tensor, + block_table: torch.Tensor, + token_map: torch.Tensor, + multipliers: torch.Tensor, + primes: torch.Tensor, + offsets: torch.Tensor, + pad_id: int, + token_mask: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The hash ids, from the checkpoint's ``NgramHashState.forward`` in torch int64 ops. + + A documented exception to a pure-argument reference, like the paged-cache + entries: a row's predecessors before its run are what the pool holds, and + the arguments only say where to look. ``token_pool`` is read at + ``block_table[b, q // P], q % P`` for those positions (an unmapped page + or a column past the table reads DEAD) and must be passed as it was + before the call -- the op writes it. Nothing is written here; ``out`` is + ignored. The pool writes themselves are certified by the test's + cross-call model, which carries each request's history independently of + the pool. + """ + ids = input_ids.long() + t_len = ids.numel() + n_max = multipliers.shape[1] + page_size = token_pool.shape[1] + max_blocks = block_table.shape[1] + vocab = token_map.numel() + rows = torch.arange(t_len, device=ids.device) + req = req_idx_per_token.long() + off = rows - cu_seq_lens.long()[req] + pos = start_pos.long()[req] + off + table = block_table.long() + + comp = torch.where( + (ids >= 0) & (ids < vocab), token_map.long()[ids.clamp(0, vocab - 1)], -1 + ) + if token_mask is not None: + comp = torch.where(token_mask, comp, -1) + blocked = comp < 0 + toks = [torch.where(blocked, pad_id, comp)] + for s in range(1, n_max): + q = pos - s + in_run = off >= s + from_run = comp[(rows - s).clamp_min(0)] + qc = q.clamp_min(0) + blk = torch.div(qc, page_size, rounding_mode="floor") + readable = ~in_run & (q >= 0) & (blk < max_blocks) + page = torch.where(readable, table[req, blk.clamp_max(max_blocks - 1)], -1) + stored = token_pool[page.clamp_min(0), qc % page_size].long() + from_hist = torch.where(page >= 0, stored, -1) + src = torch.where(in_run, from_run, from_hist) + blocked = blocked | (q < 0) | (src < 0) + toks.append(torch.where(blocked, pad_id, src)) + tok = torch.stack(toks, dim=-1) # [T, N] + products = tok.unsqueeze(1) * multipliers.long() # [T, L, N] + rolling, hashes = products[..., 0], [] + for s in range(1, n_max): + rolling = torch.bitwise_xor(rolling, products[..., s]) + hashes.append(rolling.unsqueeze(-1) % primes.long()[:, s - 1]) + return torch.cat(hashes, dim=-1) + offsets.long() + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Bit for bit: integer ids, no rounding anywhere on either side. + + Any difference is a wrong token, a wrong predecessor, a wrong block or a + wrong product. The wrong variants the test runs through the same check -- + ``pad_id`` 2 instead of ``token_map[2]``, int32 or float64 products, + non-cumulative blocking, multipliers seeded by layer index, history not + carried, the page of ``start_pos - 1`` unmapped -- each change rows. + """ + assert out.dtype == ref.dtype == torch.int64, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + rows = (out != ref).flatten(1).any(1) + bad = int(rows.sum()) + assert bad == 0, f"{bad} of {out.shape[0]} rows differ" + + +engram_hash = EngramHash() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/gemm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/gemm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/gemm/mhc_gemm_sqrsum_fma.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/gemm/mhc_gemm_sqrsum_fma.py new file mode 100644 index 000000000000..f5478ca4f717 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/gemm/mhc_gemm_sqrsum_fma.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection mixes GEMM: true-fp32 FMA product of a bf16 stream with fp32 weights, plus sum(x^2). + +The op is `torch.ops.trtllm.mhc_gemm_sqrsum_fma`, registered in +`cpp/tensorrt_llm/thop/mhcOp.cpp`. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_HC_MULT, _HIDDEN = 4, 5120 +#: The flattened [hc_mult, hidden] stream. +_K_HC = _HC_MULT * _HIDDEN +_U = 2.0**-24 # fp32 unit roundoff +#: Median of |y - y64| / ||x * w||_2 allowed over the output. +_FIDELITY = 2.0**-17 + +#: decode (1, 8, 32), DSpark draft/verify blocks (5, 6 per request: 30, 48, 192), +#: an odd count, prefill up to max_num_tokens (4096, 8192). +_HC_MS = (1, 5, 6, 8, 30, 32, 48, 192, 333, 4096, 8192) + +#: The tile_n DeepSeek-V4.1's target passes by row count: (first M, last M, tile_n). +_BANDS = ((1, 128, 1), (129, 2048, 6), (2049, 32768, 24)) +#: Each band's lower edge, interior points and upper edge: the tile-24 band's +#: upper edge is 16384 for a step (max_num_tokens) and 32768 for the decoder pass. +_BAND_POINTS = (1, 64, 128, 129, 1000, 2048, 2049, 5000, 8192, 12000, 16384, 24000, 32768) +#: The trunk's ('mhc', T, tile) records in the core target's runtime-value dumps +#: (boot gates on tep4/dep4/tp4, engines at max_batch_size 256 / max_num_tokens +#: 16384 with CUDA-graph, eager, dep4 and tp4 legs, the 128K and 1M long-context +#: engines and the long-context envelope runs). Every dumped tile equals its band's. +_DUMPED_TRUNK_M = ( + 1, 2, 3, 4, 5, 6, 7, 8, 12, 14, 16, 18, 24, 29, 32, 34, 48, 60, 96, 108, 127, 128, 129, 134, + 144, 168, 192, 228, 240, 288, 330, 336, 384, 432, 462, 480, 528, 546, 576, 597, 610, 624, + 672, 720, 768, 780, 852, 888, 890, 1018, 1152, 1536, 1906, 1935, 2033, 2038, 2048, 2049, + 2252, 2778, 2784, 3870, 4208, 4245, 4388, 4394, 4559, 4591, 4868, 7835, 7855, 7893, 7908, + 8147, 8191, 8192, 8467, 9193, 9240, 9364, 9370, 9524, 10441, 10444, 10447, 10451, 16298, + 16308, 16316, 16320, 16322, 16327, 16328, 16329, 16333, 16340, 16378, 16384, +) # fmt: skip +#: The DSpark drafter's ('dspark.mhc', N, tile) records in the same engines (N = 5G). +_DUMPED_DRAFT_M = ( + 5, 10, 15, 20, 40, 50, 80, 90, 110, 115, 120, 125, 140, 160, 190, 200, 240, 275, 280, 320, + 360, 370, 385, 400, 440, 455, 480, 520, 560, 580, 585, 600, 640, 650, 665, 710, 715, 720, + 960, 1280, +) # fmt: skip + + +def _target_tile(m: int) -> int: + for lo, hi, tile in _BANDS: + if lo <= m <= hi: + return tile + raise ValueError(m) + + +def _band_label(m: int) -> str: + if m in _BAND_POINTS: + return "a band edge or interior point" + if m in _DUMPED_TRUNK_M: + return "a trunk runtime-dump row count" + return "a drafter runtime-dump row count (5G)" + + +def _depth(k: int) -> int: + """Longest rounding chain through the kernel's reduction tree. + + `mhcGemmSqrsumFmaKernel` gives each of 256 threads a sequential FMA chain + (4 per 1024-wide step, plus a scalar tail of up to one per 256 leftover + columns), then a 5-level warp butterfly, then an 8-warp sum. + """ + return 4 * (k // 1024) + -(-(k % 1024) // 256) + 5 + 8 + + +def _gamma(h: int) -> float: + return h * _U / (1 - h * _U) + + +class MhcGemmSqrsumFma(OpWrapper): + """`y = x @ w.T` (true fp32 FMA) and `r = (x.float() ** 2).sum(-1)` in one CUDA-core launch. + + `x` bf16 `[M, K]`, `w` fp32 `[N, K]` (the checkpoint's `hc_*_fn` as stored, + no transpose), `y` fp32 `[M, N]` and `r` fp32 `[M]`, all dense row-major. + The bf16 -> fp32 widening of `x` is exact and every term is fused into an + fp32 accumulator with one rounding (`fmaf`); no tensor cores, no TF32. + + Fusion boundary: the GEMM and the sum of squares are inside the call. The + caller owns flattening the hyper-connection stream `[T, hc_mult, hidden]` to + `[T, hc_mult * hidden]` copy-major (a view of the contiguous stream), the RMS + factor `rsqrt(r / K + norm_eps)` and everything downstream (scale, base, + sigmoid, Sinkhorn), and keeping `y[:, :hc_mult]` and `r` for a later + sublayer. Overwrites `y` and `r` completely, writes nothing else, and + returns `(y, r)`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"heuristic tile M={m} {flavor}: N=24 K=20480 hc mixes", + spec=dict(M=m, K=_K_HC, flavor=flavor, tile_n=0, tile_m=0), + ) + for m in _HC_MS + for flavor in ("gauss", "hc_like") + ) + + tuple( + Cell( + why=f"K={k} M={m}: {what}", + spec=dict(M=m, K=k, flavor="gauss", tile_n=0, tile_m=0), + ) + for k, what in ( + (999, "below 1024, the scalar tail loop only"), + (1024, "exactly one main-loop step, no tail"), + (5124, "main loop plus a 4-column tail"), + ) + for m in (1, 7) + ) + + tuple( + Cell( + why=f"M={m} tile_n={_target_tile(m)}: {_band_label(m)}", + spec=dict(M=m, K=_K_HC, flavor="hc_like", tile_n=_target_tile(m), tile_m=0), + ) + for m in sorted(set(_BAND_POINTS) | set(_DUMPED_TRUNK_M) | set(_DUMPED_DRAFT_M)) + ) + ) + + note = """ + Metadata consumed: none. Stateless. + + DeepSeek-V4.1 calls this twice per layer at N = 24 (hc_mult 4: 4 pre, 4 + post, 16 comb mixes), K = 20480, with M the rows of the pass: up to + `max_num_tokens` (16384), or up to 32768 in the decoder bounded-replay pass + of layers 21-39; the DSpark drafter calls it at M = 5G. `y * rsqrt(r / K + + 1e-20)` is the reference model's `F.linear(x.flatten(2).float(), hc_fn) * + rsqrt(x.square().mean(-1) + 1e-20)`; `pre` derived from this op's outputs + matched the golden layer dumps to at most 1.192e-06 absolute on 8 sublayers. + + Preconditions: the op checks dtypes (RuntimeError) and rejects all-CPU + arguments (NotImplementedError); a mixed-device call was not probed and, + since GB300 can address host memory, should not be expected to fail. It + checks neither layout nor sizes (see `is_valid`). Alignment: `x` and `w` + data pointers 8- and 16-byte aligned and `K % 4 == 0` whenever K >= 1024, + else a sticky CUDA misaligned-address error poisons the process; freshly + allocated tensors and row slices of a contiguous `[T, K]` buffer with + `K % 4 == 0` are aligned. `tile_n > 0` must divide N (the op raises; for + N = 24 that is 1, 2, 3, 4, 6, 8, 12, 24); `tile_m` is not read. M = 0 + returns without launching. + + Deterministic, and bit-identical across every `tile_n` and `tile_m`. The + longest rounding chain is 93 at K = 20480. Built with fast math: subnormal + `x` elements are flushed to zero. Grid `(M, N / tile_n)` of 256-thread + blocks, block `(m, j)` reads row `m` and writes row `m`; `tile_n` selects one + of eight kernels compiled ahead of time. Every prefix `x[:M]`, M = 1..32768, + at its band's `tile_n` returns the first M rows of the 32768-row result and + writes nothing past row M - 1. + + Tactic: `tile_n` changes speed only. The built-in heuristic (`tile_n <= 0`) + picks 1 for M <= 32 and 8 above, slow from M = 40 up; measured on sm_103 the + fast choice is 1 for M <= 128, 6 for 129..2048 and 24 from 4096, which is + what the target passes (24 from 2049). DeepGEMM's `tf32_hc_prenorm_gemm` + computes the same pair (5.5x faster at 4096 rows, 6.9x slower at 1) but + rounds the weight to TF32 (median error about 1000x this op's), so it does + not qualify; the + autotuner in `modules/mhc/mhc_cuda.py` can pick it, and calling this op + directly pins the fp32 FMA path. + + Capturable in a CUDA graph (no host sync, no allocation); certified at the + 114 per-rank capture row counts (G, 5G, 6G over the captured batch sizes + 1..32, 40..128 step 8, 192, 256; at most 1536, so tile 1 or 6), each + captured after an eager call at that M. + + Not certified: N other than 24; K other than 999, 1024, 5124 and 20480; + M above 32768; the built-in heuristic at M values outside its own cells. + """ + + def raw_call( + self, + x: torch.Tensor, + w: torch.Tensor, + y: torch.Tensor, + r: torch.Tensor, + M: int, + N: int, + K: int, + tile_n: int = 0, + tile_m: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + torch.ops.trtllm.mhc_gemm_sqrsum_fma(x, w, y, r, M, N, K, tile_n, tile_m) + return y, r + + def reference( + self, + x: torch.Tensor, + w: torch.Tensor, + y: torch.Tensor, + r: torch.Tensor, + M: int, + N: int, + K: int, + tile_n: int = 0, + tile_m: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The float64 product and the per-element scales `compare` gates against. + + Returns `(y64, r64, bound, l2)`, all float64: `y64 = x @ w.T`, + `r64 = sum x^2`, `bound = gamma_h * sum_k |x_k w_k|` with `h` the + kernel's longest rounding chain at this K, and `l2 = ||x * w||_2` per + element. A torch fp32 product is not an adequate stand-in even with TF32 + off: at M = 4096 and 8192 its largest error measured 10x to 40x this + op's. `y`, `r`, `M`, `N`, `tile_n` and `tile_m` are ignored; nothing is + mutated. + """ + xd, wd = x.double(), w.double() + y64 = xd @ wd.T + r64 = (xd * xd).sum(-1) + bound = (xd.abs() @ wd.abs().T) * _gamma(_depth(K)) + l2 = ((xd * xd) @ (wd * wd).T).sqrt() + return y64, r64, bound, l2 + + def is_valid( + self, + x: torch.Tensor, + w: torch.Tensor, + y: torch.Tensor, + r: torch.Tensor, + M: int, + N: int, + K: int, + tile_n: int = 0, + tile_m: int = 0, + ) -> None: + # The kernel indexes every operand as dense row-major from M, N, K alone and + # never reads a stride or a shape: a strided operand, or an M/N/K that + # disagrees with the tensors, is accepted and gives silently wrong results + # (or writes past the views it was given). + assert x.shape == (M, K) and x.is_contiguous(), "x must be a contiguous bf16 [M, K]" + assert w.shape == (N, K) and w.is_contiguous(), "w must be a contiguous fp32 [N, K]" + assert y.shape == (M, N) and y.is_contiguous(), "y must be a contiguous fp32 [M, N]" + assert r.shape == (M,) and r.is_contiguous(), "r must be a contiguous fp32 [M]" + + def compare( + self, + out: tuple[torch.Tensor, torch.Tensor], + ref: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + ) -> None: + """`y` against a rigorous bound and a fidelity gate; `r` at torch's fp32 defaults. + + * Every element satisfies `|y - y64| <= gamma_h * sum_k |x_k w_k|`, the + forward-error bound of an fp32 dot product with longest rounding chain + `h`. It cannot flake on a correct kernel and catches gross errors (a + wrong row, column, layout or operand). + * The median over elements of `|y - y64| / ||x * w||_2` is at most 2^-17. + fp32 accumulation noise sits near 2^-23; measured on sm_103 the op lands + at 8e-08 to 1.5e-07 on synthetic operands and 1.5e-07 to 2.3e-07 on the + real streams (0.0155x to 0.0193x of the gate), while TF32-rounded + weights land 17.2x to 19.1x above it and bf16-rounded weights 140x to + 158x. + * `r` is a sum of non-negative terms: `assert_close` at fp32 defaults + (largest relative error measured 3.3e-07). + """ + (y, r), (y64, r64, bound, l2) = out, ref + assert y.dtype == r.dtype == torch.float32, (y.dtype, r.dtype) + assert y.shape == y64.shape and r.shape == r64.shape, (y.shape, r.shape) + err = (y.double() - y64).abs() + bound_ratio = (err / bound).max().item() + assert bound_ratio <= 1.0, ( + f"an element exceeds the fp32 rounding-chain bound: {bound_ratio:.3f}x" + ) + fidelity = (err / l2).median().item() + assert fidelity <= _FIDELITY, f"median err/||x*w||_2 = {fidelity:.3e} > {_FIDELITY:.3e}" + torch.testing.assert_close(r, r64.float()) + + +mhc_gemm_sqrsum_fma = MhcGemmSqrsumFma() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/gemm/mxfp8_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/gemm/mxfp8_mxfp8_gemm.py new file mode 100644 index 000000000000..984352ce2bc8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/gemm/mxfp8_mxfp8_gemm.py @@ -0,0 +1,319 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 x MXFP8 dense GEMM in nn.Linear layout: act @ weight.T with 1x32 UE8M0 block scales. + +Wraps the FlashInfer Python API `flashinfer.mm_mxfp8`, not a `torch.ops.trtllm` +op: its availability is a property of the environment (`flashinfer` importable), +not of the TensorRT-LLM build. `backend` selects one of the three certified +kernel families (cuDNN / cuBLASLt, CuTe DSL, CUTLASS); exactly one runs per call. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +from flashinfer import mm_mxfp8 + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +CERTIFIED_BACKENDS = ("cudnn", "cute-dsl", "cutlass") + +_BLOCK = 32 # MXFP8: one UE8M0 scale per 32 contiguous elements along K + +#: The DeepSeek-V4.1-Flash dense FP8 linears per rank, (N, K) in nn.Linear [out, in] +#: order. wq_b is [64 / tp * 512, 1280] and wo_b [5120, 8 / tp * 1024] (tp 1 under +#: attention DP). +_TARGET_SHAPES = ( + (1792, 5120, "wq_a||wkv"), + (8192, 1280, "wq_b, attention TP4"), + (32768, 1280, "wq_b, attention DP"), + (5120, 2048, "wo_b, attention TP4"), + (5120, 8192, "wo_b, attention DP"), + (4608, 5120, "shared w1||w3"), + (5120, 2304, "shared w2"), + (4096, 1280, "indexer wq_b, wq_b TP8"), + (25600, 6144, "engram wkv"), + (16384, 1280, "wq_b, attention TP2"), + (5120, 4096, "wo_b, attention TP2"), + (5120, 1024, "wo_b, attention TP8"), + (5120, 15360, "DSpark main_proj"), +) +#: decode (1..48 incl. DSpark verify rows 6 / 48), a mid batch, a prefill chunk. +_TARGET_ROWS = (1, 6, 8, 32, 48, 192, 4096) +#: The largest prefill chunk, at the widest and the deepest K. +_PREFILL_8192 = ((25600, 6144), (5120, 15360), (1792, 5120)) + + +def _pad_up(x: int, m: int) -> int: + return (x + m - 1) // m * m + + +def _unswizzle(sf: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """The `[rows, cols]` UE8M0 bytes of a flat 128x4-swizzled scale buffer.""" + r = torch.arange(rows, device=sf.device).view(-1, 1) + c = torch.arange(cols, device=sf.device).view(1, -1) + idx = ( + (c % 4) + + (c // 4) * (4 * 128) + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * _pad_up(cols, 4)) + ) + return sf.view(torch.uint8).reshape(-1)[idx] + + +def _dequant(q: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """Exact fp32 value of an MXFP8 operand: e4m3 * 2^(byte - 127), one byte per 32 along K.""" + rows, k = q.shape + scale = torch.exp2(_unswizzle(sf, rows, k // _BLOCK).float() - 127.0) + return q.float() * scale.repeat_interleave(_BLOCK, dim=-1) + + +def _cell(why: str, n: int, k: int, m: int, backend: str, **overrides) -> Cell: + spec = dict( + N=n, K=k, M=m, backend=backend, out_dtype=torch.bfloat16, act_sf=(116, 124), out=False + ) + spec.update(overrides) + return Cell(why=why, spec=spec) + + +class Mxfp8Mxfp8Gemm(OpWrapper): + """`out = a @ b.T` over two MXFP8 operands, fp32 accumulation, rounded once to `out_dtype`. + + `a[m, k] = e4m3(act[m, k]) * 2^(act_sf(m, k // 32) - 127)` and likewise `b` + from `weight` / `weight_sf`; both scale buffers are 1-D uint8 in the 128x4 + swizzled order `torch.ops.trtllm.mxfp8_quantize(x, True)` emits and + `torch.ops.trtllm.block_scale_interleave` produces. `weight` is the + row-major `[N, K]` nn.Linear weight; FlashInfer gets its `.t()` view (the + `[K, N]` column-major operand it expects), which launches nothing. + + Fusion boundary: both operands' block-scale dequantization, the fp32 + accumulation and the output cast are inside the call. The caller owns + producing `act` / `act_sf` (an activation quantizer whose two outputs this op + takes unchanged), the load-time weight-scale preparation, and any bias, + activation or cross-rank reduction; there is no alpha, global scale or bias. + Returns the `[M, N]` result: `out` itself when given, else a fresh tensor. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + _cell(f"{be} M={m} {n}x{k}: {name}", n, k, m, be) + for n, k, name in _TARGET_SHAPES + for m in _TARGET_ROWS + for be in CERTIFIED_BACKENDS + ) + + tuple( + _cell(f"{be} M=8192 {n}x{k}: the largest prefill chunk", n, k, 8192, be) + for n, k in _PREFILL_8192 + for be in CERTIFIED_BACKENDS + ) + + tuple( + _cell(f"{be} M={m} 512x1024: act scale rows pad to 128", 512, 1024, m, be) + for m in (1, 3, 7, 9, 127, 128, 129, 1000) + for be in CERTIFIED_BACKENDS + ) + + tuple( + _cell(f"{be} M=8 160x1024: weight scale rows pad 160->256", 160, 1024, 8, be) + for be in CERTIFIED_BACKENDS + ) + + tuple( + _cell(f"{be} M={m} 5120x576: K/32 = 18 scale cols pad to 20", 5120, 576, m, be) + for m in (1, 8, 300) + for be in CERTIFIED_BACKENDS + ) + + tuple( + _cell( + f"{be} fp16 out M=48 1792x5120: act scales 2^-6 lower", + 1792, + 5120, + 48, + be, + out_dtype=torch.float16, + act_sf=(110, 118), + ) + for be in CERTIFIED_BACKENDS + ) + + tuple( + _cell(f"{be} out= M=48 1792x5120: written and returned", 1792, 5120, 48, be, out=True) + for be in CERTIFIED_BACKENDS + ) + ) + + note = """ + Metadata consumed: FlashInfer's process-global autotuner cache, separate + from TensorRT-LLM's autotuner. Nothing must be prepared: outside a + `flashinfer.autotune()` context every call takes its backend's fallback + tactic, the path certified here. `cute-dsl` JIT-compiles each variant on + first use (about 0.5 s; a swap-AB kernel with N tile 8/16/32 at M <= 32, + 128x128 above); `cudnn` builds one plan per (N, K, M bucket) on the first + call in the bucket (about 20 ms) and launches a kernel that still changes + with M inside a bucket. If anything in the process tunes `mm_mxfp8` under + `flashinfer.autotune()`, later calls take the tuned tactic: not certified. + TensorRT-LLM's engine warmup does not tune it for a modeling_v2 target. + + Checkpoint weights store one UE8M0 byte per 32x32 tile; this op wants one + byte per row and K block: `sf_rows = scale.view(torch.uint8) + .repeat_interleave(32, dim=0)[:N]`, then `block_scale_interleave(sf_rows)`. + Tiling instead (`scale.repeat(32, 1)`) is accepted and wrong. Scale buffers + are exactly `pad_up(rows, 128) * pad_up(K / 32, 4)` uint8 bytes; padding + bytes are ignored. A row-major scale buffer of the right length is accepted + and silently wrong on every backend, and no guard can see it. + + Preconditions: `act` and `weight` 2-D float8_e4m3fn, contiguous (see + `is_valid`); N >= 128 and K >= 128 (ValueError); N % 32 == 0 and K % 32 == 0 + are the certified domain; M >= 1 (M = 0 raises on every backend); `out_dtype` + bf16 or fp16 (fp32 raises; fp16 overflow is inf); `out` `[M, N]` contiguous. + Wrong scale lengths, a non-uint8 scale view and a non-e4m3 `act` raise. + + Backends: the target picks `cute-dsl` at M <= 32 and `cudnn` above (measured + on sm_103 the fastest at 38 of 40 and 40 of 40 cells respectively). + `cute-dsl` and `cutlass` are bitwise equal to each other at every certified + cell; `cudnn` is held to the bound only: it differs in the last bits at + M <= 48 for K >= 5120 (split-K plus a reduce kernel, two launches), and a + row's bits can depend on the batch's M. Every backend is deterministic. + Fused projections (`wq_a || wkv`, `w1 || w3`, the drafter's `wkv` x3) are + within the bound of the separate GEMMs and bitwise equal on `cute-dsl` and + `cutlass`; on `cudnn` they differ in the last bits in a window of M (at + K = 5120: 129..640 for `wq_a || wkv`, 49..128 for `w1 || w3`, 193..640 for + `wkv` x3). Eager host cost at M = 1 is about 187 / 93 / 82 us + (cudnn / cute-dsl / cutlass); CUDA-graph replay removes it. + + CUDA graphs: the first call a process makes on a backend must be eager + (capturing it fails); after one eager call capture works at row counts + never run eagerly, and replay is bit-identical to eager. FlashInfer keeps + one `mm_mxfp8_workspace` buffer per device; the first `cudnn` call grows it + to 33,554,688 bytes at a new address, and no measured call moves it after + that. A larger request (an M, shape or cuDNN not measured) would move it + under every captured `cudnn` graph. + + Row counts certified at the twelve core linears: every M 1..32 on + `cute-dsl` and `cudnn`; the bottom, interior and top of every FlashInfer + bucket from 33 to 32768 and every dumped row count on all three backends; + every M 33..N_max on `cudnn` on a column sample, and every `cudnn` kernel + switch at full width. N_max is 32768 at the eleven decoder-layer shapes and + 16384 at the Engram `wkv`. The drafter's `main_proj` and fused `wkv` x3 run + at every M 1..1280 on all three backends. Captured on the target's backend + at G and 6G (5G where the drafter has the shape) over the batch sizes + 1..32, 40..128 step 8, 192, 256, and at every M 1..160 for the two drafter + shapes; `cutlass` captured at M = 8 and 256 only. + + Not certified: M above N_max; on `cudnn` between 33 and N_max, errors + confined to columns outside the sample except at the full-width values; + `cute-dsl` and `cutlass` above 32 rows at M other than the bucket and dumped + values; the drafter shapes above 1280 rows other than 4096 and 8192 at + `main_proj`; capture at row counts or shapes other than those listed; + non-finite inputs; a workspace request above 33,554,688 bytes; any tactic + chosen by `flashinfer.autotune()`. + """ + + def raw_call( + self, + act: torch.Tensor, + act_sf: torch.Tensor, + weight: torch.Tensor, + weight_sf: torch.Tensor, + out_dtype: torch.dtype = torch.bfloat16, + backend: str = "cudnn", + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return mm_mxfp8( + act, + weight.t(), + act_sf, + weight_sf, + out=out, + out_dtype=out_dtype, + use_8x4_sf_layout=False, + backend=backend, + ) + + def reference( + self, + act: torch.Tensor, + act_sf: torch.Tensor, + weight: torch.Tensor, + weight_sf: torch.Tensor, + out_dtype: torch.dtype = torch.bfloat16, + backend: str = "cudnn", + out: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """The exact product of the dequantized operands in true fp32, and the bound `compare` gates on. + + Returns `(ref, bound)`, both fp32 `[M, N]`: `ref = a @ b.T` with `a` and + `b` dequantized exactly from the e4m3 data and the swizzled scale bytes, + and `bound = K * 2^-24 * sum_k |a_mk b_nk| + eps(out_dtype) * |ref|`. + The product runs with TF32 off (`true_fp32_matmul`). `backend` and `out` + are ignored; nothing is mutated. + """ + a = _dequant(act, act_sf) + b = _dequant(weight, weight_sf) + with true_fp32_matmul(): + ref = a @ b.T + l1 = a.abs() @ b.abs().T + bound = act.shape[1] * 2.0**-24 * l1 + torch.finfo(out_dtype).eps * ref.abs() + return ref, bound + + def is_valid( + self, + act: torch.Tensor, + act_sf: torch.Tensor, + weight: torch.Tensor, + weight_sf: torch.Tensor, + out_dtype: torch.dtype = torch.bfloat16, + backend: str = "cudnn", + out: Optional[torch.Tensor] = None, + ) -> None: + # Every guard is pure metadata, and every guarded violation was observed on + # this machine to be accepted and computed wrongly without raising: the + # cute-dsl backend ignores the strides of a non-contiguous act or weight; + # cute-dsl and cutlass write a non-contiguous `out` as if it were + # contiguous (landing outside the view); FlashInfer's "trtllm" backend + # reads the same operands in a shuffled layout they are not in. + assert backend in CERTIFIED_BACKENDS, ( + f"backend must be one of {CERTIFIED_BACKENDS}; got {backend!r} (the 'trtllm' " + "backend needs shuffled operands and returns wrong results on these)" + ) + assert act.is_contiguous() and weight.is_contiguous(), ( + "act [M, K] and weight [N, K] must be contiguous; the cute-dsl backend " + "ignores strides and returns wrong results" + ) + assert out is None or out.is_contiguous(), ( + "out must be contiguous; the cute-dsl and cutlass backends write it as " + "contiguous regardless of its strides" + ) + + def compare(self, out: torch.Tensor, ref: tuple[torch.Tensor, torch.Tensor]) -> None: + """Every element within the fp32-accumulation bound: `|out - ref| <= bound`. + + Both operands are exact in fp32 and every block scale is a power of two, + so the kernel's only error sources are the order of its fp32 + accumulation over K, which the tensor-core MMA fixes and the reference + cannot reproduce (bounded per element by the recursive-summation bound + `K * 2^-24 * sum_k |a_mk b_nk|`), and one rounding of the fp32 sum to + the output dtype (taken as a full `eps * |ref|`). torch's default bf16 + `assert_close` does not model the first: outputs that cancel to near + zero carry an absolute error of that size, and it failed on 63 to 65 of + 80 measured shape and row cells for every candidate GEMM. Measured on + sm_103 the largest `|err| / bound` over every certified check was + 0.4940, while a row-major weight or activation scale, a tiled 32x32 + expansion and every scale byte off by one land 101x to 5810x outside it. + A single misplaced scale byte touches one column and can hide under the + bound at M = 1. + """ + expected, bound = ref + assert out.dtype in (torch.bfloat16, torch.float16), out.dtype + assert out.shape == expected.shape, (out.shape, expected.shape) + err = (out.float() - expected).abs() + ratio = (err / bound.clamp_min(torch.finfo(torch.float32).tiny)).max().item() + assert ratio <= 1.0, f"max |out - ref| / bound = {ratio:.3g} (> 1)" + + +mxfp8_mxfp8_gemm = Mxfp8Mxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/moe/gate_forward.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/moe/gate_forward.py new file mode 100644 index 000000000000..b41332a7a70e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/moe/gate_forward.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""sqrt(softplus) MoE routing: biased top-k selection, unbiased renormalized weights. + +The op is `torch.ops.trtllm.gate_forward`, registered in +`cpp/tensorrt_llm/thop/moe/routing/moeGateOp.cpp`; this entry always runs its +bias mode (`is_hash=False`). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: (n_routed_experts, num_experts_per_tok): the V4.1 main MoE, the DSpark drafter +#: MoE, and the third pair the op dispatches. +_SURFACES = ((384, 6), (128, 3), (256, 6)) +#: Decode-like through prefill-like token counts. +_SWEEP_T = (1, 2, 3, 4, 5, 6, 7, 8, 15, 30, 40, 64, 160, 1024, 4096, 8192, 262144) +#: (logit sigma, bias centre): the checkpoint's per-layer biases sit in a narrow band +#: around a large offset (9.1..57.6 over the 40 layers); the drafter's are exactly 0. +_DISTS = ((1.0, 0.0), (4.0, 12.0), (12.0, 57.0)) +_ROUTE_SCALE = 1.5 # routed_scaling_factor + + +def _warp_sum(w: torch.Tensor) -> torch.Tensor: + """Sum of the topk weights in the kernel's order: a 32-lane xor butterfly + (offsets 16, 8, 4, 2, 1) with lane j holding slot j and the other lanes 0. + + For 6 addends this is ((w0 + w4) + w2) + ((w1 + w5) + w3); for 3 it is + (w0 + w2) + w1. + """ + lanes = [w[:, j] for j in range(w.shape[1])] + lanes += [torch.zeros_like(lanes[0])] * (32 - len(lanes)) + for offset in (16, 8, 4, 2, 1): + lanes = [lanes[j] + lanes[j ^ offset] for j in range(32)] + return lanes[0] + + +def _cell(why: str, surface: tuple[int, int], tokens: int, **overrides) -> Cell: + spec = dict( + E=surface[0], + topk=surface[1], + T=tokens, + sigma=3.0, + center=10.0, + route_scale=_ROUTE_SCALE, + out=False, + ) + spec.update(overrides) + return Cell(why=why, spec=spec) + + +class GateForward(OpWrapper): + """Route each token to `topk` experts by `sqrt(softplus(logit)) + bias`. + + For every row of `router_logits` `[T, E]`, in fp32: + `s = sqrt(L > 20 ? L : log1p(exp(L)))` (torch's softplus threshold); the + `topk` experts with the largest `s + bias`, by decreasing score, equal + scores in ascending expert id; their weights from the *unbiased* `s`, + divided by `sum + 1e-20` (the sum in a warp-butterfly order) and multiplied + by `route_scale`. Bias mode only. + + Fusion boundary: softplus, sqrt, the bias add, top-k, the gather, the + renormalization and the scaling are inside the one launch. The caller owns + the router GEMM (fp32 logits), any `bias` / `bias_vl` choice and everything + downstream. `out_weights` (fp32) and `out_indices` (int32), when given, are + written in place; otherwise both are allocated. Returns + `(topk_weights [T, topk] fp32, topk_ids [T, topk] int32)`, weights first. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + _cell( + f"{e}/{k} T={t} sigma={sigma:g} bias={center:g}: token sweep", + (e, k), + t, + sigma=sigma, + center=center, + ) + for e, k in _SURFACES + for t in _SWEEP_T + for sigma, center in _DISTS + ) + + tuple( + _cell(f"{e}/{k} route_scale={s:g}: T=512 sigma=3 bias=10", (e, k), 512, route_scale=s) + for e, k in _SURFACES[:2] + for s in (1.0, 1.5, 2.5, 0.0, -1.0) + ) + + tuple( + _cell( + f"{e}/{k} caller out buffers: T=8192, written and returned", (e, k), 8192, out=True + ) + for e, k in _SURFACES[:2] + ) + + tuple( + _cell(f"{e}/{k} T=0: two empty [0, topk] outputs", (e, k), 0) for e, k in _SURFACES[:2] + ) + ) + + note = """ + Metadata consumed: none. Stateless: no cache, workspace or tuner; runs on + the ambient stream; repeated calls and a call on another stream give + identical bits. CUDA-graph capturable (the allocation of absent outputs is + a graph-pool allocation). + + Matches the checkpoint's `Gate.forward` (`gate_temp` 1, `norm_topk_prob` + true) bit for bit on real data. The kernel uses torch's softplus threshold + (`L > 20` gives `L`, else `log1pf(expf(L))`) and divides by `sum + 1e-20`, + as the checkpoint does, so from the kernel source: every finite logit gives + a finite score (no expf overflow; the golden prompts' logits reach 52.63), + and a row whose selected scores are all exactly 0 (every logit <= ~-104, or + -inf) gets weights `0 / 1e-20 * route_scale`, i.e. 0. Only a `+inf` logit + still gives NaN (that expert, `inf / inf`) and 0 for the row's others, and + a NaN logit is selected first and makes the row's weights NaN. One + difference remains, at exact ties: equal selection scores come out in + ascending expert id where `torch.topk` emits another order; with the + checkpoint's large biases 29 of 161,360 golden rows reorder tied experts + (same set, weights within 2 ulp). + + This entry passes one bias per call: the image-span + `torch.where(image_mask, bias_vl, bias)` is two calls plus a per-token + select of both outputs (bit-exact). The op's optional `image_mask` / + `vision_bias` arguments (a one-call per-token bias) are not exposed by this + entry and not certified here. Weights + are fp32; the trtllm-gen MoE runners take bf16 `topk_weights`, so the + caller casts. The summation order (the 32-lane butterfly above) is pinned: + another order matched only 78-83% of the weights. + + Preconditions: `(E, topk)` one of `(384, 6)`, `(128, 3)`, `(256, 6)` (the + op raises otherwise); fp32 logits and bias, fp32 / int32 outputs, CUDA + tensors (the op raises on each). A CPU `bias` is not rejected and is read + from host memory: not certified. Layouts the op does not police are in + `is_valid`; an empty bias is an illegal memory access that kills the CUDA + context. T = 0 returns two `[0, topk]` tensors without a launch: the op's + grid would have zero blocks, which CUDA rejects as an invalid launch. + + Row axis: one ahead-of-time kernel per `(E, topk)`, one warp per row, the + row count only sizes the grid. At `(384, 6)`, `route_scale` 1.5, every T + 1..8192 (the target's per-call MoE chunk cap) returns the first T rows of + the 8192-row call bitwise; at `(128, 3)` with a zero bias (the drafter's + call) every T 1..1280 does. Captured at `(384, 6)` at the 167 per-rank + sizes 1..6144 (G and 6G, and the attention-DP gathered 2G/12G, 4G/24G, + 8G/48G over the batch sizes 1..32, 40..128 step 8, 192, 256), plus dep8's + 9216 and 12288 rows as two chunk calls (8192 + 1024, 8192 + 4096) in one + graph; at `(128, 3)` at the 46 sizes 5G, 5..1280. + + Not certified: hash mode; a CPU bias; at `(384, 6)` capture of one call + above 6144 rows or at a size outside the 167, and other chunk splits; at + `(128, 3)` the prefix sweep above 1280, capture outside the 46 sizes (other + than T 40) or with a nonzero bias; several calls in one graph other than + dep8's two chunks; the prefix sweep and capture list at `(256, 6)`; the + drafter's real input distribution (certified on proxy inputs). + """ + + def raw_call( + self, + router_logits: torch.Tensor, + bias: torch.Tensor, + topk: int, + route_scale: float, + out_weights: Optional[torch.Tensor] = None, + out_indices: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens = router_logits.shape[0] + if out_weights is None: + out_weights = torch.empty( + (num_tokens, topk), dtype=torch.float32, device=router_logits.device + ) + if out_indices is None: + out_indices = torch.empty( + (num_tokens, topk), dtype=torch.int32, device=router_logits.device + ) + if num_tokens == 0: + # The kernel's grid is one block per four rows, and a zero-block launch is + # cudaErrorInvalidValue; there is nothing to route, so nothing is launched. + return out_weights, out_indices + # The hash mode's two inputs; empty in bias mode (never read). + no_ids = torch.empty(0, dtype=torch.int32, device=router_logits.device) + torch.ops.trtllm.gate_forward( + router_logits, + bias, + no_ids, + no_ids, + out_weights, + out_indices, + topk, + route_scale, + False, + ) + return out_weights, out_indices + + def reference( + self, + router_logits: torch.Tensor, + bias: torch.Tensor, + topk: int, + route_scale: float, + out_weights: Optional[torch.Tensor] = None, + out_indices: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """The checkpoint's `Gate.forward` in fp32, from native torch ops. + + The kernel's numerics: `sqrt(softplus(x))` with the softplus threshold + 20 (`x > 20` gives `x`, else `log1p(exp(x))`), a *stable* descending + sort of `scores + bias` (equal selection scores in ascending expert id, + the kernel's tie-break), the gather from the bare scores, the + warp-order sum, `/ (sum + 1e-20)` and then `* route_scale`. + `out_weights` and `out_indices` are ignored and not mutated; the + expected outputs are new tensors. + """ + x = router_logits.float() + scores = F.softplus(x, beta=1.0, threshold=20.0).sqrt() + choice = scores + bias.float() + ids = torch.sort(choice, dim=-1, descending=True, stable=True).indices[:, :topk] + weights = scores.gather(1, ids) + weights = weights / (_warp_sum(weights).unsqueeze(-1) + 1e-20) + weights = weights * route_scale + return weights, ids.to(torch.int32) + + def is_valid( + self, + router_logits: torch.Tensor, + bias: torch.Tensor, + topk: int, + route_scale: float, + out_weights: Optional[torch.Tensor] = None, + out_indices: Optional[torch.Tensor] = None, + ) -> None: + # Pure-metadata guards. Each violation below was observed on sm_103 to be + # accepted by the op and to return a plausible-looking wrong answer (or to + # read/write outside the tensor) instead of raising: the kernel addresses + # every tensor as a dense row-major buffer from data_ptr(), reads E = + # size(1) bias elements and writes T * topk outputs regardless of shapes. + assert router_logits.dim() == 2, ( + f"router_logits must be 2-D [T, E], got {tuple(router_logits.shape)}; a 3-D " + "input is read as size(0) rows of size(1) experts" + ) + assert router_logits.is_contiguous(), ( + "router_logits must be contiguous; a strided view is read as a dense " + "[T, E] buffer and silently gives wrong routing" + ) + num_tokens, num_experts = router_logits.shape + assert tuple(bias.shape) == (num_experts,), ( + f"bias must be [{num_experts}], got {tuple(bias.shape)}; the kernel reads " + f"{num_experts} elements from bias.data_ptr() whatever its length" + ) + assert bias.is_contiguous(), ( + "bias must be contiguous; a strided view is read as a dense [E] buffer " + "and silently gives wrong routing" + ) + for name, out in (("out_weights", out_weights), ("out_indices", out_indices)): + if out is None: + continue # raw_call allocates it at [T, topk] + assert tuple(out.shape) == (num_tokens, topk) and out.is_contiguous(), ( + f"{name} must be contiguous [{num_tokens}, {topk}], got " + f"{tuple(out.shape)} strides {out.stride()}; the kernel writes a dense " + "[T, topk] block from data_ptr()" + ) + + def compare( + self, + out: tuple[torch.Tensor, torch.Tensor], + ref: tuple[torch.Tensor, torch.Tensor], + ) -> None: + """Bit-exact on both the ids and the weights. + + The kernel evaluates the same fp32 operations in the same order as the + reference: the softplus threshold 20 and `log1pf(expf(x))` below it + (bit-identical to torch's softplus), a correctly rounded `sqrtf`, one + fp32 add of the bias, the warp-order sum, one add of 1e-20, one + division, one multiply. Measured on sm_103 bit-exact on every element of + every certified case, so any tolerance would only hide a defect; the + wrong-reference controls (weights from the biased score, sigmoid scoring, + a missing `route_scale`) differ on 100% of the weights. + """ + (weights, ids), (ref_weights, ref_ids) = out, ref + assert weights.dtype == torch.float32 and ids.dtype == torch.int32, ( + weights.dtype, + ids.dtype, + ) + assert weights.shape == ref_weights.shape and ids.shape == ref_ids.shape, ( + weights.shape, + ids.shape, + ) + assert torch.equal(ids, ref_ids), ( + f"{int((ids != ref_ids).any(-1).sum())} of {ids.shape[0]} rows route differently" + ) + assert torch.equal(weights, ref_weights), ( + f"{int((weights != ref_weights).sum())} of {weights.numel()} weights differ, " + f"max abs {float((weights - ref_weights).abs().max()):.3e}" + ) + + +gate_forward = GateForward() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/engram_gate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/engram_gate.py new file mode 100644 index 000000000000..1a2de10b5cfd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/engram_gate.py @@ -0,0 +1,440 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Engram gated residual update, per (token, hc copy), in one launch. + +The DeepSeek-V4.1 Engram gate (``Engram.forward`` of the checkpoint's +``inference/model.py``, after the ``wkv`` projection): for every token ``t`` +and hc copy ``c`` of the residual stream ``x [T, HC, D]``, with +``key [T, HC, D]``, ``value [T, D]`` and ``w = q_weight * k_weight [HC, D]``, +all in fp32:: + + rstd = rsqrt(mean_d(x[t,c]^2) + eps) * rsqrt(mean_d(key[t,c]^2) + eps) + dot = sum_d(x[t,c] * w[c] * key[t,c]) * rstd * D**-0.5 + gate = sigmoid(copysign(sqrt(max(|dot|, clamp_value)), dot)) # 0 where token_mask is False + out[t,c] = x[t,c] + gate * value[t] # one rounding to out's dtype + +The statistics are per (token, copy) over ``D``, never joint over the copies. + +The kernel is a Triton kernel defined in this file (a port of SGLang's +``fused_engram_gate``), not a ``torch.ops.trtllm`` op: no engine op computes +this gate. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# Launch configuration: fixed, no autotune. One program per (token, hc copy) +# row, the whole row in one block of next_pow2(D) lanes, _NUM_WARPS warps. +_NUM_WARPS = 8 + + +@triton.jit +def _engram_gate_kernel( + x_ptr, + key_ptr, + value_ptr, + qw_ptr, + kw_ptr, + mask_ptr, + out_ptr, + gate_ptr, + stride_xt, + stride_xc, + stride_kt, + stride_kc, + stride_vt, + stride_qc, + stride_wc, + stride_ot, + stride_oc, + stride_mt, + stride_gt, + stride_gc, + eps, + clamp_value, + D: tl.constexpr, + HC: tl.constexpr, + BLOCK: tl.constexpr, + HAS_MASK: tl.constexpr, + WRITE_GATE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + t = row // HC + c = row - t * HC + col = tl.arange(0, BLOCK) + live = col < D + x = tl.load(x_ptr + t * stride_xt + c * stride_xc + col, mask=live, other=0.0).to(tl.float32) + key = tl.load(key_ptr + t * stride_kt + c * stride_kc + col, mask=live, other=0.0).to( + tl.float32 + ) + qw = tl.load(qw_ptr + c * stride_qc + col, mask=live, other=0.0).to(tl.float32) + kw = tl.load(kw_ptr + c * stride_wc + col, mask=live, other=0.0).to(tl.float32) + w = qw * kw + # mean as sum * fp32(1/D), rsqrt / sqrt / exp / division as the CUDA math + # library computes them (what torch's float kernels call), so the only + # difference from the torch reference is the summation order. + inv_d = 1.0 / D + ms_x = tl.sum(x * x, axis=0) * inv_d + ms_k = tl.sum(key * key, axis=0) * inv_d + rstd = libdevice.rsqrt(ms_x + eps) * libdevice.rsqrt(ms_k + eps) + dot = tl.sum((x * w) * key, axis=0) * rstd * (D**-0.5) + z = libdevice.copysign( + tl.sqrt_rn(tl.maximum(tl.abs(dot), clamp_value, propagate_nan=tl.PropagateNan.ALL)), dot + ) + gate = tl.div_rn(1.0, 1.0 + libdevice.exp(-z)) + if HAS_MASK: + keep = tl.load(mask_ptr + t * stride_mt) != 0 + gate = tl.where(keep, gate, 0.0) + if WRITE_GATE: + tl.store(gate_ptr + t * stride_gt + c * stride_gc, gate) + v = tl.load(value_ptr + t * stride_vt + col, mask=live, other=0.0).to(tl.float32) + tl.store(out_ptr + t * stride_ot + c * stride_oc + col, x + gate * v, mask=live) + + +_HC = 4 # hc_mult +_DIM = 5120 # hidden_size +_EPS = 1e-20 # rms_norm_eps +_CLAMP = 1e-6 # Engram.clamp_value +_BF16 = torch.bfloat16 +_FP32 = torch.float32 + +#: Token counts driven against the reference in the plain call form (out=None, +#: gate_out given), decode-like to the per-rank max_num_tokens ceiling. +_CALL_T = (1, 2, 5, 37, 160, 1024, 8192, 16384) + +#: The keep masks driven at T = 16384 in the multimodal target's call form. +_TARGET_MASKS = ( + "none_masked", + "all_masked", + "span_at_row0", + "span_at_end", + "one_span_mid", + "multi_span", + "alternating", +) + tuple(f"rand{k}" for k in range(9)) + + +def _call_spec(t_len, masked, x_dtype=_BF16, kv_dtype=_BF16, w_dtype=_BF16, mask_dtype=torch.bool): + return dict( + kind="call", + T=t_len, + masked=masked, + x_dtype=x_dtype, + kv_dtype=kv_dtype, + w_dtype=w_dtype, + mask_dtype=mask_dtype, + mask_stride=1, + HC=_HC, + D=_DIM, + eps=_EPS, + clamp_value=_CLAMP, + ) + + +def _dt(dtype): + return str(dtype).removeprefix("torch.") + + +class EngramGate(OpWrapper): + """Engram gate + residual update: ``out[t, c] = x[t, c] + gate[t, c] * value[t]``. + + One Triton launch over ``T * HC`` (token, copy) rows. Inside the call: both + RMS statistics, the weighted dot, the signed-sqrt sigmoid gate, the mask, + the broadcast of the one ``value`` row to every copy, and the output cast. + The caller owns the Engram lookup and the ``wkv`` projection that produce + ``key`` / ``value`` (two views of one fused buffer are fine, no copy), the + choice of layers, applying the update before the layer's ``hc_pre`` mixing, + and building ``token_mask``. + + ``x`` / ``key`` are ``[T, HC, D]``, ``value`` is ``[T, D]``, ``q_weight`` / + ``k_weight`` are ``[HC, D]``, each with unit stride along ``D`` and any other + strides. All arithmetic is fp32. ``token_mask`` (``[T]``, bool or integer, + any stride, nonzero = keep) forces ``gate = 0`` where it is False/0. The + result goes to ``out`` (``[T, HC, D]``, unit stride along ``D``; may be + ``x`` itself for an in-place update) or to a new contiguous ``x.dtype`` + tensor, rounded once from fp32. ``gate_out`` (``[T, HC]``, any strides) + receives the fp32 gate when given. Returns the output tensor (``out`` when + given); ``T = 0`` returns it without a launch. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"T={t} {'masked' if m else 'no mask'}: new out, split wkv views, bf16", + spec=_call_spec(t, m), + ) + for t in _CALL_T + for m in (False, True) + ) + + tuple( + Cell( + why=f"dtypes x={_dt(xd)} kv={_dt(kd)} w={_dt(wd)}: T=160 masked", + spec=_call_spec(160, True, xd, kd, wd), + ) + for xd, kd, wd in ((_BF16, _FP32, _FP32), (_FP32, _FP32, _FP32), (_FP32, _BF16, _BF16)) + ) + + tuple( + Cell( + why=f"mask {_dt(md)} stride 3: T=64, masked rows pass x through", + spec={**_call_spec(64, True, mask_dtype=md), "mask_stride": 3}, + ) + for md in (torch.bool, torch.int32, torch.uint8) + ) + + ( + Cell( + why="clamp and sign edges: |dot| 0 and 1e-12..1e2, both signs, zero rows", + spec=dict(kind="edges", HC=_HC, D=_DIM, eps=_EPS, clamp_value=_CLAMP), + ), + ) + + tuple( + Cell( + why=f"mm T=16384 {kind}: in place, bool keep, fused wkv", + spec=dict( + kind="target_masked", + T=16384, + mask=kind, + HC=_HC, + D=_DIM, + eps=_EPS, + clamp_value=_CLAMP, + ), + ) + for kind in _TARGET_MASKS + ) + ) + + note = """ + DeepSeek-V4.1: HC = 4 (hc_mult), D = 5120 (hidden_size), eps = 1e-20 (the + config's rms_norm_eps, passed as norm_eps), clamp_value = 1e-6 (a constant + of the reference's Engram.__init__, not a config field). x is the bf16 hc + stream entering layers 1 and 14; key / value are the first 4 * 5120 and the + last 5120 columns of the bf16 wkv output [T, 25600]; q_weight / k_weight are + the checkpoint's bf16 [4, 5120] tensors as stored; out=x updates in place. + token_mask is None for text-only batches; the multimodal target passes a + contiguous bool keep (False on image-span rows) in steps holding an image + request's context chunk, and that mask may be all True. + + Metadata consumed: none. + + Semantics the members state only by example: the statistics are per (token, + copy) over D, never joint over the copies; the sqrt is signed and |dot| is + clamped before it, so dot = 0 gives sigmoid(+sqrt(clamp_value)), not 0.5; + masked rows get gate exactly 0 and equal x bit for bit; kept rows are bit + for bit what the same call computes with token_mask=None. + + Preconditions: every tensor on the same CUDA device (a CPU tensor makes + Triton's launcher raise ValueError); unit stride along D on x, key, value, + the weights and out; out is x itself or overlaps no input (a partial + overlap was not run). The kernel reads its whole row before writing it, + which is what makes out=x safe. + + Batch invariant: one program per (token, copy) row with a fixed block, and + T only sizes the grid, so a row does not depend on T or on its position. + Measured on sm_103: every prefix r = 1..16384 of a 16384-token call is bit + identical to the first r rows of it, with no mask and with two masks. + + CUDA-graph capturable: no host synchronization, no allocation when out is + given. Certified captured at T = G and 6 G over the engine's 46 graph + batches (82 counts, 1..1536) in the core's form: no mask, no gate_out, in + place, key / value split views of one static fused buffer. The masked form + runs eagerly only (captured steps are generation-only); its capture was run + only at T = 1, 5, 160. + + Accuracy: the kernel's three sums over D run in Triton's reduction order; + the pointwise tail uses the functions torch's float kernels call. That is + the only difference from the reference, and rows whose sums are exact in + any order match it bit for bit. The launch is fixed (8 warps, fp fusion + off, no autotune); the reduction order depends on the warp count, so a + different launch gives different last-bit gates and needs its own cells. + + Not certified: other HC or D; fp16; T > 16384 in one call; capture at + another T; masked capture beyond T = 1, 5, 160; other eps / clamp_value; + an out dtype other than x.dtype; a gate_out dtype other than fp32; NaN or + inf inputs; a partial overlap of out with an input. + """ + + def raw_call( + self, + x: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float, + clamp_value: float, + token_mask: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, + gate_out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + t_len, hc, dim = x.shape + if out is None: + out = torch.empty((t_len, hc, dim), dtype=x.dtype, device=x.device) + if t_len == 0: + return out + _engram_gate_kernel[(t_len * hc,)]( + x, + key, + value, + q_weight, + k_weight, + token_mask if token_mask is not None else x, + out, + gate_out if gate_out is not None else x, + x.stride(0), + x.stride(1), + key.stride(0), + key.stride(1), + value.stride(0), + q_weight.stride(0), + k_weight.stride(0), + out.stride(0), + out.stride(1), + token_mask.stride(0) if token_mask is not None else 0, + gate_out.stride(0) if gate_out is not None else 0, + gate_out.stride(1) if gate_out is not None else 0, + eps, + clamp_value, + D=dim, + HC=hc, + BLOCK=triton.next_power_of_2(dim), + HAS_MASK=token_mask is not None, + WRITE_GATE=gate_out is not None, + num_warps=_NUM_WARPS, + enable_fp_fusion=False, + ) + return out + + def is_valid( + self, + x: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float, + clamp_value: float, + token_mask: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, + gate_out: Optional[torch.Tensor] = None, + ) -> None: + """The layouts the kernel reads or writes as if they were the expected one. + + Each was run through the launch without these checks on sm_103 and + returned without an error: a stride-2 view along ``D`` is read as if + contiguous (output bit-exact on 0.1-18% of elements), and a view with + fewer rows / columns / copies than ``x`` implies is read or written past + its end (the parent buffer's data is used, or 8 rows past it written). + Not guarded, because the kernel reads them correctly: any non-unit + stride other than the one along ``D``, integer mask dtypes, a strided + mask. + """ + # Every guard below is a layout the kernel would otherwise read or write + # as if it were the expected one, without an error (probed on sm_103). + assert x.dim() == 3 and x.stride(2) == 1, ( + f"x must be [T, HC, D] with unit stride along D, got {tuple(x.shape)} {x.stride()}" + ) + t_len, hc, dim = x.shape + assert tuple(key.shape) == (t_len, hc, dim) and key.stride(2) == 1, ( + f"key must be [{t_len}, {hc}, {dim}] with unit stride along D, got {tuple(key.shape)} {key.stride()}" + ) + assert tuple(value.shape) == (t_len, dim) and value.stride(1) == 1, ( + f"value must be [{t_len}, {dim}] with unit stride along D, got {tuple(value.shape)} {value.stride()}" + ) + for name, wt in (("q_weight", q_weight), ("k_weight", k_weight)): + assert tuple(wt.shape) == (hc, dim) and wt.stride(1) == 1, ( + f"{name} must be [{hc}, {dim}] with unit stride along D, got {tuple(wt.shape)} {wt.stride()}" + ) + if token_mask is not None: + assert tuple(token_mask.shape) == (t_len,), ( + f"token_mask must be [{t_len}], got {tuple(token_mask.shape)}" + ) + if out is not None: + assert tuple(out.shape) == (t_len, hc, dim) and out.stride(2) == 1, ( + f"out must be [{t_len}, {hc}, {dim}] with unit stride along D, got {tuple(out.shape)} {out.stride()}" + ) + if gate_out is not None: + assert tuple(gate_out.shape) == (t_len, hc), ( + f"gate_out must be [{t_len}, {hc}], got {tuple(gate_out.shape)}" + ) + + def reference( + self, + x: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float, + clamp_value: float, + token_mask: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, + gate_out: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """``(out, gate)``: the expected output and the expected fp32 gate ``[T, HC]``. + + Torch ops in the order of the checkpoint's ``Engram.forward``: per-copy + ``mean`` over ``D``, ``rsqrt``, the product sum, the signed sqrt with + the clamp, ``sigmoid``, ``masked_fill``, one cast to the output dtype. + The gate is what ``gate_out`` receives (0 on masked rows) and is + returned whether or not ``gate_out`` is given. ``x`` is the stream as it + was before the call -- for an in-place call pass a copy taken before + it. ``out`` only sets the output dtype (``x.dtype`` when None) and + ``gate_out`` is ignored; neither is written. + """ + dim = x.shape[-1] + k = key.float() + weight = q_weight.float() * k_weight.float() + h = x.float() + rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt(k.square().mean(-1) + eps) + dot = (h * weight * k).sum(-1) * rstd * dim**-0.5 + gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(clamp_value).sqrt(), dot)) + if token_mask is not None: + gate = gate.masked_fill(~token_mask.bool().unsqueeze(-1), 0) + out_dtype = x.dtype if out is None else out.dtype + expected = (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(out_dtype) + return expected, gate + + def compare( + self, + out: tuple[torch.Tensor, Optional[torch.Tensor]], + ref: tuple[torch.Tensor, torch.Tensor], + ) -> None: + """Gate ``(out, gate_out)`` -- the two tensors a call fills -- against ``reference``. + + ``gate_out`` may be None for a call made without one; then only the + output is gated. The kernel differs from the reference only in the order + of its three fp32 sums over ``D``, so: + + 1. the gate at torch's default fp32 tolerance; + 2. the output at torch's default tolerance of its dtype. A per-ulp bound + is not usable: where ``x`` and ``gate * value`` nearly cancel, one + last-bit gate difference is several ulps of the small result. + + Measured on sm_103 at T = 1..16384 and on the golden layers: gate max abs + diff 1.5e-08 to 8.6e-06, output bit-exact on 0.99996 to 1.0 of bf16 + elements (about 0.91 for fp32 x, max abs diff 1.5e-05). + The test also checks the output bit for bit against ``x + gate * value`` + formed in torch from the kernel's own gate (no FMA contraction, one + rounding), which needs ``x`` and ``value`` and so lives there. The + joint-RMS, unsigned-sqrt, unclamped and wrong-mask controls each fail + these gates on every case where they differ. + """ + (got, gate), (ref_out, ref_gate) = out, ref + assert got.dtype == ref_out.dtype, (got.dtype, ref_out.dtype) + assert got.shape == ref_out.shape, (got.shape, ref_out.shape) + if gate is not None: + torch.testing.assert_close(gate, ref_gate) + torch.testing.assert_close(got, ref_out) + + +engram_gate = EngramGate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/mhc_big_fuse.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/mhc_big_fuse.py new file mode 100644 index 000000000000..058f8f4ba85e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/mhc_big_fuse.py @@ -0,0 +1,385 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection mix coefficients via the trtllm mhc_big_fuse op. + +From the raw mixes and square sums of the hc mixes GEMM: RMS-scaled sigmoid `post`, a +Sinkhorn-normalized `comb`, and (when `hidden_size > 0`) the stream collapsed with the FRESH +`pre`. DeepSeek-V4.1 lags `pre` by one sublayer, so it consumes only `post_mix` and `comb_mix`. + +The op is `torch.ops.trtllm.mhc_big_fuse`, registered in `cpp/tensorrt_llm/thop/mhcOp.cpp`. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Residual-stream copies; compiled into the kernel. +_HC = 4 +#: Columns of a mixes row: pre 0:4 | post 4:8 | comb 8:24. +_MIX_COLS = _HC * (2 + _HC) + +#: Row counts DeepSeek-V4.1 makes its call at, from the target's runtime-value dumps: the trunk's +#: ('mhc', T, tile) records (T 1..16384, engines up to max_batch_size 256 / max_num_tokens 16384, +#: CUDA-graph and eager, dep4 and tp4, 128K and 1M max_seq_len, the long-context envelope) and +#: the DSpark drafter's ('dspark.mhc', N, tile) records (N = 5G); plus 17 and the tops of the +#: step range (16384) and of the decoder bounded-replay pass (32768). +_TARGET_ROWS = tuple( + sorted( + { + 1, 2, 3, 4, 5, 6, 7, 8, 12, 14, 16, 18, 24, 29, 32, 34, 48, 60, 96, 108, 127, 128, + 129, 134, 144, 168, 192, 228, 240, 288, 330, 336, 384, 432, 462, 480, 528, 546, 576, + 597, 610, 624, 672, 720, 768, 780, 852, 888, 890, 1018, 1152, 1536, 1906, 1935, 2033, + 2038, 2048, 2049, 2252, 2778, 2784, 3870, 4208, 4245, 4388, 4394, 4559, 4591, 4868, + 7835, 7855, 7893, 7908, 8147, 8191, 8192, 8467, 9193, 9240, 9364, 9370, 9524, 10441, + 10444, 10447, 10451, 16298, 16308, 16316, 16320, 16322, 16327, 16328, 16329, 16333, + 16340, 16378, 16384, + } + | { + 5, 10, 15, 20, 40, 50, 80, 90, 110, 115, 120, 125, 140, 160, 190, 200, 240, 275, 280, + 320, 360, 370, 385, 400, 440, 455, 480, 520, 560, 580, 585, 600, 640, 650, 665, 710, + 715, 720, 960, 1280, + } + | {17, 4095, 8191, 8193, 16383, 16385, 32767, 32768} + ) +) # fmt: skip + + +def _spec( + M: int, + hidden_size: int = 5120, + num_splits: int = 1, + block_size: int = 0, + zero_rows: int = 0, + sinkhorn_repeat: int = 20, + hc_post_mult_value: float = 2.0, + kind: str = "stream", +) -> dict: + """One cell's spec. `kind` "stream" builds the mixes from a [M, 4, 5120] stream the way the + mixes GEMM does; "extreme" feeds logits of magnitude ~1e3-1e4 at unit RMS.""" + return dict( + kind=kind, + M=M, + hidden_size=hidden_size, + num_splits=num_splits, + block_size=block_size, + zero_rows=zero_rows, + sinkhorn_repeat=sinkhorn_repeat, + hc_post_mult_value=hc_post_mult_value, + ) + + +def layer_input_error_ratio( + layer_input: torch.Tensor, exact: torch.Tensor, terms: torch.Tensor +) -> float: + """max(|layer_input - exact| / bound) under the `layer_input` gate (see `MhcBigFuse.compare`).""" + bound = 2.0**-7 * exact.abs() + 16 * 2.0**-24 * terms + bound = ( + bound + torch.finfo(torch.float64).tiny + ) # an all-zero padding row must come out exactly 0 + return ((layer_input.double() - exact).abs() / bound).max().item() + + +class MhcBigFuse(OpWrapper): + """The hyper-connection coefficient step for a 4-copy residual stream, one CTA per token. + + From the hc mixes GEMM's raw mixes `y[t, 0:24]` and square sum `r[t]` (optionally split-K + partials), for rows `t < M`:: + + m = y * rsqrt(r / K + rms_eps) # pre 0:4 | post 4:8 | comb 8:24 + pre = sigmoid(m[0:4] * hc_scale[0] + hc_base[0:4]) + hc_pre_eps (not returned) + post = sigmoid(m[4:8] * hc_scale[1] + hc_base[4:8]) * hc_post_mult_value -> post_mix + comb = Sinkhorn(m[8:24] * hc_scale[2] + hc_base[8:24]) as [4, 4] -> comb_mix + layer_input = bf16(sum_j pre[j] * residual[t, j, :]) only when hidden_size > 0 + + Sinkhorn is a row softmax plus `hc_sinkhorn_eps`, a column pass, then `sinkhorn_repeat - 1` + (row, column) passes, every denominator `+ hc_sinkhorn_eps`. + + Fusion boundary: inside the call are the split-K reduction, the RMS scaling, both sigmoids, + the whole Sinkhorn normalization and (when `hidden_size > 0`) the collapse with the FRESH + `pre`. The caller owns the mixes GEMM and square sum that produce `y_acc` / `r_acc`, the + lagged `pre` collapse DeepSeek-V4.1 actually feeds its sublayers, the RMSNorm of the sublayer + input, and hc_post. Writes rows `[0, M)` of `post_mix`, `comb_mix` and (when + `hidden_size > 0`) `layer_input` in place and returns `(post_mix, comb_mix, layer_input)`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"full call M={m}: V4.1 constants" + + (", one all-zero padding row" if m > 1 else ""), + spec=_spec(m, zero_rows=1 if m > 1 else 0), + ) + for m in (1, 5, 6, 64, 65, 1024, 8192) + ) + + tuple( + Cell( + why=f"block_size={bs} M={m}: full-call launch shape", + spec=_spec(m, block_size=bs), + ) + for m in (1, 64, 65, 300) + for bs in (0, 128, 256, 512) + ) + + tuple( + Cell( + why=f"post/comb-only block_size={bs} M={m}: hidden_size 0", + spec=_spec(m, hidden_size=0, block_size=bs), + ) + for m in (1, 6, 4096) + for bs in (0, 128, 256) + ) + + tuple( + Cell( + why=f"num_splits={s} M=33: split-K partials summed in the kernel", + spec=_spec(33, num_splits=s), + ) + for s in (1, 2, 4, 8, 16) + ) + + tuple( + Cell( + why=f"sinkhorn_repeat={n} hc_post_mult={p}: M=48 full call", + spec=_spec(48, sinkhorn_repeat=n, hc_post_mult_value=p), + ) + for n in (1, 20) + for p in (1.0, 2.0) + ) + + ( + Cell( + why="logits ~1e3-1e4 M=64: softmax and sigmoids stay finite", + spec=_spec(64, kind="extreme"), + ), + ) + + tuple( + Cell( + why=f"target call M={m}: hidden_size 0, block_size 128, one split", + spec=_spec(m, hidden_size=0, block_size=128, zero_rows=1 if m > 2 else 0), + ) + for m in _TARGET_ROWS + ) + ) + + note = """ + DeepSeek-V4.1 constants: rms_eps 1e-20, hc_pre_eps = hc_sinkhorn_eps = 1e-6, + hc_post_mult_value 2.0, sinkhorn_repeat 20, K = 4 * 5120 = 20480. With them `post` and + `comb` are exactly the checkpoint's `hc_split_sinkhorn` (same Sinkhorn order, same eps + placement). `comb_mix[t, j, k]` is row j (source copy), column k (destination copy); the + softmax runs along the row, and V4.1's hc_post sums over j, consuming the tensor unchanged. + + `layer_input` is NOT V4.1's hc_pre: the kernel collapses with the `pre` of this same call, + while V4.1 lags `pre` by one sublayer. Measured on golden layer dumps it is off by relative + L2 5.3e-2 to 9.0 from the lagged collapse. V4.1 therefore calls with hidden_size 0, + block_size 128, num_splits 1 and consumes only `post_mix` / `comb_mix`; the caller computes + the lagged collapse from the previous sublayer's saved `y[:, 0:4]` and `r`, which this call + does not return. Post/comb-only mode (hidden_size 0) neither reads `residual` nor writes + `layer_input` (both may be zero-size, `[M, 4, 0]` / `[M, 0]`); `post_mix` / `comb_mix` are + bit-identical to a hidden_size 5120 call. `K` stays the flattened stream width in that mode. + + Metadata consumed: none. Stateless: no workspace, no autotuner. + + Preconditions: fp32 `y_acc`, `r_acc`, `hc_scale`, `hc_base`, `post_mix`, `comb_mix`; bf16 + `residual`, `layer_input` (a bf16 `y_acc` and an fp32 `residual` raise RuntimeError, the + other dtypes were not probed). num_splits in {1, 2, 4, 8, 16}, else the op raises + RuntimeError. With num_splits > 1 split s of token t is read at flat row s * M + t, so + `y_acc` / `r_acc` must be exactly `[num_splits, M, ...]`. Shapes are trusted beyond dtype; + hidden_size is a multiple of 8. `sinkhorn_repeat <= 0` computes exactly what 1 does. + An all-zero token row (a CUDA-graph padding row) gives finite outputs and a zero + `layer_input` row. Rows at and beyond M are never written; M = 0 touches nothing. + + block_size: 0 or negative = auto (128 for M <= 16, else 256); 1-255 -> 128, 256-511 -> 256, + >= 512 -> 512; all bit-identical. 512 with num_splits 1 and M <= 64 stages the residual + through a bulk async copy; that path was not run in post/comb-only mode, use 0, 128 or 256 + there. The V4.1 call launches exactly one kernel at every M, + `mhcBigFuseKernel<1, 128, false, false, false>`, compiled ahead of time. Measured on sm_103, + post/comb-only at block_size 128 is the fastest or within 0.1 us of the fastest candidate at + every M (11.3 us at M = 8192, against 69-75 us for any call that also collapses). + + Certified for every M in 1..32768 of the V4.1 call: each prefix of a 32768-row call is + bitwise equal to its first M rows, and the call captured into a CUDA graph at every + decode / DSpark-verify / drafter row count up to 1536 replays bitwise equal to eager. + + The registered op also takes `norm_weight`, `norm_eps`, `pre_mix_ext`, `pre_mix_out` (a + fused next-layer RMSNorm and a lagged-`pre` collapse); this entry leaves them at their + defaults and certifies nothing about them. The kernel is launched with PDL when + TRTLLM_ENABLE_PDL is on. No fake/meta implementation; torch.compile not exercised. + + Not certified: hidden_size other than 5120 and 0; K other than 20480; block_size >= 512 + with hidden_size 0; num_splits > 1 at M other than 33 or with more than M rows per split; + M > 32768; eps other than V4.1's; sinkhorn_repeat other than 1 and 20; hc_post_mult_value + other than 1.0 and 2.0; non-contiguous `hc_scale` / `hc_base`. + """ + + def raw_call( + self, + y_acc: torch.Tensor, + r_acc: torch.Tensor, + residual: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + layer_input: torch.Tensor, + M: int, + K: int, + hidden_size: int, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + num_splits: int, + block_size: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + torch.ops.trtllm.mhc_big_fuse( + y_acc, + r_acc, + residual, + hc_scale, + hc_base, + post_mix, + comb_mix, + layer_input, + M, + K, + hidden_size, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + num_splits, + block_size, + ) + return post_mix, comb_mix, layer_input + + def is_valid( + self, + y_acc: torch.Tensor, + r_acc: torch.Tensor, + residual: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + layer_input: torch.Tensor, + M: int, + K: int, + hidden_size: int, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + num_splits: int, + block_size: int, + ) -> None: + # The kernel addresses every one of these as a dense row-major buffer. A strided view of any + # of the six was accepted and read or wrote the wrong elements with no error. + assert y_acc.is_contiguous() and r_acc.is_contiguous(), "y_acc and r_acc must be contiguous" + assert residual.is_contiguous(), "residual must be contiguous [M, 4, hidden_size]" + assert post_mix.is_contiguous() and comb_mix.is_contiguous(), ( + "post_mix and comb_mix must be contiguous" + ) + assert layer_input.is_contiguous(), "layer_input must be contiguous" + + def reference( + self, + y_acc: torch.Tensor, + r_acc: torch.Tensor, + residual: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + layer_input: torch.Tensor, + M: int, + K: int, + hidden_size: int, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + num_splits: int, + block_size: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """`(post_mix, comb_mix, layer_input, layer_input_terms)` expected after the call. + + fp32 torch ops in the kernel's operation order for `pre` / `post` / `comb`; split partials + are summed in split order, read at flat row `s * M + t` like the kernel. `layer_input` is + returned as the EXACT fp64 collapse of that fp32 `pre` (not rounded to bf16: the gate in + `compare` carries the final rounding), and `layer_input_terms` is `sum_j |pre_j * x_j|` + in fp64, the scale of the kernel's fp32 work, which the gate needs. Each expected output + is a copy of the passed buffer with rows `[0, M)` replaced, so rows the op must not + write keep their contents; nothing passed in is mutated. `layer_input` (and + `layer_input_terms`) rows are only replaced when `hidden_size > 0`. `block_size` is + ignored (it picks the launch shape only); `residual` is only read when + `hidden_size > 0`. + """ + y_rows = y_acc.reshape(-1, _MIX_COLS).float() + r_rows = r_acc.reshape(-1).float() + y, r = y_rows[:M].clone(), r_rows[:M].clone() + for s in range(1, num_splits): + y += y_rows[s * M : (s + 1) * M] + r += r_rows[s * M : (s + 1) * M] + eps = hc_sinkhorn_eps + mixes = y * torch.rsqrt(r / K + rms_eps)[:, None] + pre = torch.sigmoid(mixes[:, :_HC] * hc_scale[0] + hc_base[:_HC]) + hc_pre_eps + post = ( + torch.sigmoid(mixes[:, _HC : 2 * _HC] * hc_scale[1] + hc_base[_HC : 2 * _HC]) + * hc_post_mult_value + ) + comb = (mixes[:, 2 * _HC :] * hc_scale[2] + hc_base[2 * _HC :]).view(-1, _HC, _HC) + comb = comb.softmax(-1) + eps + comb = comb / (comb.sum(-2, keepdim=True) + eps) + for _ in range(sinkhorn_repeat - 1): + comb = comb / (comb.sum(-1, keepdim=True) + eps) + comb = comb / (comb.sum(-2, keepdim=True) + eps) + + exp_post = post_mix.clone() + exp_post.view(-1, _HC)[:M] = post + exp_comb = comb_mix.clone() + exp_comb.view(-1, _HC * _HC)[:M] = comb.reshape(-1, _HC * _HC) + exp_li = layer_input.to(torch.float64, copy=True) + terms_sum = torch.zeros_like(exp_li) + if hidden_size > 0: + x = residual.reshape(-1, _HC, hidden_size)[:M].double() + terms = pre.double()[:, :, None] * x + exp_li.view(-1, hidden_size)[:M] = terms.sum(1) + terms_sum.view(-1, hidden_size)[:M] = terms.abs().sum(1) + return exp_post, exp_comb, exp_li, terms_sum + + def compare( + self, + out: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ref: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + ) -> None: + """`post_mix` / `comb_mix` at fp32 `assert_close` defaults; `layer_input` per element. + + `post` / `comb`: the kernel is built with fast math (`expf`, reciprocals), which + accounts for last-ulp differences; rtol 1.3e-6, atol 1e-5 held in every certified cell. + + `layer_input` against the exact fp64 collapse: one bf16 ulp of the exact value (the final + rounding, which a few-ulp fp32 difference can push across a rounding midpoint) plus 16 + fp32 ulps of S = sum_j |pre_j * x_j|. The second term is the kernel's fp32 work: a 4-term + FMA chain rounds each running sum (<= 4 ulps of S), and its fast-math sigmoid differs + from torch's by a few ulps of each pre_j. A row that cancels to near zero keeps that + absolute error, which a purely relative bf16 gate misreads: the default bf16 + `assert_close` failed on 2 of 41943040 such elements at M = 8192 (greatest relative + difference 0.625). Observed max error / bound under this gate: 0.498. Controls at + M = 256: a 19-pass Sinkhorn lands 350x the fp32 atol away, a transposed comb 97000x, + and the lagged-`pre` collapse 1.2e6x the `layer_input` bound. + """ + (post, comb, li), (r_post, r_comb, r_li, r_terms) = out, ref + torch.testing.assert_close(post, r_post) + torch.testing.assert_close(comb, r_comb) + assert li.shape == r_li.shape, (li.shape, r_li.shape) + if li.numel(): + worst = layer_input_error_ratio(li, r_li, r_terms) + assert worst <= 1.0, f"layer_input outside the gate: max err/bound = {worst:.3f}" + + +mhc_big_fuse = MhcBigFuse() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/mhc_hc_head_apply.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/mhc_hc_head_apply.py new file mode 100644 index 000000000000..d220a8ab4107 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/mhc_hc_head_apply.py @@ -0,0 +1,301 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection collapse with caller-supplied mixes via the mHC hc_head_apply kernel. + +The mixes and the RMS statistic are inputs, not derived from ``x``: they may come +from a different (earlier) stream than the one being collapsed, which is how +DeepSeek-V4.1's lagged ``hc_pre`` uses this op. + +The op is `torch.ops.trtllm.mhc_hc_head_apply`, registered in `cpp/tensorrt_llm/thop/mhcOp.cpp`. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The op is built with fast math and is not bit-identical to torch; measured on sm_103 between +#: 0.999863 and 1.0 of elements are, in every certified cell. Every wrong variant the test +#: rejects keeps at most 0.40. +_MIN_BITEXACT = 0.999 + +#: Row counts DeepSeek-V4.1 makes its call at, from the target's runtime-value dumps: the trunk's +#: ('mhc', T, tile) records (T 1..16384, engines up to max_batch_size 256 / max_num_tokens 16384, +#: CUDA-graph and eager, dep4 and tp4, 128K and 1M max_seq_len, the long-context envelope) and +#: the DSpark drafter's ('dspark.mhc', N, tile) records (N = 5G); plus the tops of the step range +#: (16384) and of the decoder bounded-replay pass (32768). +_TARGET_ROWS = tuple( + sorted( + { + 1, 2, 3, 4, 5, 6, 7, 8, 12, 14, 16, 18, 24, 29, 32, 34, 48, 60, 96, 108, 127, 128, + 129, 134, 144, 168, 192, 228, 240, 288, 330, 336, 384, 432, 462, 480, 528, 546, 576, + 597, 610, 624, 672, 720, 768, 780, 852, 888, 890, 1018, 1152, 1536, 1906, 1935, 2033, + 2038, 2048, 2049, 2252, 2778, 2784, 3870, 4208, 4245, 4388, 4394, 4559, 4591, 4868, + 7835, 7855, 7893, 7908, 8147, 8191, 8192, 8467, 9193, 9240, 9364, 9370, 9524, 10441, + 10444, 10447, 10451, 16298, 16308, 16316, 16320, 16322, 16327, 16328, 16329, 16333, + 16340, 16378, 16384, + } + | { + 5, 10, 15, 20, 40, 50, 80, 90, 110, 115, 120, 125, 140, 160, 190, 200, 240, 275, 280, + 320, 360, 370, 385, 400, 440, 455, 480, 520, 560, 580, 585, 600, 640, 650, 665, 710, + 715, 720, 960, 1280, + } + | {4095, 8191, 8193, 16383, 16385, 32767, 32768} + ) +) # fmt: skip + + +def pre_weights( + mixes: torch.Tensor, + sqrsum: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + mult: int, + K: int, + norm_eps: float, + eps: float, +) -> torch.Tensor: + """The collapse weights `pre [rows, mult]` the op forms from its statistics, in fp32 torch.""" + rstd = torch.rsqrt(sqrsum / K + norm_eps) + return torch.sigmoid(mixes * rstd[:, None] * scale[0] + base[:mult]) + eps + + +def collapse_error_ratio(out: torch.Tensor, exact: torch.Tensor, terms: torch.Tensor) -> float: + """max(|out - exact| / bound) under the per-element gate of `MhcHcHeadApply.compare`. + + An element that is NaN in both `out` and `exact` (a row the op does not write, NaN-filled + by the caller) counts as exact. + """ + if not out.numel(): + return 0.0 + bound = 2.0**-7 * exact.abs() + 16 * 2.0**-24 * terms + torch.finfo(torch.float64).tiny + err = (out.double() - exact).abs() + err = torch.where(out.isnan() & exact.isnan(), torch.zeros_like(err), err) + return (err / bound).max().item() + + +class MhcHcHeadApply(OpWrapper): + """Collapse a `mult`-copy residual stream into one hidden vector per token, from given mixes. + + For every token `t < M`:: + + rstd[t] = rsqrt(sqrsum[t] / K + norm_eps) + pre[t, j] = sigmoid(mixes[t, j] * rstd[t] * scale[0] + base[j]) + eps j < mult + out[t, :] = bf16( sum_{j < mult} pre[t, j] * x[t, j, :] ) fp32 accumulation, j in order + + `pre` is a function of `mixes[t]`, `sqrsum[t]`, `scale[0]`, `base[:mult]`, `K`, `norm_eps` + and `eps` only -- `x` is never read to form it, so the statistics may describe a different + (earlier) stream than the one being collapsed. That cross-stream use is what is certified. + + Fusion boundary: the call does the RMS scaling, the sigmoid, the `+ eps`, the weighted sum and + the bf16 cast. The caller owns the mixes GEMM and sum of squares (and keeping them until the + consuming sublayer), and any norm applied to the collapsed vector afterwards; `post` and + `comb` are not computed here. Writes rows `[0, M)` of `out` in place and returns `out`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"lagged collapse M={m}: independent producer and consumer streams", + spec=dict(M=m, eps=1e-6, saturated=False), + ) + for m in (1, 5, 6, 32, 1000, 8192) + ) + + ( + Cell( + why="saturated sigmoid M=256: pre ~ eps and ~ 1 + eps on two copies", + spec=dict(M=256, eps=1e-6, saturated=True), + ), + Cell( + why="eps=0.25 M=64: the + eps after the sigmoid is visible in bf16", + spec=dict(M=64, eps=0.25, saturated=False), + ), + ) + + tuple( + Cell( + why=f"target call M={m}: mult 4, hidden 5120, K 20480, eps 1e-6", + spec=dict(M=m, eps=1e-6, saturated=False), + ) + for m in _TARGET_ROWS + ) + ) + + note = """ + Every cell: mult 4, hidden_size 5120, K 20480, norm_eps 1e-20, fresh producer and consumer + streams drawn independently, and (M >= 3) an all-zero padded row (sqrsum 0, so rstd 1e10 + and pre = sigmoid(base) + eps, finite) and an attention-sink row of stream norm ~2e6. + + DeepSeek-V4.1's lagged hc_pre. Pass the statistics of the producer stream through the + PRODUCER sublayer's `hc_*_fn`, that sublayer's `scale` / `base`, and the current stream as + `x`: the attention input of layer L >= 1 uses the stream entering layer L-1's FFN with + `hc_ffn_*` of L-1 (x: the stream entering layer L, after the Engram add on 1 and 14); the + FFN input of layer L uses the stream entering layer L's attention with `hc_attn_*` of L (x: + the stream after L's attention hc_post); the final-norm input uses the stream entering layer + 39's FFN with `hc_ffn_*` of 39 (x: layer 39's output; the checkpoint has no hc_head + tensors). Layer 0's attention input and each DSpark stage's first collapse use the one-hot + pre [1, 0, 0, 0], which this op cannot produce: that is copy 0 of the stream, not a call. + Each site was checked against the checkpoint's golden layer dumps. Wiring the consuming + sublayer's `scale` / `base` instead gives row-relative error median 0.02 to 0.81. + + `mixes[t, j]` is the RAW product of the flattened producer row with row j of `hc_*_fn` (no + rsqrt applied) and `sqrsum[t]` the SUM (not mean) of its squares: `y[:, :mult]` and `r` of + `torch.ops.trtllm.mhc_gemm_sqrsum_fma`, with `y[:, :mult]` copied to a contiguous buffer. + Sibling ops that collapse with the fresh `pre` of the same stream + (`torch.ops.trtllm.mhc_big_fuse`'s `layer_input`, `torch.ops.trtllm.mhc_fused_hc`) compute + the wrong collapse under V4.1's lag. + + Metadata consumed: none. Stateless; no allocation, no host synchronization; capturable. + + Preconditions: fp32 `mixes`, `sqrsum`, `scale`, `base`; bf16 `x`, `out` -- a wrong dtype + raises RuntimeError before launch. The op reads no shape or stride; the layouts it silently + miscomputes are in `is_valid`. A contiguous view with a storage offset is fine. Only + `scale[0]` and `base[:mult]` are read (the checkpoint's [3] and [24] pass as is). Rows >= M + of `out` are not written, no input is modified, and M = 0 touches nothing. CUDA tensors + (a CPU `mixes` was accepted without error on GB300; not certified). hidden_size % 8 == 0 + with 16-byte-aligned rows (8-element vector accesses, not probed); mult <= 8 (an 8-entry + shared array), only 4 certified. + + Arithmetic: one CTA of 256 threads per token, built with fast math (approximate exp, + reciprocal, rsqrt), a 4-term fp32 sum in copy order, one RNE rounding to bf16. At V4.1's + eps 1e-6 the `+ eps` sits below bf16 resolution (omitting it keeps 0.998 of elements + bit-identical), so no output-level check confirms a caller passed it; its handling is + certified at eps 0.25, where omitting it is rejected. + + Every M in 1..32768 of the V4.1 call: each prefix of a 32768-row call is bitwise equal to its + first M rows; the call launches exactly one `mhcHcHeadApplyKernel`, compiled ahead of time, + at every M; captured into a CUDA graph at every decode / DSpark-verify / drafter row count + up to 1536 it replays bitwise equal to eager. + + Not certified: M > 32768; mult != 4, hidden_size != 5120, K != 20480, norm_eps != 1e-20; + CPU inputs; torch.compile (no fake implementation is registered). + """ + + def raw_call( + self, + mixes: torch.Tensor, + sqrsum: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + M: int, + mult: int, + hidden_size: int, + K: int, + norm_eps: float, + eps: float, + ) -> torch.Tensor: + torch.ops.trtllm.mhc_hc_head_apply( + mixes, sqrsum, x, out, scale, base, M, mult, hidden_size, K, norm_eps, eps + ) + return out + + def is_valid( + self, + mixes: torch.Tensor, + sqrsum: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + M: int, + mult: int, + hidden_size: int, + K: int, + norm_eps: float, + eps: float, + ) -> None: + # The op addresses every tensor by raw pointer at the geometry the integer + # arguments describe and never reads a shape or a stride. Each layout + # rejected here was measured on sm_103 to return a wrong result without an + # error (and, for a strided or short `out`, to write outside it). + assert mixes.dim() == 2 and mixes.is_contiguous(), ( + "mixes must be a contiguous [>=M, mult] tensor" + ) + assert mixes.shape[1] == mult and mixes.shape[0] >= M, ( + "mixes must be a contiguous [>=M, mult] tensor" + ) + assert sqrsum.dim() == 1 and sqrsum.is_contiguous() and sqrsum.shape[0] >= M, ( + "sqrsum must be a contiguous [>=M] tensor" + ) + assert x.dim() == 3 and x.is_contiguous(), ( + "x must be a contiguous [>=M, mult, hidden_size] tensor" + ) + assert tuple(x.shape[1:]) == (mult, hidden_size) and x.shape[0] >= M, ( + "x must be a contiguous [>=M, mult, hidden_size] tensor" + ) + assert out.dim() == 2 and out.is_contiguous(), ( + "out must be a contiguous [>=M, hidden_size] tensor" + ) + assert out.shape[1] == hidden_size and out.shape[0] >= M, ( + "out must be a contiguous [>=M, hidden_size] tensor" + ) + + def reference( + self, + mixes: torch.Tensor, + sqrsum: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + M: int, + mult: int, + hidden_size: int, + K: int, + norm_eps: float, + eps: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """`(out, exact, terms)` expected after the call, each the shape of `out`. + + `pre` in fp32 torch (`pre_weights`). `out` is the torch collapse: fp32 weighted sum over + the copies, rounded once to bf16. `exact` is the same collapse in fp64, unrounded, and + `terms` is `sum_j |pre_j * x_j|` in fp64: the scale of the kernel's fp32 work, which the + gate in `compare` needs. Rows `>= M` keep the passed `out`'s contents (`terms` 0 there); + nothing passed in is mutated. `hidden_size` is implied by `x` and not read. + """ + pre = pre_weights(mixes[:M], sqrsum[:M], scale, base, mult, K, norm_eps, eps) + expected = out.clone() + expected[:M] = torch.sum(pre[:, :, None] * x[:M].float(), dim=1).to(torch.bfloat16) + terms = pre.double()[:, :, None] * x[:M].double() + exact = out.to(torch.float64, copy=True) + exact[:M] = terms.sum(1) + term_scale = torch.zeros_like(exact) + term_scale[:M] = terms.abs().sum(1) + return expected, exact, term_scale + + def compare( + self, + out: torch.Tensor, + ref: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ) -> None: + """A per-element term bound, plus at least 0.999 of elements bit-identical to torch. + + Per element: one bf16 ulp of the exact value (the final rounding, which a few-ulp fp32 + difference can push across a rounding midpoint) plus 16 fp32 ulps of + S = sum_j |pre_j * x_j|. The kernel forms `pre` with fast-math exp, reciprocal and rsqrt, + a few fp32 ulps from torch's, and sums the four products in fp32: both err in proportion + to the terms, not to a result that cancels to near zero. The attention-sink row does + cancel like that on a few elements, where torch's default bf16 relative tolerance + misreads a correct result (at M = 15, 20 and 180 one element is 7.3e-4 to 9.8e-4 from + the torch collapse, relative 0.018 to 0.041). An all-zero row must come out exactly 0. + Measured on sm_103 the op's largest error is 0.498 of the bound; on 2048 rows the wrong + variants land at 7.3e5 (another sublayer's scale / base) to 2.1e7 (the consumer's RMS + statistic) times it. + """ + expected, exact, terms = ref + assert out.dtype == expected.dtype == torch.bfloat16, (out.dtype, expected.dtype) + assert out.shape == expected.shape, (out.shape, expected.shape) + ratio = collapse_error_ratio(out, exact, terms) + assert ratio <= 1.0, f"an element exceeds the term bound: max err / bound = {ratio:.3f}" + if out.numel(): + same = (out == expected) | (out.isnan() & expected.isnan()) + frac = same.float().mean().item() + assert frac >= _MIN_BITEXACT, f"bit-exact fraction {frac} < {_MIN_BITEXACT}" + + +mhc_hc_head_apply = MhcHcHeadApply() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/mhc_post_mapping.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/mhc_post_mapping.py new file mode 100644 index 000000000000..58ac5677ed63 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/norm/mhc_post_mapping.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection post mapping (hc_post) via the trtllm mhc_post_mapping kernel. + +The op is `torch.ops.trtllm.mhc_post_mapping`, registered in `cpp/tensorrt_llm/thop/mhcOp.cpp`. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_HC = 4 +_H_CKPT = 5120 + +#: Row counts DeepSeek-V4.1 makes its call at, from the target's runtime-value dumps: the trunk's +#: ('mhc', T, tile) records (T 1..16384, engines up to max_batch_size 256 / max_num_tokens 16384, +#: CUDA-graph and eager, dep4 and tp4, 128K and 1M max_seq_len, the long-context envelope) and +#: the DSpark drafter's ('dspark.mhc', N, tile) records (N = 5G); plus the tops of the step range +#: (16384) and of the decoder bounded-replay pass (32768). +_TARGET_ROWS = tuple( + sorted( + { + 1, 2, 3, 4, 5, 6, 7, 8, 12, 14, 16, 18, 24, 29, 32, 34, 48, 60, 96, 108, 127, 128, + 129, 134, 144, 168, 192, 228, 240, 288, 330, 336, 384, 432, 462, 480, 528, 546, 576, + 597, 610, 624, 672, 720, 768, 780, 852, 888, 890, 1018, 1152, 1536, 1906, 1935, 2033, + 2038, 2048, 2049, 2252, 2778, 2784, 3870, 4208, 4245, 4388, 4394, 4559, 4591, 4868, + 7835, 7855, 7893, 7908, 8147, 8191, 8192, 8467, 9193, 9240, 9364, 9370, 9524, 10441, + 10444, 10447, 10451, 16298, 16308, 16316, 16320, 16322, 16327, 16328, 16329, 16333, + 16340, 16378, 16384, + } + | { + 5, 10, 15, 20, 40, 50, 80, 90, 110, 115, 120, 125, 140, 160, 190, 200, 240, 275, 280, + 320, 360, 370, 385, 400, 440, 455, 480, 520, 560, 580, 585, 600, 640, 650, 665, 710, + 715, 720, 960, 1280, + } + | {4095, 8191, 8193, 16383, 16385, 32767, 32768} + ) +) # fmt: skip + + +def _spec(T: int, H: int = _H_CKPT, out: str = "none", flat: bool = False) -> dict: + """`out`: "none" lets the op allocate, "alias" passes `out=residual` (the target's form), + "view" passes a row-prefix view of a longer buffer. `flat` passes `post_mix` as [T, 4, 1] + and `comb_mix` as [T, 16].""" + return dict(T=T, H=H, out=out, flat=flat) + + +class MhcPostMapping(OpWrapper): + """The hyper-connection post step (hc_post) for a 4-copy residual stream. + + Expands a sublayer's output back to the 4 copies and mixes the old copies in through the + per-token `comb` matrix:: + + out[t, j, h] = post_mix[t, j] * x[t, h] + sum_{k=0..3} comb_mix[t, k, j] * residual[t, k, h] + + `comb_mix[t, k, j]` is [source copy k][destination copy j]: the residual is mixed through + `comb` transposed relative to `comb @ residual`, exactly the orientation of the DeepSeek-V4.1 + checkpoint's `Block.hc_post`, whose Sinkhorn output passes unchanged. Per element the kernel + computes `acc = post[j] * x` (one fp32 multiply), then for k = 0..3 in order + `acc = fma(comb[k, j], residual[k], acc)`, then one RNE rounding to bf16. + + Fusion boundary: the expansion, the comb mix, the sum and the bf16 cast -- nothing else. The + caller owns computing `post_mix` / `comb_mix` (the hc mixes GEMM and the Sinkhorn step), + producing `x` in bf16 (an fp32 all-reduce result must be cast first), and collapsing the new + stream for the next sublayer. Writes `out` (a fresh contiguous `[T, 4, H]` bf16 tensor when + None) and returns it; `out` may be `residual` itself. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell(why=f"checkpoint width T={t}: H 5120, op-allocated out", spec=_spec(t)) + for t in (1, 5, 6, 32, 4096, 8192) + ) + + tuple( + Cell(why=f"width H={h} T=32: op-allocated out", spec=_spec(32, H=h)) + for h in (1032, 4096, 7168) + ) + + ( + Cell( + why="flat mixes T=16: post [T, 4, 1], comb [T, 16]", + spec=_spec(16, flat=True), + ), + Cell( + why="out view T=32: row prefix of a longer buffer", + spec=_spec(32, out="view"), + ), + ) + + tuple( + Cell(why=f"in place T={t}: out is residual, H 5120", spec=_spec(t, out="alias")) + for t in (1, 32, 4096) + ) + + tuple( + Cell( + why=f"target call T={t}: aliased, H 5120, sink row", + spec=_spec(t, out="alias"), + ) + for t in _TARGET_ROWS + ) + ) + + note = """ + Every cell carries a doubly-stochastic comb built the way the checkpoint builds it and + (T > 1) a 1e5-scale attention-sink row (the checkpoint's BOS stream reaches |x| ~ 3.5e5). + The other widths cover a whole number of 2048-element sweeps (4096), another model's width + (7168), and a width that leaves most threads idle on its last sweep (1032 = 8 * 129). + + The op is bit-exact to its own FMA order in every cell. Against the checkpoint's expression + (products rounded separately, then summed) it rounds differently on about 5 in 100,000 + elements, and default bf16 `assert_close` against that expression is not a safe gate where + the sink row cancels to a small output: it rejects 1 element in 8 of 254 target-call and + capture cases (max absolute difference 1.17e-2, relative 0.185, all in the sink row). Every + element is within half a bf16 ulp of the exact sum plus 16 fp32 ulps of the term + magnitudes (max 0.996 of that bound). + + Rows are independent: a NaN / inf in one row's `post_mix`, `comb_mix`, `x` or `residual` + leaves every other row bit-identical, so CUDA-graph padding rows cannot contaminate real + rows. Real-data check on the checkpoint's golden hc_post inputs (BOS rows to 3.5e5): the + golden stream is reproduced bit-exactly on all but about 2 in 100,000 elements; with comb + transposed only 8% are. + + Metadata consumed: none. Stateless; capturable in a CUDA graph. + + Preconditions: hc width exactly 4 (compiled in). H % 8 == 0 and 16-byte-aligned base + pointers (8-element vector accesses); a violation is loud but fatal to the CUDA context + (H = 12 died with "misaligned address") and is not checked. Dtypes are fixed (bf16 + `residual`, `x`, `out`; fp32 `post_mix`, `comb_mix`); a wrong dtype raises RuntimeError + before launch. The op reads every tensor as a dense buffer of the size T and H imply; the + layouts it silently miscomputes or writes out of bounds for are in `is_valid`. `out` may + alias `residual` exactly (the kernel declares its pointers `__restrict__`; exact aliasing is + a property of the compiled kernel the test re-checks, not one its source promises); partial + overlap of `out` with any input was not tested. T = 0 returns an empty `[0, 4, H]` tensor. + + DeepSeek-V4.1 calls `(h, sublayer_out, post, comb, out=h)` twice per layer at H = 5120. + Certified for every T in 1..32768 of that aliased call: each prefix of a 32768-row call is + bitwise equal to its first T rows; the call launches exactly one `mhcPostMappingKernel`, + compiled ahead of time, at every T; captured into a CUDA graph at every decode / + DSpark-verify / drafter row count up to 1536 it replays bitwise equal to eager. + + Alternative measured on sm_103: `flashinfer.mhc.mhc_post` is bit-identical; this op is + within 0.2 us at decode sizes, faster at prefill sizes (59.0 vs 70.9 us at T = 4096, 112.6 + vs 136.6 us at T = 8192), and about 3x faster eager. + + Not certified: T > 32768; H other than 5120 (every T) and 1032, 4096, 7168 (T = 32); partial + aliasing of `out`; torch.compile. + """ + + def raw_call( + self, + residual: torch.Tensor, + x: torch.Tensor, + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + num_tokens, _, hidden_size = residual.shape + if out is None: + out = torch.empty_like(residual, memory_format=torch.contiguous_format) + torch.ops.trtllm.mhc_post_mapping( + residual, x, post_mix, comb_mix, out, num_tokens, hidden_size + ) + return out + + def is_valid( + self, + residual: torch.Tensor, + x: torch.Tensor, + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + out: Optional[torch.Tensor] = None, + ) -> None: + # The op takes T and H as integers and reads every tensor as a dense + # row-major buffer of the size those imply: a non-contiguous view, a wrong + # hc width, or a short tensor is read (or, for `out`, written) past its + # bounds with no error. Each of these was measured to return silently wrong + # results on sm_103, so they are checked here. + assert residual.dim() == 3 and residual.shape[1] == 4, "residual must be [T, 4, H]" + num_tokens, _, hidden_size = residual.shape + assert x.shape == (num_tokens, hidden_size), "x must be [T, H] matching residual" + assert post_mix.numel() == num_tokens * 4, ( + "post_mix must hold T*4 values ([T, 4] or [T, 4, 1])" + ) + assert comb_mix.numel() == num_tokens * 16, ( + "comb_mix must hold T*16 values ([T, 4, 4] or [T, 16])" + ) + # An absent `out` is allocated contiguous and shaped like `residual` by the call. + if out is not None: + assert out.shape == residual.shape, "out must be [T, 4, H] like residual" + for name, t in ( + ("residual", residual), + ("x", x), + ("post_mix", post_mix), + ("comb_mix", comb_mix), + ("out", out), + ): + if t is not None: + assert t.is_contiguous(), f"{name} must be contiguous" + + def reference( + self, + residual: torch.Tensor, + x: torch.Tensor, + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + out: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """`(out, exact, terms)`, each `[T, 4, H]`. + + `out` is the kernel's contracted order in native torch: `acc = fp32(post[j] * x)`, then + for k = 0..3 `acc = fp32(comb[k, j] * r[k] + acc)` with the product and the add rounded + once (the fp64 sum stands in for the fused multiply-add: the fp32 x fp32 product is + exact in fp64), then one RNE rounding to bf16. `exact` is the five-term sum in fp64 and + `terms` is `|post[j] * x| + sum_k |comb[k, j] * residual[k]|` in fp64, which the gate in + `compare` needs. `out` (the argument) is ignored: the expected values do not depend on + where they are written. Reads `residual`, so when the call aliases `out` to it, take + the reference before the call. + """ + t = residual.shape[0] + post = post_mix.reshape(t, _HC).float() + comb = comb_mix.reshape(t, _HC, _HC).float() + acc = post.unsqueeze(-1) * x.float().unsqueeze(-2) + terms = post.double().unsqueeze(-1) * x.double().unsqueeze(-2) + exact = terms.clone() + term_scale = terms.abs() + for k in range(_HC): + prod = comb[:, k, :].double().unsqueeze(-1) * residual[:, k, :].double().unsqueeze(1) + acc = (prod + acc.double()).float() + exact += prod + term_scale += prod.abs() + return acc.to(torch.bfloat16), exact, term_scale + + def compare( + self, + out: torch.Tensor, + ref: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ) -> None: + """Equal to the kernel-order reference on every element, and within the term bound. + + The op matched its own FMA order on every element of every certified cell, so nothing + is left to round differently. The term bound is half a bf16 ulp of the exact value (the + final rounding) plus 16 fp32 ulps of S = |post[j] * x| + sum_k |comb[k, j] * r[k]|: + any fp32 order errs by a few fp32 ulps of the TERMS, which on a cancelling sink row is + large next to the result. Max 0.996 of the bound measured on sm_103. Controls: comb + applied untransposed, and `post` rolled by one copy, both fail. + """ + expected, exact, terms = ref + assert out.dtype == expected.dtype == torch.bfloat16, (out.dtype, expected.dtype) + assert out.shape == expected.shape, (out.shape, expected.shape) + bad = (out != expected).sum().item() + assert bad == 0, f"{bad} of {out.numel()} elements differ from the kernel-order reference" + if out.numel(): + bound = 2.0**-8 * exact.abs() + 16 * 2.0**-24 * terms + torch.finfo(torch.float64).tiny + ratio = ((out.double() - exact).abs() / bound).max().item() + assert ratio <= 1.0, f"max err / bound against the exact sum = {ratio:.3f}" + + +mhc_post_mapping = MhcPostMapping() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/e4m3_quantize_pow2_dynamic.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/e4m3_quantize_pow2_dynamic.py new file mode 100644 index 000000000000..079bcf813ccc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/e4m3_quantize_pow2_dynamic.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamic per-tensor E4M3 quantization with a power-of-two scale, via the trtllm Triton op. + +The op is `torch.ops.trtllm.e4m3_quantize_pow2_dynamic`, registered in +`tensorrt_llm/_torch/custom_ops/e4m3_quantize_pow2_dynamic.py`. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops.e4m3_quantize_pow2_dynamic # noqa: F401 — registers the op +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_E4M3_MAX = 448.0 +_AMAX_FLOOR = 1e-4 + +#: Row counts the DeepSeek-V4.1 engine hands the attention query at: decode at +#: 1..32 requests x next_n 1 or 6 (DSpark verify), and context chunks up to 8192. +_QUERY_ROWS = (1, 2, 6, 8, 32, 48, 192, 300, 2048, 4097, 8192) + + +class E4m3QuantizePow2Dynamic(OpWrapper): + """Quantize a whole tensor to E4M3 with one power-of-two scale from its amax. + + One call, three launches, no host synchronization:: + + amax = max(max|x|, 1e-4) # fp32, whole tensor + s = 2^ceil(log2(amax * fp32(1/448))) + out = e4m3_rne(clamp(x * (1/s), -448, 448)) # exact scaling + scale_out[0] = s * scale_mult # one fp32 multiply + + Fusion boundary: the amax reduction, the scale and the quantization are + inside the call; the consumer of `scale_out` (an attention's `bmm1_scale`, + with `scale_mult = softmax_scale * kv_scale`) is not. Writes `out` and + element 0 of `scale_out` in place and returns `(out, scale_out)`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"query T={t} H={h}: [T, H, 512] bf16 at an engine row count", + spec=dict( + kind="query", T=t, H=h, dtype=torch.bfloat16, scale_mult=512**-0.5 * 2.0**-7 + ), + ) + for h in (16, 64) + for t in _QUERY_ROWS + ) + + tuple( + Cell( + why=f"input dtype {str(dt).removeprefix('torch.')}: [33, 16, 512]", + spec=dict(kind="query", T=33, H=16, dtype=dt, scale_mult=1.0), + ) + for dt in (torch.bfloat16, torch.float16, torch.float32) + ) + + tuple( + Cell( + why=f"scale edge amax={amax!r}: {reason}", + spec=dict(kind="edge", amax=amax, dtype=torch.float32, scale_mult=1.0), + ) + for amax, reason in ( + (448.0, "fp32(1/448) rounds 448 * 2^k up to 2^(k+1)"), + (448.0 * 2**-9, "the same rounding nine binades down"), + (449.0, "just above 448, scale 2"), + (3.0e4, "far above 448"), + (1e-6, "below the 1e-4 amax floor"), + (0.0, "an all-zero tensor, the floor alone sets the scale"), + ) + ) + ) + + note = """ + H is 16 for tep4/tp4 and 64 for dep4. Per-(token, head) magnitudes in the + cells span 2^-6 .. 2^6, the range the golden queries reach. + + `amax * fp32(1/448)` is formed in fp32 like DeepSeek-V4.1's MXFP8 + activation scale; because `fp32(1/448)` is slightly above `1/448`, an amax + of exactly `448 * 2^k` gives `s = 2^(k+1)`. No value saturates: + `max|x| / s` lies in `[224, 448]`, and `out * s` is `x` rounded to E4M3's + 3-bit mantissa (relative round-trip error <= 2^-4 in E4M3's normal range). + + Metadata consumed: none. + + Preconditions: `x` contiguous bf16/fp16/fp32 holding finite values (a NaN + makes the amax NaN); `out` contiguous float8_e4m3fn with `x.numel()` + elements; `scale_out` float32 with at least one element. The op raises + ValueError on a wrong dtype, device, contiguity or numel. The one scale + spans the whole tensor, so an outlier coarsens every other value's + subnormal threshold. + + Allocates one float of scratch per call; inside CUDA-graph capture it comes + from the graph's pool, and a replay recomputes the scale from the data. + + Not certified: tensors above 8192 x 64 x 512 elements. + """ + + def raw_call( + self, + x: torch.Tensor, + out: torch.Tensor, + scale_out: torch.Tensor, + scale_mult: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + torch.ops.trtllm.e4m3_quantize_pow2_dynamic(x, out, scale_out, scale_mult) + return out, scale_out + + def reference( + self, + x: torch.Tensor, + out: torch.Tensor, + scale_out: torch.Tensor, + scale_mult: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Codes and the scale output, derived independently of the kernel. + + The scale exponent comes from `torch.frexp` in fp64 rather than from + the kernel's IEEE-bit derivation; the codes are torch's RNE cast of + `x * 2^-e`, an exact scaling. `out` and `scale_out` are named to mirror + `raw_call` and are not mutated: the expected `out` is returned as a new + tensor of its shape, and the expected `scale_out` is a copy with + element 0 replaced. + """ + amax = max(x.float().abs().max().item(), _AMAX_FLOOR) if x.numel() else _AMAX_FLOOR + scaled = torch.tensor(amax, dtype=torch.float32) * torch.tensor( + 1.0 / _E4M3_MAX, dtype=torch.float32 + ) + mantissa, exponent = torch.frexp(scaled.double()) + e = int(exponent.item()) - (1 if float(mantissa.item()) == 0.5 else 0) + codes = (x.float() * (2.0**-e)).clamp(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn) + expected_scale = scale_out.clone() + expected_scale[0] = ( + torch.tensor(2.0**e, dtype=torch.float32) + * torch.tensor(scale_mult, dtype=torch.float32) + ).item() + return codes.reshape(out.shape), expected_scale + + def compare( + self, + out: tuple[torch.Tensor, torch.Tensor], + ref: tuple[torch.Tensor, torch.Tensor], + ) -> None: + """Bitwise on the codes and on `scale_out[0]`. + + Kernel and reference make the same exact power-of-two scaling and the + same RNE cast, so there is nothing to round differently: any code that + differs is a wrong scale. Controls that use a non-power-of-two scale or + round the scale down change more than 30% of the codes. + """ + (codes, scale_out), (ref_codes, ref_scale) = out, ref + assert codes.dtype == ref_codes.dtype == torch.float8_e4m3fn, (codes.dtype, ref_codes.dtype) + assert codes.shape == ref_codes.shape, (codes.shape, ref_codes.shape) + bad = (codes.view(torch.uint8) != ref_codes.view(torch.uint8)).sum().item() + assert bad == 0, f"{bad} of {codes.numel()} codes differ" + assert scale_out[0].item() == ref_scale[0].item(), ( + scale_out[0].item(), + ref_scale[0].item(), + ) + + +e4m3_quantize_pow2_dynamic = E4m3QuantizePow2Dynamic() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/mxfp8_dequantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/mxfp8_dequantize.py new file mode 100644 index 000000000000..21079d6ae6bd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/mxfp8_dequantize.py @@ -0,0 +1,270 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gather + dequantize native MXFP8 rows (E4M3 codes x UE8M0 per 32) from a flat or paged pool into +bf16, via the Triton op `torch.ops.trtllm.mxfp8_dequantize` +(`tensorrt_llm/_torch/custom_ops/mxfp8_dequantize.py`).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops.mxfp8_dequantize # noqa: F401 — registers torch.ops.trtllm.mxfp8_dequantize +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_BLOCK = 32 +_W_KV = 512 # DeepSeek-V4.1 window-KV head_dim (448 nope + 64 RoPE, all quantized) + +#: Rows of the full-byte-range identity read: decode-sized through a prefill-sized pool. +_IDENTITY_ROWS = (1, 6, 128, 4096) +#: Widths other than the window KV's 512 (448 is 14 blocks, not a power of two). +_OTHER_WIDTHS = (32, 128, 448, 1024) +#: Token counts whose checkpoint-recipe rows round-trip to the bf16 container. +_RECIPE_ROWS = (1, 6, 32, 4096, 8192) +_GATHER_PAGE_SIZES = (1, 64, 128, 256) +_GATHER_SLOT_DTYPES = (torch.int32, torch.int64) +_GATHER_ROWS = (1, 128, 1024, 8192) + + +def _num_slots(t: torch.Tensor) -> int: + return t.shape[0] if t.dim() == 2 else t.shape[0] * t.shape[1] + + +def _slot_rows(t: torch.Tensor, slots: torch.Tensor) -> torch.Tensor: + """Rows of a [slots, C] or [pages, page_size, C] tensor at the given slots.""" + if t.dim() == 2: + return t[slots] + return t[slots // t.shape[1], slots % t.shape[1]] + + +class Mxfp8Dequantize(OpWrapper): + """Read native MXFP8 rows out of a pool, decode them to bf16, and write each one to a slot of `out`. + + The pool is E4M3 codes plus one UE8M0 byte per 32 codes: DeepSeek-V4.1's SWA window-KV layout, + 512 + 16 = 528 B per slot. For output row `r`:: + + s = src_slots[r] if src_slots is not None else r + d = dst_slots[r] if dst_slots is not None else r + if d < 0 or d >= out_slots: nothing is written for row r + elif s < 0 or s >= slots: out[d, :] = +0.0 + else: out[d, i] = bf16_rne(e4m3(codes[s, i]) * 2^(scales[s, i // 32] - 127)) + + `rows` is `len(src_slots)`, else `len(dst_slots)`, else the pool's slot count. Slot `s` of a + paged `[pages, page_size, W]` tensor is `(s // page_size, s % page_size)`, for the pool and + for `out` alike; leading strides are free, so a 528 B layout is two views of one uint8 buffer. + + Fusion boundary: the slot gather, the E4M3 x UE8M0 decode, the bf16 rounding and the optional + slot scatter are inside the call. Deriving the slot lists from block tables, remapping the + sparse-attention index table, the attention itself, and writing the pool + (`mxfp8_fake_quantize`'s native form) are not. Writes `out` in place, only at the + destination slots of the rows, and returns `out`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"identity {n} rows x 512: every e4m3 code x UE8M0 byte 0..254" + + (", bytes 0 / 1 / 254 pinned" if n >= 16 else ""), + spec=dict(kind="identity", rows=n, width=_W_KV, pin_extremes=n >= 16), + ) + for n in _IDENTITY_ROWS + ) + + tuple( + Cell( + why=f"identity 257 rows x {w}: width {w}, every code x byte 0..254", + spec=dict(kind="identity", rows=257, width=w, pin_extremes=False), + ) + for w in _OTHER_WIDTHS + ) + + tuple( + Cell( + why=f"recipe rows T={t}: equals the bf16 container" + + (", the 1e-4 floor byte 105 present" if t >= 32 else ""), + spec=dict(kind="recipe", rows=t), + ) + for t in _RECIPE_ROWS + ) + + tuple( + Cell( + why=f"gather {rows} rows, page_size {ps}, {str(sd).removeprefix('torch.')} slots" + + (": 1/16 invalid" if rows >= 16 else "") + + ", from a 40960-slot pool", + spec=dict(kind="gather", rows=rows, page_size=ps, slot_dtype=sd), + ) + for ps in _GATHER_PAGE_SIZES + for sd in _GATHER_SLOT_DTYPES + for rows in _GATHER_ROWS + ) + + tuple( + Cell( + why=f"528 B rows {'paged' if paged else 'flat'}, {str(sd).removeprefix('torch.')}" + " scales: two strided views of one buffer, 4096 rows", + spec=dict(kind="interleaved", paged=paged, scale_dtype=sd), + ) + for paged in (False, True) + for sd in (torch.uint8, torch.float8_e8m0fnu) + ) + + tuple( + Cell( + why=f"page footer, page_size {ps}: code rows then scale rows per page", + spec=dict(kind="footer", page_size=ps), + ) + for ps in (128, 256) + ) + + ( + Cell( + why="strided src_slots [::2], out the left half of [777, 1024]", + spec=dict(kind="strided"), + ), + Cell( + why="shadow pool: src = dst over 512 slots, -1 skipped, the rest kept", + spec=dict(kind="shadow"), + ), + Cell( + why="scatter: invalid dst skipped, invalid src -> zero row, int64/int32", + spec=dict(kind="scatter"), + ), + Cell( + why="dst_slots only into paged [16, 128, 512]: row r reads slot r", + spec=dict(kind="dst_only"), + ), + Cell( + why="zero rows: an empty src_slots launches nothing, out untouched", + spec=dict(kind="zero_rows"), + ), + ) + ) + + note = """ + Exactness: `e4m3(code) * 2^k` has at most 4 significant bits and is exact in fp32 for every + byte 0..254; the one rounding is the bf16 RNE conversion, exact whenever the product is a bf16 + normal (every byte the fake-quantize recipe produces, 105..247). Byte 0 is the fp32-subnormal + scale 2^-127, not zero. Byte 255 makes the whole 32-block NaN; codes 0x7F / 0xFF are NaN for + that element; code 0x80 is -0.0. Dequantizing `mxfp8_fake_quantize`'s native output + reproduces its bf16-container output bit for bit. + + Call forms the target uses: compact gather (`out` [rows, 512], `src_slots` the referenced ring + slots); shadow pool (`out` a bf16 pool of the native pool's geometry, `src_slots = dst_slots`; + slots not named keep their bytes); fixed-size window buffer (`out` one persistent zero-filled + [max_num_tokens * 128, 512] buffer, `src_slots` the window half of the attention's int32 + [T, 128] index table flattened, `dst_slots` None; rows at and past 128 * T keep what an + earlier call left). + + Metadata consumed: none. The slot addresses come only from `src_slots`, `dst_slots` and the + tensors' geometry. + + Preconditions the op raises ValueError on, leaving `out` untouched: codes CUDA, + float8_e4m3fn or uint8, 2-D or 3-D, last dim a positive multiple of 32 with unit stride; + scales uint8 or float8_e8m0fnu, exactly `codes.shape[:-1] + (W // 32,)` with unit last stride, + on the codes device; out bfloat16 with last dim W at unit stride, on the codes device; slot + lists int32/int64, 1-D, on the codes device, equal lengths when both given; without + `dst_slots`, `out` needs at least `rows` slots. Unchecked: two rows with the same valid + destination and different sources race (identical sources are fine and certified); `out` + must not overlap the pool. + + CUDA graphs: make one eager call of the same variant before capturing (Triton compiles on + first use). A variant is W, whether each slot list is given, and the class of `out`'s slot + count (1, a multiple of 16, any other), which Triton specializes on; rows are grid axis 0 and + never a kernel argument. A shadow or a fixed-size `out` therefore compiles one variant. + + Certified by the rest of this entry's test, beyond the cells, on sm_103: native pools of the + live DeepSeek-V4.1 cache manager (`TRTLLM_DSV41_KV_STORAGE=native`, `get_v41_native_kv_views`, + `[B, 128, 528]` interleaved, strides (67584, 528, 1)) up to B = 1,648,320 pages read compact + (the measured ceiling over the scanned engine configurations including PP2 stages, at 0.5 x + total memory; past byte offset 2^36) and up to 1,098,720 pages into a shadow (past 2^37), + with a wrap control showing a mod-2^32 offset would be seen; every row count 1..16,384 compact + and into a shadow; per-query window reads up to 1,048,576 rows; the fixed-size window buffer + at 262,144 / 524,288 / 1,048,576 / 2,097,152 rows (max_num_tokens 2048 / 4096 / 8192 / + 16,384; independent of max_batch_size and max_seq_len), every row count up to the buffer + (sampled at 2,097,152); renumbered reads over random multi-request steps on live managers at + every max_batch_size 1..32 and at (16,384, 256); CUDA-graph capture at every graph batch + 1..32, 40..128 step 8, 192 and 256 with next_n 1 and 6, adding no compiled variant after one + eager warm-up. + + Not certified: duplicate destinations with different sources; `out` aliasing the pool; a + first-ever call inside a capture; torch.compile (a fake is registered, not exercised); a + native pool above 1,648,320 pages or a shadow above 1,098,720; more than 16,384 rows into a + rows-sized out or a shadow other than the per-query counts; live-manager steps at positions + above 32,781; a fixed window buffer of another size; page sizes other than 1, 64, 128, 256; + widths other than 32, 128, 448, 512, 1024. + """ + + def raw_call( + self, + codes: torch.Tensor, + scales: torch.Tensor, + out: torch.Tensor, + src_slots: Optional[torch.Tensor] = None, + dst_slots: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + torch.ops.trtllm.mxfp8_dequantize(codes, scales, out, src_slots, dst_slots) + return out + + def reference( + self, + codes: torch.Tensor, + scales: torch.Tensor, + out: torch.Tensor, + src_slots: Optional[torch.Tensor] = None, + dst_slots: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The expected `out` after the call: a copy of `out` with the rows' destinations filled. + + `out` is read for its prior contents (every slot the call does not write keeps them) and is + not mutated; pass a snapshot taken before the op ran. The decode is torch's + float8_e4m3fn -> fp64 cast and `torch.ldexp(code, byte - 127)` in fp64, then fp64 -> fp32 + (exact) -> bf16 RNE: independent of the kernel's exponent-bit construction. Byte 255 + decodes to NaN; NaN payloads are not modelled, so cells carry no NaN encodings. + """ + w = codes.shape[-1] + num_slots = _num_slots(codes) + expected = out.clone() + out_slots = _num_slots(expected) + if src_slots is not None: + rows = src_slots.shape[0] + elif dst_slots is not None: + rows = dst_slots.shape[0] + else: + rows = num_slots + dev = codes.device + src = src_slots.long() if src_slots is not None else torch.arange(rows, device=dev) + dst = dst_slots.long() if dst_slots is not None else torch.arange(rows, device=dev) + written = (dst >= 0) & (dst < out_slots) + src, dst = src[written], dst[written] + valid = (src >= 0) & (src < num_slots) + s = src.clamp(0, max(num_slots - 1, 0)) + c = _slot_rows(codes.view(torch.uint8), s).contiguous().view(torch.float8_e4m3fn) + b = _slot_rows(scales.view(torch.uint8), s).long() + v = c.float().double().view(-1, w // _BLOCK, _BLOCK) + vals = torch.ldexp(v, (b - 127).double().unsqueeze(-1)) + vals = torch.where((b == 255).unsqueeze(-1), torch.nan, vals) + vals = vals.view(-1, w).float().to(torch.bfloat16) + vals[~valid] = 0.0 + if expected.dim() == 2: + expected[dst] = vals + else: + expected[dst // expected.shape[1], dst % expected.shape[1]] = vals + return expected + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise, bf16 compared as int16, over all of `out`. + + Each element is one e4m3 value times a power of two rounded once to bf16, so there is no + tolerance to choose; untouched slots must keep their prior bits. On T = 4096 recipe rows + the gate fails for each deliberately wrong decode of the same bytes: exponent bias 126 or + 128, an e5m2 decode, one scale byte per 64 codes, a slot off by one. + """ + assert out.dtype == ref.dtype == torch.bfloat16, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + a = out.reshape(-1).view(torch.int16) + e = ref.reshape(-1).view(torch.int16) + bad = torch.zeros((), dtype=torch.int64, device=a.device) + for i in range(0, a.numel(), 1 << 26): # chunked: a full bool reduction copies to int64 + bad += (a[i : i + (1 << 26)] != e[i : i + (1 << 26)]).sum() + assert int(bad) == 0, f"{int(bad)} of {out.numel()} elements differ bitwise" + + +mxfp8_dequantize = Mxfp8Dequantize() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/mxfp8_fake_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/mxfp8_fake_quantize.py new file mode 100644 index 000000000000..446882487f09 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/mxfp8_fake_quantize.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 fake quantization (E4M3 x UE8M0 per 32, amax floor 1e-4) written into slots of a +bf16 container pool, or as native e4m3 codes + UE8M0 bytes, via the Triton op +`torch.ops.trtllm.mxfp8_fake_quantize` (`tensorrt_llm/_torch/custom_ops/mxfp8_fake_quantize.py`).""" + +from __future__ import annotations + +from typing import Optional, Union + +import torch + +import tensorrt_llm._torch.custom_ops.mxfp8_fake_quantize # noqa: F401 — registers torch.ops.trtllm.mxfp8_fake_quantize +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_BLOCK = 32 +_E4M3_MAX = 448.0 +_AMAX_FLOOR = 1e-4 +_W_KV = 512 # DeepSeek-V4.1 window-KV head_dim (448 nope + 64 RoPE, all quantized) + +#: Decode (1, 5 / 6 = DSpark draft block / verify), batch, window and prefill row counts. +_CONTAINER_ROWS = (1, 5, 6, 32, 128, 4096, 8192) +#: Widths other than the window KV's 512 (448 is 14 blocks, not a power of two). +_OTHER_WIDTHS = (32, 128, 448, 1024) +#: The two tokens_per_block values the DeepSeek-V4 cache manager accepts, and their pool pages. +_POOL_PAGES = {128: 300, 256: 150} +_SCATTER_ROWS = (1, 6, 32, 4096, 8192) +_SLOT_DTYPES = (torch.int32, torch.int64) +_NATIVE_ROWS = (1, 6, 4096) + + +def _num_slots(t: torch.Tensor) -> int: + return t.shape[0] if t.dim() == 2 else t.shape[0] * t.shape[1] + + +def _write_slots(t: torch.Tensor, slots: torch.Tensor, rows: torch.Tensor) -> None: + """Write `rows` to the given slots of a [slots, C] or [pages, page_size, C] tensor.""" + if t.dim() == 2: + t[slots] = rows + else: + t[slots // t.shape[1], slots % t.shape[1]] = rows + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.view(torch.int16) if t.dtype == torch.bfloat16 else t.view(torch.uint8) + + +class Mxfp8FakeQuantize(OpWrapper): + """MXFP8-fake-quantize every bf16 row of `x` and write each row to one slot of `out`. + + For row `r` and 32-element block `b` along the last dim:: + + amax = max(max |x[r, 32b : 32b+32]|, 1e-4) # fp32; the floor is fp32(1e-4) + e = ceil(log2(amax / 448)) # exact; scale 2^e >= 2^-22 + code[r,i] = e4m3_rne(clamp(x[r,i] / 2^e, -448, 448)) # the one rounding + byte[r,b] = e + 127 # UE8M0 + value = e4m3(code) * 2^e # exact in bf16 + + This is the checkpoint reference's `act_quant(x, 32, "ue8m0", torch.float8_e8m0fnu, + inplace=True)` on the post-RoPE window KV. The dtype of `out` selects the form: a bfloat16 + `out` (the bf16 container) receives the values; a float8_e4m3fn / uint8 `out` (native) + receives the codes and `out_scale` (uint8 / float8_e8m0fnu, column b = block b) the bytes. + Both forms hold the same value set. Row r goes to slot `slot_mapping[r]` (skipped when + negative or >= slots), or to slot r when `slot_mapping` is None; slot `s` of a paged + `[pages, page_size, W]` tensor is `(s // page_size, s % page_size)`. + + Fusion boundary: the per-block amax, the scale, the e4m3 rounding, the dequantization (bf16 + form) and the slot scatter are inside the call. kv_norm, RoPE, deriving `slot_mapping` from + block tables, and any read-back of a native pool (`mxfp8_dequantize`) are not. + + Writes in place and returns what it wrote: `out` in the bf16-container form (`out_scale` + None), `(out, out_scale)` in the native form. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"bf16 container T={t} x 512, identity: decode to prefill rows", + spec=dict(kind="container", rows=t, width=_W_KV), + ) + for t in _CONTAINER_ROWS + ) + + tuple( + Cell( + why=f"bf16 container 257 x {w}, identity: width {w}", + spec=dict(kind="container", rows=257, width=w), + ) + for w in _OTHER_WIDTHS + ) + + ( + Cell( + why="in place, out is x (T=1000): rows read whole before written", + spec=dict(kind="in_place", rows=1000), + ), + Cell( + why="x a column slice of [300, 1024]: row-strided input", + spec=dict(kind="x_column_slice", rows=300), + ), + Cell( + why="x [300, 1, 512]: leading dims viewed as one row dim", + spec=dict(kind="x_3d", rows=300), + ), + ) + + tuple( + Cell( + why=f"scatter T={t}, {str(sd).removeprefix('torch.')} slots, into bf16 " + f"[{_POOL_PAGES[ps]}, {ps}, 512]" + (": 1/8 of rows -1" if t >= 8 else ""), + spec=dict(kind="scatter", rows=t, page_size=ps, slot_dtype=sd), + ) + for ps in _POOL_PAGES + for t in _SCATTER_ROWS + for sd in _SLOT_DTYPES + ) + + ( + Cell( + why="pool the left half of [40, 128, 1024]: strided paged out", + spec=dict(kind="pool_left_half", rows=777), + ), + Cell( + why="flat pool [5120, 640][:, 64:576], int32 slots: column slice", + spec=dict(kind="pool_column_slice", rows=777), + ), + Cell( + why="out-of-range slots 256, 257, 1e9, -1, -1e9 skipped", + spec=dict(kind="out_of_range", rows=8), + ), + ) + + tuple( + Cell( + why=f"native 528 B interleaved rows, T={t}: codes then UE8M0 bytes", + spec=dict(kind="native_interleaved", rows=t), + ) + for t in _NATIVE_ROWS + ) + + tuple( + Cell( + why=f"native page footer, page_size {ps}, T=2000: fp8 codes, e8m0 scales", + spec=dict(kind="native_footer", rows=2000, page_size=ps), + ) + for ps in (128, 256) + ) + + ( + Cell( + why="native identity T=4096, fp8 codes + uint8 scales: floor byte 105", + spec=dict(kind="native_identity", rows=4096), + ), + Cell( + why="zero rows, mapped: no launch, the pool untouched", + spec=dict(kind="zero_rows", mapped=True), + ), + Cell( + why="zero rows, identity: no launch, the pool untouched", + spec=dict(kind="zero_rows", mapped=False), + ), + ) + ) + + note = """ + The kernel's scale exponent is a bit-level ceil-log2 of `fp32(amax * fp32(1/448))`; for every + bf16 amax and at the floor it equals the mathematical `ceil(log2(amax / 448))`, including + `amax = 7 * 2^k` where `amax / 448` is an exact power of two. (A float64 `torch.log2` is not a + valid way to compute it: `ceil(log2(56 / 448))` came out -2, one too high.) The scale is + rounded up, so nothing saturates; the floor block's byte is 105 (2^-22). Native dequantization + `e4m3(code) * 2^(byte - 127)` is bit-identical to the bf16 container. + + Metadata consumed: none. Slot addresses come only from `slot_mapping` and `out`'s geometry. + + Preconditions the op raises ValueError on: `x` not bf16, not CUDA, rank < 2, last dim not a + positive multiple of 32, non-unit inner stride; `out` last dim != W or non-unit stride, not + 2-D/3-D, a dtype other than bf16 / float8_e4m3fn / uint8, or on another device; a native + `out` without `out_scale`, a bf16 `out` with one, `out_scale` of the wrong dtype or shape; + `slot_mapping` not int32/int64 or not one entry per row; without `slot_mapping`, `out` with + fewer slots than rows. `x`'s leading dims must be viewable as one row dim (row strides + free). Unchecked: `x` must be finite (a non-finite row changes no other row, but what is + written for it is not certified); two rows mapped to one slot race; `out` may be `x` itself + only with `slot_mapping=None`, any other overlap is a race. + + CUDA graphs: make one eager call of the same variant before capturing (Triton compiles on + first use). A variant is the width, identity vs mapped, bf16 vs native, whether the slot + map's stride is 1 and its pointer 16-byte aligned, and, in the identity form only, the class + of the row count (1, a multiple of 16, any other) since there the slot count is the row count. + The mapped form compiles one variant whatever T: rows are grid axis 0, never an argument. + + Certified by the rest of this entry's test, beyond the cells, on sm_103: every row count + 1..8192 in the DSpark drafter's identity call (`out` the first T rows of a longer buffer, the + rest untouched) and CUDA-graph capture at every draft block N = 5G; every row count 1..32768 + in the paged-scatter call; bf16 SWA pools `[B, 128, 512]` up to B = 958,720 pages and native + `[B, 128, 528]` pools of the live DeepSeek-V4.1 cache manager (native KV storage, + `get_v41_native_kv_views`, strides (67584, 528, 1)) up to B = 1,648,320 pages (the + measured maxima over the scanned engine configurations including PP2 stages, at 0.5 x total + memory), with far-end slots at every byte offset 2^31..2^36, pool-size invariance and a wrap + control; CUDA-graph capture at T = G and 6G for every graph batch G in 1..32, 40..128 step 8, + 192, 256; the DSpark drafter's stride-3 stage-ring writes (captured at M = G and 5G, eager + at every M 1..1280, seeding n 1..128); native writes over live single- and multi-request + steps (up to 16384 tokens and 256 requests, positions up to 32,781), round-tripped through + `torch.ops.trtllm.mxfp8_dequantize`. + + Not certified: non-finite values; duplicate slots; a first-ever call inside a capture; a + captured row count outside those listed; an identity-form capture whose `out` has more slots + than T; more than 32768 rows in one call; torch.compile (a fake is registered, not + exercised); a native pool above 1,648,320 pages or a bf16 pool above 958,720; page sizes + other than 1, 128, 256; widths other than 32, 128, 448, 512, 1024. + """ + + def raw_call( + self, + x: torch.Tensor, + out: torch.Tensor, + slot_mapping: Optional[torch.Tensor] = None, + out_scale: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + torch.ops.trtllm.mxfp8_fake_quantize(x, out, slot_mapping, out_scale) + return out if out_scale is None else (out, out_scale) + + def reference( + self, + x: torch.Tensor, + out: torch.Tensor, + slot_mapping: Optional[torch.Tensor] = None, + out_scale: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + """The expected `out` (and `out_scale`) after the call: copies with the rows' slots filled. + + `out` and `out_scale` are read for their prior contents (every slot the call does not write + keeps them) and are not mutated; pass snapshots taken before the op ran (for the in-place + form, a snapshot of `x` as both). The scale exponent comes from `torch.frexp` of the fp64 + quotient `amax / 448` (mantissa 0.5 means an exact power of two), not from the kernel's + IEEE-bit trick; the codes are torch's float8_e4m3fn RNE cast of an exact power-of-two + scaling. + """ + w = x.shape[-1] + xf = x.reshape(-1, w).float().view(-1, w // _BLOCK, _BLOCK) + rows = xf.shape[0] + m, e = torch.frexp(xf.abs().amax(-1).clamp_min(_AMAX_FLOOR).double() / _E4M3_MAX) + e = torch.where(m == 0.5, e - 1, e).to(torch.int32) + s = torch.exp2(e.double()).float().unsqueeze(-1) + q = (xf / s).clamp(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn) + if slot_mapping is not None: + slots = slot_mapping.long() + else: + slots = torch.arange(rows, device=x.device) + live = (slots >= 0) & (slots < _num_slots(out)) + dst = slots[live] + if out.dtype == torch.bfloat16: + expected = out.clone() + _write_slots(expected, dst, (q.float() * s).view(rows, w).to(torch.bfloat16)[live]) + return expected + expected = out.view(torch.uint8).clone() + _write_slots(expected, dst, q.view(rows, w).view(torch.uint8)[live]) + expected_scale = out_scale.view(torch.uint8).clone() + _write_slots(expected_scale, dst, (e + 127).to(torch.uint8)[live]) + return expected.view(out.dtype), expected_scale.view(out_scale.dtype) + + def compare( + self, + out: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]], + ref: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]], + ) -> None: + """Bitwise over all of `out` (and `out_scale`): bf16 as int16, codes and bytes as uint8. + + The op computes one exact value per element, so there is no tolerance to choose; untouched + slots must keep their prior bits. On T = 4096 rows covering every path of the recipe the + gate sees each plausible wrong recipe differ from the op: no 1e-4 floor, `floor` instead of + `ceil` for the exponent, block 64, and the scale floored at 2^-13 instead of amax at 1e-4 + (SGLang's V4.1 store; on real window KV only the synthetic data discriminates the floors). + """ + outs = out if isinstance(out, tuple) else (out,) + refs = ref if isinstance(ref, tuple) else (ref,) + assert len(outs) == len(refs), (len(outs), len(refs)) + for name, o, r in zip(("out", "out_scale"), outs, refs): + assert o.dtype == r.dtype, (name, o.dtype, r.dtype) + assert o.shape == r.shape, (name, o.shape, r.shape) + a, b = _bits(o).reshape(-1), _bits(r).reshape(-1) + bad = torch.zeros((), dtype=torch.int64, device=a.device) + for i in range(0, a.numel(), 1 << 26): # chunked: a full bool reduction copies to int64 + bad += (a[i : i + (1 << 26)] != b[i : i + (1 << 26)]).sum() + assert int(bad) == 0, f"{name}: {int(bad)} of {o.numel()} elements differ bitwise" + + +mxfp8_fake_quantize = Mxfp8FakeQuantize() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/mxfp8_to_e4m3.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/mxfp8_to_e4m3.py new file mode 100644 index 000000000000..3b566f4a974f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/mxfp8_to_e4m3.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gather native MXFP8 rows (E4M3 codes x UE8M0 per 32) from a flat or paged pool and transcode them to +E4M3 with one per-tensor scale, via the Triton op `torch.ops.trtllm.mxfp8_to_e4m3` +(`tensorrt_llm/_torch/custom_ops/mxfp8_to_e4m3.py`).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops.mxfp8_to_e4m3 # noqa: F401 — registers torch.ops.trtllm.mxfp8_to_e4m3 +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_BLOCK = 32 +_E4M3_MAX = 448.0 + +#: Per-tensor factors of the compact gather: powers of two around the target's bound, and two +#: non-powers of two (one extra fp32 rounding before the E4M3 rounding). +_COMPACT_INV_SCALES = tuple(2.0**k for k in (-2, 0, 3, 6, 7, 8)) + (0.75, 1.0 / 3.0) +#: Query-row counts of the target's window read on the engine-sized pool. +_ENGINE_QUERIES = (1, 32, 192, 2048, 8192) +#: SWA pool pages of an engine runtime-value dump. +_ENGINE_POOL_PAGES = 188160 + + +def _num_slots(t: torch.Tensor) -> int: + return t.shape[0] if t.dim() == 2 else t.shape[0] * t.shape[1] + + +def _slot_rows(t: torch.Tensor, slots: torch.Tensor) -> torch.Tensor: + """Rows of a [slots, C] or [pages, page_size, C] tensor at the given slots.""" + if t.dim() == 2: + return t[slots] + return t[slots // t.shape[1], slots % t.shape[1]] + + +class Mxfp8ToE4m3(OpWrapper): + """Gather native MXFP8 rows from a pool and write them as E4M3 with ONE per-tensor factor. + + For output row `r`:: + + s = src_slots[r] if src_slots is not None else r + d = dst_slots[r] if dst_slots is not None else r + v[i] = e4m3(codes[s, i]) * 2 ^ (scales[s, i // 32] - 127) # exact in fp32 + out[d, i] = e4m3_rne(clamp(v[i] * inv_scale, -448, 448)) # NaN stays NaN + + A source slot that is negative or `>= slots` gives a row of `+0.0` codes; a destination slot + that is negative or `>= out slots` is skipped. Slot `s` of a paged `[pages, page_size, W]` + tensor is `(s // page_size, s % page_size)`; leading strides are free, so DeepSeek-V4.1's + 528 B window row is passed as two views (the first 512 bytes, the last 16). + + Fusion boundary: the gather, the dequantization, the rescale and the E4M3 cast are one launch. + The caller picks `inv_scale` (`1 / kv_scale`); nothing here computes an amax. Element for + element this is `e4m3(mxfp8_dequantize(...) * inv_scale)` and writes half the bytes. Writes + `out` in place, only at the destination slots of the rows, and returns `out`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"compact gather inv_scale {s:g}, " + + ("paged [320, 128, 528]" if paged else "flat [40960, 528]") + + ": 300 queries x 128 window columns, -1 tails", + spec=dict(kind="compact", inv_scale=s, paged=paged, queries=300), + ) + for s in _COMPACT_INV_SCALES + for paged in (True, False) + ) + + tuple( + Cell( + why=f"engine pool T={t}: [188160, 128, 528] first and last pages into " + "[8192 * 128, 512]", + spec=dict( + kind="engine_pool", queries=t, inv_scale=128.0, pool_pages=_ENGINE_POOL_PAGES + ), + ) + for t in _ENGINE_QUERIES + ) + + tuple( + Cell( + why=f"identity inv_scale {s:g}: no slot lists, normal range exact", + spec=dict(kind="identity", inv_scale=s, pool_pages=64), + ) + for s in (1.0, 16.0, 128.0) + ) + + ( + Cell( + why="dst_slots scatter, inv_scale 64: -1 and >= out_slots skipped", + spec=dict(kind="scatter", inv_scale=64.0), + ), + ) + ) + + note = """ + The target's call is the compact window read: `src_slots` the window half of the attention's + index table flattened (T queries x 128 columns, -1 for an empty column), `out` the first + `128 * T` rows of a `[max_num_tokens * 128, 512]` float8_e4m3fn buffer, `dst_slots` None. + + A power-of-two `inv_scale` is exact in E4M3's normal range: `out / inv_scale == v` for every + result with magnitude in [2^-6, 448] (measured on sm_103: none inexact at `inv_scale` 1, 16, + 128). Results below 2^-6 round to E4M3 subnormals (error <= 2^-10 / inv_scale); results above + 448 saturate to +-448. A non-power-of-two `inv_scale` adds one fp32 rounding (the multiply). + UE8M0 byte 0 is 2^-127 (fp32 subnormal); byte 255 makes the whole 32-block NaN codes; e4m3 NaN + codes stay NaN; signed zeros are kept. + + Metadata consumed: none. The caller supplies the slot lists. + + Preconditions: the op raises ValueError on a wrong dtype, shape, stride or device, and on a + non-positive `inv_scale`. `codes` float8_e4m3fn / uint8, `[slots, W]` or + `[pages, page_size, W]`, W % 32 == 0, unit last-dim stride; `scales` uint8 / float8_e8m0fnu, + `codes.shape[:-1] + (W / 32,)`; `out` float8_e4m3fn with unit last-dim stride; slot lists + int32 / int64 of one length; without `dst_slots`, `out` needs at least `rows` slots. The + caller chooses `inv_scale` so the values it needs do not saturate (for DeepSeek-V4.1's normed + window KV the target derives a power-of-two bound from the norm weights, + |x| <= sqrt(512) * max|w|). Graph-capturable: no host synchronization, no allocation; one + Triton program per output row, no autotuning. + + Not certified: W other than 512; page geometries other than the cells'; `inv_scale` values that + make the fp32 product overflow or underflow fp32 (not reachable from a UE8M0 byte and a scale + in 2^-8 .. 2^8); duplicate valid destinations with different sources. + """ + + def raw_call( + self, + codes: torch.Tensor, + scales: torch.Tensor, + out: torch.Tensor, + inv_scale: float, + src_slots: Optional[torch.Tensor] = None, + dst_slots: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + torch.ops.trtllm.mxfp8_to_e4m3(codes, scales, out, inv_scale, src_slots, dst_slots) + return out + + def reference( + self, + codes: torch.Tensor, + scales: torch.Tensor, + out: torch.Tensor, + inv_scale: float, + src_slots: Optional[torch.Tensor] = None, + dst_slots: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The expected `out` after the call: a copy of `out` with the rows' destinations filled. + + `out` is read for its prior contents (every slot the call does not write keeps them) and is + not mutated; pass a snapshot taken before the op ran. The decode is the float8_e4m3fn -> + fp64 cast and `torch.ldexp(code, byte - 127)` in fp64, exact in fp32; the one fp32 rounding + is the multiply by the fp32 `inv_scale` (none for a power of two), then the clamp to +-448 + and torch's RNE cast to E4M3 -- independent of the kernel's exponent-bit construction. + Byte 255 decodes to NaN; NaN code payloads are not modelled, so cells carry no NaN + encodings. + """ + w = codes.shape[-1] + num_slots = _num_slots(codes) + expected = out.view(torch.uint8).clone() + out_slots = _num_slots(expected) + if src_slots is not None: + rows = src_slots.shape[0] + elif dst_slots is not None: + rows = dst_slots.shape[0] + else: + rows = num_slots + dev = codes.device + src = src_slots.long() if src_slots is not None else torch.arange(rows, device=dev) + dst = dst_slots.long() if dst_slots is not None else torch.arange(rows, device=dev) + written = (dst >= 0) & (dst < out_slots) + src, dst = src[written], dst[written] + valid = (src >= 0) & (src < num_slots) + s = src.clamp(0, max(num_slots - 1, 0)) + c = _slot_rows(codes.view(torch.uint8), s).contiguous().view(torch.float8_e4m3fn) + b = _slot_rows(scales.view(torch.uint8), s).long() + v = torch.ldexp( + c.float().double().view(-1, w // _BLOCK, _BLOCK), (b - 127).double().unsqueeze(-1) + ) + v = torch.where((b == 255).unsqueeze(-1), torch.nan, v) + y = v.float() * torch.tensor(inv_scale, dtype=torch.float32, device=dev) + q = y.clamp(-_E4M3_MAX, _E4M3_MAX).view(-1, w).to(torch.float8_e4m3fn).view(torch.uint8) + q[~valid] = 0 + if expected.dim() == 2: + expected[dst] = q + else: + expected[dst // expected.shape[1], dst % expected.shape[1]] = q + return expected.view(out.dtype) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise on the uint8 codes over all of `out`. + + Every output is one E4M3 rounding of an exact fp32 product (two roundings for a + non-power-of-two `inv_scale`, the same two on both sides), so there is no tolerance to + choose; untouched slots must keep their prior bytes. Controls on recipe rows each change + more than 30% of the codes: UE8M0 bias 126, `inv_scale` ignored, codes read as E5M2. + """ + assert out.dtype == ref.dtype == torch.float8_e4m3fn, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + a = out.view(torch.uint8).reshape(-1) + e = ref.view(torch.uint8).reshape(-1) + bad = torch.zeros((), dtype=torch.int64, device=a.device) + for i in range(0, a.numel(), 1 << 26): # chunked: a full bool reduction copies to int64 + bad += (a[i : i + (1 << 26)] != e[i : i + (1 << 26)]).sum() + assert int(bad) == 0, f"{int(bad)} of {out.numel()} codes differ" + + +mxfp8_to_e4m3 = Mxfp8ToE4m3() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/nvfp4_fake_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/nvfp4_fake_quantize.py new file mode 100644 index 000000000000..a3d9702a9f78 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/nvfp4_fake_quantize.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 fake quantization (E2M1 x E4M3 per 16, amax floor 6*2^-9, no global scale) written into +slots of a bf16 container pool, or as native packed E2M1 bytes + E4M3 block-scale bytes, via the +trtllm Triton op. + +The op is `torch.ops.trtllm.nvfp4_fake_quantize`, registered in +`tensorrt_llm/_torch/custom_ops/nvfp4_fake_quantize.py`. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops.nvfp4_fake_quantize # noqa: F401 — registers torch.ops.trtllm.nvfp4_fake_quantize +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_BLOCK = 16 +_FLOOR = 6 * 2.0**-9 # the model's amax floor for E4M3-scaled FP4 +_E4M3_MAX = 448.0 + +_FORMS = ("core", "native") +_ROWS = (1, 5, 6, 64, 1024, 8192) +_WIDTHS = (16, 48, 128, 1024) + + +class Nvfp4FakeQuantize(OpWrapper): + """Fake-quantize bf16 rows to NVFP4 and write each row to a slot of a pool, in one launch. + + For row `r` and block `b` (elements `16b .. 16b+15` of the last dim):: + + amax = max(max_i |x[r, 16b+i]|, 6 * 2^-9) # fp32 + scale = e4m3_rne(min(amax / 6, 448)) # IEEE fp32 division + y_i = x[r, 16b+i] / scale # IEEE fp32 division + q_i = sign(y_i) * E2M1_rne(|y_i|) # ties to even, saturates at 6 + + The dtype of `out` picks the form. Core (bf16 `out` `[..., W]`): `q_i * scale`, exact + in bf16. Native (uint8 `out` `[..., W/2]` plus `out_scale` `[..., W/16]`, float8_e4m3fn + or uint8): the packed codes (`code[2j] | code[2j+1] << 4`, sign in bit 3, kept when + `q_i` is zero) and the scale bytes. Row `r` goes to slot `slot_mapping[r]`, skipped + when negative or `>= slots`, or to slot `r` when `slot_mapping` is None; a paged `out` + `[pages, page_size, C]` maps slot `s` to `(s // page_size, s % page_size)`. + + Fusion boundary: the amax, the scale, the quantization, the dequantization (core form) + and the scatter into slots are inside the call; RoPE (applied before it), the compressor + and any read of the pool are not. Writes `out` (and `out_scale`) in place and returns + `out` in the core form, `(out, out_scale)` in the native form. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"{form} random rows t={t}: [t, 512] dense, slot r = row r", + spec=dict(kind="random", form=form, t=t, width=512), + ) + for form in _FORMS + for t in _ROWS + ) + + tuple( + Cell(why=f"{form} {why}", spec=dict(kind=kind, form=form)) + for form in _FORMS + for kind, why in ( + ("ties", "exact midpoints: every E4M3 scale, +-1 bf16 ulp"), + ("floor", "amax floor: blocks at/below 6*2^-9, zero, -0.0"), + ("every_amax", "every bf16 amax: one block each, 448 saturates"), + ) + ) + + tuple( + Cell( + why=f"{form} width {w}: [33, {w}] random multi-scale rows", + spec=dict(kind="random", form=form, t=33, width=w), + ) + for form in _FORMS + for w in _WIDTHS + ) + + ( + Cell( + why="core leading dims: x [96, 1, 512] into out [96, 512]", + spec=dict(kind="leading_dims", form="core"), + ), + Cell( + why="core row stride: x a [96, 640][:, :512] column slice", + spec=dict(kind="row_stride", form="core"), + ), + ) + + tuple( + Cell( + why=f"bf16 pool {layout} {idx}: 37 rows into 4096 slots, 3 out of range", + spec=dict(kind="bf16_pool", form="core", layout=layout, idx_dtype=idx), + ) + for layout in ("flat", "paged64", "paged128") + for idx in ("int32", "int64") + ) + + tuple( + Cell( + why=f"native pool {layout}: 37 rows into 4096 slots, 3 out of range", + spec=dict(kind="native_pool", form="native", layout=layout, idx_dtype="int32"), + ) + for layout in ("split", "row288", "paged_row288") + ) + + ( + Cell( + why="zero rows: an empty batch and slot map write nothing", + spec=dict(kind="zero_rows", form="core"), + ), + ) + ) + + note = """ + Cells. Dense calls with slot r = row r in both forms over random multi-scale rows (the + low end below the floor), every E4M3 scale at exact E2M1 midpoints and one bf16 ulp + either side, the floor region and every positive finite bf16 amax; widths 16, 48, 128, + 1024; `x` with leading dims and as a column slice; 37 rows into 4096-slot pools (bf16 + flat / 64-row / 128-row pages with int32 and int64 slots; native as split pools, one + 288-byte row, and pages of 288-byte rows), three of them at -1, slots + 5 and -2^30 or + -2^40; an empty batch. + + Value set: the DeepSeek-V4.1 compressed KV, the model's + `fp4_act_quant(x, 16, inplace, scale_dtype=float8_e4m3fn)`, measured bit-exact on sm_103 + against the checkpoint's own kernel (every E4M3 scale at its midpoints, the floor region, + every bf16 amax) and against the golden compressed KV (5,158,912 elements), both forms. + Scale byte 0x01 at or below the floor, 0x7E (448) for amax >= 2688, never NaN. FlashInfer's + NVFP4 KV quantize and `trtllm.fp4_quantize` are a different value set: they multiply by a + reciprocal (3,608 golden elements differ, all exact midpoints) and flush blocks with + amax <= 6*2^-10 to zero. FlashInfer's `nvfp4_kv_dequantize` rejects the strided 288-byte + row views (it needs contiguous inputs). + + Metadata consumed: none; the slot map is an argument. DeepSeek-V4.1's target writes the + native form as `(latent, codes, md.v41_compress_slot_mapping[r][:n], scales)` with `codes` + / `scales` the cache manager's `[..., :256]` / `[..., 256:]` views of its 288-byte rows. + + Preconditions: `x` bf16, W a positive multiple of 16, unit last-dim stride, leading dims + viewable as rows; `out` `[slots, C]` or `[pages, page_size, C]`, unit last-dim stride; + `slot_mapping` int32/int64 `[rows]`; `out_scale` exactly with a uint8 `out`; without a + slot map `out` has at least `rows` slots. The op raises ValueError otherwise. Use -1 for + "no slot". `out` may be `x` itself when `slot_mapping` is None (certified at 256 rows). + + CUDA graphs: capture after one eager call has compiled the variant -- (W, slot map or not, + form), and in the slot-r form the class of the row count (1, a multiple of 16, other). + Capture is certified at every graph batch G of 1..32, 40..128 step 8, 192, 256 and at 6G + (ratio 1) / 3G (ratio 2) rows, replay bitwise equal to eager; a first-ever call captured + is not. + + Envelope, in the test's far-end and sweep cases: bf16 views up to `[2118144, 64, 512]` / + `[1059584, 128, 512]` and native views up to `[6480780, 64, 288]` / `[3245959, 128, 288]` + (the PP2 stage-alone cache managers at 0.5 x total memory), rows on both sides of every + 2^31..2^36 offset, nothing else written; the result does not depend on the view's page + count; every row count 1..16384; multi-request steps on live managers up to 256 requests + and position 32,781. All offsets are 64-bit. + + Not certified: non-finite input; two rows naming one slot; overlap of `x` and `out` other + than exact in place; widths other than 16, 48, 128, 512, 1024; more than 16384 rows; a + capture at a row count outside the list above; views larger than the envelope (a KV quota + above 0.5 x total memory, pipeline splits other than PP2's even split, unsampled + max_batch_size); page sizes other than 1, 64, 128; torch.compile. + """ + + def raw_call( + self, + x: torch.Tensor, + out: torch.Tensor, + slot_mapping: Optional[torch.Tensor] = None, + out_scale: Optional[torch.Tensor] = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + torch.ops.trtllm.nvfp4_fake_quantize(x, out, slot_mapping, out_scale) + return out if out_scale is None else (out, out_scale) + + def reference( + self, + x: torch.Tensor, + out: torch.Tensor, + slot_mapping: Optional[torch.Tensor] = None, + out_scale: Optional[torch.Tensor] = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """The whole of `out` (and `out_scale`) after the call: quantized rows in the slots + the op writes, the prior contents everywhere else. + + The model's recipe in torch ops: fp32 amax clamped to the floor, `amax / 6` cast by + torch to float8_e4m3fn (round-to-nearest-even, clamped to 448 first), `x / scale` in + fp32, E2M1 rounding by `torch.round` (half to even) on the three binades, the sign of + `y` kept. Nothing is mutated: the expected buffers are contiguous copies of `out` / + `out_scale` with the written slots replaced. + """ + width = x.shape[-1] + xf = x.reshape(-1, width).float().unflatten(-1, (-1, _BLOCK)) + amax = xf.abs().amax(-1, keepdim=True).clamp_min(_FLOOR) + s8 = (amax / 6.0).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) + s = s8.float() + y = (xf / s).clamp(-6.0, 6.0) + a = y.abs() + r = torch.where( + a < 2, + torch.round(a * 2) / 2, + torch.where(a < 4, torch.round(a), torch.round(a / 2) * 2), + ) + + rows = xf.shape[0] + flat_out = out.clone(memory_format=torch.contiguous_format) + dest = flat_out.view(-1, out.shape[-1]) + if slot_mapping is None: + src = torch.arange(rows, device=x.device) + slots = src + else: + slots = slot_mapping.long() + src = torch.nonzero((slots >= 0) & (slots < dest.shape[0])).flatten() + slots = slots[src] + + if out_scale is None: + deq = (torch.copysign(r, y) * s).flatten(-2).to(torch.bfloat16) + dest[slots] = deq[src] + return flat_out + + code = torch.where(r < 2, r * 2, torch.where(r == 6, 7.0, r + 2)).to(torch.int32) + nib = (code | (torch.signbit(y).to(torch.int32) << 3)).flatten(-2) + packed = (nib[..., 0::2] | (nib[..., 1::2] << 4)).to(torch.uint8) + flat_scale = out_scale.clone(memory_format=torch.contiguous_format) + dest[slots] = packed[src] + scale_bytes = s8.squeeze(-1).view(torch.uint8) + flat_scale.view(-1, out_scale.shape[-1])[slots] = scale_bytes[src].view(out_scale.dtype) + return flat_out, flat_scale + + def compare( + self, + out: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + ref: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + ) -> None: + """Bitwise on every element of `out` (and `out_scale`), written or not. + + `E2M1 x E4M3` is exact in bf16 and both sides make the same IEEE divisions and the + same round-to-nearest-even, so a differing element is a wrong scale, tie, nibble or + slot, and a differing element outside the written slots is a write the op must not + make. Controls on 1024 random rows plus the tie and floor sets (538,624 elements): + a power-of-two scale changes 233,906 of them, no amax floor 42,579, ties toward zero + 2,919, `x * (1 / scale)` 190. + """ + assert isinstance(out, tuple) == isinstance(ref, tuple), (type(out), type(ref)) + pairs = zip(out, ref) if isinstance(out, tuple) else ((out, ref),) + for got, want in pairs: + assert got.dtype == want.dtype, (got.dtype, want.dtype) + assert got.shape == want.shape, (got.shape, want.shape) + raw = torch.int16 if got.element_size() == 2 else torch.uint8 + bad = (got.view(raw) != want.view(raw)).sum().item() + assert bad == 0, f"{bad} of {got.numel()} elements of {tuple(got.shape)} differ" + + +nvfp4_fake_quantize = Nvfp4FakeQuantize() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/nvfp4_kv_to_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/nvfp4_kv_to_bf16.py new file mode 100644 index 000000000000..f3255825ddea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/nvfp4_kv_to_bf16.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Paged NVFP4 KV dequantize-gather to bf16 (packed E2M1 x one E4M3 scale per 16 x a global scale), +the bitwise twin of FlashInfer's `nvfp4_kv_dequantize_paged` in one launch, via the trtllm Triton op. + +The op is `torch.ops.trtllm.nvfp4_kv_to_bf16`, registered in +`tensorrt_llm/_torch/custom_ops/nvfp4_kv_to_bf16.py`. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops.nvfp4_kv_to_bf16 # noqa: F401 — registers torch.ops.trtllm.nvfp4_kv_to_bf16 +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +_BLOCK = 16 + +_TABLE_DTYPES = ("int32", "int64") + + +class Nvfp4KvToBf16(OpWrapper): + """Dequantize the NVFP4 pool rows a table names into contiguous bf16 rows, in one launch. + + For every request `b` and every position `t < min(seq_lens[b], L)`, with + `page = block_tables[b, t // page_size]`:: + + out[b, t, 0, i] = bf16( + (E2M1[code_i] * e4m3(sf[i // 16])) * global_scale + ) # only if 0 <= page < pages + + `kv_cache` `[pages, page_size, 1, W/2]` uint8 holds the packed codes (even element in the low + nibble) and `kv_cache_sf` `[pages, page_size, 1, W/16]` the E4M3 scale bytes, read through any + strides with a unit last-dim stride: the two halves of DeepSeek-V4.1's 288-byte row in place. + Positions at or past `seq_lens[b]` (or `L`) and entries whose page id is negative or `>= pages` + are left unwritten. The signature, operand views and results are those of FlashInfer's + `nvfp4_kv_dequantize_paged` as the DeepSeek-V4.1 target calls it. + + Fusion boundary: the page lookup, the E2M1 decode, the E4M3 scale, the global scale and the + bf16 cast are inside the call; the pool and its writes, the table and the remapping of indices + for the attention that reads `out` are the caller's. Writes `out` in place and returns it. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"per query R={r} {dt}: [R, 512] top-k table over pages of one row", + spec=dict(kind="topk", R=r, table_dtype=dt), + ) + for r in (1, 6, 32, 64, 192, 384, 1536, 8192) + for dt in _TABLE_DTYPES + ) + + tuple( + Cell( + why=f"per request page {ps} {dt}: lengths 0..8192, a hole, an id past", + spec=dict(kind="request", page_size=ps, table_dtype=dt), + ) + for ps in (64, 128) + for dt in _TABLE_DTYPES + ) + + ( + Cell( + why="prefill 16384 entries at page 128: the whole pool", + spec=dict(kind="prefill_16k"), + ), + ) + + tuple( + Cell( + why=f"compact r={ratio} B={b}: [B, 132*128/r, 1, 512], random lengths", + spec=dict(kind="compact", ratio=ratio, B=b), + ) + for ratio in (1, 2) + for b in (1, 8, 32) + ) + + tuple( + Cell( + why=f"engine pool {rows} rows {form}: ends of the pool named", + spec=dict(kind="engine_pool", rows=rows, page_size=ps, form=form), + ) + for rows, ps in ((453120, 64), (151040, 128)) + for form in ("per_request", "per_query") + ) + ) + + note = """ + Cells. W = 512 over pools of random bytes (every code byte, scale bytes 0x00..0x7E) read as the + two strided halves of 288-byte rows, `global_scale` 1.0. Per query: a 16,384-row pool viewed as + pages of one row, `[R, 512]` tables with -1 padding, ids one past and far past the pool, a row + that ends 37 entries early. Per request: pages of 64 and 128, lengths 0, 1, page - 1, page + 1, + 700, 4097, 8192, a table hole and a page id past the pool; one 16,384-entry request; the + target's compact views `[B, 132 * 128 / r, 1, 512]`; pools of 453,120 x 64 / 151,040 x 128 rows + with both ends named. + + Exact for `global_scale` 1.0: `E2M1 x E4M3` has at most 6 significant bits, so the fp32 product + and its bf16 cast are exact. Code 0x8 gives -0.0; a NaN scale byte gives NaN. Native rows written + by `nvfp4_fake_quantize` read back bit-equal to its core output. Every certified call is also + bitwise equal to FlashInfer's `nvfp4_kv_dequantize_paged` on the same operands, in one launch + instead of two: measured on sm_103 1.4x (one 512-row gather) to 14.7x (256 requests of 768 + entries) faster. + + Metadata consumed: none. DeepSeek-V4.1's target uses two forms. Per request: + `compress_block_tables[r][:B]`, `compressed_kv_lens_cuda[r][:B]` over the native views at page + size 128 / r, into `comp_out[r][:B]`, the `[max_batch_size, cols * 128 / r, 1, 512]` view of one + buffer. Per query: the views re-viewed with page size 1, the index op's `[T, 512]` int32 pool + slots (-1 after each query's valid entries), `seq_lens` 512, into the `[T, 512, 1, 512]` view + of a `[max_num_tokens * 512, 512]` buffer. + + Preconditions: every `t < min(seq_lens[b], L)` has `t // page_size < max_pages` (the table is + read there, unchecked); `global_scale` a CUDA float32 `[1]`; `block_tables` int32/int64 with a + unit last-dim stride; `seq_lens` int32; `out` contiguous bf16 `[B, L, 1, W]`, W a power of two. + The op raises ValueError on a wrong dtype, shape, stride or device. Unwritten entries keep what + `out` held: the target masks the -1 columns by count and never reads them. + + No host synchronization and no allocation: graph-capturable after one eager call of the same + form (a first-ever call inside a capture is not certified). Triton specializes the kernel on its + integer arguments (strides, page count, page size, table stride, L); measured, nothing compiles + after the first eager call of a (form, pool view, out width) over every B 1..256, every length, + every R 1..16,384 and every capture. + + Envelope, in the test's cache-manager cases: native views up to `[6480780, 64, 288]` / + `[3245959, 128, 288]` (PP2 stage-alone managers at 0.5 x total memory), reads past byte offset + 2^36, the result independent of the view's page count; per query every R 1..16,384 and, under + capture, R = G and 6G; per request at table widths 132, 164, 260, 516, 1028, 2052, 8196 with up + to 256 requests (128 at width 8196), every length, captures at every graph batch G of 1..32, + 40..128 step 8, 192, 256 with next_n 1 and 6; live multi-request steps up to 256 requests. + + Not certified: W other than 512; `global_scale` other than 1.0 beyond 0.5, 2 and 1/3 (held + only to the FlashInfer twin); views beyond the ones above; a per-query call of more than 16,384 + rows; the per-request form above 256 requests, above 128 at width 8196 (1 GiB per request), or + at another table width; torch.compile. + """ + + def raw_call( + self, + kv_cache: torch.Tensor, + kv_cache_sf: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + global_scale: torch.Tensor, + out: torch.Tensor, + ) -> torch.Tensor: + torch.ops.trtllm.nvfp4_kv_to_bf16( + kv_cache, kv_cache_sf, block_tables, seq_lens, global_scale, out + ) + return out + + def reference( + self, + kv_cache: torch.Tensor, + kv_cache_sf: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + global_scale: torch.Tensor, + out: torch.Tensor, + ) -> torch.Tensor: + """The whole of `out` after the call: dequantized rows where the op writes, `out`'s prior + bits everywhere else. + + An E2M1 lookup table times the E4M3 scale in fp32 (exact), times `global_scale` (one fp32 + rounding, none for 1.0), cast to bf16 once. `out` is not mutated: the expected buffer is a + copy of it. + """ + dev = out.device + num_requests, length = out.shape[0], out.shape[1] + pages, page_size = kv_cache.shape[0], kv_cache.shape[1] + t = torch.arange(length, device=dev) + col = (t // page_size).clamp_max(max(block_tables.shape[1] - 1, 0)) + page = block_tables[:, col].long() + lens = seq_lens.long().clamp_max(length)[:, None] + written = (t[None, :] < lens) & (page >= 0) & (page < pages) + page = page.clamp(0, pages - 1)[written] + row = (t % page_size)[None, :].expand(num_requests, length)[written] + codes = kv_cache[page, row, 0] + scales = kv_cache_sf[page, row, 0].contiguous().view(torch.float8_e4m3fn).float() + lut = torch.tensor(_E2M1 + tuple(-v for v in _E2M1), dtype=torch.float32, device=dev) + values = torch.stack([lut[(codes & 0xF).long()], lut[(codes >> 4).long()]], -1).flatten(-2) + values = (values * scales.repeat_interleave(_BLOCK, -1)) * global_scale.float() + expected = out.clone() + expected[:, :, 0][written] = values.to(out.dtype) + return expected + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise on every element of `out`, written or not. + + With `global_scale` 1.0 the products are exact, so a differing element is a wrong row, + scale or nibble, and a differing element outside the written set (the test pre-fills a NaN + sentinel) is a write the op must not make. Controls (nibbles swapped, one scale per 32 + instead of 16) each change more than 30% of the values. + """ + assert out.dtype == ref.dtype == torch.bfloat16, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + bad = (out.view(torch.int16) != ref.view(torch.int16)).sum().item() + assert bad == 0, f"{bad} of {out.numel()} elements differ" + + +nvfp4_kv_to_bf16 = Nvfp4KvToBf16() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/nvfp4_kv_to_e4m3.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/nvfp4_kv_to_e4m3.py new file mode 100644 index 000000000000..6d02d87ffa7b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/ops/quantization/nvfp4_kv_to_e4m3.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gather paged NVFP4 KV rows and transcode them to E4M3 with one per-tensor scale, via the trtllm +Triton op. + +The op is `torch.ops.trtllm.nvfp4_kv_to_e4m3`, registered in +`tensorrt_llm/_torch/custom_ops/nvfp4_kv_to_e4m3.py`. +""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops.nvfp4_kv_to_e4m3 # noqa: F401 — registers torch.ops.trtllm.nvfp4_kv_to_e4m3 +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +_E4M3_MAX = 448.0 +_BLOCK = 16 + + +class Nvfp4KvToE4m3(OpWrapper): + """Transcode the NVFP4 pool rows each request reads into an E4M3 buffer, in one launch. + + For every request `b` and every position `t < min(seq_lens[b], L)`, with + `page = block_tables[b, t // page_size]`:: + + v[i] = E2M1[code_i] * e4m3(sf[i // 16]) # exact in fp32 + out[b, t, 0, i] = e4m3_rne(clamp(v[i] * inv_scale, -448, 448)) # only if 0 <= page < pages + + The pool is DeepSeek-V4.1's 288-byte compressed-KV row read in place as two + views: `kv_cache` `[pages, page_size, 1, W/2]` uint8 (packed E2M1, even + element in the low nibble) and `kv_cache_sf` `[pages, page_size, 1, W/16]` + E4M3 scale bytes; there is no global scale. Positions at or past + `seq_lens[b]` (or `L`) and rows whose page id is negative or `>= pages` are + left unwritten. + + Fusion boundary: page lookup, E2M1 decode, E4M3 scale, rescale and E4M3 cast + are inside the call; choosing `inv_scale` is the caller's. Writes `out` in + place and returns it. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"compact r={ratio} inv_scale={inv}: whole-context [8, 132*128/r, 1, 512]", + spec=dict(kind="compact", ratio=ratio, inv_scale=inv), + ) + for ratio in (2, 1) + for inv in (1.0, 16.0, 64.0, 128.0, 0.25, 0.8) + ) + + tuple( + Cell( + why=f"engine pool r={ratio} mbs={mbs}: engine page count, ends named", + spec=dict(kind="engine", ratio=ratio, pages=pages, mbs=mbs, inv_scale=64.0), + ) + for mbs in (1, 8, 32) + for ratio, pages in ((2, 453120), (1, 151040)) + ) + + tuple( + Cell( + why=f"per-query top-k r={ratio} T={t}: one-row pages, -1 padded", + spec=dict(kind="per_query", ratio=ratio, T=t, inv_scale=2.0), + ) + for ratio in (2, 1) + for t in (1, 6, 192, 2048, 8192) + ) + ) + + note = """ + Cells. The target's compact compressed buffer `[B, 132 * 128 / r, 1, 512]` + at compress ratio 2 (64-row pages) and 1 (128-row pages): 8 requests with + lengths 0, 1, page-1, page+1, 700/r, 4097/r, 16398/r and `L`, one + page-table hole. The engine's pool page counts (453120 x 64 and + 151040 x 128 rows) with pages named at both ends, into the + `max_batch_size` views 1, 8, 32. The per-query top-k form the target reads + when it is smaller than the per-request one: the pool viewed as pages of + one row (`[pages * bk, 1, 1, C]`, a slot number is its own page id), a + `[T, 512]` table of selected slots with `-1` after each row's valid + entries, `seq_lens` 512 for every row, `out` `[T, 512, 1, 512]`. + + Lossy by construction: `E2M1 x E4M3` values carry up to 6 significant bits + and E4M3 keeps 4, so the result rounds by at most `2^-4` relative in + E4M3's normal range (43% of values inexact on random rows, about 23% on + DeepSeek-V4.1's real compressed KV). Code `0x8` (-0.0) keeps its sign; + `6 * 448` saturates to 448; a NaN scale byte gives NaN. + + Metadata consumed: none. DeepSeek-V4.1's target passes + `compress_block_tables[r][:B]` and `compressed_kv_lens_cuda[r][:B]`. + + Preconditions: every `t < min(seq_lens[b], L)` has + `t // page_size < max_pages` (the table is read there); the caller picks an + `inv_scale` > 0 under which the values it needs do not saturate; + `block_tables` int32/int64 and `seq_lens` int32; `out` contiguous + float8_e4m3fn with `W` a power of two. The op raises ValueError on a wrong + dtype, shape, stride, device or a non-positive `inv_scale`. No host + synchronization and no allocation: graph-capturable. Grid + `(ceil(L / 8), B)`, eight rows per program, early exit past `seq_lens[b]`. + + Not certified: `W` other than 512; `max_batch_size` other than 1, 8, 32 in + the engine-pool case (8 in the others); CUDA-graph capture of the per-query + form specifically (capture is certified on the whole-context form, the same + kernel). + """ + + def raw_call( + self, + kv_cache: torch.Tensor, + kv_cache_sf: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + inv_scale: float, + ) -> torch.Tensor: + torch.ops.trtllm.nvfp4_kv_to_e4m3( + kv_cache, kv_cache_sf, block_tables, seq_lens, out, inv_scale + ) + return out + + def reference( + self, + kv_cache: torch.Tensor, + kv_cache_sf: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + inv_scale: float, + ) -> torch.Tensor: + """The whole of `out` after the call: transcoded rows where the op writes, `out`'s + prior bytes everywhere else. + + An E2M1 lookup table times the E4M3 scale in fp32 (exact: at most 6 + significant bits), times the fp32 `inv_scale` (one rounding, none for a + power of two), clamped to +-448 and cast by torch to float8_e4m3fn + (round-to-nearest-even). `out` is not mutated: the expected buffer is a + copy of it. + """ + dev = out.device + num_requests, length = out.shape[0], out.shape[1] + pages, page_size = kv_cache.shape[0], kv_cache.shape[1] + t = torch.arange(length, device=dev) + col = (t // page_size).clamp_max(block_tables.shape[1] - 1) + page = block_tables[:, col].long() + lens = seq_lens.long().clamp_max(length)[:, None] + written = (t[None, :] < lens) & (page >= 0) & (page < pages) + page = page.clamp(0, pages - 1)[written] + row = (t % page_size)[None, :].expand(num_requests, length)[written] + codes = kv_cache[page, row, 0] + scales = kv_cache_sf[page, row, 0].contiguous().view(torch.float8_e4m3fn).float() + lut = torch.tensor(_E2M1 + tuple(-v for v in _E2M1), dtype=torch.float32, device=dev) + values = torch.stack([lut[(codes & 0xF).long()], lut[(codes >> 4).long()]], -1).flatten(-2) + values = values * scales.repeat_interleave(_BLOCK, -1) + y = values * torch.tensor(inv_scale, dtype=torch.float32, device=dev) + y = y.clamp(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn) + expected = out.view(torch.uint8).clone() + expected[:, :, 0][written] = y.view(torch.uint8) + return expected.view(out.dtype) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise on every byte of `out`, written or not. + + The products are exact in fp32 and the cast is torch's RNE on both + sides, so a differing code is a wrong row, scale or nibble order; a + differing byte outside the written set is a write the op must not make. + Controls (nibbles swapped, one scale per 32 instead of 16, `inv_scale` + ignored) each change more than 30% of the codes. + """ + assert out.dtype == ref.dtype == torch.float8_e4m3fn, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + bad = (out.view(torch.uint8) != ref.view(torch.uint8)).sum().item() + assert bad == 0, f"{bad} of {out.numel()} bytes differ" + + +nvfp4_kv_to_e4m3 = Nvfp4KvToE4m3() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/weights.py new file mode 100644 index 000000000000..fd06dd181862 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/weights.py @@ -0,0 +1,1430 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight tables: deepseek-v4.1-flash-nvfp4 / sm_103 / unified. + +Two tables, one per weight pass of the engine: + +* ``MODEL_WEIGHTS`` -- the trunk, the vision tower and the aligner, declared into + the core's ``w`` and filled by the shell's ``load_weights``; +* ``DRAFT_WEIGHTS`` -- the three DSpark stages (``mtp.*``), declared into the + core's ``w_ds`` only when the engine was built with a ``speculative_config``, + and filled by the loader's second, lazily mapped pass over the same + checkpoint (``load_draft_weights``). + +Each states every weight once, as a ``W`` row: shape, checkpoint source, dtype, +load-time transform and the layers that carry it. ``dims`` is the one place a +per-rank width is derived, from the mapping alone. + +**Per-rank sources.** The member (tep4 / dep4 / tp4, and tep8 / dep8 over two +nodes) decides the slices, read off the mapping, never off the path segment: + +* attention heads: TP members hold heads ``[h r, h r + h)``, ``h = 64 / tp`` + (16 at tp4, 8 at tp8) -- ``wq_b`` rows, ``attn_sink``, ``wo_a`` groups + ``[g r, g r + g)``, ``g = 8 / tp``, and ``wo_b`` columns split on the same + boundary; attention DP holds all 64; +* routed experts: EP members hold experts ``[e r, e r + e)`` whole, ``e = 384 / tp`` + (96 at EP4, 48 at EP8); tp4 holds all 384 at intermediate ``[576r, 576r+576)``; +* Engram tables (``layers.{1,14}.engram.embed.{weight,scale}``): with + host_offload on (the default) they are no parameter at all -- loaded whole, + once per node, into the shared host copy of the sibling ``engram_host_table`` + (``core.attach_engram_host_table``) and checked against the checkpoint here, + never copied to a device; with ``TRTLLM_DSV41_ENGRAM_HOST_OFFLOAD=0``, rows + ``[ceil(rows/tp) r, ...)`` on the device of every member; +* everything else is replicated: hyper-connections, norms, ``wq_a``/``wkv``, + compressor, indexer, router (``gate.bias`` and ``gate.bias_vl``), shared + expert, embedding, and the whole vision tower + aligner. + +A source that is a per-rank part of a checkpoint tensor is a ``Slice``: the +part is read **before** any transform, because a relayout is a function of the +per-rank row count, and a lazily mapped safetensors entry then reads only this +rank's bytes off disk. + +Formats. Dense FP8 weights are E4M3 in 32x32 blocks with one E8M0 ``.scale`` +per block: exactly MXFP8 once each block scale is repeated over its 32 rows +(``repeat_interleave`` -- tiling is silently wrong) and 128x4-swizzled, which +is what ``mxfp8_mxfp8_gemm`` reads. Fused projections expand each part's scales +and then concatenate (``wq_a || wkv``; the shared expert's ``[w1 ; w3]`` = +``[gate ; up]``, the half order ``silu_and_mul`` reads). ``wo_a`` is +dequantized to bf16 exactly as the checkpoint's own converter does. Routed +NVFP4 experts are packed for the trtllm-gen block-scale runner: FC1 = +``[w3 ; w1]`` = ``[up ; gate]`` (the OPPOSITE half order), interleaved, 32-row +block shuffled (weight and scale bytes), scales 128x4 swizzled; FC2 shuffled, +scales swizzled. The per-expert fp32 scalars are loaded whole (the layer-level +maxima need all 384) and folded in the core's ``post_load_weights``. + +The vision tower (``vision.*``), the aligner (``aligner.*``) and the three span +delimiters are bf16 in the checkpoint and load as they are, with two +relayouts: ``aligner.w1``'s columns go from the reference unfold order +``(channel, kh, kw)`` to ``(kh, kw, channel)``, so the aligner input is a +plain row gather of the ViT output; ``image_start`` / ``image_end`` / +``image_newline`` are stacked into one ``[3, 5120]`` table in that order. + +Predicted non-load sets of ``MODEL_WEIGHTS.load`` (asserted exactly, never relaxed): + +* ``mtp.*`` (2401 keys) -- the three DSpark stages, ``DRAFT_WEIGHTS``' rows; +* EP members: the weight and weight_scale keys of every routed expert outside + this rank's window (6 per expert per layer); +* pipeline stages: every key of another stage's layers, except the previous + stage's last ``hc_ffn`` triple, and ``head.weight`` outside the last stage. + +**The DSpark drafter** (``DRAFT_WEIGHTS``, layer ``s`` = stage ``mtp.{s}``): +replicated on every rank of every member. Dense stage weights are MXFP8 +exactly like the trunk's; the 128 routed experts per stage are MXFP4 (E2M1 x +UE8M0/32) packed for the MXFP4 trtllm-gen runner with the NVFP4 runner's +byte-generic relayouts (FC1 = ``[w3 ; w1]`` interleaved and 32-row shuffled, +FC2 shuffled, scales 128x4 swizzled); ``wkv`` of the three stages is also +stacked into one ``[1536, 5120]`` operand for the ring write. Its predicted +non-load set is the confidence head (the worker drafts the whole block, +``confidence_threshold`` 0) and the all-zero ``gate.bias_vl``. +""" + +from dataclasses import dataclass +from types import SimpleNamespace + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2._weights import ModelWeights, W + +from . import engram_host_table as _engram_host + +SF16 = 16 # NVFP4: one e4m3 scale per 16 along K +MX = 32 # MXFP8 / checkpoint FP8 block + +# Hyper-connection mixes: 4 pre + 4 post + 16 comb per copy count of 4. +HC = 4 +MIX = (2 + HC) * HC +IDX_HEADS = 32 +IDX_DIM = 128 +VIT_DIM = 1024 +VIT_LAYERS = 32 +VIT_HEADS = 16 +VIT_INTER = 2816 +VIT_PATCH_IN = 3 * 14 * 14 # a patch flattened (channel, kh, kw) +VIT_RATIO = 3 # the aligner's 3x3 block +_AL_TAPS = VIT_RATIO * VIT_RATIO +# The checkpoint's Engram layers (1 and 14). One row pair per slot, because the +# two tables have different row counts and a row's shape cannot depend on its +# layer index. +_ENGRAM_SLOTS = 2 + +# Parameter-key suffixes after which the load releases transform scratch. +_RELEASE_AFTER = ("_fc1_w", "_fc1_s", "_fc2_w", "_fc2_s", "_eng_table") + +# Copy threshold above which an untransformed row is streamed in row chunks, +# so a 25 GB Engram shard never sits in host RAM at once. +_CHUNK_ELEMS = 1 << 28 + + +@dataclass(frozen=True) +class Slice: + """This rank's part of one checkpoint tensor: ``weights[key][rows]``. + + ``rows`` is an index tuple over the tensor's leading dimensions, applied + before any transform. + """ + + key: str + rows: tuple + + +def _split_source(ckpt) -> tuple[tuple[str, ...], tuple, bool]: + """``(keys, per-key index tuples or None, multi)`` of one manifest source.""" + multi = isinstance(ckpt, tuple) + parts = ckpt if multi else (ckpt,) + keys = tuple(p.key if isinstance(p, Slice) else p for p in parts) + srcs = tuple(p.rows if isinstance(p, Slice) else None for p in parts) + return keys, srcs, multi + + +def _is_lazy(v) -> bool: + return not isinstance(v, torch.Tensor) + + +def _materialize(v, src=None) -> torch.Tensor: + """Realize a checkpoint entry (a lazy safetensors slice or a tensor), + reading only ``src`` of it. 0-dim entries (NVFP4 scalars) reject ``[:]``.""" + if _is_lazy(v): + if len(v.get_shape()) == 0: + return v[()] + return v[src] if src is not None else v[:] + if v.ndim == 0: + return v + return v[src] if src is not None else v + + +def _shape(v) -> tuple: + return tuple(v.get_shape()) if _is_lazy(v) else tuple(v.shape) + + +def mx_sf_numel(n: int, k: int) -> int: + """Bytes of the 128x4-swizzled UE8M0 scales of an [n, k] MXFP8 weight.""" + return ((n + 127) // 128 * 128) * ((k // MX + 3) // 4 * 4) + + +# ---- transforms (run on the destination device) --------------------------- + + +def _swizzle(x: torch.Tensor) -> torch.Tensor: + """128x4 block-scale swizzle of a [M, C] uint8 scale matrix, flat.""" + return torch.ops.trtllm.block_scale_interleave(x.contiguous()).reshape(-1) + + +def _expand_rows(scale: torch.Tensor, rows: int) -> torch.Tensor: + """[ceil(N/32), K/32] E8M0 block scales -> [N, K/32] uint8 per-row scales.""" + return scale.view(torch.uint8).repeat_interleave(MX, dim=0)[:rows] + + +def _mx_sf(t, rows) -> torch.Tensor: + """One fused MXFP8 weight's 128x4-swizzled scales: each part's 32x32 block + scales expanded to its ``rows``, concatenated, then swizzled once.""" + scales = t if isinstance(t, tuple) else (t,) + assert len(scales) == len(rows), (len(scales), rows) + parts = [_expand_rows(s, n) for s, n in zip(scales, rows)] + return _swizzle(torch.cat(parts, dim=0)) + + +def _wqkv_a_sf(t, core) -> torch.Tensor: + return _mx_sf(t, (core.q_lora, core.head_dim)) + + +def _wq_b_sf(t, core) -> torch.Tensor: + return _mx_sf(t, (core.h_loc * core.head_dim,)) + + +def _hidden_sf(t, core) -> torch.Tensor: + """Scales of a weight with ``hidden`` output rows (``wo_b``, ``sh_w2``).""" + return _mx_sf(t, (core.hidden,)) + + +def _idx_wq_b_sf(t, core) -> torch.Tensor: + return _mx_sf(t, (IDX_HEADS * IDX_DIM,)) + + +def _sh_w13_sf(t, core) -> torch.Tensor: + return _mx_sf(t, (core.moe_inter, core.moe_inter)) + + +def _eng_wkv_sf(t, core) -> torch.Tensor: + return _mx_sf(t, (core.eng_wkv_n,)) + + +def _cat_rows_f32(t, core) -> torch.Tensor: + return torch.cat([x.float() for x in t], dim=0).contiguous() + + +def _to_f32(t, core) -> torch.Tensor: + return t.float().contiguous() + + +def _u8(t, core) -> torch.Tensor: + return t.view(torch.uint8).contiguous() + + +def _idx_wproj(t, core) -> torch.Tensor: + """weights_proj with the indexer's 128^-0.5 * 32^-0.5 = 2^-6 folded in + (a power of two: exact in bf16).""" + return (t.float() * 2.0**-6).to(torch.bfloat16).contiguous() + + +def _dequant_wo_a(t, core) -> torch.Tensor: + """FP8 E4M3 x E8M0 32x32 -> bf16, the checkpoint converter's own recipe.""" + w, s = t + out = w.float().unflatten(0, (-1, MX)).unflatten(-1, (-1, MX)) * s.float()[:, None, :, None] + return out.flatten(2, 3).flatten(0, 1).to(torch.bfloat16).contiguous() + + +def _block32_perm(rows: int, device) -> torch.Tensor: + """Gather index of the 32-row block shuffle: inside each block of 32 rows, + source row ``4u + v`` lands at destination row ``8v + u``.""" + assert rows % 32 == 0, rows + j = torch.arange(32) + dst = (j % 4) * 8 + j // 4 + idx = torch.empty(32, dtype=torch.long) + idx[dst] = j + blocks = rows // 32 + return (idx.repeat(blocks) + torch.arange(blocks).repeat_interleave(32) * 32).to(device) + + +def _interleave_perm(rows: int, device) -> torch.Tensor: + """[up | gate] -> up0, gate0, up1, gate1, ...""" + p = torch.empty(rows, dtype=torch.long) + p[0::2] = torch.arange(0, rows // 2) + p[1::2] = torch.arange(rows // 2, rows) + return p.to(device) + + +def _fc1_perm(rows: int, device) -> torch.Tensor: + return _interleave_perm(rows, device)[_block32_perm(rows, device)] + + +def _expert_fc1_w(t, core) -> torch.Tensor: + """(w3 = up, w1 = gate) packed nibbles -> one expert's FC1 operand.""" + up, gate = t + w = torch.cat([up, gate], dim=0) + return torch.index_select(w, 0, _fc1_perm(w.shape[0], w.device)).contiguous() + + +def _expert_fc1_s(t, core) -> torch.Tensor: + up, gate = t + s = torch.cat([up.view(torch.uint8), gate.view(torch.uint8)], dim=0) + rows, cols = s.shape + s = torch.index_select(s, 0, _fc1_perm(rows, s.device)) + return _swizzle(s).reshape(rows, cols) + + +def _expert_fc2_w(t, core) -> torch.Tensor: + return torch.index_select(t, 0, _block32_perm(t.shape[0], t.device)).contiguous() + + +def _expert_fc2_s(t, core) -> torch.Tensor: + s = t.view(torch.uint8) + rows, cols = s.shape + s = torch.index_select(s, 0, _block32_perm(rows, s.device)) + return _swizzle(s).reshape(rows, cols) + + +def _aligner_w1(t, core) -> torch.Tensor: + """[5120, 9216] columns ``c * 9 + (kh * 3 + kw)`` -> ``(kh * 3 + kw) * 1024 + c``.""" + n, k = t.shape + assert k == VIT_DIM * _AL_TAPS, t.shape + return t.view(n, VIT_DIM, _AL_TAPS).transpose(1, 2).reshape(n, k).contiguous() + + +def _stack_rows(t, core) -> torch.Tensor: + return torch.stack(t, dim=0).contiguous() + + +# ---- sources ------------------------------------------------------------------ + + +def _one(template): + """A single checkpoint key, no index into the destination.""" + return lambda d, i, t=template: [(t.format(i=i), None)] + + +def _rows(template, slice_fn): + """This rank's rows of one checkpoint tensor.""" + return lambda d, i, t=template: [(Slice(t.format(i=i), slice_fn(d)), None)] + + +def _pair(a, b): + """Two checkpoint tensors handed to the transform together.""" + return lambda d, i, x=a, y=b: [((x.format(i=i), y.format(i=i)), None)] + + +def _fused(*parts): + """One parameter assembled from several checkpoint tensors, each written to + its own destination rows: ``parts`` are ``(template, d -> dst slice)``.""" + return lambda d, i, p=parts: [(t.format(i=i), fn(d)) for t, fn in p] + + +def _wo_a_src(d, i): + a = f"layers.{i}.attn" + return [ + ((Slice(f"{a}.wo_a.weight", (d.o_rows,)), Slice(f"{a}.wo_a.scale", (d.o_sf_rows,))), None) + ] + + +def _expert_window(names, transform_pair: bool, cols=None): + """This rank's routed experts, one source per local expert at ``(local,)``. + + ``names`` is ``(up, gate)`` for an FC1 row (w3, w1: the up half first) or + ``(down,)`` for an FC2 row. Every source is this rank's part of the expert: + its intermediate rows (FC1) or columns (FC2) under MoE TP, all of it under EP. + """ + + def build(d, i): + out = [] + for local in range(d.e_loc): + q = f"layers.{i}.ffn.experts.{d.e_off + local}" + if transform_pair: + rows = (slice(d.i_off, d.i_off + d.i_loc),) + src = tuple(Slice(f"{q}.{n}", rows) for n in names) + else: + src = Slice(f"{q}.{names[0]}", cols(d)) + out.append((src, (local,))) + return out + + return build + + +def _expert_scalars(proj, suffix): + """A per-expert fp32 scalar over the whole routing space, not windowed: the + layer-level maxima ``post_load_weights`` folds need all 384 on every rank.""" + return lambda d, i: [ + (f"layers.{i}.ffn.experts.{e}.{proj}.{suffix}", (e,)) for e in range(d.num_experts) + ] + + +def _engram_shard(name): + return lambda d, i: [ + ( + Slice( + f"layers.{i}.engram.embed.{name}", + ( + slice( + d.eng_start[d.engram_layers.index(i)], d.eng_end[d.engram_layers.index(i)] + ), + ), + ), + slice(0, d.eng_end[d.engram_layers.index(i)] - d.eng_start[d.engram_layers.index(i)]), + ) + ] + + +# ---- layer sets --------------------------------------------------------------- + + +def _stage(d): + """This pipeline stage's layers.""" + return d.pp_layer_ids + + +def _hc_ffn_layers(d): + """The stage's layers, plus the previous stage's last layer: its FFN mixes + collapse this stage's first attention input.""" + return list(d.pp_layer_ids) + ([d.pp_prev] if d.pp_prev is not None else []) + + +def _kv_sources_of_ratio(r): + return lambda d: [i for i in d.pp_layer_ids if i in d.kv_sources and d.ratios[i] == r] + + +def _kv_sources(d): + return [i for i in d.pp_layer_ids if i in d.kv_sources] + + +def _index_sources(d): + return [i for i in d.pp_layer_ids if i in d.index_sources] + + +def _engram_layers(d): + return [i for i in d.pp_layer_ids if i in d.engram_layers] + + +def _engram_device_slot(li): + """Engram slot ``li``'s layer, when its table is a device row shard.""" + return lambda d: ( + [d.engram_layers[li]] + if not d.eng_host and li < len(d.engram_layers) and d.engram_layers[li] in d.local + else [] + ) + + +def _mlp_scalar_rows(): + return tuple( + W( + f"e_{name}", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_expert_scalars(proj, suffix), + layers=_stage, + ) + for name, (proj, suffix) in { + "isc1": ("w1", "input_scale"), + "isc3": ("w3", "input_scale"), + "isc2": ("w2", "input_scale"), + "ws1": ("w1", "weight_scale_2"), + "ws3": ("w3", "weight_scale_2"), + "ws2": ("w2", "weight_scale_2"), + }.items() + ) + + +def _hc_rows(sublayer, layers): + return ( + W( + f"hc_{sublayer}_fn", + shape=lambda d: (MIX, d.hc_k), + dtype=torch.float32, + src=_one(f"layers.{{i}}.hc_{sublayer}_fn"), + layers=layers, + ), + W( + f"hc_{sublayer}_base", + shape=lambda d: (MIX,), + dtype=torch.float32, + src=_one(f"layers.{{i}}.hc_{sublayer}_base"), + layers=layers, + ), + W( + f"hc_{sublayer}_scale", + shape=lambda d: (3,), + dtype=torch.float32, + src=_one(f"layers.{{i}}.hc_{sublayer}_scale"), + layers=layers, + ), + ) + + +def _vision_rows(): + """The vision tower, aligner and span delimiters (replicated, bf16).""" + rows = [ + W( + "vis_patch_w", + shape=lambda d: (VIT_DIM, VIT_PATCH_IN), + src="vision.patch_embed.proj.weight", + layers=None, + ), + W( + "vis_patch_b", + shape=lambda d: (VIT_DIM,), + src="vision.patch_embed.proj.bias", + layers=None, + ), + W("vis_norm", shape=lambda d: (VIT_DIM,), src="vision.norm.weight", layers=None), + W( + "vis_al_w1", + shape=lambda d: (d.hidden, VIT_DIM * _AL_TAPS), + src="aligner.w1.weight", + transform=_aligner_w1, + layers=None, + ), + W("vis_al_b1", shape=lambda d: (d.hidden,), src="aligner.w1.bias", layers=None), + W("vis_al_w2", shape=lambda d: (d.hidden, d.hidden), src="aligner.w2.weight", layers=None), + W("vis_al_b2", shape=lambda d: (d.hidden,), src="aligner.w2.bias", layers=None), + W( + "vis_span_tail", + shape=lambda d: (3, d.hidden), + src=lambda d, i: [(("image_start", "image_end", "image_newline"), None)], + transform=_stack_rows, + layers=None, + ), + ] + for L in range(VIT_LAYERS): + p, v = f"vis{L}_", f"vision.blocks.{L}." + rows += [ + W(p + "norm1", shape=lambda d: (VIT_DIM,), src=v + "norm1.weight", layers=None), + W( + p + "wqkv_w", + shape=lambda d: (3 * VIT_DIM, VIT_DIM), + src=v + "attn.wqkv.weight", + layers=None, + ), + W(p + "wqkv_b", shape=lambda d: (3 * VIT_DIM,), src=v + "attn.wqkv.bias", layers=None), + W( + p + "wo_w", + shape=lambda d: (VIT_DIM, VIT_DIM), + src=v + "attn.wo.weight", + layers=None, + ), + W(p + "wo_b", shape=lambda d: (VIT_DIM,), src=v + "attn.wo.bias", layers=None), + W(p + "norm2", shape=lambda d: (VIT_DIM,), src=v + "norm2.weight", layers=None), + W( + p + "w1", + shape=lambda d: (2 * VIT_INTER, VIT_DIM), + src=v + "mlp.w1.weight", + layers=None, + ), + W(p + "w2", shape=lambda d: (VIT_DIM, VIT_INTER), src=v + "mlp.w2.weight", layers=None), + ] + return tuple(rows) + + +def _engram_vocab_size(tc) -> int: + """The Engram bucket vocabulary as one number, from either form of the config: + the checkpoint's integer, or a per-n-gram list of the same value.""" + bucket = getattr(tc, "engram_bucket_vocab_size", None) + if bucket: + return int(bucket) + v = tc.engram_vocab_size + if isinstance(v, (list, tuple)): + assert len(set(v)) == 1, v + v = v[0] + return int(v) + + +def _text_config(core): + cfg = core.model_config.pretrained_config + tc = getattr(cfg, "text_config", None) + return cfg if tc is None else tc + + +def _layout(core) -> SimpleNamespace: + """The per-rank partition of the trunk, from the configuration and mapping alone. + + Shared by ``dims`` and the core: the core reads its head, group, expert and + Engram windows from here rather than deriving them a second time. + """ + tc = _text_config(core) + mapping = core.model_config.mapping + tp = mapping.tp_size + dp = bool(mapping.enable_attention_dp) + attn_tp = 1 if dp else tp + attn_rank = 0 if dp else mapping.tp_rank + heads, head_dim, o_groups, o_lora = ( + tc.num_attention_heads, + tc.head_dim, + tc.o_groups, + tc.o_lora_rank, + ) + h_loc = heads // attn_tp + g_loc = o_groups // attn_tp + h0 = attn_rank * h_loc + g0 = attn_rank * g_loc + q_rows = slice(h0 * head_dim, (h0 + h_loc) * head_dim) + o_rows = slice(g0 * o_lora, (g0 + g_loc) * o_lora) + num_experts, moe_inter = tc.n_routed_experts, tc.moe_intermediate_size + if mapping.moe_ep_size == tp: + e_loc = num_experts // tp + e_off = e_loc * mapping.moe_ep_rank + i_loc, i_off = moe_inter, 0 + else: + e_loc, e_off = num_experts, 0 + i_loc = moe_inter // tp + i_off = i_loc * mapping.moe_tp_rank + num_layers = tc.num_hidden_layers + pp_layer_ids = [int(i) for i in mapping.pp_layers(num_layers)] + eng_rows = [int(n) for n in tc.engram_num_embeddings] + eng_host = _engram_host.host_offload_enabled() + if eng_host: + eng_start, eng_end = [0] * len(eng_rows), list(eng_rows) + eng_part = list(eng_rows) + else: + eng_part = [(n + tp - 1) // tp for n in eng_rows] + eng_start = [p * mapping.tp_rank for p in eng_part] + eng_end = [min(s + p, n) for s, p, n in zip(eng_start, eng_part, eng_rows)] + eng_cols = (tc.engram_max_ngram_size - 1) * tc.engram_n_heads + return SimpleNamespace( + num_layers=num_layers, + hidden=tc.hidden_size, + vocab=tc.vocab_size, + heads=heads, + head_dim=head_dim, + q_lora=tc.q_lora_rank, + o_groups=o_groups, + o_lora=o_lora, + group_in=heads * head_dim // o_groups, + attn_tp=attn_tp, + attn_rank=attn_rank, + h_loc=h_loc, + g_loc=g_loc, + q_rows=q_rows, + q_sf_rows=slice(q_rows.start // MX, q_rows.stop // MX), + sink_rows=slice(h0, h0 + h_loc), + o_rows=o_rows, + o_sf_rows=slice(o_rows.start // MX, o_rows.stop // MX), + num_experts=num_experts, + moe_inter=moe_inter, + e_loc=e_loc, + e_off=e_off, + i_loc=i_loc, + i_off=i_off, + hc_k=HC * tc.hidden_size, + ratios=list(tc.compress_ratios)[:num_layers], + kv_sources=tuple(tc.kv_source_layer_ids), + index_sources=tuple(tc.index_source_layer_ids), + engram_layers=tuple(tc.engram_layer_ids), + eng_rows=eng_rows, + eng_part=eng_part, + eng_start=eng_start, + eng_end=eng_end, + eng_host=eng_host, + eng_dim=tc.engram_head_dim, + eng_cols=eng_cols, + eng_in=eng_cols * tc.engram_head_dim, + eng_wkv_n=(HC + 1) * tc.hidden_size, + pp_layer_ids=pp_layer_ids, + local=frozenset(pp_layer_ids), + pp_prev=pp_layer_ids[0] - 1 if pp_layer_ids[0] > 0 else None, + last_pp=mapping.is_last_pp_rank(), + dtype=core.model_config.torch_dtype, + ) + + +class DeepseekV41UnifiedWeights(ModelWeights): + """deepseek-v4.1-flash-nvfp4 / sm_103 / unified: the trunk, vision tower and aligner.""" + + LAYER_PREFIX = "layers.{i}" + + WEIGHTS: tuple[W, ...] = ( + W("embed", shape=lambda d: (d.vocab, d.hidden), src="embed.weight", layers=None), + W("final_norm", shape=lambda d: (d.hidden,), src="norm.weight", layers=None), + # --- hyper-connections: the stage's layers (and the previous stage's last FFN mixes) --- + *_hc_rows("attn", _stage), + *_hc_rows("ffn", _hc_ffn_layers), + W("attn_norm", shape=lambda d: (d.hidden,), src="{p}.attn_norm.weight", layers=_stage), + W("ffn_norm", shape=lambda d: (d.hidden,), src="{p}.ffn_norm.weight", layers=_stage), + # --- attention: the fused low-rank query + window latent, MXFP8 --- + W( + "wqkv_a", + shape=lambda d: (d.q_lora + d.head_dim, d.hidden), + dtype=torch.float8_e4m3fn, + src=_fused( + ("layers.{i}.attn.wq_a.weight", lambda d: slice(0, d.q_lora)), + ("layers.{i}.attn.wkv.weight", lambda d: slice(d.q_lora, d.q_lora + d.head_dim)), + ), + layers=_stage, + ), + W( + "wqkv_a_sf", + shape=lambda d: (mx_sf_numel(d.q_lora + d.head_dim, d.hidden),), + dtype=torch.uint8, + src=_pair("layers.{i}.attn.wq_a.scale", "layers.{i}.attn.wkv.scale"), + transform=_wqkv_a_sf, + layers=_stage, + ), + W("q_norm", shape=lambda d: (d.q_lora,), src="{p}.attn.q_norm.weight", layers=_stage), + W("kv_norm", shape=lambda d: (d.head_dim,), src="{p}.attn.kv_norm.weight", layers=_stage), + W( + "wq_b", + shape=lambda d: (d.h_loc * d.head_dim, d.q_lora), + dtype=torch.float8_e4m3fn, + src=_rows("layers.{i}.attn.wq_b.weight", lambda d: (d.q_rows,)), + layers=_stage, + ), + W( + "wq_b_sf", + shape=lambda d: (mx_sf_numel(d.h_loc * d.head_dim, d.q_lora),), + dtype=torch.uint8, + src=_rows("layers.{i}.attn.wq_b.scale", lambda d: (d.q_sf_rows,)), + transform=_wq_b_sf, + layers=_stage, + ), + W( + "sink", + shape=lambda d: (d.h_loc,), + dtype=torch.float32, + src=_rows("layers.{i}.attn.attn_sink", lambda d: (d.sink_rows,)), + layers=_stage, + ), + W( + "wo_a", + shape=lambda d: (d.g_loc * d.o_lora, d.group_in), + src=_wo_a_src, + transform=_dequant_wo_a, + layers=_stage, + ), + W( + "wo_b", + shape=lambda d: (d.hidden, d.g_loc * d.o_lora), + dtype=torch.float8_e4m3fn, + src=_rows("layers.{i}.attn.wo_b.weight", lambda d: (slice(None), d.o_rows)), + layers=_stage, + ), + W( + "wo_b_sf", + shape=lambda d: (mx_sf_numel(d.hidden, d.g_loc * d.o_lora),), + dtype=torch.uint8, + src=_rows("layers.{i}.attn.wo_b.scale", lambda d: (slice(None), d.o_sf_rows)), + transform=_hidden_sf, + layers=_stage, + ), + # --- compressor and indexer K: the KV sources. A ratio-2 compressor is + # [wkv ; wgate] promoted to fp32 (the reference's fp32 pooling path), a + # ratio-1 one stays the checkpoint's bf16: two rows over disjoint layers. --- + W( + "comp_w", + shape=lambda d: (2 * d.head_dim, d.hidden), + dtype=torch.float32, + src=_pair( + "layers.{i}.attn.compressor.wkv.weight", "layers.{i}.attn.compressor.wgate.weight" + ), + transform=_cat_rows_f32, + layers=_kv_sources_of_ratio(2), + ), + W( + "comp_w", + shape=lambda d: (d.head_dim, d.hidden), + src="{p}.attn.compressor.wkv.weight", + layers=_kv_sources_of_ratio(1), + ), + W( + "comp_norm", + shape=lambda d: (d.head_dim,), + src="{p}.attn.compressor.norm.weight", + layers=_kv_sources, + ), + W( + "idx_wk", + shape=lambda d: (IDX_DIM, d.head_dim), + src="{p}.attn.indexer.wk.weight", + layers=_kv_sources, + ), + W( + "idx_k_norm", + shape=lambda d: (IDX_DIM,), + src="{p}.attn.indexer.k_norm.weight", + layers=_kv_sources, + ), + # --- indexer query: the index sources --- + W( + "idx_wq_b", + shape=lambda d: (IDX_HEADS * IDX_DIM, d.q_lora), + dtype=torch.float8_e4m3fn, + src="{p}.attn.indexer.wq_b.weight", + layers=_index_sources, + ), + W( + "idx_wq_b_sf", + shape=lambda d: (mx_sf_numel(IDX_HEADS * IDX_DIM, d.q_lora),), + dtype=torch.uint8, + src="{p}.attn.indexer.wq_b.scale", + transform=_idx_wq_b_sf, + layers=_index_sources, + ), + W( + "idx_wproj", + shape=lambda d: (IDX_HEADS, d.hidden), + src="{p}.attn.indexer.weights_proj.weight", + transform=_idx_wproj, + layers=_index_sources, + ), + # --- MoE router: the checkpoint's bf16 gate weight, its bias, and the + # bias that selects experts for image-span tokens --- + W( + "gate_w", + shape=lambda d: (d.num_experts, d.hidden), + src="{p}.ffn.gate.weight", + layers=_stage, + ), + W( + "gate_bias", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src="{p}.ffn.gate.bias", + layers=_stage, + ), + W( + "gate_bias_vl", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src="{p}.ffn.gate.bias_vl", + layers=_stage, + ), + # --- this rank's routed experts, NVFP4 --- + W( + "fc1_w", + shape=lambda d: (d.e_loc, 2 * d.i_loc, d.hidden // 2), + dtype=torch.uint8, + src=_expert_window(("w3.weight", "w1.weight"), True), + transform=_expert_fc1_w, + layers=_stage, + ), + W( + "fc1_s", + shape=lambda d: (d.e_loc, 2 * d.i_loc, d.hidden // SF16), + dtype=torch.uint8, + src=_expert_window(("w3.weight_scale", "w1.weight_scale"), True), + transform=_expert_fc1_s, + layers=_stage, + ), + W( + "fc2_w", + shape=lambda d: (d.e_loc, d.hidden, d.i_loc // 2), + dtype=torch.uint8, + src=_expert_window( + ("w2.weight",), + False, + cols=lambda d: (slice(None), slice(d.i_off // 2, (d.i_off + d.i_loc) // 2)), + ), + transform=_expert_fc2_w, + layers=_stage, + ), + W( + "fc2_s", + shape=lambda d: (d.e_loc, d.hidden, d.i_loc // SF16), + dtype=torch.uint8, + src=_expert_window( + ("w2.weight_scale",), + False, + cols=lambda d: (slice(None), slice(d.i_off // SF16, (d.i_off + d.i_loc) // SF16)), + ), + transform=_expert_fc2_s, + layers=_stage, + ), + *_mlp_scalar_rows(), + # --- the shared expert, MXFP8, [gate ; up] packed --- + W( + "sh_w13", + shape=lambda d: (2 * d.moe_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_fused( + ("layers.{i}.ffn.shared_experts.w1.weight", lambda d: slice(0, d.moe_inter)), + ( + "layers.{i}.ffn.shared_experts.w3.weight", + lambda d: slice(d.moe_inter, 2 * d.moe_inter), + ), + ), + layers=_stage, + ), + W( + "sh_w13_sf", + shape=lambda d: (mx_sf_numel(2 * d.moe_inter, d.hidden),), + dtype=torch.uint8, + src=_pair( + "layers.{i}.ffn.shared_experts.w1.scale", "layers.{i}.ffn.shared_experts.w3.scale" + ), + transform=_sh_w13_sf, + layers=_stage, + ), + W( + "sh_w2", + shape=lambda d: (d.hidden, d.moe_inter), + dtype=torch.float8_e4m3fn, + src="{p}.ffn.shared_experts.w2.weight", + layers=_stage, + ), + W( + "sh_w2_sf", + shape=lambda d: (mx_sf_numel(d.hidden, d.moe_inter),), + dtype=torch.uint8, + src="{p}.ffn.shared_experts.w2.scale", + transform=_hidden_sf, + layers=_stage, + ), + # --- Engram: the tables as a device row shard (host_offload off), and + # the projections every member holds --- + *( + row + for li in range(_ENGRAM_SLOTS) + for row in ( + W( + "eng_table", + shape=lambda d, li=li: (d.eng_part[li], d.eng_dim), + dtype=torch.float8_e4m3fn, + src=_engram_shard("weight"), + layers=_engram_device_slot(li), + ), + W( + "eng_scale", + shape=lambda d, li=li: (d.eng_part[li], d.eng_dim // MX), + dtype=torch.uint8, + src=_engram_shard("scale"), + transform=_u8, + layers=_engram_device_slot(li), + ), + ) + ), + W( + "eng_wkv", + shape=lambda d: (d.eng_wkv_n, d.eng_in), + dtype=torch.float8_e4m3fn, + src="{p}.engram.wkv.weight", + layers=_engram_layers, + ), + W( + "eng_wkv_sf", + shape=lambda d: (mx_sf_numel(d.eng_wkv_n, d.eng_in),), + dtype=torch.uint8, + src="{p}.engram.wkv.scale", + transform=_eng_wkv_sf, + layers=_engram_layers, + ), + W( + "eng_q", + shape=lambda d: (HC, d.hidden), + src="{p}.engram.q_weight", + layers=_engram_layers, + ), + W( + "eng_k", + shape=lambda d: (HC, d.hidden), + src="{p}.engram.k_weight", + layers=_engram_layers, + ), + # --- the vision tower, aligner and span delimiters --- + *_vision_rows(), + ) + + def dims(self, core) -> SimpleNamespace: + return _layout(core) + + def host_keys(self, d) -> list[str]: + """Checkpoint keys of the Engram tables held in the node's shared host copy.""" + if not d.eng_host: + return [] + return [ + f"layers.{i}.engram.embed.{name}" + for i in _engram_layers(d) + for name in ("weight", "scale") + ] + + def expected_unconsumed(self, core, weights) -> set: + """Exactly the keys this rank does not consume (module docstring).""" + keys = set(weights.keys()) + out = {k for k in keys if k.startswith("mtp.")} + prev_hc = {f"layers.{core.pp_prev}.hc_ffn_{f}" for f in ("fn", "base", "scale")} + out |= { + k + for k in keys + if (i := _layer_of(k)) is not None and i not in core.local and k not in prev_hc + } + if not core.last_pp: + out.add("head.weight") + if core.e_loc != core.num_experts: + lo, hi = core.e_off, core.e_off + core.e_loc + for i in core.pp_layer_ids: + for e in range(core.num_experts): + if lo <= e < hi: + continue + for w_ in ("w1", "w2", "w3"): + out.add(f"layers.{i}.ffn.experts.{e}.{w_}.weight") + out.add(f"layers.{i}.ffn.experts.{e}.{w_}.weight_scale") + return out + + def load(self, model, weights) -> None: + """Overrides the base: sources are per-rank parts of checkpoint tensors read + before their transform, the per-tensor scalars are 0-dim, the Engram tables + may live in the node's shared host copy, a huge untransformed row streams in + chunks, and the vocab-parallel ``lm_head`` is a slice.""" + core = model.model + d = self.dims(core) + manifest = self.manifest(core) + consumed: set = set() + ckpt_dir = getattr(weights, "checkpoint_dir", None) + if ckpt_dir is not None: + core._ckpt_dir = ckpt_dir + + def fill(param, ckpt, dst_index, transform) -> None: + keys, srcs, multi = _split_source(ckpt) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + dst = param.data if dst_index is None else param.data[dst_index] + if transform is None and len(keys) == 1: + v = weights[keys[0]] + s0 = srcs[0] + shp = _shape(v) + n_elems = 1 + for d_ in shp: + n_elems *= d_ + if len(shp) >= 1 and n_elems > _CHUNK_ELEMS and (s0 is None or len(s0) == 1): + # Stream a huge row range in chunks. + r0 = 0 if s0 is None or s0[0].start is None else s0[0].start + r1 = shp[0] if s0 is None or s0[0].stop is None else s0[0].stop + assert dst.shape[0] == r1 - r0, (keys, tuple(dst.shape), r0, r1) + row_elems = n_elems // shp[0] + step = max(1, _CHUNK_ELEMS // row_elems) + for a in range(r0, r1, step): + b = min(a + step, r1) + part = _materialize(v, (slice(a, b),)) + assert part.dtype == dst.dtype, (keys, part.dtype, dst.dtype) + dst[a - r0 : b - r0].copy_(part.to(dst.device)) + consumed.update(keys) + return + src_all = tuple( + _materialize(weights[k], s).to(dst.device, non_blocking=True) + for k, s in zip(keys, srcs) + ) + if transform is not None: + out = transform(src_all if multi else src_all[0], core) + else: + assert len(src_all) == 1, "a multi-key row needs a transform" + out = src_all[0] + assert tuple(dst.shape) == tuple(out.shape), (keys, tuple(dst.shape), tuple(out.shape)) + assert out.dtype == dst.dtype, (keys, out.dtype, dst.dtype) + dst.copy_(out, non_blocking=True) + consumed.update(keys) + + assert set(manifest.keys()) == set(core.w.keys()), ( + "manifest/parameter drift", + sorted(set(manifest.keys()) ^ set(core.w.keys()))[:8], + ) + if core.eng_attach: + # Collective over the node's ranks (local rank 0 fills the shared + # file), and first: it fails fast if the node lacks the host memory. + core.attach_engram_host_table(core._ckpt_dir) + host = self.host_keys(d) + if host: + held = {t["key"]: tuple(t["shape"]) for t in core._eng_table.manifest["tensors"]} + for key in host: + assert key in weights, f"checkpoint key missing: {key}" + assert held.get(key) == _shape(weights[key]), ( + key, + held.get(key), + _shape(weights[key]), + ) + consumed.add(key) + for param_key, sources in manifest.items(): + for ckpt, dst_index, transform in sources: + fill(core.w[param_key], ckpt, dst_index, transform) + # The expert relayouts allocate per-expert scratch e_loc times per + # layer and an Engram shard is tens of GB: release before the next + # parameter so peak load memory stays one parameter deep. + if param_key.endswith(_RELEASE_AFTER): + torch.cuda.synchronize() + torch.cuda.empty_cache() + + # Shell-registered exception: lm_head. The shell builds it vocab-parallel + # ([V/tp, H], rows of this TP rank) on the TP members and whole under + # attention DP; asserted, not adapted. Outside the last pipeline stage the + # engine skips the lm_head and drops its weight. + if core.last_pp: + lm = model.lm_head.weight + if tuple(lm.shape) == (core.vocab, core.hidden): + assert core.dp, "a replicated lm_head on a head-split member" + fill(lm, "head.weight", None, None) + else: + part = core.vocab // core.tp + assert tuple(lm.shape) == (part, core.hidden) and not core.dp, tuple(lm.shape) + fill( + lm, + Slice("head.weight", (slice(core.tp_rank * part, (core.tp_rank + 1) * part),)), + None, + None, + ) + shell = { + n for n, _ in model.named_parameters() if not n.startswith(("model.w.", "model.w_ds.")) + } + want = {"lm_head.weight"} if core.last_pp else set() + assert shell == want, f"unfed shell parameters: {sorted(shell)}" + + torch.cuda.synchronize() + leftover = set(weights.keys()) - consumed + expected = self.expected_unconsumed(core, weights) + assert leftover == expected, ( + "checkpoint coverage: leftover keys are not exactly the predicted non-load set; " + f"unexpected {sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" + ) + + +def _layer_of(key: str): + """The backbone layer of a ``layers.{i}.`` checkpoint key, else None.""" + if not key.startswith("layers."): + return None + return int(key.split(".", 2)[1]) + + +MODEL_WEIGHTS = DeepseekV41UnifiedWeights() + + +# ---- the DSpark drafter (the ``load_draft_weights`` pass) -------------------- + + +def _fc1_w_u8(t, core) -> torch.Tensor: + """MXFP4 (w3 = up, w1 = gate) packed nibbles -> one expert's FC1 operand + (the NVFP4 relayout is byte-generic).""" + up, gate = t + return _expert_fc1_w((up.view(torch.uint8), gate.view(torch.uint8)), core) + + +def _fc2_w_u8(t, core) -> torch.Tensor: + return _expert_fc2_w(t.view(torch.uint8), core) + + +def _ds_wqkv_a_sf(t, core) -> torch.Tensor: + return _mx_sf(t, (core.ds.q_lora, core.ds.head_dim)) + + +def _ds_wq_b_sf(t, core) -> torch.Tensor: + return _mx_sf(t, (core.ds.heads * core.ds.head_dim,)) + + +def _ds_hidden_sf(t, core) -> torch.Tensor: + return _mx_sf(t, (core.ds.hidden,)) + + +def _ds_sh_w13_sf(t, core) -> torch.Tensor: + return _mx_sf(t, (core.ds.inter, core.ds.inter)) + + +def _ds_wkv3_sf(t, core) -> torch.Tensor: + return _mx_sf(t, (core.ds.head_dim,) * core.ds.stages) + + +def _ds_experts(names, transform_pair: bool): + """All of a stage's routed experts, replicated, one source per expert at ``(e,)``.""" + + def build(d, s): + out = [] + for e in range(d.experts): + q = f"mtp.{s}.ffn.experts.{e}" + src = tuple(f"{q}.{n}" for n in names) if transform_pair else f"{q}.{names[0]}" + out.append((src, (e,))) + return out + + return build + + +def _ds_hc_rows(sublayer): + return ( + W( + f"hc_{sublayer}_fn", + shape=lambda d: (MIX, HC * d.hidden), + dtype=torch.float32, + src=f"{{p}}.hc_{sublayer}_fn", + ), + W( + f"hc_{sublayer}_base", + shape=lambda d: (MIX,), + dtype=torch.float32, + src=f"{{p}}.hc_{sublayer}_base", + ), + W( + f"hc_{sublayer}_scale", + shape=lambda d: (3,), + dtype=torch.float32, + src=f"{{p}}.hc_{sublayer}_scale", + ), + ) + + +class DeepseekV41DSparkWeights(ModelWeights): + """The three DSpark stages of deepseek-v4.1-flash-nvfp4, replicated on every rank. + + Layer ``s`` is stage ``mtp.{s}``; ``d.num_layers`` is the stage count. + """ + + LAYER_PREFIX = "mtp.{i}" + + WEIGHTS: tuple[W, ...] = ( + *_ds_hc_rows("attn"), + *_ds_hc_rows("ffn"), + W("attn_norm", shape=lambda d: (d.hidden,), src="{p}.attn_norm.weight"), + W("ffn_norm", shape=lambda d: (d.hidden,), src="{p}.ffn_norm.weight"), + W( + "wqkv_a", + shape=lambda d: (d.q_lora + d.head_dim, d.hidden), + dtype=torch.float8_e4m3fn, + src=( + ("{p}.attn.wq_a.weight", lambda d: slice(0, d.q_lora)), + ("{p}.attn.wkv.weight", lambda d: slice(d.q_lora, d.q_lora + d.head_dim)), + ), + ), + W( + "wqkv_a_sf", + shape=lambda d: (mx_sf_numel(d.q_lora + d.head_dim, d.hidden),), + dtype=torch.uint8, + src=lambda d, s: [((f"mtp.{s}.attn.wq_a.scale", f"mtp.{s}.attn.wkv.scale"), None)], + transform=_ds_wqkv_a_sf, + ), + W("q_norm", shape=lambda d: (d.q_lora,), src="{p}.attn.q_norm.weight"), + W("kv_norm", shape=lambda d: (d.head_dim,), src="{p}.attn.kv_norm.weight"), + W( + "wq_b", + shape=lambda d: (d.heads * d.head_dim, d.q_lora), + dtype=torch.float8_e4m3fn, + src="{p}.attn.wq_b.weight", + ), + W( + "wq_b_sf", + shape=lambda d: (mx_sf_numel(d.heads * d.head_dim, d.q_lora),), + dtype=torch.uint8, + src="{p}.attn.wq_b.scale", + transform=_ds_wq_b_sf, + ), + W("sink", shape=lambda d: (d.heads,), dtype=torch.float32, src="{p}.attn.attn_sink"), + W( + "wo_a", + shape=lambda d: (d.o_groups * d.o_lora, d.heads * d.head_dim // d.o_groups), + src=lambda d, s: [((f"mtp.{s}.attn.wo_a.weight", f"mtp.{s}.attn.wo_a.scale"), None)], + transform=_dequant_wo_a, + ), + W( + "wo_b", + shape=lambda d: (d.hidden, d.o_groups * d.o_lora), + dtype=torch.float8_e4m3fn, + src="{p}.attn.wo_b.weight", + ), + W( + "wo_b_sf", + shape=lambda d: (mx_sf_numel(d.hidden, d.o_groups * d.o_lora),), + dtype=torch.uint8, + src="{p}.attn.wo_b.scale", + transform=_ds_hidden_sf, + ), + W( + "gate_w", + shape=lambda d: (d.experts, d.hidden), + dtype=torch.float32, + src="{p}.ffn.gate.weight", + transform=_to_f32, + ), + W("gate_bias", shape=lambda d: (d.experts,), dtype=torch.float32, src="{p}.ffn.gate.bias"), + W( + "fc1_w", + shape=lambda d: (d.experts, 2 * d.inter, d.hidden // 2), + dtype=torch.uint8, + src=_ds_experts(("w3.weight", "w1.weight"), True), + transform=_fc1_w_u8, + ), + W( + "fc1_s", + shape=lambda d: (d.experts, 2 * d.inter, d.hidden // MX), + dtype=torch.uint8, + src=_ds_experts(("w3.scale", "w1.scale"), True), + transform=_expert_fc1_s, + ), + W( + "fc2_w", + shape=lambda d: (d.experts, d.hidden, d.inter // 2), + dtype=torch.uint8, + src=_ds_experts(("w2.weight",), False), + transform=_fc2_w_u8, + ), + W( + "fc2_s", + shape=lambda d: (d.experts, d.hidden, d.inter // MX), + dtype=torch.uint8, + src=_ds_experts(("w2.scale",), False), + transform=_expert_fc2_s, + ), + W( + "sh_w13", + shape=lambda d: (2 * d.inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=( + ("{p}.ffn.shared_experts.w1.weight", lambda d: slice(0, d.inter)), + ("{p}.ffn.shared_experts.w3.weight", lambda d: slice(d.inter, 2 * d.inter)), + ), + ), + W( + "sh_w13_sf", + shape=lambda d: (mx_sf_numel(2 * d.inter, d.hidden),), + dtype=torch.uint8, + src=lambda d, s: [ + ( + ( + f"mtp.{s}.ffn.shared_experts.w1.scale", + f"mtp.{s}.ffn.shared_experts.w3.scale", + ), + None, + ) + ], + transform=_ds_sh_w13_sf, + ), + W( + "sh_w2", + shape=lambda d: (d.hidden, d.inter), + dtype=torch.float8_e4m3fn, + src="{p}.ffn.shared_experts.w2.weight", + ), + W( + "sh_w2_sf", + shape=lambda d: (mx_sf_numel(d.hidden, d.inter),), + dtype=torch.uint8, + src="{p}.ffn.shared_experts.w2.scale", + transform=_ds_hidden_sf, + ), + # --- shared by the stages: stage 0's main projection, the stacked ring + # write operand, and the last stage's norm and Markov head --- + W( + "main_proj", + shape=lambda d: (d.hidden, d.targets * d.hidden), + dtype=torch.float8_e4m3fn, + src="mtp.0.main_proj.weight", + layers=None, + ), + W( + "main_proj_sf", + shape=lambda d: (mx_sf_numel(d.hidden, d.targets * d.hidden),), + dtype=torch.uint8, + src="mtp.0.main_proj.scale", + transform=_ds_hidden_sf, + layers=None, + ), + W("main_norm", shape=lambda d: (d.hidden,), src="mtp.0.main_norm.weight", layers=None), + W( + "wkv3", + shape=lambda d: (d.num_layers * d.head_dim, d.hidden), + dtype=torch.float8_e4m3fn, + src=lambda d, i: [ + (f"mtp.{s}.attn.wkv.weight", slice(s * d.head_dim, (s + 1) * d.head_dim)) + for s in range(d.num_layers) + ], + layers=None, + ), + W( + "wkv3_sf", + shape=lambda d: (mx_sf_numel(d.num_layers * d.head_dim, d.hidden),), + dtype=torch.uint8, + src=lambda d, i: [ + (tuple(f"mtp.{s}.attn.wkv.scale" for s in range(d.num_layers)), None) + ], + transform=_ds_wkv3_sf, + layers=None, + ), + W( + "norm", + shape=lambda d: (d.hidden,), + src=lambda d, i: [(f"mtp.{d.num_layers - 1}.norm.weight", None)], + layers=None, + ), + W( + "mk_embed", + shape=lambda d: (d.vocab, d.markov_rank), + src=lambda d, i: [(f"mtp.{d.num_layers - 1}.markov_head.embed.weight", None)], + layers=None, + ), + W( + "mk_head", + shape=lambda d: (d.vocab, d.markov_rank), + dtype=torch.float32, + src=lambda d, i: [(f"mtp.{d.num_layers - 1}.markov_head.head.weight", None)], + transform=_to_f32, + layers=None, + ), + ) + + def dims(self, core) -> SimpleNamespace: + tc = _text_config(core) + return SimpleNamespace( + num_layers=int(tc.num_nextn_predict_layers), + hidden=int(tc.hidden_size), + vocab=int(tc.vocab_size), + heads=int(tc.num_attention_heads), + head_dim=int(tc.head_dim), + q_lora=int(tc.q_lora_rank), + o_groups=int(tc.o_groups), + o_lora=int(tc.o_lora_rank), + inter=int(tc.moe_intermediate_size), + experts=int(tc.dspark_n_routed_experts), + markov_rank=int(tc.dspark_markov_rank), + targets=len(tc.dspark_target_layer_ids), + dtype=core.model_config.torch_dtype, + ) + + def expected_unconsumed(self, core, weights) -> set: + """``mtp.*`` keys the drafter does not consume (module docstring).""" + mtp = [k for k in weights.keys() if k.startswith("mtp.")] + out = {k for k in mtp if ".confidence_head." in k} + out |= {k for k in mtp if k.endswith(".ffn.gate.bias_vl")} + return out + + def load(self, model, weights) -> None: + """Fill the core's ``w_ds`` from the second weight session; asserts ``mtp.*`` + coverage exactly.""" + core = model.model + assert core.ds is not None, "load_draft on a core built without the DSpark drafter" + manifest = self.manifest(core) + assert set(manifest.keys()) == set(core.w_ds.keys()), ( + "draft manifest/parameter drift", + sorted(set(manifest.keys()) ^ set(core.w_ds.keys()))[:8], + ) + consumed: set = set() + + def fill(param, ckpt, dst_index, transform) -> None: + keys, srcs, multi = _split_source(ckpt) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + dst = param.data if dst_index is None else param.data[dst_index] + src_all = tuple( + _materialize(weights[k], s).to(dst.device, non_blocking=True) + for k, s in zip(keys, srcs) + ) + if transform is not None: + out = transform(src_all if multi else src_all[0], core) + else: + assert len(src_all) == 1, "a multi-key row needs a transform" + out = src_all[0] + assert tuple(dst.shape) == tuple(out.shape), (keys, tuple(dst.shape), tuple(out.shape)) + assert out.dtype == dst.dtype, (keys, out.dtype, dst.dtype) + dst.copy_(out, non_blocking=True) + consumed.update(keys) + + for param_key, sources in manifest.items(): + for ckpt, dst_index, transform in sources: + fill(core.w_ds[param_key], ckpt, dst_index, transform) + if param_key.endswith("_fc2_s"): + torch.cuda.synchronize() + torch.cuda.synchronize() + mtp = {k for k in weights.keys() if k.startswith("mtp.")} + leftover = mtp - consumed + expected = self.expected_unconsumed(core, weights) + assert leftover == expected, ( + "mtp coverage: leftover keys are not exactly the predicted non-load set; " + f"unexpected {sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" + ) + + +DRAFT_WEIGHTS = DeepseekV41DSparkWeights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/routing.py new file mode 100644 index 000000000000..f93920c367b5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v41/routing.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Where a DeepseekV41ForConditionalGeneration config lands. Read this file and you know. + +The engine names the DeepSeek-V4.1 checkpoint that carries its vision tower +``DeepseekV41ForConditionalGeneration``; every target here runs that tower, so +a text-only V4.1 config (``DeepseekV41ForCausalLM``) is not routed here. + +One forward-reading decision tree per architecture family: the criteria are +evaluated in the order a reader would ask them, and every branch that does not +end in a target returns None (in ``auto`` the engine then uses the built-in +DeepSeek-V4.1 implementation; in ``require`` it raises, quoting the trace below). +""" + +from __future__ import annotations + +from typing import Optional + +from ..._router_index import NULL_TRACE, ModelingV2Context, Trace + +# The one GPU architecture these targets are written for. sm is part of a +# target's identity, not a knob: a different SM is a different target. +_SM = (10, 3) + +# Config-shape fingerprint -> checkpoint identity; see the note in the gpt_oss +# routing module for what shape-sniffing does and does not pin. +# +# (num_hidden_layers, hidden_size, n_routed_experts, q_lora_rank, head_dim, +# o_groups, index_n_heads, num_nextn_predict_layers). Every field changes an +# operand shape the target declares: the expert stack, the q-LoRA pair, the +# 512-wide shared latent, the grouped output projection, the indexer heads and +# the DSpark stage count the weight table predicts. The attributes are read +# off the top-level config, which answers them from the checkpoint's +# `text_config`. +_CHECKPOINTS = { + (40, 5120, 384, 1280, 512, 8, 32, 3): "deepseek_v4_1_flash_nvfp4", +} + +_TARGETS = { + ("deepseek_v4_1_flash_nvfp4", "unified"): "ModelingV2DeepseekV41FlashNvfp4Sm103Unified", +} + +# Synthetic architecture name -> the module whose import registers it. +TARGET_MODULES = { + "ModelingV2DeepseekV41FlashNvfp4Sm103Unified": ( + "models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.modeling" + ), +} + +# Checkpoint identity -> the module carrying the input processor its targets +# read their multimodal inputs from (``INPUT_PROCESSOR``, ``MODEL_TYPE``, +# ``PLACEHOLDER_METADATA``); ``modeling_v2_input_processor`` registers it. +INPUT_PROCESSORS = { + "deepseek_v4_1_flash_nvfp4": ( + "models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.input_processor" + ), +} + +# The global quant_algo the engine reads for this checkpoint's export -- fp8 +# 32x32 dense weights plus NVFP4 routed experts. A modelopt reading of the +# export gives MIXED_PRECISION; the engine's DeepSeek-V4.1 loader reads the +# dense FP8 blocks as FP8_BLOCK_SCALES and keeps the experts' NVFP4 per layer. +# The target loads the checkpoint's own tensors under either reading, and a +# checkpoint re-exported another way fails the weight table's dtype asserts. +_QUANT = frozenset({"MIXED_PRECISION", "FP8_BLOCK_SCALES"}) + + +def _quant(q) -> Optional[str]: + """The global quant_algo as a string, or None for an unquantized config.""" + algo = getattr(q, "quant_algo", None) + return None if algo is None else str(getattr(algo, "value", algo)) + + +def _parallel(m) -> Optional[str]: + """Name the parallel segment, or None if no target implements it. + + The unified target folds its members at runtime from the mapping: + tep4 (attention TP4 + MoE EP4), dep4 (attention DP4 + MoE EP4) and tp4 + (attention TP4 + MoE TP4) on one node, tep8 / dep8 (EP8) over two NVL72 + trays, and any of them per pipeline stage (pp2 x tep2 on one node, pp2 x + tep4 over two trays). Without pipeline parallelism the TP group is 4 or 8 + ranks; with it, 2 or 4 per stage. Whether a pipeline split is valid is the + target's to say (a stage must hold every layer its layers read from). + """ + tp, pp = m.tp_size, m.pp_size + if m.world_size != tp * pp or tp not in ((4, 8) if pp == 1 else (2, 4)): + return None + ep, mtp = m.moe_ep_size, m.moe_tp_size + if ep == tp and mtp == 1: + return "unified" # tep, or dep with attention DP + if tp in (2, 4) and ep == 1 and mtp == tp and not m.enable_attention_dp: + return "unified" # tp (MoE TP) + return None + + +def _shape(c) -> tuple: + """The config-shape fingerprint ``_CHECKPOINTS`` is keyed by.""" + return ( + c.num_hidden_layers, + c.hidden_size, + c.n_routed_experts, + c.q_lora_rank, + c.head_dim, + c.o_groups, + c.index_n_heads, + c.num_nextn_predict_layers, + ) + + +def input_processor(c) -> Optional[str]: + """The input-processor module for this checkpoint's targets, or None. + + Decided by the checkpoint alone: the input processor is built where the GPU + and the parallel mapping are not known, so this cannot replay ``route``. + """ + return INPUT_PROCESSORS.get(_CHECKPOINTS.get(_shape(c))) + + +def route(ctx: ModelingV2Context, trace: Trace = NULL_TRACE) -> Optional[str]: + c, m = ctx.pretrained_config, ctx.mapping + + if not trace.check("sm", ctx.sm, ctx.sm == _SM): + return None + + shape = _shape(c) + ckpt = trace.resolve("shape", shape, _CHECKPOINTS.get(shape)) + if ckpt is None: + return None + + quant = _quant(ctx.quant_config) + if not trace.check("quant", quant, quant in _QUANT): + return None + + parallel = trace.resolve( + "parallel", + f"ws={m.world_size} tp={m.tp_size} ep={m.moe_ep_size} " + f"moe_tp={m.moe_tp_size} attention_dp={m.enable_attention_dp}", + _parallel(m), + ) + if parallel is None: + return None + + return _TARGETS.get((ckpt, parallel)) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/modeling.py index 6adc9e2b0798..c810d0d02ad7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/modeling.py @@ -2,38 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """ModelingV2 target: gpt-oss-120b / sm_103 / tp1 — self-contained modeling code. -Flat single-entry forward assembled from catalog entries only; every call -that creates or transforms a tensor is a catalog entry, everything else is -tensor-metadata reads and Python control flow. Attention consumes runtime -state fully explicitly through thop_attention: per-step arguments are -projected from the engine-prepared TrtllmAttentionMetadata once per forward -in _build_step_args and shared by all layers. Two attention arguments are -per-layer here rather than per-step: the fp32 sink logits (one extra softmax -denominator column per query head) and attention_window_size — this -checkpoint alternates sliding_attention (window 128) and full_attention -layers, and the window is a pure mask, so one shared pool and one -block-offset table serve both kinds. - -Every layer's MLP is an MXFP4 sparse mixture of experts (128 experts, top-4, -renormalized, clamped GLU) run W4A8: mxfp8_quantize turns the bf16 hidden -states into e4m3 data plus per-32 UE8M0 block scales — widening hidden 2880 -to the FC1 K alignment 3072 inside that call — and one -mxe4m3_mxe2m1_block_scale_moe_runner call per layer covers routing, both -grouped GEMMs, the clamped activation, the MXFP8 requantization between them -and the combine. The router bias is folded into the router GEMM's fused bias -because the MoE op silently ignores routing_bias on this routing method. The -expert weights are declared in the kernel-ready padded/shuffled/swizzled -layout — identical for the W4A16 and W4A8 members of this kernel family — and -the manifest loop in weights.py transforms the checkpoint's block/scale -tensors into it at load time. - -Weights are target-owned: a flat ParameterDict declared here (HF [out, in] -storage so checkpoint rows copy in unchanged), loaded by the manifest loop -in the sibling weights.py, with column-major GEMM views derived once after -load. The registration shell inherits DecoderModelForCausalLM for lm_head, -packed-batch logits gathering, and the meta-init/load/post-load hooks. - -The import-time and first-forward contract checks below fail fast on drift. +Flat single-entry forward assembled from catalog entries only, running the +MXFP4 sparse MoE W4A8 on a single GPU (tp1). Weights are target-owned and +declared in the sibling weights.py. Each target is self-contained: it binds +every op it calls against the real weights in its own `__init__` and shares +no code with the other target. """ import math @@ -42,28 +15,24 @@ from torch import nn from transformers import PretrainedConfig -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( - fused_qk_norm_rope, -) +from tensorrt_llm._torch._experimental.modeling_v2._target import Phase, Target, phase_of +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import advance_step_generation from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( - thop_attention, + ThopAttention, ) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import CublasMm from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.mxe4m3_mxe2m1_block_scale_moe_runner import ( # noqa: E501 - mxe4m3_mxe2m1_block_scale_moe_runner, + Mxe4m3Mxe2m1BlockScaleMoeRunner, ) from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 - flashinfer_fused_add_rmsnorm, + FlashinferFusedAddRmsnorm, ) from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( - flashinfer_rmsnorm, + FlashinferRmsnorm, ) from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.mxfp8_quantize import ( - mxfp8_quantize, + Mxfp8Quantize, ) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.embedding import embedding -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.empty import empty -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.reshape import reshape from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.model_config import ModelConfig @@ -74,131 +43,8 @@ ) from . import weights as _weights +from .ops.attention.fused_qk_norm_rope import FusedQkNormRope -# The GPU architecture this target IS. Routing will not send another one here, -# but a direct instantiation could, and the certification is per arch: this -# assert is what the version pin used to be. In-tree the version moves with -# the code, so pinning it is meaningless; the architecture does not. -_SM = (10, 3) - - -#: Every trtllm op this target reaches for, in its forward and in the weight -#: load. Declared here, asserted in tests/unittest/_torch/modeling_v2: a symbol -#: that does not exist is a fact of the build, and the place to find that out -#: is a machine with the extension built rather than every import of this -#: module. -REQUIRED_TRTLLM_OPS = ( - "fused_qk_norm_rope", - "flashinfer_rmsnorm", - "flashinfer_fused_add_rmsnorm", - "cublas_mm", - "mxfp8_quantize", - "mxe4m3_mxe2m1_block_scale_moe_runner", -) - -# Metadata fields consumed each step (sourcing mirrors the in-tree -# FallbackFmha for this trtllm version; existence checked at first forward). -_STEP_FIELDS = ( - "kv_lens_cuda_runtime", - "kv_lens_runtime", - "host_total_kv_lens", - "prompt_lens_cuda_runtime", - "prompt_lens_cpu_runtime", - "host_request_types_runtime", - "kv_cache_block_offsets", - "host_kv_cache_pool_pointers", - "host_kv_cache_pool_mapping", - "effective_workspace", - "tokens_per_block", - "max_num_requests", - "max_context_length", - "max_seq_len", - "num_contexts", - "num_ctx_tokens", - "trtllm_gen_jit_warmup", - "use_paged_context_fmha", - "effective_beam_width", - "cache_indirection", - "block_ids_per_seq", - "max_context_q_len_override", - "is_cross", - "is_spec_decoding_enabled", - "use_spec_decoding", - "is_spec_dec_tree", - "spec_decoding_generation_lengths", - "spec_decoding_position_offsets_for_cpp", - "spec_decoding_packed_mask", - "spec_decoding_bl_tree_mask_offset", - "spec_decoding_bl_tree_mask", - "max_total_draft_tokens", - "spec_bl_tree_first_sparse_mask_offset_kv", - "num_sparse_topk", - "flash_mla_tile_scheduler_metadata", - "flash_mla_num_splits", - # Both are engine-prepared per-instance constants (max_num_sequences - # defaults to max_num_requests; the tree-mask flag is set from - # is_spec_dec_dynamic_tree), so they project like the rest. - "max_num_sequences", - "force_prepare_spec_dec_tree_mask", -) - - -def _build_step_args(md: TrtllmAttentionMetadata) -> dict: - """Project the prepared metadata onto thop_attention's explicit batch - state, once per forward; every runtime-owned value passes through as - the engine prepared it. CUDA-graph classes: tensors are engine-owned - persistent buffers refreshed in place (reference class); Python ints - are per-capture constants (host-derived class — a decode-only graph - always sees num_contexts == 0). attention_window_size is absent here - on purpose: it varies per layer and is passed at the call site.""" - return dict( - sequence_length=md.kv_lens_cuda_runtime, - host_past_key_value_lengths=md.kv_lens_runtime, - host_total_kv_lens=md.host_total_kv_lens, - context_lengths=md.prompt_lens_cuda_runtime, - host_context_lengths=md.prompt_lens_cpu_runtime, - host_request_types=md.host_request_types_runtime, - kv_cache_block_offsets=md.kv_cache_block_offsets, - host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, - workspace_=md.effective_workspace, - tokens_per_block=md.tokens_per_block, - max_num_requests=md.max_num_requests, - max_context_length=md.max_context_length, - max_seq_len=md.max_seq_len, - num_contexts=md.num_contexts, - num_ctx_tokens=md.num_ctx_tokens, - trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup, - use_paged_context_fmha=md.use_paged_context_fmha, - beam_width=md.effective_beam_width, - cache_indirection=md.cache_indirection, - block_ids_per_seq=md.block_ids_per_seq, - max_context_q_len_override=md.max_context_q_len_override, - is_cross=md.is_cross, - is_spec_decoding_enabled=md.is_spec_decoding_enabled, - use_spec_decoding=md.use_spec_decoding, - is_spec_dec_tree=md.is_spec_dec_tree, - spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, - spec_decoding_position_offsets_for_cpp=md.spec_decoding_position_offsets_for_cpp, - spec_decoding_packed_mask=md.spec_decoding_packed_mask, - spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, - spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, - spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, - spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, - num_sparse_topk=md.num_sparse_topk, - flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, - flash_mla_num_splits=md.flash_mla_num_splits, - max_num_sequences=md.max_num_sequences, - force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, - ) - - -# Per-call constants of this target's call shape — the values the in-tree -# path sources from the attention module and forward args: packed-QKV -# causal GQA, RoPE applied outside (fused_qk_norm_rope), bf16 activations -# over a bf16 KV pool, no MLA / mRoPE / cross / relative-bias / sparse -# features. attention_sinks and attention_window_size are per-layer, not -# constants, and are passed at the call site. _CALL_CONSTANTS = dict( output_sf=None, k=None, @@ -213,9 +59,7 @@ def _build_step_args(md: TrtllmAttentionMetadata) -> dict: quant_mode=0, kv_scale_orig_quant=None, kv_scale_quant_orig=None, - # In-kernel RoPE is disabled (position_embedding_type=0): rotation - # happens outside in fused_qk_norm_rope. The rope_* values below are - # the contract's inert placeholders, not this model's rope config. + # In-kernel RoPE is disabled here; the rope_* values below are inert placeholders, not this model's rope config. position_embedding_type=0, rotary_inv_freq=None, rotary_cos_sin=None, @@ -251,108 +95,49 @@ def _build_step_args(md: TrtllmAttentionMetadata) -> dict: cross_kv=None, relative_attention_bias=None, relative_attention_max_distance=0, - # Held at the op's defaults. All three are MLA-only surface that this - # target does not use -- skip_correction is forced to 0.0 for a non-MLA layer by - # the engine's own resolver, and kv_norm_* folds an MLA kv_a_layernorm - # that does not exist here. kv_norm_weight=None, kv_norm_eps=1e-6, skip_correction_threshold=0.0, ) -# The MoE op's clamped gated activation: (up + beta) * gate * -# sigmoid(alpha * gate) after gate.clamp(max=limit), up.clamp(+-limit) — -# exactly this checkpoint's expert activation. alpha/beta are fixed in the -# HF reference; limit is config (swiglu_limit). -_GLU_ALPHA = 1.702 -_GLU_BETA = 1.0 -# Renormalize routing (top-k first, then fp32 softmax over the selected -# logits) and the only gated-activation kernel in this dtype family. -_ROUTING_METHOD_RENORMALIZE = 1 -_ACT_TYPE_SWIGLU = 0 -# FC1's K alignment for the trtllm-gen MXFP4 weight family: it sizes the -# declared expert operands and is the alignment mxfp8_quantize pads the -# hidden states up to, so the two always agree. -_FC1_K_ALIGN = 512 -# The MoE op reads the activation scales as a linear (row-major) buffer. -# The 128x4 swizzled order has the same byte count whenever num_tokens is a -# multiple of 128 — which every decode CUDA-graph batch of 128 or 256 is — -# and is then accepted silently as a wrong answer, so this is spelled out -# rather than left to the quantizer's default. -_LINEAR_SCALE_LAYOUT = False - - -def _pad_up(x: int, align: int) -> int: - return (x + align - 1) // align * align - - -class ModelingV2Core(DecoderModel): +_FC1_K_ALIGN = _weights.FC1_K_ALIGN + + +class GptOssTp1ModelingV2(DecoderModel): def __init__(self, model_config: ModelConfig): super().__init__(model_config) cfg = model_config.pretrained_config - assert cfg is not None - - # This target IS this geometry — assert, never adapt. - assert torch.cuda.get_device_capability() == _SM, ( - f"target certified on sm_{_SM[0]}{_SM[1]}, running on " - f"sm_{''.join(map(str, torch.cuda.get_device_capability()))}" - ) - assert model_config.mapping.tp_size == 1, "tp1 target" - assert model_config.mapping.pp_size == 1, "tp1 target" - # This checkpoint's config.json declares no dtype at all, so - # `pretrained_config.torch_dtype` is None. The engine resolves bf16 - # regardless (ModelConfig.torch_dtype defaults it), but the shell - # sizes lm_head from the *pretrained* value and would materialize it - # in the torch default fp32 — two layers away from where it is read. - # The shell normalizes that before super().__init__ (see below), so - # by the time the core is built both surfaces agree, and both are - # asserted: the declaration the shell and the KV pool are sized from, - # and the value quant_mode=0 must agree with. - dt = model_config.torch_dtype - assert cfg.torch_dtype == torch.bfloat16, cfg.torch_dtype - assert dt == torch.bfloat16, f"bf16 target, engine resolved {dt}" - assert not cfg.tie_word_embeddings, "untied lm_head" - assert cfg.attention_bias, "q/k/v/o carry bias" - assert cfg.hidden_act == "silu", "clamped GLU over a silu-shaped gate" - - # bf16 KV pool only: an fp8 pool needs quant_mode=128 plus the two - # fp32 scale tensors, which this target does not build. The expert - # quantization (mxfp4 blocks + e8m0 scales) is consumed directly by - # the MoE op, so quant_config.quant_algo — which the engine reads - # off quantization_config as W4A8_MXFP4_MXFP8, an mxfp8-activation - # recipe this target does not implement — is inert here. - kv_algo = model_config.quant_config.kv_cache_quant_algo - assert kv_algo is None, f"unsupported KV algo {kv_algo}; bf16 pool only" - quant = getattr(cfg, "quantization_config", None) - assert isinstance(quant, dict) and quant.get("quant_method") == "mxfp4", ( - "expert weights must be mxfp4 blocks + e8m0 scales" - ) - self.num_layers = cfg.num_hidden_layers - self.hidden = cfg.hidden_size - self.heads_q = cfg.num_attention_heads - self.heads_kv = cfg.num_key_value_heads - self.head_dim = cfg.head_dim - self.eps = cfg.rms_norm_eps - - # RoPE: YaRN over the full head_dim, half-split (neox) pairs. The - # engine hands this checkpoint the transformers-5.x migrated rope - # dict (rope_theta lives inside it and cfg.rope_theta is absent); - # a checkpoint written before that migration keeps the flat field, - # so both shapes are resolved here and every scalar the ramp - # depends on is asserted rather than defaulted. rope = getattr(cfg, "rope_scaling", None) or getattr(cfg, "rope_parameters", None) - assert isinstance(rope, dict), "rope parameters must be a dict" - assert rope.get("rope_type") == "yarn", "YaRN rope scaling" - assert rope.get("partial_rotary_factor", 1.0) == 1.0, "full-width rotation" - assert rope.get("attention_factor") is None, "attention_factor is derived" - assert rope.get("mscale") is None and rope.get("mscale_all_dim") is None, ( - "the derived attention_factor assumes no mscale pair" - ) self.theta = float(rope.get("rope_theta", getattr(cfg, "rope_theta", 0.0))) - assert self.theta > 0.0, "rope theta" - self.rotary_dim = self.head_dim - self.yarn_factor = float(rope["factor"]) + + layer_types = list(cfg.layer_types) + self.sliding = [t == "sliding_attention" for t in layer_types] + self.window = cfg.sliding_window + + d = _weights.MODEL_WEIGHTS.dims(self) + self.inter_pad = d.inter_pad + self.fc1_k_pad = d.fc1_k_pad + self.fc2_rows_pad = d.fc2_rows_pad + self.w = _weights.MODEL_WEIGHTS.declare(d) + + self._prefill: Target | None = None + self._decode: Target | None = None + + def post_load_weights(self) -> None: + """Construct the two targets, now that the weights are real. + + Must only be called after meta init is over, never from `__init__`, where the + shell's containers are still meta. + + Named for the shell hook that drives it. Not an override: that hook lives on + DecoderModelForCausalLM, and a core derives from DecoderModel, which has no + such method -- so there is no super() to call here. + """ + cfg = self.model_config.pretrained_config + rope = getattr(cfg, "rope_scaling", None) or getattr(cfg, "rope_parameters", None) + rotary_dim = cfg.head_dim + yarn_factor = float(rope["factor"]) beta_fast = float(rope["beta_fast"]) beta_slow = float(rope["beta_slow"]) orig_max = float(rope["original_max_position_embeddings"]) @@ -360,7 +145,7 @@ def __init__(self, model_config: ModelConfig): def correction_dim(rotations: float) -> float: return ( - self.rotary_dim + rotary_dim * math.log(orig_max / (rotations * 2.0 * math.pi)) / (2.0 * math.log(self.theta)) ) @@ -369,161 +154,405 @@ def correction_dim(rotations: float) -> float: high = correction_dim(beta_slow) if truncate: low, high = math.floor(low), math.ceil(high) - self.yarn_low = max(low, 0.0) - self.yarn_high = min(high, self.rotary_dim - 1.0) - self.yarn_attn_factor = ( - 0.1 * math.log(self.yarn_factor) + 1.0 if self.yarn_factor > 1.0 else 1.0 + yarn_low = max(low, 0.0) + yarn_high = min(high, rotary_dim - 1.0) + yarn_attn_factor = 0.1 * math.log(yarn_factor) + 1.0 if yarn_factor > 1.0 else 1.0 + + self._prefill = PrefillTarget( + self, + yarn_factor=yarn_factor, + yarn_low=yarn_low, + yarn_high=yarn_high, + yarn_attn_factor=yarn_attn_factor, + ) + self._decode = DecodeTarget( + self, + yarn_factor=yarn_factor, + yarn_low=yarn_low, + yarn_high=yarn_high, + yarn_attn_factor=yarn_attn_factor, ) - # Sliding window: half the layers mask to the newest `window` keys, - # the rest are plain causal. The window is a mask only — the op - # appends at absolute positions — so both kinds share one pool, one - # layer->pool mapping and one block-offset table. - layer_types = list(cfg.layer_types) - assert len(layer_types) == self.num_layers - assert set(layer_types) <= {"sliding_attention", "full_attention"}, layer_types - self.sliding = [t == "sliding_attention" for t in layer_types] - self.window = cfg.sliding_window - assert any(self.sliding) and self.window > 0, "alternating sliding window" - - # Sparse MoE on every layer; no dense MLP branch exists. - self.num_experts = cfg.num_local_experts - self.topk = cfg.num_experts_per_tok - self.inter = cfg.intermediate_size - self.glu_limit = float(cfg.swiglu_limit) - assert 0 < self.topk < self.num_experts, "MoE top-k bound" - - # Kernel bounds the geometry must fit: fused_qk_norm_rope's head_dim - # set, thop_attention's GQA rule, and the MoE op's padded operand - # widths (the padded intermediate is what `intermediate_size` means - # to that call, and hidden_states reach it widened to fc1_k_pad by - # the quantizer). - assert self.head_dim in (64, 128, 256), "fused_qk_norm_rope head_dim set" - assert self.heads_q % self.heads_kv == 0, "thop_attention GQA rule" - self.inter_pad = _pad_up(self.inter, 128) - self.fc1_k_pad = _pad_up(self.hidden, _FC1_K_ALIGN) - self.fc2_rows_pad = _pad_up(self.hidden, 128) - assert self.hidden % 32 == 0, "valid_hidden_size must be a multiple of 32" - assert self.inter % 32 == 0, "valid_intermediate_size must be a multiple of 32" - - q_width = self.heads_q * self.head_dim - kv_width = self.heads_kv * self.head_dim - - # Weight declaration: HF [out, in] row-major storage so checkpoint - # rows copy in unchanged; GEMM consumes .t() column-major views - # built after load. The expert tensors are the exception — they are - # declared in the MoE op's kernel-ready layout (padded, row-permuted - # weights/scales/biases), which weights.py builds from the - # checkpoint's block/scale tensors during the load. Meta-init - # intercepts torch.empty here — real CUDA storage arrives when the - # engine materializes the registry. - def P(*shape, dtype=dt): - return nn.Parameter(torch.empty(*shape, dtype=dtype), requires_grad=False) - - fc1_rows = 2 * self.inter_pad - w = nn.ParameterDict() - for i in range(self.num_layers): - w[f"l{i}_norm1"] = P(self.hidden) - w[f"l{i}_qkv"] = P(q_width + 2 * kv_width, self.hidden) - w[f"l{i}_qkv_bias"] = P(q_width + 2 * kv_width) - w[f"l{i}_sinks"] = P(self.heads_q, dtype=torch.float32) - w[f"l{i}_o"] = P(self.hidden, q_width) - w[f"l{i}_o_bias"] = P(self.hidden) - w[f"l{i}_norm2"] = P(self.hidden) - w[f"l{i}_router"] = P(self.num_experts, self.hidden) - w[f"l{i}_router_bias"] = P(self.num_experts) - w[f"l{i}_fc1_w"] = P(self.num_experts, fc1_rows, self.fc1_k_pad // 2, dtype=torch.uint8) - w[f"l{i}_fc1_s"] = P( - self.num_experts, fc1_rows, self.fc1_k_pad // 32, dtype=torch.uint8 + def _select_target(self, attn_metadata) -> Target: + """Which target runs this step.""" + return self._decode if phase_of(attn_metadata) is Phase.DECODE else self._prefill + + def forward( + self, + attn_metadata: AttentionMetadata, + *args, + **kwargs, + ) -> torch.Tensor: + """Route the step, and nothing else.""" + advance_step_generation() + return self._select_target(attn_metadata).forward(attn_metadata, *args, **kwargs) + + +class PrefillTarget(Target): + """The general case: context rows, and possibly generation rows beside them.""" + + def __init__( + self, + core: GptOssTp1ModelingV2, + *, + yarn_factor: float, + yarn_low: float, + yarn_high: float, + yarn_attn_factor: float, + ) -> None: + """Bind every op this target calls, once, against the real weights.""" + super().__init__(core) + cfg = core.model_config.pretrained_config + w = core.w + device = w["final_norm"].device + dtype = w["final_norm"].dtype + n = cfg.num_hidden_layers + + self._norm0 = FlashinferRmsnorm() + self._norm0.bind_const(weight=w["l0_norm1"], eps=cfg.rms_norm_eps) + + self._qkv = CublasMm() + for i in range(n): + self._qkv.bind_layered(i, mat_b=w[f"l{i}_qkv"].t(), bias=w[f"l{i}_qkv_bias"]) + + # fused_qk_norm_rope requires valid q/k norm weight tensors even when is_qk_norm=False, where + # the values are unused. + no_qk_norm = torch.zeros(cfg.head_dim, dtype=dtype, device=device) + self._qk_rope = FusedQkNormRope() + self._qk_rope.bind_const( + num_heads_q=cfg.num_attention_heads, + num_heads_k=cfg.num_key_value_heads, + num_heads_v=cfg.num_key_value_heads, + head_dim=cfg.head_dim, + rotary_dim=cfg.head_dim, + eps=cfg.rms_norm_eps, + q_weight=no_qk_norm, + k_weight=no_qk_norm, + base=core.theta, + is_neox=True, + factor=yarn_factor, + low=yarn_low, + high=yarn_high, + attention_factor=yarn_attn_factor, + is_qk_norm=False, + ) + + self._attn = ThopAttention() + for i in range(n): + self._attn.bind_layered(i, attention_sinks=w[f"l{i}_sinks"], local_layer_idx=i) + self._attn.bind_const( + num_heads=cfg.num_attention_heads, + num_kv_heads=cfg.num_key_value_heads, + head_size=cfg.head_dim, + **_CALL_CONSTANTS, + ) + + self._o_proj = CublasMm() + for i in range(n): + self._o_proj.bind_layered(i, mat_b=w[f"l{i}_o"].t(), bias=w[f"l{i}_o_bias"]) + + self._norm2 = FlashinferFusedAddRmsnorm() + for i in range(n): + self._norm2.bind_layered(i, weight=w[f"l{i}_norm2"]) + self._norm2.bind_const(eps=cfg.rms_norm_eps) + + self._router = CublasMm() + for i in range(n): + self._router.bind_layered(i, mat_b=w[f"l{i}_router"].t(), bias=w[f"l{i}_router_bias"]) + + self._quant = Mxfp8Quantize() + # The swizzled 128x4 order has the same byte count whenever num_tokens % 128 == 0, which every + # decode CUDA-graph batch is, so a wrong value here is accepted silently as a wrong answer. + self._quant.bind_const(swizzled_layout=False, alignment=_FC1_K_ALIGN) + + self._moe = Mxe4m3Mxe2m1BlockScaleMoeRunner() + for i in range(n): + self._moe.bind_layered( + i, + gemm1_weights=w[f"l{i}_fc1_w"], + gemm1_weights_scale=w[f"l{i}_fc1_s"], + gemm1_bias=w[f"l{i}_fc1_b"], + gemm2_weights=w[f"l{i}_fc2_w"], + gemm2_weights_scale=w[f"l{i}_fc2_s"], + gemm2_bias=w[f"l{i}_fc2_b"], ) - w[f"l{i}_fc1_b"] = P(self.num_experts, fc1_rows, dtype=torch.float32) - w[f"l{i}_fc2_w"] = P( - self.num_experts, - self.fc2_rows_pad, - self.inter_pad // 2, - dtype=torch.uint8, + self._moe.bind_glu(cfg.num_local_experts, cfg.swiglu_limit, device) + self._moe.bind_const( + # The MoE op silently ignores routing_bias on this routing method; the bias rides the + # router GEMM epilogue instead. + routing_bias=None, + num_experts=cfg.num_local_experts, + top_k=cfg.num_experts_per_tok, + n_group=None, + topk_group=None, + intermediate_size=core.inter_pad, + valid_hidden_size=cfg.hidden_size, + valid_intermediate_size=cfg.intermediate_size, + local_expert_offset=0, + local_num_experts=cfg.num_local_experts, + routed_scaling_factor=None, + routing_method_type=1, + act_type=0, + ) + + self._norm_next = FlashinferFusedAddRmsnorm() + for i in range(n): + next_w = w[f"l{i + 1}_norm1"] if i + 1 < n else w["final_norm"] + self._norm_next.bind_layered(i, weight=next_w) + self._norm_next.bind_const(eps=cfg.rms_norm_eps) + + def step_args(self, md: TrtllmAttentionMetadata) -> dict: + """Project the prepared metadata onto thop_attention's explicit batch state, once per forward.""" + return dict( + sequence_length=md.kv_lens_cuda_runtime, + host_past_key_value_lengths=md.kv_lens_runtime, + host_total_kv_lens=md.host_total_kv_lens, + context_lengths=md.prompt_lens_cuda_runtime, + host_context_lengths=md.prompt_lens_cpu_runtime, + host_request_types=md.host_request_types_runtime, + kv_cache_block_offsets=md.kv_cache_block_offsets, + host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, + workspace_=md.effective_workspace, + tokens_per_block=md.tokens_per_block, + max_num_requests=md.max_num_requests, + max_context_length=md.max_context_length, + max_seq_len=md.max_seq_len, + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup, + use_paged_context_fmha=md.use_paged_context_fmha, + beam_width=md.effective_beam_width, + cache_indirection=md.cache_indirection, + block_ids_per_seq=md.block_ids_per_seq, + max_context_q_len_override=md.max_context_q_len_override, + is_cross=md.is_cross, + is_spec_decoding_enabled=md.is_spec_decoding_enabled, + use_spec_decoding=md.use_spec_decoding, + is_spec_dec_tree=md.is_spec_dec_tree, + spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, + spec_decoding_position_offsets_for_cpp=md.spec_decoding_position_offsets_for_cpp, + spec_decoding_packed_mask=md.spec_decoding_packed_mask, + spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, + spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, + spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, + spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, + num_sparse_topk=md.num_sparse_topk, + flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, + flash_mla_num_splits=md.flash_mla_num_splits, + max_num_sequences=md.max_num_sequences, + force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, + ) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None = None, + position_ids: torch.IntTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + lora_params: dict | None = None, + **kwargs, + ) -> torch.Tensor: + core = self.core + cfg = core.model_config.pretrained_config + + self._check_step_contract(attn_metadata) + self._attn.bind_const(**self.step_args(attn_metadata)) + + full_window = attn_metadata.max_seq_len + for i, sliding in enumerate(core.sliding): + self._attn.bind_layered( + i, attention_window_size=core.window if sliding else full_window ) - w[f"l{i}_fc2_s"] = P( - self.num_experts, - self.fc2_rows_pad, - self.inter_pad // 32, - dtype=torch.uint8, + + pos = torch.reshape(position_ids, [-1]) + + if inputs_embeds is None: + h = nn.functional.embedding(input_ids, core.w["embed"]) + else: + h = inputs_embeds + num_tokens = h.shape[0] + dt = h.dtype + attn_out = torch.empty( + [num_tokens, cfg.num_attention_heads * cfg.head_dim], dtype=dt, device=h.device + ) + + x = self._norm0(h) + residual = h + for i in range(cfg.num_hidden_layers): + qkv = self._qkv(x, layer=i) + qkv = self._qk_rope(qkv, position_ids=pos) + attn_out = self._attn(q=qkv, output=attn_out, layer=i) + o = self._o_proj(attn_out, layer=i) + o, residual = self._norm2(o, residual, layer=i) + router_logits = self._router(o, layer=i) + hidden_fp8, hidden_sf = self._quant(o) + moe = self._moe( + routing_logits=router_logits, + hidden_states=hidden_fp8, + hidden_states_scale=hidden_sf, + layer=i, ) - w[f"l{i}_fc2_b"] = P(self.num_experts, self.fc2_rows_pad, dtype=torch.float32) - w["final_norm"] = P(self.hidden) - w["embed"] = P(cfg.vocab_size, self.hidden) - self.w = w - - self._layers: list | None = None - self._call_tensors: dict | None = None - self._step_contract_checked = False - - def build_layer_views(self) -> None: - """Post-load derivation: per-layer tuples of column-major GEMM views, - norm weights and expert operands (kills hot-path dict lookups; .t() - is zero-copy), plus the constant fp32 call tensors the MoE and RoPE - calls require. Meta is over here, so real tensors may be created.""" - w = self.w + moe, residual = self._norm_next(moe, residual, layer=i) + x = moe + return x + + +class DecodeTarget(Target): + """The specialization: routed to only when there are no context rows.""" + + def __init__( + self, + core: GptOssTp1ModelingV2, + *, + yarn_factor: float, + yarn_low: float, + yarn_high: float, + yarn_attn_factor: float, + ) -> None: + """Bind every op this target calls, once, against the real weights.""" + super().__init__(core) + cfg = core.model_config.pretrained_config + w = core.w device = w["final_norm"].device - # Per-expert activation scalars, sized by local_num_experts. - self._call_tensors = { - "alpha": torch.full( - (self.num_experts,), _GLU_ALPHA, dtype=torch.float32, device=device - ), - "beta": torch.full((self.num_experts,), _GLU_BETA, dtype=torch.float32, device=device), - "limit": torch.full( - (self.num_experts,), self.glu_limit, dtype=torch.float32, device=device - ), - # fused_qk_norm_rope requires valid q/k norm weight tensors even - # with is_qk_norm=False, where their values are unused. - "no_qk_norm": torch.zeros(self.head_dim, dtype=w["final_norm"].dtype, device=device), - } - layers = [] - for i in range(self.num_layers): - next_norm = w[f"l{i + 1}_norm1"] if i + 1 < self.num_layers else w["final_norm"] - layers.append( - ( - w[f"l{i}_qkv"].t(), - w[f"l{i}_qkv_bias"], - w[f"l{i}_sinks"], - w[f"l{i}_o"].t(), - w[f"l{i}_o_bias"], - w[f"l{i}_norm2"], - w[f"l{i}_router"].t(), - w[f"l{i}_router_bias"], - w[f"l{i}_fc1_w"], - w[f"l{i}_fc1_s"], - w[f"l{i}_fc1_b"], - w[f"l{i}_fc2_w"], - w[f"l{i}_fc2_s"], - w[f"l{i}_fc2_b"], - next_norm, - ) + dtype = w["final_norm"].dtype + n = cfg.num_hidden_layers + + self._norm0 = FlashinferRmsnorm() + self._norm0.bind_const(weight=w["l0_norm1"], eps=cfg.rms_norm_eps) + + self._qkv = CublasMm() + for i in range(n): + self._qkv.bind_layered(i, mat_b=w[f"l{i}_qkv"].t(), bias=w[f"l{i}_qkv_bias"]) + + # fused_qk_norm_rope requires valid q/k norm weight tensors even when is_qk_norm=False, where + # the values are unused. + no_qk_norm = torch.zeros(cfg.head_dim, dtype=dtype, device=device) + self._qk_rope = FusedQkNormRope() + self._qk_rope.bind_const( + num_heads_q=cfg.num_attention_heads, + num_heads_k=cfg.num_key_value_heads, + num_heads_v=cfg.num_key_value_heads, + head_dim=cfg.head_dim, + rotary_dim=cfg.head_dim, + eps=cfg.rms_norm_eps, + q_weight=no_qk_norm, + k_weight=no_qk_norm, + base=core.theta, + is_neox=True, + factor=yarn_factor, + low=yarn_low, + high=yarn_high, + attention_factor=yarn_attn_factor, + is_qk_norm=False, + ) + + self._attn = ThopAttention() + for i in range(n): + self._attn.bind_layered(i, attention_sinks=w[f"l{i}_sinks"], local_layer_idx=i) + self._attn.bind_const( + num_heads=cfg.num_attention_heads, + num_kv_heads=cfg.num_key_value_heads, + head_size=cfg.head_dim, + **_CALL_CONSTANTS, + ) + + self._o_proj = CublasMm() + for i in range(n): + self._o_proj.bind_layered(i, mat_b=w[f"l{i}_o"].t(), bias=w[f"l{i}_o_bias"]) + + self._norm2 = FlashinferFusedAddRmsnorm() + for i in range(n): + self._norm2.bind_layered(i, weight=w[f"l{i}_norm2"]) + self._norm2.bind_const(eps=cfg.rms_norm_eps) + + self._router = CublasMm() + for i in range(n): + self._router.bind_layered(i, mat_b=w[f"l{i}_router"].t(), bias=w[f"l{i}_router_bias"]) + + self._quant = Mxfp8Quantize() + # The swizzled 128x4 order has the same byte count whenever num_tokens % 128 == 0, which every + # decode CUDA-graph batch is, so a wrong value here is accepted silently as a wrong answer. + self._quant.bind_const(swizzled_layout=False, alignment=_FC1_K_ALIGN) + + self._moe = Mxe4m3Mxe2m1BlockScaleMoeRunner() + for i in range(n): + self._moe.bind_layered( + i, + gemm1_weights=w[f"l{i}_fc1_w"], + gemm1_weights_scale=w[f"l{i}_fc1_s"], + gemm1_bias=w[f"l{i}_fc1_b"], + gemm2_weights=w[f"l{i}_fc2_w"], + gemm2_weights_scale=w[f"l{i}_fc2_s"], + gemm2_bias=w[f"l{i}_fc2_b"], ) - self._layers = layers - - def _check_step_contract(self, md, position_ids) -> None: - """First-forward fail-fast: the metadata fields this target consumes - must exist (they are private trtllm surface), and the KV pool layout - must be one `thop_attention` is certified over. Everything checked is - fixed at engine construction — once per model instance is sound. - - This checkpoint alternates sliding-window and full attention, and - KVCacheManagerV2 gives each attention-window class its own layer - group, so its mapping carries two pool ids rather than one. The bound - is the catalog's: `thop_attention.md` certifies one pool and two, and - this forward passes the mapping through untouched -- the op reads the - pool column itself. A third pool would be outside what was measured. - """ - missing = [name for name in _STEP_FIELDS if not hasattr(md, name)] - assert not missing, f"metadata fields missing: {missing}" - assert position_ids.dtype == torch.int32 - pools = {row[0] for row in md.host_kv_cache_pool_mapping.tolist()} - assert pools <= {0, 1}, ( - f"KV addressing over {len(pools)} pools is not certified " - f"(thop_attention.md certifies one and two); layer->pool ids {sorted(pools)}" + self._moe.bind_glu(cfg.num_local_experts, cfg.swiglu_limit, device) + self._moe.bind_const( + # The MoE op silently ignores routing_bias on this routing method; the bias rides the + # router GEMM epilogue instead. + routing_bias=None, + num_experts=cfg.num_local_experts, + top_k=cfg.num_experts_per_tok, + n_group=None, + topk_group=None, + intermediate_size=core.inter_pad, + valid_hidden_size=cfg.hidden_size, + valid_intermediate_size=cfg.intermediate_size, + local_expert_offset=0, + local_num_experts=cfg.num_local_experts, + routed_scaling_factor=None, + routing_method_type=1, + act_type=0, + ) + + self._norm_next = FlashinferFusedAddRmsnorm() + for i in range(n): + next_w = w[f"l{i + 1}_norm1"] if i + 1 < n else w["final_norm"] + self._norm_next.bind_layered(i, weight=next_w) + self._norm_next.bind_const(eps=cfg.rms_norm_eps) + + def step_args(self, md: TrtllmAttentionMetadata) -> dict: + """Project the prepared metadata onto thop_attention's explicit batch state, once per forward.""" + return dict( + sequence_length=md.kv_lens_cuda_runtime, + host_past_key_value_lengths=md.kv_lens_runtime, + host_total_kv_lens=md.host_total_kv_lens, + context_lengths=md.prompt_lens_cuda_runtime, + host_context_lengths=md.prompt_lens_cpu_runtime, + host_request_types=md.host_request_types_runtime, + kv_cache_block_offsets=md.kv_cache_block_offsets, + host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, + workspace_=md.effective_workspace, + tokens_per_block=md.tokens_per_block, + max_num_requests=md.max_num_requests, + max_context_length=md.max_context_length, + max_seq_len=md.max_seq_len, + num_contexts=0, + num_ctx_tokens=0, + trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup, + use_paged_context_fmha=md.use_paged_context_fmha, + beam_width=md.effective_beam_width, + cache_indirection=md.cache_indirection, + block_ids_per_seq=md.block_ids_per_seq, + max_context_q_len_override=md.max_context_q_len_override, + is_cross=md.is_cross, + is_spec_decoding_enabled=md.is_spec_decoding_enabled, + use_spec_decoding=md.use_spec_decoding, + is_spec_dec_tree=md.is_spec_dec_tree, + spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, + spec_decoding_position_offsets_for_cpp=md.spec_decoding_position_offsets_for_cpp, + spec_decoding_packed_mask=md.spec_decoding_packed_mask, + spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, + spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, + spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, + spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, + num_sparse_topk=md.num_sparse_topk, + flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, + flash_mla_num_splits=md.flash_mla_num_splits, + max_num_sequences=md.max_num_sequences, + force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, ) - self._step_contract_checked = True def forward( self, @@ -534,154 +563,69 @@ def forward( lora_params: dict | None = None, **kwargs, ) -> torch.Tensor: - assert self._layers is not None and self._call_tensors is not None, ( - "load_weights must run before forward" - ) - assert position_ids is not None - assert isinstance(attn_metadata, TrtllmAttentionMetadata) - # Inputs this target does not implement must fail loudly, not be - # silently dropped (unlike runtime-owned features, which pass through). - assert lora_params is None, "LoRA is not implemented by this target" - assert kwargs.get("spec_metadata") is None, ( - "speculative decoding is not implemented by this target" - ) - if not self._step_contract_checked: - self._check_step_contract(attn_metadata, position_ids) + core = self.core + cfg = core.model_config.pretrained_config + + self._check_step_contract(attn_metadata) + self._attn.bind_const(**self.step_args(attn_metadata)) - step = _build_step_args(attn_metadata) - # Full-attention layers take the no-window value; sliding layers take - # the checkpoint's window. Both are host-derived per-capture constants - # (max_seq_len is an engine-construction constant). full_window = attn_metadata.max_seq_len - const = self._call_tensors - no_qk_norm = const["no_qk_norm"] - pos = reshape(position_ids, [-1]) + for i, sliding in enumerate(core.sliding): + self._attn.bind_layered( + i, attention_window_size=core.window if sliding else full_window + ) + + pos = torch.reshape(position_ids, [-1]) if inputs_embeds is None: - assert input_ids is not None - h = embedding(input_ids, self.w["embed"]) + h = nn.functional.embedding(input_ids, core.w["embed"]) else: h = inputs_embeds num_tokens = h.shape[0] dt = h.dtype - attn_out = empty([num_tokens, self.heads_q * self.head_dim], dt, h.device) + attn_out = torch.empty( + [num_tokens, cfg.num_attention_heads * cfg.head_dim], dtype=dt, device=h.device + ) - x = flashinfer_rmsnorm(h, self.w["l0_norm1"], self.eps) + x = self._norm0(h) residual = h - for i in range(self.num_layers): - ( - w_qkv, - b_qkv, - sinks, - w_o, - b_o, - w_n2, - w_rt, - b_rt, - fc1_w, - fc1_s, - fc1_b, - fc2_w, - fc2_s, - fc2_b, - w_next, - ) = self._layers[i] - qkv = cublas_mm(x, w_qkv, b_qkv) - fused_qk_norm_rope( - qkv, - num_heads_q=self.heads_q, - num_heads_k=self.heads_kv, - num_heads_v=self.heads_kv, - head_dim=self.head_dim, - rotary_dim=self.rotary_dim, - eps=self.eps, - q_weight=no_qk_norm, - k_weight=no_qk_norm, - base=self.theta, - is_neox=True, - position_ids=pos, - factor=self.yarn_factor, - low=self.yarn_low, - high=self.yarn_high, - attention_factor=self.yarn_attn_factor, - is_qk_norm=False, - ) - thop_attention( - q=qkv, - output=attn_out, - local_layer_idx=i, - num_heads=self.heads_q, - num_kv_heads=self.heads_kv, - head_size=self.head_dim, - attention_sinks=sinks, - attention_window_size=self.window if self.sliding[i] else full_window, - **step, - **_CALL_CONSTANTS, - ) - o = cublas_mm(attn_out, w_o, b_o) - flashinfer_fused_add_rmsnorm(o, residual, w_n2, self.eps) - # The router bias rides the GEMM epilogue: the MoE op silently - # ignores routing_bias on this routing method. - router_logits = cublas_mm(o, w_rt, b_rt) - # The quantizer owns the hidden widening 2880 -> fc1_k_pad: it - # zero-fills the padded columns and their scale bytes, which - # multiply zero-valued padded weights. - hidden_fp8, hidden_sf = mxfp8_quantize(o, _LINEAR_SCALE_LAYOUT, _FC1_K_ALIGN) - moe = mxe4m3_mxe2m1_block_scale_moe_runner( - router_logits, - None, - hidden_fp8, - hidden_sf, - fc1_w, - fc1_s, - fc1_b, - const["alpha"], - const["beta"], - const["limit"], - fc2_w, - fc2_s, - fc2_b, - self.num_experts, - self.topk, - None, - None, - self.inter_pad, - self.hidden, - self.inter, - 0, - self.num_experts, - None, - _ROUTING_METHOD_RENORMALIZE, - _ACT_TYPE_SWIGLU, + for i in range(cfg.num_hidden_layers): + qkv = self._qkv(x, layer=i) + qkv = self._qk_rope(qkv, position_ids=pos) + attn_out = self._attn(q=qkv, output=attn_out, layer=i) + o = self._o_proj(attn_out, layer=i) + o, residual = self._norm2(o, residual, layer=i) + router_logits = self._router(o, layer=i) + hidden_fp8, hidden_sf = self._quant(o) + moe = self._moe( + routing_logits=router_logits, + hidden_states=hidden_fp8, + hidden_states_scale=hidden_sf, + layer=i, ) - flashinfer_fused_add_rmsnorm(moe, residual, w_next, self.eps) + moe, residual = self._norm_next(moe, residual, layer=i) x = moe return x @register_auto_model("ModelingV2GptOss120bSm103Tp1") -class ModelingV2GptOss120bSm103Tp1(DecoderModelForCausalLM[ModelingV2Core, PretrainedConfig]): +class ModelingV2GptOss120bSm103Tp1(DecoderModelForCausalLM[GptOssTp1ModelingV2, PretrainedConfig]): def __init__(self, model_config: ModelConfig): cfg = model_config.pretrained_config - assert cfg is not None - # DecoderModelForCausalLM sizes lm_head from the *pretrained* dtype, - # which this checkpoint's config.json does not declare — leaving - # lm_head.weight fp32 while every other tensor is bf16, and failing - # two layers away from here. The engine has already resolved the - # dtype it will run at; adopt it, rather than let the default of a - # missing field decide. Only fills the gap: a declared dtype wins. + # This checkpoint's config.json declares no dtype, so lm_head would be sized fp32 while every + # other tensor is bf16, failing two layers away. if cfg.torch_dtype is None: cfg.torch_dtype = model_config.torch_dtype super().__init__( - ModelingV2Core(model_config), + GptOssTp1ModelingV2(model_config), config=model_config, hidden_size=cfg.hidden_size, vocab_size=cfg.vocab_size, ) def load_weights(self, weights, *args, **kwargs): - _weights.load(self, weights) + _weights.MODEL_WEIGHTS.load(self, weights) def post_load_weights(self): super().post_load_weights() - self.model.build_layer_views() + self.model.post_load_weights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/__init__.py new file mode 100644 index 000000000000..1059445198db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/attention/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/attention/fused_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/attention/fused_qk_norm_rope.py new file mode 100644 index 000000000000..59ec9a6d395a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/attention/fused_qk_norm_rope.py @@ -0,0 +1,268 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-head QK RMS norm followed by RoPE, in place on a packed qkv buffer.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + assert_within_ulp, +) + + +class FusedQkNormRope(OpWrapper): + """In place on `qkv`: RMS-norm the q and k heads, then rotate them. + Returns `qkv` -- the same buffer, mutated and handed back so a call site + can show what it touched without allocating. + + The v heads are not touched. `qkv` is one `[num_tokens, (Hq + Hk + Hv) * + head_dim]` buffer, which is the layout the fused projection already + produces, so nothing is split or copied. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # gpt-oss-120b is the only caller, and it asks for the RoPE half only: + # `is_qk_norm=False`, neox pairing, and YaRN-blended frequencies. 64 query + # heads, 8 key and 8 value, head_dim 64, all of which the cells carry + # because the op reads them as scalars rather than deriving them. + _GPT_OSS = dict( + num_heads_q=64, + num_heads_k=8, + num_heads_v=8, + head_dim=64, + rotary_dim=64, + base=150000.0, + is_neox=True, + is_qk_norm=False, + # The checkpoint's YaRN correction range, from beta_fast=32 / + # beta_slow=1 over its original 4096-token context. + factor=32.0, + low=8.0, + high=18.0, + attention_factor=1.0, + ) + + CELLS: tuple[Cell, ...] = ( + Cell(why="decode: one token", spec=dict(tokens=1, **_GPT_OSS)), + Cell(why="prefill-shaped batch", spec=dict(tokens=2048, **_GPT_OSS)), + Cell( + why="a token count that is not a multiple of any tile", + spec=dict(tokens=37, **_GPT_OSS), + ), + ) + + note = """ + Three of this op's paths are uncertified because no shipped target reaches + them, and the first is the one its own name advertises: + + * `is_qk_norm=True` -- gpt-oss passes False, so the norm half runs in + `reference` but never against the kernel; + * `is_neox=False`, the interleaved pairing; + * `use_mrope`, `use_gemma`. + + Adding a target that wants any of them means adding the cell, not assuming + the path works because the argument exists. + """ + + def raw_call( + self, + qkv: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + is_neox: bool, + position_ids: torch.Tensor, + factor: float = 1.0, + low: float = 0.0, + high: float = 0.0, + attention_factor: float = 1.0, + is_qk_norm: bool = True, + use_gemma: bool = False, + use_mrope: bool = False, + mrope_section1: int = 0, + mrope_section2: int = 0, + ) -> torch.Tensor: + """Writes `qkv` in place and returns it -- the op destroys the q/k + heads' prior contents, so a caller that wants them has to have kept + its own copy beforehand.""" + torch.ops.trtllm.fused_qk_norm_rope( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + return qkv + + def reference( + self, + qkv: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + is_neox: bool, + position_ids: torch.Tensor, + factor: float = 1.0, + low: float = 0.0, + high: float = 0.0, + attention_factor: float = 1.0, + is_qk_norm: bool = True, + use_gemma: bool = False, + use_mrope: bool = False, + mrope_section1: int = 0, + mrope_section2: int = 0, + ) -> torch.Tensor: + """The buffer the op writes back, fp32 throughout. + + Returns rather than mutates so a cell can drive both sides from the + same input; the in-place write is what `raw_call` does. + """ + num_heads = num_heads_q + num_heads_k + num_heads_v + x = qkv.float().view(-1, num_heads, head_dim).clone() + half = rotary_dim // 2 + inv_freq = self._inv_freq(rotary_dim, base, factor, low, high, qkv.device) + + if use_mrope: + angle = position_ids.float()[:, :, None] * inv_freq # [3, T, half] + cos, sin = angle.cos(), angle.sin() + + def pick(c: torch.Tensor) -> torch.Tensor: + out = c[0].clone() + out[:, 1 : mrope_section1 * 3 : 3] = c[1][:, 1 : mrope_section1 * 3 : 3] + out[:, 2 : mrope_section2 * 3 : 3] = c[2][:, 2 : mrope_section2 * 3 : 3] + return out + + cos, sin = pick(cos), pick(sin) + else: + angle = position_ids.float()[:, None] * inv_freq # [T, half] + cos, sin = angle.cos(), angle.sin() + cos = cos[:, None, :] * attention_factor + sin = sin[:, None, :] * attention_factor + + for start, count, weight in ( + (0, num_heads_q, q_weight.float()), + (num_heads_q, num_heads_k, k_weight.float()), + ): + h = x[:, start : start + count, :] + if is_qk_norm: + h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + eps) + h = h * ((1.0 + weight) if use_gemma else weight) + r = h[..., :rotary_dim] + if is_neox: + x1, x2 = r[..., :half], r[..., half:] + rotated = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) + else: + x1, x2 = r[..., ::2], r[..., 1::2] + rotated = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten( + -2 + ) + x[:, start : start + count, :] = torch.cat([rotated, h[..., rotary_dim:]], -1) + return x.view(qkv.shape[0], -1).to(qkv.dtype) + + @staticmethod + def _inv_freq( + rotary_dim: int, base: float, factor: float, low: float, high: float, device + ) -> torch.Tensor: + """YaRN-blended inverse frequencies; plain RoPE when factor == 1.""" + j = torch.arange(rotary_dim // 2, dtype=torch.float32, device=device) + pos_freqs = base ** (2.0 * j / rotary_dim) + inv_extrapolation = 1.0 / pos_freqs + inv_interpolation = 1.0 / (factor * pos_freqs) + if high == low: + high = high + 0.001 # the kernel guards the ramp singularity the same way + ramp = ((j - low) / (high - low)).clamp(0.0, 1.0) + extrapolation = 1.0 - ramp + return inv_interpolation * (1.0 - extrapolation) + inv_extrapolation * extrapolation + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """A rotation is a two-term combination, and two terms can cancel. + + `x1 * cos - x2 * sin` goes to zero wherever the two products nearly + match, and there torch's per-element rtol is meaningless: at 2048 tokens + roughly 2e-4 of the outputs land close enough that a 7.3e-4 absolute + difference reads as a relative difference of 23. The decode and 37-token + cells pass under torch's default band only because they are too small to + sample one. + + So: 2 ulp element-wise against the row scale, 1 of relative RMS -- tight, + because one rotation is the shortest chain in this catalog. Measured + worst case across these cells is 7.3e-4 absolute at a row scale of ~4, + about 0.05 ulp. + """ + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=1.0) + + def is_valid( + self, + qkv: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + is_neox: bool, + position_ids: torch.Tensor, + *unused_args: object, + **unused_kwargs: object, + ) -> None: + # `eps`, `q_weight`, `k_weight`, `base` and `is_neox` are named though + # nothing here reads them: they mirror `raw_call`, and gpt-oss passes + # them by keyword. + # + # The op reads the head counts as scalars and addresses qkv from + # data_ptr(); a width that disagrees with them is read as whatever the + # counts say and the heads silently come out shifted. + expected = (num_heads_q + num_heads_k + num_heads_v) * head_dim + assert qkv.shape[-1] == expected, ( + f"qkv's last dim must be (Hq + Hk + Hv) * head_dim = {expected}; got " + f"{qkv.shape[-1]}. The op trusts the counts and reads the buffer flat" + ) + assert qkv.is_contiguous(), "qkv must be contiguous; the op writes it through data_ptr()" + assert rotary_dim <= head_dim, ( + f"rotary_dim ({rotary_dim}) exceeds head_dim ({head_dim}); the rotation " + "would read past the end of each head" + ) + assert rotary_dim % 2 == 0, f"rotary_dim must be even to pair; got {rotary_dim}" + assert position_ids.shape[-1] == qkv.shape[0], ( + f"one position per token: {position_ids.shape[-1]} ids for {qkv.shape[0]} tokens" + ) + + +fused_qk_norm_rope = FusedQkNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/weights.py index 23b87584d6cd..cb58413cbbb1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/weights.py @@ -1,45 +1,36 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Weight manifest and loader: gpt-oss-120b / sm_103 / tp1. - -MANIFEST is a data table: target parameter -> the checkpoint keys that fill -it, each with an optional destination index into the parameter and an -optional source transform. tp1: no sharding — every checkpoint tensor -reaches exactly one parameter whole. - -Storage layout is HF [out, in] row-major, so every attention/router/embed -copy is layout-preserving; the GEMM-side column-major views are derived -after load in modeling.build_layer_views(). Two families need a transform: - -* **fp32 promotions.** The expert biases and the attention sink logits are - bf16 on disk, but the MoE op rejects a bf16 bias and thop_attention reads - the sink buffer as raw fp32. - -* **The MXFP4 expert operands.** The checkpoint stores each layer's experts - as `gate_up_proj_blocks` [E, 2I, H/32, 16] / `_scales` [E, 2I, H/32] and - `down_proj_blocks` [E, H, I/32, 16] / `_scales` [E, H, I/32] — E2M1 codes - two per byte (low nibble = even K index) with one E8M0 exponent per 32 K - elements, already in [out, in] orientation. The MoE op wants them padded, - row-permuted and (scales only) swizzled; `_prep_fc1` / `_prep_fc2` below - are that recipe, applied on device, one layer at a time. - - The parity trap: this checkpoint's `2I` axis runs (gate, up, gate, up, - ...) — HF reads `gate = gate_up[..., ::2]`, `up = gate_up[..., 1::2]` — - while the kernel's interleave wants destination row `2i` = **up** `i` and - `2i+1` = gate `i`. The halves are therefore split by parity and - re-concatenated as [up ; gate] before the permutation. Getting it - backwards is finite, plausibly scaled and invisible to a boot check. - -Loading contract: `load(model, weights)` consumes the engine-provided -per-rank checkpoint dict (safetensors lazy slices), fills every declared -parameter exactly once, and asserts full bidirectional coverage — every -target parameter written, every checkpoint key consumed. +"""Weight table: gpt-oss-120b / sm_103 / tp1. + +One `W` per weight role, carrying its shape, checkpoint source, dtype and +load-time transform. tp1: no sharding -- every checkpoint tensor reaches +exactly one parameter whole. Storage layout is HF [out, in] row-major, so the +attention, router and embed copies are layout-preserving. + +The parity trap in the MXFP4 expert operands: this checkpoint's `2I` axis +runs (gate, up, gate, up, ...) -- HF reads `gate = gate_up[..., ::2]`, `up = +gate_up[..., 1::2]` -- while the kernel's interleave wants destination row +`2i` = **up** `i` and `2i+1` = gate `i`. Getting it backwards is finite, +plausibly scaled and invisible to a boot check. """ +from types import SimpleNamespace + import torch +from tensorrt_llm._torch._experimental.modeling_v2._weights import ModelWeights, W + SCALE_BLOCK = 32 # mxfp4: one E8M0 exponent per 32 elements along K +# FC1's K alignment for the trtllm-gen MXFP4 weight family. It sizes the +# declared expert operands here and is the alignment mxfp8_quantize pads the +# hidden states up to in modeling.py; the two must be the same number. +FC1_K_ALIGN = 512 + + +def _pad_up(x: int, align: int) -> int: + return (x + align - 1) // align * align + def _pad_rows_cols(t: torch.Tensor, rows: int, cols: int) -> torch.Tensor: """Zero-pad a `[E, R, C]` stack to `[E, rows, cols]`.""" @@ -99,7 +90,8 @@ def _fc1_perm(rows: int, device) -> torch.Tensor: def _prep_fc1_weight(blocks: torch.Tensor, core) -> torch.Tensor: """[E, 2I, H/32, 16] uint8 blocks -> kernel FC1 operand.""" - b = blocks.reshape(core.num_experts, 2 * core.inter, core.hidden // 2) + cfg = core.model_config.pretrained_config + b = blocks.reshape(cfg.num_local_experts, 2 * cfg.intermediate_size, cfg.hidden_size // 2) up, gate = _split_gate_up(b) cols = core.fc1_k_pad // 2 w = torch.cat( @@ -135,8 +127,9 @@ def _prep_fc1_bias(bias: torch.Tensor, core) -> torch.Tensor: def _prep_fc2_weight(blocks: torch.Tensor, core) -> torch.Tensor: """[E, H, I/32, 16] uint8 blocks -> kernel FC2 operand (no interleave).""" + cfg = core.model_config.pretrained_config w = _pad_rows_cols( - blocks.reshape(core.num_experts, core.hidden, core.inter // 2), + blocks.reshape(cfg.num_local_experts, cfg.hidden_size, cfg.intermediate_size // 2), core.fc2_rows_pad, core.inter_pad // 2, ) @@ -160,87 +153,123 @@ def _to_fp32(t: torch.Tensor, core) -> torch.Tensor: return t.float() -def _manifest(core) -> dict: - """target param key -> list of (ckpt key, index into the param | None, - source transform | None).""" - q_width = core.heads_q * core.head_dim - kv_width = core.heads_kv * core.head_dim - rows: dict = {} - for i in range(core.num_layers): - p = f"model.layers.{i}" - rows[f"l{i}_norm1"] = [(f"{p}.input_layernorm.weight", None, None)] - rows[f"l{i}_qkv"] = [ - (f"{p}.self_attn.q_proj.weight", (slice(0, q_width),), None), - ( - f"{p}.self_attn.k_proj.weight", - (slice(q_width, q_width + kv_width),), - None, - ), - ( - f"{p}.self_attn.v_proj.weight", - (slice(q_width + kv_width, q_width + 2 * kv_width),), - None, +class GptOssWeights(ModelWeights): + """gpt-oss-120b / sm_103 / tp1.""" + + WEIGHTS: tuple[W, ...] = ( + W("norm1", shape=lambda d: (d.hidden,), src="{p}.input_layernorm.weight"), + W( + "qkv", + shape=lambda d: (d.q_width + 2 * d.kv_width, d.hidden), + src=( + ("{p}.self_attn.q_proj.weight", lambda d: (slice(0, d.q_width),)), + ( + "{p}.self_attn.k_proj.weight", + lambda d: (slice(d.q_width, d.q_width + d.kv_width),), + ), + ( + "{p}.self_attn.v_proj.weight", + lambda d: (slice(d.q_width + d.kv_width, d.q_width + 2 * d.kv_width),), + ), ), - ] - rows[f"l{i}_qkv_bias"] = [ - (f"{p}.self_attn.q_proj.bias", (slice(0, q_width),), None), - (f"{p}.self_attn.k_proj.bias", (slice(q_width, q_width + kv_width),), None), - ( - f"{p}.self_attn.v_proj.bias", - (slice(q_width + kv_width, q_width + 2 * kv_width),), - None, + ), + W( + "qkv_bias", + shape=lambda d: (d.q_width + 2 * d.kv_width,), + src=( + ("{p}.self_attn.q_proj.bias", lambda d: (slice(0, d.q_width),)), + ( + "{p}.self_attn.k_proj.bias", + lambda d: (slice(d.q_width, d.q_width + d.kv_width),), + ), + ( + "{p}.self_attn.v_proj.bias", + lambda d: (slice(d.q_width + d.kv_width, d.q_width + 2 * d.kv_width),), + ), ), - ] - rows[f"l{i}_sinks"] = [(f"{p}.self_attn.sinks", None, _to_fp32)] - rows[f"l{i}_o"] = [(f"{p}.self_attn.o_proj.weight", None, None)] - rows[f"l{i}_o_bias"] = [(f"{p}.self_attn.o_proj.bias", None, None)] - rows[f"l{i}_norm2"] = [(f"{p}.post_attention_layernorm.weight", None, None)] - rows[f"l{i}_router"] = [(f"{p}.mlp.router.weight", None, None)] - rows[f"l{i}_router_bias"] = [(f"{p}.mlp.router.bias", None, None)] - rows[f"l{i}_fc1_w"] = [(f"{p}.mlp.experts.gate_up_proj_blocks", None, _prep_fc1_weight)] - rows[f"l{i}_fc1_s"] = [(f"{p}.mlp.experts.gate_up_proj_scales", None, _prep_fc1_scale)] - rows[f"l{i}_fc1_b"] = [(f"{p}.mlp.experts.gate_up_proj_bias", None, _prep_fc1_bias)] - rows[f"l{i}_fc2_w"] = [(f"{p}.mlp.experts.down_proj_blocks", None, _prep_fc2_weight)] - rows[f"l{i}_fc2_s"] = [(f"{p}.mlp.experts.down_proj_scales", None, _prep_fc2_scale)] - rows[f"l{i}_fc2_b"] = [(f"{p}.mlp.experts.down_proj_bias", None, _prep_fc2_bias)] - rows["final_norm"] = [("model.norm.weight", None, None)] - rows["embed"] = [("model.embed_tokens.weight", None, None)] - return rows - - -def load(model, weights) -> None: - core = model.model - manifest = _manifest(core) - consumed: set = set() - - def fill(param: torch.nn.Parameter, ckpt_key: str, index, transform) -> None: - assert ckpt_key in weights, f"checkpoint key missing: {ckpt_key}" - src = weights[ckpt_key][:] # materialize the lazy slice - dst = param.data if index is None else param.data[index] - if transform is not None: - # The expert transforms are heavy row gathers over ~1 GB of - # blocks; run them where the destination lives. - src = transform(src.to(dst.device, non_blocking=True), core) - assert dst.shape == src.shape, (ckpt_key, tuple(dst.shape), tuple(src.shape)) - assert src.dtype == dst.dtype, (ckpt_key, src.dtype, dst.dtype) - dst.copy_(src, non_blocking=True) - consumed.add(ckpt_key) - - assert set(manifest.keys()) == set(core.w.keys()), ( - "manifest/parameter drift", - set(manifest.keys()) ^ set(core.w.keys()), + ), + W( + "sinks", + shape=lambda d: (d.heads_q,), + dtype=torch.float32, + src="{p}.self_attn.sinks", + transform=_to_fp32, + ), + W("o", shape=lambda d: (d.hidden, d.q_width), src="{p}.self_attn.o_proj.weight"), + W("o_bias", shape=lambda d: (d.hidden,), src="{p}.self_attn.o_proj.bias"), + W("norm2", shape=lambda d: (d.hidden,), src="{p}.post_attention_layernorm.weight"), + W("router", shape=lambda d: (d.num_experts, d.hidden), src="{p}.mlp.router.weight"), + W("router_bias", shape=lambda d: (d.num_experts,), src="{p}.mlp.router.bias"), + W( + "fc1_w", + shape=lambda d: (d.num_experts, d.fc1_rows, d.fc1_k_pad // 2), + dtype=torch.uint8, + src="{p}.mlp.experts.gate_up_proj_blocks", + transform=_prep_fc1_weight, + ), + W( + "fc1_s", + shape=lambda d: (d.num_experts, d.fc1_rows, d.fc1_k_pad // 32), + dtype=torch.uint8, + src="{p}.mlp.experts.gate_up_proj_scales", + transform=_prep_fc1_scale, + ), + W( + "fc1_b", + shape=lambda d: (d.num_experts, d.fc1_rows), + dtype=torch.float32, + src="{p}.mlp.experts.gate_up_proj_bias", + transform=_prep_fc1_bias, + ), + W( + "fc2_w", + shape=lambda d: (d.num_experts, d.fc2_rows_pad, d.inter_pad // 2), + dtype=torch.uint8, + src="{p}.mlp.experts.down_proj_blocks", + transform=_prep_fc2_weight, + ), + W( + "fc2_s", + shape=lambda d: (d.num_experts, d.fc2_rows_pad, d.inter_pad // 32), + dtype=torch.uint8, + src="{p}.mlp.experts.down_proj_scales", + transform=_prep_fc2_scale, + ), + W( + "fc2_b", + shape=lambda d: (d.num_experts, d.fc2_rows_pad), + dtype=torch.float32, + src="{p}.mlp.experts.down_proj_bias", + transform=_prep_fc2_bias, + ), + W("final_norm", shape=lambda d: (d.hidden,), src="model.norm.weight", layers=None), + W( + "embed", + shape=lambda d: (d.vocab, d.hidden), + src="model.embed_tokens.weight", + layers=None, + ), ) - for param_key, sources in manifest.items(): - for ckpt_key, index, transform in sources: - fill(core.w[param_key], ckpt_key, index, transform) - # The transforms allocate several GB of scratch per layer; release it - # before the next one so peak load memory stays one layer deep. - if param_key.endswith("_fc2_b"): - torch.cuda.empty_cache() - - # Shell-registered exception: the base class owns lm_head (untied). - fill(model.lm_head.weight, "lm_head.weight", None, None) - - torch.cuda.synchronize() - leftover = set(weights.keys()) - consumed - assert not leftover, f"unconsumed checkpoint keys: {sorted(leftover)[:8]}" + + RELEASE_AFTER = "_fc2_b" + + def dims(self, core) -> SimpleNamespace: + cfg = core.model_config.pretrained_config + inter_pad = _pad_up(cfg.intermediate_size, 128) + return SimpleNamespace( + num_layers=cfg.num_hidden_layers, + hidden=cfg.hidden_size, + q_width=cfg.num_attention_heads * cfg.head_dim, + kv_width=cfg.num_key_value_heads * cfg.head_dim, + heads_q=cfg.num_attention_heads, + num_experts=cfg.num_local_experts, + vocab=cfg.vocab_size, + fc1_rows=2 * inter_pad, + fc1_k_pad=_pad_up(cfg.hidden_size, FC1_K_ALIGN), + inter_pad=inter_pad, + fc2_rows_pad=_pad_up(cfg.hidden_size, 128), + dtype=core.model_config.torch_dtype, + ) + + +MODEL_WEIGHTS = GptOssWeights() diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 6a52657739df..c41144be691a 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -150,6 +150,14 @@ class AttentionMetadata: # The shape is (batch_size) if provided. prompt_lens: Optional[List[int]] = None + # Per context request, in batch order: whether this step's chunk is the + # last chunk of the request's context phase (``prompt_lens`` holds the + # chunk length, so a chunk alone cannot tell), and whether the request + # asked for context logits (prompt logprobs imply them). Set by the + # PyTorch decoder runner's input preparation; None where it is not set. + context_last_chunk: Optional[List[bool]] = None + context_return_logits: Optional[List[bool]] = None + # Explicit query/KV sequence boundaries for attention kernels that operate # on packed varlen context inputs. These are sequence/segment boundaries, # not necessarily request boundaries. diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py index 9879c0e40a4b..3b93d07fef43 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py @@ -14,9 +14,9 @@ # limitations under the License. from collections import defaultdict -from dataclasses import replace +from dataclasses import dataclass, replace from math import gcd -from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Tuple +from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Sequence, Tuple import torch @@ -47,6 +47,7 @@ from tensorrt_llm.runtime.kv_cache_manager_v2 import ( BAD_PAGE_INDEX, AttentionLayerConfig, + AttentionReusePolicy, BatchDesc, BufferConfig, BufferId, @@ -68,18 +69,137 @@ ) from .offload import SparseOffloadLayerDescriptor, SparseOffloadState from .params import ( + DEEPSEEK_V4_ATTENTION, DEEPSEEK_V4_NON_SLIDING_ATTENTION, DEEPSEEK_V4_SLIDING_ATTENTION, DEEPSEEK_V4_SPARSE_RATIO, DeepseekV4AttentionType, compress_ratio_has_attention, is_overlap_compressor, + is_sliding_attention_type, +) +from .v41_layout import ( + DEEPSEEK_V41_ATTENTION, + DEEPSEEK_V41_DECODER_REPLAY_ENV, + DEEPSEEK_V41_DEFAULT_ENCODER_REPLAY, + DEEPSEEK_V41_ENCODER_REPLAY_ENV, + DEEPSEEK_V41_KV_STORAGE_ENV, + DEEPSEEK_V41_KV_STORAGE_NATIVE, + DEEPSEEK_V41_NON_SLIDING_ATTENTION, + DEEPSEEK_V41_SLIDING_ATTENTION, + DEEPSEEK_V41_SLIDING_TABLE_COLUMNS, + DeepseekV41Layout, + deepseek_v41_native_entry_split, + deepseek_v41_token_bytes, + get_deepseek_v41_decoder_replay_setting, + get_deepseek_v41_encoder_replay_setting, + get_deepseek_v41_layout, ) if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata COMPRESS_BLOCK_SCALE_ROLE = DataRole("deepseek_v4_compress_block_scale") + + +def _v41_encoder_replay_refusal( + layout: DeepseekV41Layout, mapping: Mapping, max_num_tokens: Optional[int], kwargs: dict +) -> Optional[Tuple[type, str]]: + """Why encoder SWA bounded replay cannot run on this engine, as (error type, reason), + or None when it can. Attention TP and attention DP both run it: under attention DP + each rank's manager replays the hits of the requests that rank holds.""" + if kwargs.get("is_disagg", False): + return ValueError, ( + "disaggregated serving: a generation engine receives the context's windows and " + "has nothing to replay" + ) + if kwargs.get("joint_kv_cache_reuse", False): + return ValueError, "a separately reused draft KV pool" + if kwargs.get("kv_connector_manager") is not None: + return ValueError, ( + "a KV cache connector: a connector-served prefix brings no window to replay from" + ) + if max_num_tokens is not None and max_num_tokens <= layout.replay_window: + return ValueError, ( + f"max_num_tokens <= {layout.replay_window} (got {max_num_tokens}): the first chunk " + "after a hit holds every replay row" + ) + return None + + +_V41_ENCODER_REPLAY_ANNOUNCED = set() + + +def _announce_v41_encoder_replay_off(mapping: Mapping, reason: str) -> None: + """One line per process on rank 0 (on stdout: worker ranks log below the info level).""" + if mapping.rank != 0 or reason in _V41_ENCODER_REPLAY_ANNOUNCED: + return + _V41_ENCODER_REPLAY_ANNOUNCED.add(reason) + print( + f"[TRT-LLM] [I] {DEEPSEEK_V41_ENCODER_REPLAY_ENV} unset: DeepSeek-V4.1 encoder SWA " + f"bounded replay resolves off, it does not support {reason}; it runs aggregated " + "(tep4, dep4, tp4), without a KV connector", + flush=True, + ) + + +def _v41_decoder_replay_refusal( + layout: DeepseekV41Layout, enable_block_reuse: bool, encoder_replay_off: Optional[str] +) -> Optional[str]: + """Why decoder SWA bounded replay cannot run beside this engine's prefix cache, or None + when it can. ``encoder_replay_off`` is why encoder replay resolved off (None when on). + + Decoder replay leaves the sliding-window KV of the layers after the last KV-source + layer holding each prompt's last ``replay_window`` positions only. A prefix-cache hit + inside that tail is served correctly only when encoder replay re-forwards the tail's + cached rows; otherwise the request reads window KV that none of its forwards wrote. + """ + if enable_block_reuse and encoder_replay_off is not None: + return ( + f"block reuse while encoder SWA bounded replay is off ({encoder_replay_off}): a " + f"prefix-cache hit inside a prompt's last {layout.replay_window} positions would " + "read sliding-window KV that no forward of the request wrote" + ) + return None + + +_V41_DECODER_REPLAY_ANNOUNCED = set() + + +def _announce_v41_decoder_replay_off(mapping: Mapping, reason: str) -> None: + """One line per process on rank 0 (on stdout: worker ranks log below the info level).""" + if mapping.rank != 0 or reason in _V41_DECODER_REPLAY_ANNOUNCED: + return + _V41_DECODER_REPLAY_ANNOUNCED.add(reason) + print( + f"[TRT-LLM] [I] {DEEPSEEK_V41_DECODER_REPLAY_ENV} unset: DeepSeek-V4.1 decoder SWA " + f"bounded replay resolves off, it does not support {reason}; it runs with encoder " + "replay on or with block reuse off", + flush=True, + ) + + +@dataclass(frozen=True) +class DeepseekV41ReplayPlan: + """Replay rows of one request's first context chunk after a prefix-cache hit. + + The chunk starts at ``start`` = ``reused - num_tokens``: its first ``num_tokens`` + rows re-forward cached tokens (their main KV, index K and compressor state are + cached and stay as they are), the rest is the uncached suffix. + ``history_token_ids`` are the prompt ids of positions + ``start - len(history_token_ids) .. start - 1`` (the Engram history the first + replay rows hash, which no cached page holds). + """ + + start: int + num_tokens: int + history_token_ids: Tuple[int, ...] + + @property + def reused(self) -> int: + return self.start + self.num_tokens + + KV_CACHE_COPY_ALIGNMENT = 16 NVFP4_VECTOR_SIZE = 16 @@ -102,6 +222,8 @@ def get_attn_dim( return state_factor * index_head_dim if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE: return state_factor * index_head_dim + if attn_type == DeepseekV4AttentionType.ENGRAM_TOKEN: + return 1 raise ValueError(f"Unsupported DeepSeek-V4 attention type: {attn_type}") @@ -278,8 +400,55 @@ def _get_attn_bytes_per_token( return token_bytes +def _estimate_v41_swa_cache_size( + layout: DeepseekV41Layout, + layers: List[int], + tokens_per_block: int, + max_draft_len: int, + *, + context: bool, + scratch: bool, +) -> Tuple[int, int]: + """DeepSeek-V4.1 counterpart of ``_estimate_swa_cache_size`` over owned windowed roles. + + Scratch keys are (role, compress ratio), as in DeepSeek-V4. A window that + does not start on a block boundary spans one more block than its + rounded-up length, so each windowed layer-role is charged one extra block + per request. + """ + tokens_per_block = int(tokens_per_block) + size_per_token = 0 + size_per_request = 0 + scratch_keys = set() + for layer in layers: + for attn_type in DEEPSEEK_V41_SLIDING_ATTENTION: + if not layout.has_attention(layer, attn_type): + continue + window_size = layout.window(attn_type, max_draft_len) + if window_size is None or window_size <= 0: + continue + window_tokens = ( + (int(window_size) + tokens_per_block - 1) // tokens_per_block + 1 + ) * tokens_per_block + token_bytes = deepseek_v41_token_bytes( + attn_type, layout.head_dim, layout.index_head_dim, layout.kv_storage + ) + if not context: + size_per_request += window_tokens * token_bytes + elif not scratch: + size_per_token += token_bytes + else: + scratch_key = (attn_type, layout.compress_ratios[layer]) + if scratch_key in scratch_keys: + size_per_request += window_tokens * token_bytes + else: + scratch_keys.add(scratch_key) + size_per_token += token_bytes + return size_per_token, size_per_request + + def _get_index_mode(attn_type: DeepseekV4AttentionType) -> PageIndexMode: - if attn_type in DEEPSEEK_V4_SLIDING_ATTENTION: + if is_sliding_attention_type(attn_type): return PageIndexMode.PER_LAYER else: return PageIndexMode.SHARED @@ -299,6 +468,13 @@ class DeepseekV4CacheManager(KVCacheManagerV2): # The block size of the (indexer) compressed cache. # For other attention types, block size is tokens_per_block. compressed_block_sizes: List[int] + # DeepSeek-V4.1 layout mode, selected by the checkpoint config passed as + # ``pretrained_config``; None keeps the DeepSeek-V4 layout. + v41_layout: Optional[DeepseekV41Layout] = None + # Why encoder SWA bounded replay resolved off (None while it is on), and whether + # decoder SWA bounded replay runs on this engine; both resolved at construction. + v41_encoder_replay_off_reason: Optional[str] = None + v41_decoder_replay: bool = False def _get_typical_seq_len(self, kv_cache_config: KvCacheConfig) -> int: """Retain DeepSeek-V4's max-length pool-sizing model by default.""" @@ -347,6 +523,28 @@ def __init__( "and validated fetch addressing, scratch lifetime, and transfer ordering." ) + self.v41_layout = get_deepseek_v41_layout(kwargs.get("pretrained_config")) + if self.v41_layout is not None: + self._init_deepseek_v41( + kv_cache_config, + kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + max_batch_size=max_batch_size, + max_beam_width=max_beam_width, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + vocab_size=vocab_size, + mapping=mapping, + dtype=dtype, + compressor_dtype=compressor_dtype, + sparse_attn_config=sparse_attn_config, + max_input_len=max_input_len, + max_num_tokens=max_num_tokens, + **kwargs, + ) + return + # DeepSeek-V4 specific attributes initialization assert kv_cache_type == CacheTypeCpp.SELFKONLY, "DeepSeek-V4 only supports SELFKONLY" assert num_kv_heads == 1, "DeepSeek-V4 only supports num_kv_heads == 1" @@ -476,6 +674,342 @@ def __init__( PageIndexMode.SHARED, ) + def _init_deepseek_v41( + self, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: int, + max_batch_size: int, + max_beam_width: int, + tokens_per_block: int, + max_seq_len: int, + vocab_size: int, + mapping: Mapping, + dtype: DataType, + compressor_dtype: DataType, + sparse_attn_config: DeepSeekV4SparseAttentionConfig, + max_input_len: Optional[int], + max_num_tokens: Optional[int], + **kwargs, + ) -> None: + """Initialize the DeepSeek-V4.1 layout mode. + + Pools follow ``self.v41_layout``: SWA on every backbone layer; main KV + (``COMPRESS``) and native MXFP4 index K (``INDEXER_COMPRESS``) on the + KV-source layers only; the compressor pending pair on ratio-2 KV sources; + the Engram token history on the first Engram layer. SWA and main KV + hold the reference's fake-quantized values, stored as selected by + ``TRTLLM_DSV41_KV_STORAGE``: bf16 containers, or native uint8 entries + (``get_v41_native_kv_views``). + """ + layout = self.v41_layout + assert kv_cache_type == CacheTypeCpp.SELFKONLY, "DeepSeek-V4.1 only supports SELFKONLY" + assert num_kv_heads == 1, "DeepSeek-V4.1 only supports num_kv_heads == 1" + assert compressor_dtype == DataType.FLOAT, ( + f"Unsupported compressor dtype: {compressor_dtype}, only support FP32/TF32" + ) + if dtype != DataType.BF16: + raise ValueError( + "DeepSeek-V4.1 selects its SWA / main-KV storage with " + f"{DEEPSEEK_V41_KV_STORAGE_ENV} (bf16 containers or native), not with the KV " + f"cache dtype; got KV cache dtype {dtype}. Leave the KV cache quantization unset." + ) + if kv_cache_config.dtype == "fp8_ds_mla": + raise ValueError("DeepSeek-V4.1 does not use the fp8_ds_mla KV cache layout.") + if tokens_per_block != 128: + raise ValueError( + "DeepSeek-V4.1 requires tokens_per_block == 128: a ratio-1 index-K page holds " + "tokens_per_block entries and the paged indexer logits kernel accepts at most " + f"128. Got {tokens_per_block}." + ) + if num_layers > layout.num_layers: + raise ValueError( + f"DeepSeek-V4.1 cache manager got {num_layers} layers, but the checkpoint " + f"layout has {layout.num_layers} backbone layers." + ) + if sparse_attn_config.window_size != layout.window_size: + raise ValueError( + f"sparse_attention_config.window_size ({sparse_attn_config.window_size}) must " + f"equal the checkpoint's sliding_window ({layout.window_size})." + ) + if sparse_attn_config.index_head_dim != layout.index_head_dim: + raise ValueError( + f"sparse_attention_config.index_head_dim ({sparse_attn_config.index_head_dim}) " + f"must equal the checkpoint's index_head_dim ({layout.index_head_dim})." + ) + # Pipeline parallelism: a stage holds only its own layers' pools, so no + # layer may read a pool of another stage. + stage_error = layout.stage_dependency_error(mapping.pp_layers(layout.num_layers)) + if stage_error is not None: + raise ValueError( + f"DeepSeek-V4.1 pipeline stage {mapping.pp_rank}: {stage_error}. This " + f"checkpoint's stage boundaries can only be at layers " + f"{list(layout.pipeline_boundaries())} (set pp_partition)." + ) + + self.use_fp8_ds_mla = False + self.index_head_dim = layout.index_head_dim + # Per backbone layer, 0 / 1 / 2 exactly as in the checkpoint. + self._compress_ratios = list(layout.compress_ratios) + spec_config = kwargs.get("spec_config", None) + self._max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 + self._swa_window_size = layout.window_size + self._compressor_dtype = compressor_dtype + self._use_nvfp4_compress = False + self._nvfp4_residual_dim = 0 + self.compressed_block_sizes = [ + tokens_per_block // ratio if ratio > 0 else tokens_per_block + for ratio in self._compress_ratios + ] + self._init_indexer_dtype(sparse_attn_config) + if self._indexer_k_dtype != "fp4": + raise ValueError( + "DeepSeek-V4.1 stores index K natively as MXFP4 (indexer_k_dtype='fp4'); " + f"got {self._indexer_k_dtype!r}." + ) + self._max_input_len = max_input_len + self._max_num_tokens = max_num_tokens + self._v41_replay_plans: Dict[int, DeepseekV41ReplayPlan] = {} + layout = self._resolve_v41_encoder_replay(mapping, max_num_tokens, kwargs) + self.v41_decoder_replay = self._resolve_v41_decoder_replay( + mapping, kv_cache_config.enable_block_reuse + ) + + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + max_batch_size=max_batch_size, + max_beam_width=max_beam_width, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + vocab_size=vocab_size, + mapping=mapping, + dtype=dtype, + max_num_tokens=max_num_tokens, + **kwargs, + ) + self.is_vswa = True + self._assert_layer_pool_scale() + + first_pp_layer = self.pp_layers[0] + self.swa_pool_ptr = self.impl.get_mem_pool_base_address( + self._layer_attn_to_layer_id[first_pp_layer, DeepseekV4AttentionType.SWA], + DeepseekV4AttentionType.SWA.role, + PageIndexMode.PER_LAYER, + ) + # Base of the first local KV source of each ratio. Every KV source has + # its own base (``get_buffers(owner, COMPRESS)``); the per-ratio page + # table is shared by all owners of that ratio. + self.compress_pool_ptrs = {} + self.compress_scale_pool_ptrs = {} + for ratio in layout.compressed_ratios: + owner = next( + ( + layer + for layer in self.pp_layers + if layout.is_kv_source(layer) and self._compress_ratios[layer] == ratio + ), + None, + ) + if owner is None: + continue + self.compress_pool_ptrs[ratio] = self.impl.get_mem_pool_base_address( + self._layer_attn_to_layer_id[owner, DeepseekV4AttentionType.COMPRESS], + DeepseekV4AttentionType.COMPRESS.role, + PageIndexMode.SHARED, + ) + entry_bytes = { + attn_type.name: deepseek_v41_token_bytes( + attn_type, self.head_dim, self.index_head_dim, layout.kv_storage + ) + for attn_type in ( + DeepseekV4AttentionType.SWA, + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ) + } + logger.info( + f"DeepSeek-V4.1 KV storage {layout.kv_storage!r}: bytes per entry {entry_bytes}, " + f"persistent KV bytes per token {self.get_cache_bytes_per_token()}" + ) + if layout.encoder_replay: + logger.info( + "DeepSeek-V4.1 encoder SWA bounded replay on: the prefix cache holds main KV " + "and index K only; a hit of r tokens (r a multiple of " + f"{layout.reuse_alignment}) re-forwards the last min({layout.replay_window}, r) " + f"cached tokens (block reuse {'on' if self.enable_block_reuse else 'off'})" + ) + + def _resolve_v41_encoder_replay( + self, mapping: Mapping, max_num_tokens: Optional[int], kwargs: dict + ) -> DeepseekV41Layout: + """Resolve encoder SWA bounded replay against this engine and return the layout. + + An explicit ``TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY=1`` on a configuration + replay does not support raises; an explicit 0 turns it off. Unset, it is + ``DEEPSEEK_V41_DEFAULT_ENCODER_REPLAY`` where supported (aggregated, no KV + connector; attention TP or DP) and off otherwise, which rank 0 states once. + """ + layout = self.v41_layout + setting = get_deepseek_v41_encoder_replay_setting() + refusal = _v41_encoder_replay_refusal(layout, mapping, max_num_tokens, kwargs) + if setting is None: + enabled = DEEPSEEK_V41_DEFAULT_ENCODER_REPLAY and refusal is None + if DEEPSEEK_V41_DEFAULT_ENCODER_REPLAY and refusal is not None: + _announce_v41_encoder_replay_off(mapping, refusal[1]) + off_reason = refusal[1].split(": ", 1)[0] + else: + off_reason = None if enabled else f"{DEEPSEEK_V41_ENCODER_REPLAY_ENV} unset" + elif setting: + if refusal is not None: + error, reason = refusal + raise error( + f"{DEEPSEEK_V41_ENCODER_REPLAY_ENV}=1 (DeepSeek-V4.1 encoder SWA bounded " + f"replay) does not support {reason}. Set {DEEPSEEK_V41_ENCODER_REPLAY_ENV}=0." + ) + enabled, off_reason = True, None + else: + enabled, off_reason = False, f"{DEEPSEEK_V41_ENCODER_REPLAY_ENV}=0" + self.v41_encoder_replay_off_reason = off_reason + if enabled != layout.encoder_replay: + layout = replace(layout, encoder_replay=enabled) + self.v41_layout = layout + return layout + + def _resolve_v41_decoder_replay(self, mapping: Mapping, enable_block_reuse: bool) -> bool: + """Resolve decoder SWA bounded replay against this engine (ruling R21). + + The model runs it; the manager decides where it may run, after encoder replay: + with block reuse on it needs encoder replay (``_v41_decoder_replay_refusal``). + Unset, it is on (attention TP or DP) unless that refusal applies, which rank 0 + states once; an explicit ``TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY=1`` there + raises. An explicit 0 is always off. + """ + setting = get_deepseek_v41_decoder_replay_setting() + if setting is False: + return False + refusal = _v41_decoder_replay_refusal( + self.v41_layout, enable_block_reuse, self.v41_encoder_replay_off_reason + ) + if refusal is None: + return True + if setting is None: + _announce_v41_decoder_replay_off(mapping, refusal) + return False + raise ValueError( + f"{DEEPSEEK_V41_DECODER_REPLAY_ENV}=1 (DeepSeek-V4.1 decoder SWA bounded replay) " + f"does not support {refusal}. Set {DEEPSEEK_V41_DECODER_REPLAY_ENV}=0, or run it " + "with encoder replay on or with block reuse off (kv_cache_config.enable_block_reuse)." + ) + + def _v41_has_local_source(self, compress_ratio: int) -> bool: + """Whether this pipeline stage holds a KV source of ``compress_ratio``.""" + layout = self.v41_layout + return any( + layout.is_kv_source(layer) and layout.compress_ratios[layer] == compress_ratio + for layer in self.pp_layers + ) + + def _is_v41_native_kv(self, attn_type: DeepseekV4AttentionType) -> bool: + """Whether ``attn_type``'s pool holds native uint8 entries (V4.1 native storage).""" + return ( + self.v41_layout is not None + and self.v41_layout.kv_storage == DEEPSEEK_V41_KV_STORAGE_NATIVE + and attn_type in (DeepseekV4AttentionType.SWA, DeepseekV4AttentionType.COMPRESS) + ) + + @property + def num_sliding_attention_columns(self) -> int: + """Width of the per-layer sliding block table; column = ``attn_type.value``.""" + if self.v41_layout is not None: + return DEEPSEEK_V41_SLIDING_TABLE_COLUMNS + return len(DEEPSEEK_V4_SLIDING_ATTENTION) + + @property + def _sliding_attention_types(self) -> Tuple[DeepseekV4AttentionType, ...]: + if self.v41_layout is not None: + return DEEPSEEK_V41_SLIDING_ATTENTION + return DEEPSEEK_V4_SLIDING_ATTENTION + + @property + def _attention_types(self) -> Tuple[DeepseekV4AttentionType, ...]: + if self.v41_layout is not None: + return DEEPSEEK_V41_ATTENTION + return DEEPSEEK_V4_ATTENTION + + @property + def disagg_num_attn_types(self) -> int: + """Width of the (model layer, role) encoding of the disaggregated synthetic layer ids. + + ``kv_extractor`` encodes a manager layer as ``model_layer * width + role.value``. + The width counts the roles this manager's layout can allocate, so it is the + same on every pipeline stage: 7 for DeepSeek-V4, whose roles end before + ``ENGRAM_TOKEN``, and 8 in the DeepSeek-V4.1 layout mode, which owns an + ``ENGRAM_TOKEN`` pool. + """ + return max(attn_type.value for attn_type in self._attention_types) + 1 + + def _layer_has_attention(self, layer_idx: int, attn_type: DeepseekV4AttentionType) -> bool: + """Whether ``layer_idx`` owns a pool of ``attn_type``.""" + if self.v41_layout is not None: + return self.v41_layout.has_attention(layer_idx, attn_type) + return compress_ratio_has_attention(self._compress_ratios[layer_idx], attn_type) + + def _compress_table_ratios(self) -> List[int]: + """Ratios that have a shared COMPRESS page table.""" + if self.v41_layout is not None: + return list(self.v41_layout.compressed_ratios) + return [ + compress_ratio + for compress_ratio in set(self._compress_ratios) + if compress_ratio_has_attention(compress_ratio, DeepseekV4AttentionType.COMPRESS) + ] + + def _non_sliding_bytes_per_token(self) -> int: + """Persistent (non-windowed) cache bytes per token of the local layers.""" + if self.v41_layout is not None: + return self.v41_layout.persistent_bytes_per_token(self.pp_layers) + compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] + return _estimate_non_sliding_attn_size_per_token( + self.head_dim, + self.index_head_dim, + compress_ratios, + self.dtype == DataType.FP8, + indexer_k_dtype=self._indexer_k_dtype, + use_fp8_ds_mla=self.use_fp8_ds_mla, + has_nvfp4_compress=self._use_nvfp4_compress, + nvfp4_residual_dim=self.nvfp4_residual_dim, + ) + + def _swa_cache_size(self, *, context: bool, scratch: bool) -> Tuple[int, int]: + """Windowed-cache (bytes per token, bytes per request) of the local layers.""" + if self.v41_layout is not None: + return _estimate_v41_swa_cache_size( + self.v41_layout, + list(self.pp_layers), + self.tokens_per_block, + self._max_draft_len, + context=context, + scratch=scratch, + ) + compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] + return _estimate_swa_cache_size( + self.head_dim, + self.index_head_dim, + compress_ratios, + self.dtype == DataType.FP8, + self.tokens_per_block, + self._swa_window_size, + context=context, + scratch=scratch, + indexer_k_dtype=self._indexer_k_dtype, + use_fp8_ds_mla=self.use_fp8_ds_mla, + ) + @property def nvfp4_residual_dim(self) -> int: """Residual dimensions fixed when the NVFP4 cache is constructed.""" @@ -543,6 +1077,8 @@ def get_buffers(self, layer_idx: int, attn_type: DeepseekV4AttentionType) -> tor Returns: The buffer tensor (shape: [num_blocks, tokens_per_block, attn_dim]) For blockwise FP8 layers, shape is [num_blocks, tokens_per_block, attn_dim + scale_size] + For DeepSeek-V4.1 native SWA / main KV, uint8 [num_blocks, block, entry bytes] + (528 / 288); see ``get_v41_native_kv_views``. """ layer_id = self._layer_attn_to_layer_id[(layer_idx, attn_type)] data_role = attn_type.role @@ -563,9 +1099,13 @@ def get_buffers(self, layer_idx: int, attn_type: DeepseekV4AttentionType) -> tor DeepseekV4AttentionType.SWA, DeepseekV4AttentionType.COMPRESS, ) + v41_native = self._is_v41_native_kv(attn_type) if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: # Indexer always pack data + per-block scales into the same row. dim_per_token = self._indexer_data_size + self._indexer_scale_size + elif v41_native: + # One contiguous row per entry: codes, then block scales. + dim_per_token = sum(deepseek_v41_native_entry_split(attn_type, self.head_dim)) elif footer_scale: # Footer-scale pages are byte-packed rather than token-major. from . import footer_scale_kv @@ -595,13 +1135,43 @@ def get_buffers(self, layer_idx: int, attn_type: DeepseekV4AttentionType) -> tor dtype = self._compressor_dtype elif attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: dtype = self._indexer_dtype - elif footer_scale or ( - attn_type == DeepseekV4AttentionType.COMPRESS and self._use_nvfp4_compress + elif attn_type == DeepseekV4AttentionType.ENGRAM_TOKEN: + dtype = DataType.INT32 + elif ( + footer_scale + or v41_native + or (attn_type == DeepseekV4AttentionType.COMPRESS and self._use_nvfp4_compress) ): dtype = DataType.UINT8 return convert_to_torch_tensor(TensorWrapper(addr, dtype, shape)) + def get_v41_native_kv_views( + self, layer_idx: int, attn_type: DeepseekV4AttentionType + ) -> Tuple[torch.Tensor, torch.Tensor]: + """(codes, scales) views of a DeepSeek-V4.1 native SWA or main-KV pool. + + Both are uint8 views of ``get_buffers(layer_idx, attn_type)``, page + geometry ``[num_blocks, block]`` and slot numbering unchanged: + + - SWA: codes ``[..., 512]`` E4M3, scales ``[..., 16]`` UE8M0 (one per 32); + - COMPRESS: codes ``[..., 256]`` packed E2M1 (even element in the low + nibble), scales ``[..., 32]`` E4M3 (one per 16, no global scale). + + These are the ``out`` / ``out_scale`` of the native writers + (``mxfp8_fake_quantize`` / ``nvfp4_fake_quantize``) and the + ``codes`` / ``scales`` (``kv_cache`` / ``kv_cache_sf`` after + ``unsqueeze(2)``) of the dequantizing readers. + """ + if not self._is_v41_native_kv(attn_type): + raise RuntimeError( + f"{attn_type.name} of layer {layer_idx} is not a DeepSeek-V4.1 native pool " + f"(set {DEEPSEEK_V41_KV_STORAGE_ENV}={DEEPSEEK_V41_KV_STORAGE_NATIVE})." + ) + code_bytes, _ = deepseek_v41_native_entry_split(attn_type, self.head_dim) + buffer = self.get_buffers(layer_idx, attn_type) + return buffer[..., :code_bytes], buffer[..., code_bytes:] + def get_compress_scale_buffers(self, layer_idx: int) -> torch.Tensor: """Return E4M3 per-16 scales for an NVFP4 COMPRESS layer.""" if not self._use_nvfp4_compress: @@ -717,34 +1287,39 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: pin_memory=prefer_pinned(), device="cpu", ) - for compress_ratio in set(self._compress_ratios) - if compress_ratio_has_attention(compress_ratio, DeepseekV4AttentionType.COMPRESS) + for compress_ratio in self._compress_table_ratios() } + # DeepSeek-V4.1: (role, ratio) -> (pool id, page-index scale) of the page + # table shared by the KV sources of that ratio, for COMPRESS and + # INDEXER_COMPRESS. + self._v41_shared_table_params: Dict[ + Tuple[DeepseekV4AttentionType, int], Tuple[int, int] + ] = {} - # layer offsets per layer and attn, shape [num_local_layers, len(DEEPSEEK_V4_SLIDING_ATTENTION)]. + # layer offsets per layer and attn, shape [num_local_layers, num_sliding_attention_columns]. self._layer_offsets = torch.full( - (self.num_local_layers, len(DEEPSEEK_V4_SLIDING_ATTENTION)), + (self.num_local_layers, self.num_sliding_attention_columns), -1, dtype=torch.int32, device="cpu", ) # Pool ids per layer and sliding attention type, shape [num_local_layers, num_sliding_attention_types]. self._layer_attn_pool_ids = torch.full( - (self.num_local_layers, len(DEEPSEEK_V4_SLIDING_ATTENTION)), + (self.num_local_layers, self.num_sliding_attention_columns), -1, dtype=torch.int32, device="cpu", ) # Scales per layer and sliding attention type, shape [num_local_layers, num_sliding_attention_types]. self._layer_attn_scales = torch.ones( - (self.num_local_layers, len(DEEPSEEK_V4_SLIDING_ATTENTION)), + (self.num_local_layers, self.num_sliding_attention_columns), dtype=torch.int32, device="cpu", ) # Scratch pages per block per layer and sliding attention type, shape # [num_local_layers, num_sliding_attention_types]. self._scratch_pages = torch.zeros( - (self.num_local_layers, len(DEEPSEEK_V4_SLIDING_ATTENTION)), + (self.num_local_layers, self.num_sliding_attention_columns), dtype=torch.int32, device="cpu", ) @@ -786,9 +1361,20 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: self._hca_compress_pool_id = self.layer_to_pool_mapping_dict[compress_layer_id] self._hca_compress_scale = int(compress_converter.scale) + if self.v41_layout is not None and self.v41_layout.is_kv_source(layer_idx): + for attn_type in DEEPSEEK_V41_NON_SLIDING_ATTENTION: + shared_layer_id = self._layer_attn_to_layer_id[layer_idx, attn_type] + shared_converter = self.impl.get_page_index_converter( + shared_layer_id, attn_type.role + ) + self._v41_shared_table_params[attn_type, compress_ratio] = ( + self.layer_to_pool_mapping_dict[shared_layer_id], + int(shared_converter.scale), + ) + local_layer_idx = self.layer_offsets[layer_idx] - for attn_type in DEEPSEEK_V4_SLIDING_ATTENTION: - if not compress_ratio_has_attention(self._compress_ratios[layer_idx], attn_type): + for attn_type in self._sliding_attention_types: + if not self._layer_has_attention(layer_idx, attn_type): continue layer_id = self._layer_attn_to_layer_id[layer_idx, attn_type] pool_id = self.layer_to_pool_mapping_dict[layer_id] @@ -808,7 +1394,7 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: self._precomputed_sliding_block_tables = torch.empty( ( self.num_local_layers, - len(DEEPSEEK_V4_SLIDING_ATTENTION), + self.num_sliding_attention_columns, self.host_kv_cache_block_offsets.size(1), self.max_blocks_per_seq, ), @@ -946,51 +1532,18 @@ def get_cache_indices( def _get_extra_quota_padding(self) -> int: """Ensure each attention type has minimal space when max_tokens is small.""" - return len(DeepseekV4AttentionType) * (2 << 20) + return len(self._attention_types) * (2 << 20) def _get_quota_from_max_tokens(self, max_tokens: int) -> int: - compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] - has_fp8_kv_cache = self.dtype == DataType.FP8 - non_sliding_attn_size_per_token = _estimate_non_sliding_attn_size_per_token( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - has_nvfp4_compress=self._use_nvfp4_compress, - nvfp4_residual_dim=self.nvfp4_residual_dim, - ) + non_sliding_attn_size_per_token = self._non_sliding_bytes_per_token() ( context_swa_size_per_token, _, - ) = _estimate_swa_cache_size( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - self.tokens_per_block, - self._swa_window_size, - context=True, - scratch=self.enable_swa_scratch_reuse, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - ) + ) = self._swa_cache_size(context=True, scratch=self.enable_swa_scratch_reuse) ( generation_swa_size_per_token, generation_swa_size_per_request, - ) = _estimate_swa_cache_size( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - self.tokens_per_block, - self._swa_window_size, - context=False, - scratch=False, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - ) + ) = self._swa_cache_size(context=False, scratch=False) max_context_tokens = ( self._max_num_tokens if self._max_num_tokens is not None else max_tokens ) @@ -1006,45 +1559,14 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: return int(generation_quota + context_extra_quota + padding) def _get_max_tokens_from_quota(self, quota: int) -> float: - compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] - has_fp8_kv_cache = self.dtype == DataType.FP8 - non_sliding_attn_size_per_token = _estimate_non_sliding_attn_size_per_token( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - has_nvfp4_compress=self._use_nvfp4_compress, - nvfp4_residual_dim=self.nvfp4_residual_dim, - ) - context_swa_size_per_token, _ = _estimate_swa_cache_size( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - self.tokens_per_block, - self._swa_window_size, - context=True, - scratch=self.enable_swa_scratch_reuse, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, + non_sliding_attn_size_per_token = self._non_sliding_bytes_per_token() + context_swa_size_per_token, _ = self._swa_cache_size( + context=True, scratch=self.enable_swa_scratch_reuse ) ( generation_swa_size_per_token, generation_swa_size_per_request, - ) = _estimate_swa_cache_size( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - self.tokens_per_block, - self._swa_window_size, - context=False, - scratch=False, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - ) + ) = self._swa_cache_size(context=False, scratch=False) padding = self._get_extra_quota_padding() size_per_batch = self.max_batch_size * generation_swa_size_per_request + padding if quota < size_per_batch: @@ -1066,6 +1588,8 @@ def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerC """ Add DeepSeek-V4 layers and warmup constraints to the cache config. """ + if self.v41_layout is not None: + return self._build_v41_cache_config(config) layers: List[AttentionLayerConfig] = [] layer_attn_to_layer_id: Dict[Tuple[int, DeepseekV4AttentionType], LayerId] = {} manager_layer_id_to_layer_attn: Dict[ @@ -1215,6 +1739,20 @@ def _add_layer( # number of layers in the KVCacheManagerPy self._num_manager_layers = len(layers) + return replace( + config, + layers=layers, + constraints=self._decode_capacity_constraints(config), + ) + + def _decode_capacity_constraints(self, config: KVCacheManagerConfigPy) -> List[BatchDesc]: + """``config.constraints`` plus the longest decode request beside short ones. + + The windowed and compressed pools must also hold one request at + ``max_seq_len`` alongside minimal decode requests filling the resident + capacity. Added only while the pool split is derived + (``kv_cache_config.pool_ratio`` unset); used by both layouts. + """ constraints = list(config.constraints) # _build_base_config copies kv_cache_config.pool_ratio to initial_pool_ratio. if config.initial_pool_ratio is None: @@ -1242,10 +1780,104 @@ def _add_layer( * (constraint_batch_size - 1) ) ) + return constraints + + def _build_v41_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: + """Add the DeepSeek-V4.1 owner-only pools to the cache config. + + One manager layer per (model layer, lifecycle): SWA on every layer; the + persistent COMPRESS + INDEXER_COMPRESS pair on KV sources; the windowed + compressor pending pair on ratio-2 KV sources; the windowed Engram token + history on its owner. Reader layers get no compressed entry here, so + every ``(layer, role)`` key names exactly one owner. + """ + layout = self.v41_layout + layers: List[AttentionLayerConfig] = [] + layer_attn_to_layer_id: Dict[Tuple[int, DeepseekV4AttentionType], LayerId] = {} + manager_layer_id_to_layer_attn: Dict[ + Tuple[LayerId, DataRole], Tuple[int, DeepseekV4AttentionType] + ] = {} + + def _add_layer( + layer_idx: int, + attention_types: List[DeepseekV4AttentionType], + sliding_window_size: int | None, + ) -> None: + layer_id = LayerId(len(layers)) + for attn_type in attention_types: + layer_attn_to_layer_id[layer_idx, attn_type] = layer_id + manager_layer_id_to_layer_attn[layer_id, attn_type.role] = (layer_idx, attn_type) + # Encoder SWA bounded replay: the windowed roles (SWA, compressor state, + # Engram history) stay private to each request and out of the prefix + # cache; a reuse hit gets fresh pages for its live window, which the + # replay rows rebuild. Only main KV and index K are shared. + reuse_policy = ( + AttentionReusePolicy.PRIVATE + if layout.encoder_replay and sliding_window_size is not None + else AttentionReusePolicy.REQUIRED + ) + layers.append( + AttentionLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig( + role=attn_type.role, + size=self._get_attn_bytes_per_block(attn_type, layer_idx), + ) + for attn_type in attention_types + ], + sliding_window_size=sliding_window_size, + num_sink_tokens=None, + reuse_policy=reuse_policy, + ) + ) + + draft_len = self._max_draft_len + for layer in self.pp_layers: + if not 0 <= layer < layout.num_layers: + raise NotImplementedError( + f"Layer {layer} is outside the {layout.num_layers} DeepSeek-V4.1 backbone " + "layers; speculative layers appended to the target cache are not supported " + "(DSpark keeps its stage windows in the drafter)." + ) + _add_layer( + layer, + [DeepseekV4AttentionType.SWA], + layout.window(DeepseekV4AttentionType.SWA, draft_len), + ) + if layout.has_attention(layer, DeepseekV4AttentionType.COMPRESS): + _add_layer( + layer, + [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ], + None, + ) + if layout.has_attention(layer, DeepseekV4AttentionType.COMPRESSOR_KV): + _add_layer( + layer, + [ + DeepseekV4AttentionType.COMPRESSOR_KV, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + ], + layout.window(DeepseekV4AttentionType.COMPRESSOR_KV, draft_len), + ) + if layout.has_attention(layer, DeepseekV4AttentionType.ENGRAM_TOKEN): + _add_layer( + layer, + [DeepseekV4AttentionType.ENGRAM_TOKEN], + layout.window(DeepseekV4AttentionType.ENGRAM_TOKEN, draft_len), + ) + + self._layer_attn_to_layer_id = layer_attn_to_layer_id + self._manager_layer_id_to_layer_attn = manager_layer_id_to_layer_attn + self._disagg_ignored_roles = frozenset() + self._num_manager_layers = len(layers) return replace( config, layers=layers, - constraints=constraints, + constraints=self._decode_capacity_constraints(config), ) def _init_indexer_dtype(self, sparse_attn_config: DeepSeekV4SparseAttentionConfig) -> None: @@ -1299,7 +1931,7 @@ def _assert_layer_pool_scale(self) -> None: (attn_type, layer_idx) for attn_type in DeepseekV4AttentionType for layer_idx in self.pp_layers - if compress_ratio_has_attention(self._compress_ratios[layer_idx], attn_type) + if self._layer_has_attention(layer_idx, attn_type) ] for attn_type, layer_idx in comb: compress_ratio = self._compress_ratios[layer_idx] @@ -1359,6 +1991,13 @@ def _get_attn_bytes_per_block( """ Get the cache bytes per block for a specific attention type and layer. """ + if self.v41_layout is not None: + block_size = self.tokens_per_block + if attn_type in DEEPSEEK_V41_NON_SLIDING_ATTENTION: + block_size = self.compressed_block_sizes[layer_idx] + return block_size * deepseek_v41_token_bytes( + attn_type, self.head_dim, self.index_head_dim, self.v41_layout.kv_storage + ) has_fp8_kv_cache = self.dtype == DataType.FP8 token_bytes = get_token_bytes( self.head_dim, @@ -1386,18 +2025,7 @@ def _get_attn_bytes_per_block( def get_cache_bytes_per_token(self) -> int: """Get the average cache bytes per token for DeepSeek-V4.""" - has_fp8_kv_cache = self.dtype == DataType.FP8 - compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] - return _estimate_non_sliding_attn_size_per_token( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - has_nvfp4_compress=self._use_nvfp4_compress, - nvfp4_residual_dim=self.nvfp4_residual_dim, - ) + return self._non_sliding_bytes_per_token() def get_max_resource_count(self) -> int: # Keep scheduler capacity tied to physical GPU KV quota in bytes. @@ -1432,29 +2060,9 @@ def _get_generation_bytes(self, request: llm_request.LlmRequest) -> int: return self._get_cache_bytes_for_tokens(total_tokens, context=False) def _get_cache_bytes_for_tokens(self, total_tokens: int, *, context: bool) -> int: - has_fp8_kv_cache = self.dtype == DataType.FP8 - compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] - non_sliding_attn_size_per_token = _estimate_non_sliding_attn_size_per_token( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - has_nvfp4_compress=self._use_nvfp4_compress, - nvfp4_residual_dim=self.nvfp4_residual_dim, - ) - swa_size_per_token, swa_size_per_request = _estimate_swa_cache_size( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - self.tokens_per_block, - self._swa_window_size, - context=context, - scratch=self.enable_swa_scratch_reuse, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, + non_sliding_attn_size_per_token = self._non_sliding_bytes_per_token() + swa_size_per_token, swa_size_per_request = self._swa_cache_size( + context=context, scratch=self.enable_swa_scratch_reuse ) return int( total_tokens * (non_sliding_attn_size_per_token + swa_size_per_token) @@ -1481,6 +2089,138 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: """Exclude storage-only NVFP4 alignment pages from KV transfer.""" return self._disagg_ignored_roles + # ------------------------------------------------------------------ + # DeepSeek-V4.1 encoder SWA bounded replay: scheduling of replay rows. + # ------------------------------------------------------------------ + + @property + def v41_encoder_replay(self) -> bool: + return self.v41_layout is not None and self.v41_layout.encoder_replay + + def validate_request(self, request: llm_request.LlmRequest) -> None: + """Refuse, at admission, what DeepSeek-V4.1 decoder SWA bounded replay cannot serve. + + Under decoder replay the layers after the last KV-source layer run over each + prompt's last ``replay_window`` positions only, so context logits, and the prompt + logprobs computed from them, do not exist. The model refuses them inside the + forward, which fails the whole executor step; here only the request fails. A + generation-only request (disaggregated serving) ran its context elsewhere. + + Nothing here depends on the prompt length, so a request paused for recompute, + whose prompt then holds its generated tokens too, needs no second check. + """ + super().validate_request(request) + if not self.v41_decoder_replay or request.is_generation_only_request: + return + if getattr(request, "py_return_context_logits", False): + raise ValueError( + "DeepSeek-V4.1 decoder SWA bounded replay runs the layers after the last " + f"KV-source layer over each prompt's last {self.v41_layout.replay_window} " + "positions only, so return_context_logits / prompt_logprobs do not exist. " + "Drop these options from the request, or serve them with " + f"{DEEPSEEK_V41_DECODER_REPLAY_ENV}=0." + ) + + @staticmethod + def _v41_replayable(request: llm_request.LlmRequest) -> bool: + """Whether a request's cached rows can be re-forwarded from their token ids alone. + + Multimodal rows need their embeddings, and context logits would include the + replay rows' approximate logits: such requests take no prefix reuse instead. + """ + return request.multimodal_hashes is None and not getattr( + request, "py_return_context_logits", False + ) + + def prepare_context_cache( + self, req: llm_request.LlmRequest, reuse_limit: int | None = None + ) -> int | None: + """Claim the prefix cache; under encoder replay, return the replay start. + + Encoder SWA bounded replay keeps no sliding-window state in the prefix cache, so + a hit of ``reused`` tokens (a multiple of ``reuse_alignment``) re-forwards the + last ``min(replay_window, reused)`` of them. Returning the replay start makes the + caller settle the context cursor there, so the replay rows are scheduled at the + head of the first chunk and charged to the token budget like any context row. + """ + if not self.v41_encoder_replay or not self.enable_block_reuse: + return super().prepare_context_cache(req, reuse_limit) + if not req.is_first_context_chunk: + return super().prepare_context_cache(req, reuse_limit) + layout = self.v41_layout + align = layout.reuse_alignment + if not self._v41_replayable(req): + reuse_limit = 0 + elif req.py_request_id not in self.kv_cache_map: + probed = self.probe_context_reuse(req) + if probed is not None: + aligned = probed - probed % align + reuse_limit = aligned if reuse_limit is None else min(reuse_limit, aligned) + reused = super().prepare_context_cache(req, reuse_limit) + if reused is not None and reused % align: + # A partially covered page can end the claimed match off the alignment. + # Claim again below it; the tree still covers that shorter prefix. + for limit in (reused - reused % align, 0): + self.free_resources(req) + reused = super().prepare_context_cache(req, limit) + if reused is None or reused % align == 0: + break + if reused is not None and reused % align: + raise RuntimeError( + f"request {req.py_request_id}: no reuse boundary aligned to {align} " + f"(claimed {reused} tokens)" + ) + if reused is None: + return None + num_replay = min(layout.replay_window, reused) + start = reused - num_replay + history = layout.engram_history_len if layout.engram_token_owner is not None else 0 + history_ids = ( + tuple(req.get_tokens_range(0, max(0, start - history), start)) if start > 0 else () + ) + self._v41_replay_plans[req.py_request_id] = DeepseekV41ReplayPlan( + start, num_replay, history_ids + ) + return start + + def resize_context(self, req: llm_request.LlmRequest, num_tokens: int) -> bool: + plan = self._v41_replay_plans.get(req.py_request_id) if self.v41_encoder_replay else None + if ( + plan is not None + and req.context_current_position == plan.start + and num_tokens < plan.num_tokens + ): + # Every replay row belongs to the first chunk: a later chunk would start + # inside the cached prefix, below the cache's history. + return False + return super().resize_context(req, num_tokens) + + def free_resources(self, request: llm_request.LlmRequest, pin_on_release: bool = False): + if self.v41_layout is not None: + self._v41_replay_plans.pop(request.py_request_id, None) + return super().free_resources(request, pin_on_release) + + def get_v41_replay_plan( + self, request_id: int, context_start: int + ) -> Optional[DeepseekV41ReplayPlan]: + """The replay plan of a context chunk starting at ``context_start``, if it has one. + + Only the chunk starting at the replay start carries replay rows; later chunks + of the same request are ordinary. + """ + plan = self._v41_replay_plans.get(request_id) + if plan is None or plan.num_tokens == 0 or plan.start != context_start: + return None + return plan + + def get_v41_replay_plans( + self, request_ids: Sequence[int], context_starts: Sequence[int] + ) -> List[Optional[DeepseekV41ReplayPlan]]: + return [ + self.get_v41_replay_plan(request_id, int(start)) + for request_id, start in zip(request_ids, context_starts) + ] + def get_indexer_k_cache_buffers(self, layer_idx: int) -> torch.Tensor: """ Get the buffers for the indexer k cache for a specific layer. @@ -1976,7 +2716,16 @@ def copy_batch_compress_block_tables( assert beam_width == 1, "DSV4 only supports beam width 1 now" copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) staging = self._host_compress_block_tables_staging[compress_ratio] - if compress_ratio == 4: + if self.v41_layout is not None: + pool_id, scale = self._v41_shared_table_params.get( + (DeepseekV4AttentionType.COMPRESS, compress_ratio), (None, None) + ) + if pool_id is None and not self._v41_has_local_source(compress_ratio): + # A pipeline stage without a KV source of this ratio: no layer of + # the stage reads the table. + dst_tensor[:num_seqs].fill_(BAD_PAGE_INDEX) + return + elif compress_ratio == 4: pool_id = self._csa_compress_pool_id scale = self._csa_compress_scale elif compress_ratio == 128: @@ -2002,10 +2751,28 @@ def copy_batch_indexer_compress_block_tables( beam_width: int, num_contexts: int, num_seqs: int, + compress_ratio: Optional[int] = None, ) -> None: - """Build the shared INDEXER_COMPRESS compatibility block table.""" + """Build the shared INDEXER_COMPRESS compatibility block table. + + ``compress_ratio`` selects the per-ratio table in the DeepSeek-V4.1 + layout and is ignored by DeepSeek-V4, which has one indexer ratio. + """ assert beam_width == 1, "DSV4 only supports beam width 1 now" copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) + if self.v41_layout is not None: + pool_id, scale = self._v41_shared_table_params.get( + (DeepseekV4AttentionType.INDEXER_COMPRESS, compress_ratio), (None, None) + ) + if pool_id is None and not self._v41_has_local_source(compress_ratio): + host_block_table[:num_seqs] = BAD_PAGE_INDEX + return + if pool_id is None or scale is None: + raise RuntimeError( + f"Missing INDEXER_COMPRESS pool metadata for compress ratio {compress_ratio}" + ) + host_block_table[:num_seqs] = self._compute_shared_block_table(pool_id, scale, copy_idx) + return pool_id = self._csa_indexer_compress_pool_id scale = self._csa_indexer_compress_scale if pool_id is None or scale is None: @@ -2015,6 +2782,33 @@ def copy_batch_indexer_compress_block_tables( @staticmethod def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping, **kwargs): config = model_config.pretrained_config + v41_layout = get_deepseek_v41_layout(config) + if v41_layout is not None: + if ( + v41_layout.encoder_replay + and get_deepseek_v41_encoder_replay_setting() is None + and _v41_encoder_replay_refusal( + v41_layout, mapping, kwargs.get("max_num_tokens"), kwargs + ) + is not None + ): + v41_layout = replace(v41_layout, encoder_replay=False) + pp_layers = mapping.pp_layers(v41_layout.num_layers) + spec_config = getattr(model_config, "spec_config", None) + max_draft_len = getattr(spec_config, "max_draft_len", 0) or 0 + swa_size_per_token, swa_size_per_request = _estimate_v41_swa_cache_size( + v41_layout, + pp_layers, + kwargs["tokens_per_block"], + max_draft_len, + context=False, + scratch=False, + ) + max_batch_size = int(kwargs.get("max_batch_size") or 0) + return ( + v41_layout.persistent_bytes_per_token(pp_layers) + swa_size_per_token, + swa_size_per_request * max_batch_size, + ) head_dim = config.kv_lora_rank + config.qk_rope_head_dim index_head_dim = model_config.sparse_attention_config.index_head_dim pp_layers = mapping.pp_layers(model_config.get_num_attention_layers()) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py index 0e64ab823449..8a36800bb6ff 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/kernels.py @@ -850,3 +850,179 @@ def prepare_sparse_write_table( PAGE_SCALE=page_scale, BLOCK=256, ) + + +######################################################## +# torch custom op over the same kernel +######################################################## + +_INT32_MAX = 2**31 - 1 + + +def _is_power_of_two(n: int) -> bool: + return n > 0 and (n & (n - 1)) == 0 + + +@torch.library.custom_op( + "trtllm::dsv4_local_to_global_indices", + mutates_args=(), + schema=( + "(Tensor req_id, Tensor swa_block_table, Tensor swa_local_indices, " + "int swa_tokens_per_block, int swa_row_offset, Tensor? compressed_block_table, " + "Tensor? compressed_local_indices, int compressed_tokens_per_block, " + "int compressed_row_offset, bool enable_pdl) -> Tensor" + ), +) +def dsv4_local_to_global_indices( + req_id: torch.Tensor, + swa_block_table: torch.Tensor, + swa_local_indices: torch.Tensor, + swa_tokens_per_block: int, + swa_row_offset: int, + compressed_block_table: torch.Tensor | None, + compressed_local_indices: torch.Tensor | None, + compressed_tokens_per_block: int, + compressed_row_offset: int, + enable_pdl: bool, +) -> torch.Tensor: + """Map per-token local SWA and compressed indices to pool rows, one kernel launch. + + Same kernel as ``deepseek_v4_local_to_global_indices`` with each pool's page size + passed explicitly instead of derived from a compress ratio, so a compressed + segment exists whenever its two tensors are given -- including a 1:1 compressed + pool. For token ``t`` of request row ``r = req_id[t]`` and a local index ``l`` in + a segment with block table ``bt``, page size ``tpb`` and row offset ``off``:: + + out = off + bt[r, l // tpb] * tpb + l % tpb if 0 <= l and l // tpb < bt.shape[1] + out = -1 otherwise + + Returns int32 ``[T, W_swa + W_compressed]``: the SWA segment in its own columns + first, then the compressed segment (width 0 when the compressed tensors are None). + """ + if (compressed_block_table is None) != (compressed_local_indices is None): + raise ValueError( + "compressed_block_table and compressed_local_indices are given together or not at all" + ) + has_compressed = compressed_block_table is not None + int_args = (req_id, swa_block_table, swa_local_indices) + if has_compressed: + int_args += (compressed_block_table, compressed_local_indices) + for t in int_args: + if t.dtype != torch.int32: + raise ValueError( + f"req_id, block tables and local indices must be int32 tensors, got {t.dtype}" + ) + if req_id.dim() != 1 or req_id.stride(0) != 1: + raise ValueError("req_id must be a contiguous 1-D tensor") + num_tokens = req_id.shape[0] + if swa_local_indices.dim() != 2 or swa_local_indices.shape[0] != num_tokens: + raise ValueError("swa_local_indices must be [num_tokens, W_swa]") + if swa_block_table.dim() != 2: + raise ValueError("swa_block_table must be [rows, max_blocks]") + num_swa = swa_local_indices.shape[1] + if not _is_power_of_two(num_swa): + raise ValueError(f"W_swa must be a power of two, got {num_swa}") + if swa_tokens_per_block <= 0: + raise ValueError(f"swa_tokens_per_block must be positive, got {swa_tokens_per_block}") + if not 0 <= swa_row_offset <= _INT32_MAX: + raise ValueError(f"swa_row_offset must be in [0, 2**31), got {swa_row_offset}") + + if has_compressed: + if compressed_local_indices.dim() != 2 or compressed_local_indices.shape[0] != num_tokens: + raise ValueError("compressed_local_indices must be [num_tokens, W_compressed]") + if compressed_block_table.dim() != 2: + raise ValueError("compressed_block_table must be [rows, max_blocks]") + num_compressed = compressed_local_indices.shape[1] + if not _is_power_of_two(num_compressed): + raise ValueError(f"W_compressed must be a power of two, got {num_compressed}") + if compressed_tokens_per_block <= 0: + raise ValueError( + f"compressed_tokens_per_block must be positive, got {compressed_tokens_per_block}" + ) + if not 0 <= compressed_row_offset <= _INT32_MAX: + raise ValueError( + f"compressed_row_offset must be in [0, 2**31), got {compressed_row_offset}" + ) + bt_c, idx_c = compressed_block_table, compressed_local_indices + tpb_c = compressed_tokens_per_block + max_blocks_c = compressed_block_table.shape[1] + offset_c = compressed_row_offset + else: + # Never read: has_compressed is a constexpr that compiles the segment out. + num_compressed = 0 + bt_c, idx_c = swa_block_table, swa_local_indices + tpb_c = swa_tokens_per_block + max_blocks_c = 1 + offset_c = 0 + + out = torch.empty( + (num_tokens, num_swa + num_compressed), dtype=torch.int32, device=req_id.device + ) + if num_tokens == 0: + return out + + _deepseek_v4_local_to_global_kernel[(num_tokens,)]( + req_id, + swa_block_table, + bt_c, + swa_local_indices, + idx_c, + out, + out, + swa_row_offset, + offset_c, + swa_tokens_per_block, + tpb_c, + swa_block_table.shape[1], + max_blocks_c, + num_swa, + num_compressed, + num_swa + num_compressed, + has_compressed, + swa_block_table.stride(0), + swa_block_table.stride(1), + bt_c.stride(0), + bt_c.stride(1), + swa_local_indices.stride(0), + swa_local_indices.stride(1), + idx_c.stride(0), + idx_c.stride(1), + out.stride(0), + out.stride(1), + None, + None, + None, + None, + None, + None, + 1.0, + WRITE_FMHA_SCHEDULER=False, + out_extra_stride0=out.stride(0), + out_extra_stride1=out.stride(1), + SPLIT_EXTRA=False, + LAUNCH_WITH_PDL=enable_pdl, + active_request_count_ptr=None, + MASK_DECODE_PADDING=False, + MASK_INVALID_PAGES=False, + launch_pdl=enable_pdl, + ) + return out + + +@dsv4_local_to_global_indices.register_fake +def _( + req_id, + swa_block_table, + swa_local_indices, + swa_tokens_per_block, + swa_row_offset, + compressed_block_table, + compressed_local_indices, + compressed_tokens_per_block, + compressed_row_offset, + enable_pdl, +): + width = swa_local_indices.shape[1] + if compressed_local_indices is not None: + width += compressed_local_indices.shape[1] + return req_id.new_empty((req_id.shape[0], width), dtype=torch.int32) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py index 316897a406a6..05525b4d47a6 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py @@ -5,7 +5,7 @@ from __future__ import annotations import math -from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple import torch @@ -23,11 +23,22 @@ DeepSeekV4MetadataParams, is_compress_layer, ) +from .v41_layout import DeepseekV41Layout if TYPE_CHECKING: from .cache_manager import DeepseekV4CacheManager +def _v41_layout_of(manager) -> Optional[DeepseekV41Layout]: + """The DeepSeek-V4.1 layout a cache manager runs in, or None for DeepSeek-V4.""" + layout = getattr(manager, "v41_layout", None) + return layout if isinstance(layout, DeepseekV41Layout) else None + + +def _num_sliding_table_columns(manager) -> int: + return getattr(manager, "num_sliding_attention_columns", len(DEEPSEEK_V4_SLIDING_ATTENTION)) + + class DeepseekV4TrtllmAttentionMetadata(DSAtrtllmAttentionMetadata): # The set of compress ratios for the layers compress_ratio_set: Set[int] @@ -37,6 +48,9 @@ class DeepseekV4TrtllmAttentionMetadata(DSAtrtllmAttentionMetadata): num_total_compressed_tokens: Dict[int, int] # The max number of context compressed tokens for each compress ratio max_ctx_compressed_tokens: Dict[int, int] + # DeepSeek-V4.1 layout of the cache manager, or None (DeepSeek-V4). A class + # default, so an instance that skipped __post_init__ reads DeepSeek-V4. + v41_layout: Optional["DeepseekV41Layout"] = None def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -46,6 +60,10 @@ def _invalidate_mla_scheduler_buffers(self) -> None: super()._invalidate_mla_scheduler_buffers() self._fused_q_gen_cu_seqlens_valid = False + @property + def is_deepseek_v41(self) -> bool: + return self.v41_layout is not None + def __post_init__(self): super().__post_init__() self.num_total_compressed_tokens = {} @@ -62,39 +80,15 @@ def __post_init__(self): f"TileSizeKV of the FMHA kernel. (got {window_size})." ) capture_graph = self.is_cuda_graph - self.compress_ratio_set = set(self.compress_ratios) - # Cache a sorted list for deterministic iteration order across - # torch.compile traces. Avoids repeated set→list conversion and - # guarantees stable specialization keys. - self._compress_ratios_sorted = sorted(self.compress_ratio_set) - _supported_ratios = {1, 4, 128} - _unsupported = self.compress_ratio_set - _supported_ratios - assert not _unsupported, ( - f"Unsupported compress ratios {_unsupported}. Only {_supported_ratios} are supported." - ) - - attention_types = [] - for compress_ratio in self.compress_ratio_set: - if compress_ratio == 1: - attention_types.append((1, DeepseekV4AttentionType.SWA)) - elif compress_ratio == 4: - attention_types.append((1, DeepseekV4AttentionType.SWA)) - attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESS)) - attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESSOR_KV)) - attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESSOR_SCORE)) - attention_types.append((compress_ratio, DeepseekV4AttentionType.INDEXER_COMPRESS)) - attention_types.append( - (compress_ratio, DeepseekV4AttentionType.INDEXER_COMPRESSOR_KV) - ) - attention_types.append( - (compress_ratio, DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE) - ) - else: - attention_types.append((1, DeepseekV4AttentionType.SWA)) - attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESS)) - attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESSOR_KV)) - attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESSOR_SCORE)) - self.attention_type_set = set(attention_types) + # DeepSeek-V4.1 layout mode follows the cache manager, which selects it + # from the checkpoint config. The layer ratios then come from the + # checkpoint (0 = window only, 1 / 2 = compressed); the sparse config's + # normalized compress_ratios carry no layout. + self.v41_layout = _v41_layout_of(self.kv_cache_manager) + if self.v41_layout is not None: + self._init_v41_ratio_sets(self.v41_layout) + else: + self._init_v4_ratio_sets() # Create buffers for the compressor # cu_seq_lens_cuda is the cumulative sequence lengths for the requests @@ -232,26 +226,35 @@ def __post_init__(self): capture_graph=capture_graph, ) - # Compute max_compressed_indices for CUDA graph compatibility. - # For ratio=4, the indexer selects index_topk compressed tokens. - # For ratio=128, use max_seq_len / 128 rounded up to next power of 2 - raw_128 = math.ceil(self.max_seq_len / 128) - po2_128 = 1 << (raw_128 - 1).bit_length() if raw_128 > 0 else 1 - self.max_compressed_indices = { - 1: 0, # No compressed indices - 4: self.sparse_mla_topk, # index_topk from indexer - 128: po2_128, # All compressed tokens, rounded to power of 2 - } + if self.v41_layout is not None: + # Every compressed V4.1 layer attends the window plus the index + # top-k; the compressed local indices come from the indexer. + self.max_compressed_indices = { + ratio: (self.sparse_mla_topk if ratio > 0 else 0) for ratio in self._topk_len_ratios + } + self.compressed_local_indices_cuda = None + self._init_v41_buffers(capture_graph) + else: + # Compute max_compressed_indices for CUDA graph compatibility. + # For ratio=4, the indexer selects index_topk compressed tokens. + # For ratio=128, use max_seq_len / 128 rounded up to next power of 2 + raw_128 = math.ceil(self.max_seq_len / 128) + po2_128 = 1 << (raw_128 - 1).bit_length() if raw_128 > 0 else 1 + self.max_compressed_indices = { + 1: 0, # No compressed indices + 4: self.sparse_mla_topk, # index_topk from indexer + 128: po2_128, # All compressed tokens, rounded to power of 2 + } - # Compressed local indices for compress_ratio=128 - # Note: ratio=4 uses dynamic topk_indices from indexer, so we only pre-allocate for ratio=128 - self.compressed_local_indices_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens, self.max_compressed_indices[128]), - cache_name="compressed_local_indices_cuda", - dtype=torch.int32, - capture_graph=capture_graph, - ) + # Compressed local indices for compress_ratio=128 + # Note: ratio=4 uses dynamic topk_indices from indexer, so we only pre-allocate for ratio=128 + self.compressed_local_indices_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens, self.max_compressed_indices[128]), + cache_name="compressed_local_indices_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) # Sliding-window caches use per-layer page indices and are indexed by # [local_layer_idx, attention_type.value, sequence, block]. COMPRESS @@ -259,7 +262,7 @@ def __post_init__(self): # compatibility table for the generic DSA indexer path. block_table_shape = ( self.kv_cache_manager.num_local_layers, - len(DEEPSEEK_V4_SLIDING_ATTENTION), + _num_sliding_table_columns(self.kv_cache_manager), self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq, ) @@ -283,8 +286,7 @@ def __post_init__(self): dtype=torch.int32, capture_graph=capture_graph, ) - for compress_ratio in self._compress_ratios_sorted - if is_compress_layer(compress_ratio) + for compress_ratio in self._compress_table_ratios } self._init_sparse_offload_state() @@ -298,7 +300,7 @@ def __post_init__(self): dtype=torch.int32, capture_graph=capture_graph, ) - for compress_ratio in self.compress_ratio_set + for compress_ratio in self._topk_len_ratios } # cached_token_lens_cuda: number of tokens already cached per request @@ -424,6 +426,380 @@ def apply_device_ragged_layout( # therefore captured once after the layout update and replayed in order. self._invalidate_mla_scheduler_buffers() + def _init_v4_ratio_sets(self) -> None: + self.compress_ratio_set = set(self.compress_ratios) + # Cache a sorted list for deterministic iteration order across + # torch.compile traces. Avoids repeated set→list conversion and + # guarantees stable specialization keys. + self._compress_ratios_sorted = sorted(self.compress_ratio_set) + _supported_ratios = {1, 4, 128} + _unsupported = self.compress_ratio_set - _supported_ratios + assert not _unsupported, ( + f"Unsupported compress ratios {_unsupported}. Only {_supported_ratios} are supported." + ) + + attention_types = [] + for compress_ratio in self.compress_ratio_set: + if compress_ratio == 1: + attention_types.append((1, DeepseekV4AttentionType.SWA)) + elif compress_ratio == 4: + attention_types.append((1, DeepseekV4AttentionType.SWA)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESS)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESSOR_KV)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESSOR_SCORE)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.INDEXER_COMPRESS)) + attention_types.append( + (compress_ratio, DeepseekV4AttentionType.INDEXER_COMPRESSOR_KV) + ) + attention_types.append( + (compress_ratio, DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE) + ) + else: + attention_types.append((1, DeepseekV4AttentionType.SWA)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESS)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESSOR_KV)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESSOR_SCORE)) + self.attention_type_set = set(attention_types) + # sparse_mla_topk_lens keys: every layer ratio. + self._topk_len_ratios = self.compress_ratio_set + # Ratios with a COMPRESS page table. + self._compress_table_ratios = [ + compress_ratio + for compress_ratio in self._compress_ratios_sorted + if is_compress_layer(compress_ratio) + ] + + def _init_v41_ratio_sets(self, layout: "DeepseekV41Layout") -> None: + """Ratio key spaces of the DeepSeek-V4.1 layout mode. + + Per-ratio compressed metadata (lengths, positions, masks, page tables) + is keyed by the ratios that have a compressed stream, {1, 2}; + ``sparse_mla_topk_lens`` is keyed by every layer ratio, {0, 1, 2}. + """ + self.compress_ratios = list(layout.compress_ratios) + self.compress_ratio_set = set(layout.compressed_ratios) + self._compress_ratios_sorted = sorted(self.compress_ratio_set) + self._topk_len_ratios = list(layout.layer_ratios) + self._compress_table_ratios = list(self._compress_ratios_sorted) + attention_types = {(0, DeepseekV4AttentionType.SWA)} + for layer in range(layout.num_layers): + ratio = layout.compress_ratios[layer] + for attn_type in layout.attention_types(layer): + attention_types.add((ratio, attn_type)) + self.attention_type_set = attention_types + + def _init_v41_buffers(self, capture_graph: bool) -> None: + """Buffers only the DeepSeek-V4.1 layout mode fills.""" + # Per-ratio INDEXER_COMPRESS (native MXFP4 index K) page tables, one + # row per sequence of the batch, in pages of the owner's index-K pool + # (tokens_per_block // ratio entries per page). + self.indexer_compress_block_tables = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq), + cache_name=f"indexer_compress_block_tables_{compress_ratio}", + dtype=torch.int32, + capture_graph=capture_graph, + ) + for compress_ratio in self._compress_ratios_sorted + } + self.host_indexer_compress_block_tables = { + compress_ratio: torch.empty_like( + self.indexer_compress_block_tables[compress_ratio], + device="cpu", + pin_memory=prefer_pinned(), + ) + for compress_ratio in self._compress_ratios_sorted + } + # Per token: compressed entries visible at ratio r, (position + 1) // r. + # Its generation rows, viewed [num_generations, tokens_per_generation], + # are the paged indexer's context lengths. + self.compressed_visible_lens_cuda = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name=f"compressed_visible_lens_cuda_{compress_ratio}", + dtype=torch.int32, + capture_graph=capture_graph, + ) + for compress_ratio in self._compress_ratios_sorted + } + max_positions = self.max_seq_len + max(self.max_draft_tokens, 0) + # Static upper bound of compressed_visible_lens per ratio, e.g. the + # paged indexer's max_context_len inside a captured step. + self.max_compressed_kv_lens = { + compress_ratio: max_positions // compress_ratio + for compress_ratio in self._compress_ratios_sorted + } + + # Address book: where this step's KV rows are written. Filled by + # _prepare_v41_address_book() in prepare() and on_update_kv_lens(); + # -1 marks a row that must not be written. Rows at or past the step's + # token / compressed-row count are not written. + manager = self.kv_cache_manager + # MXFP4 index-K entry: two E2M1 codes per byte, then one UE8M0 byte + # per quant block (64 + 4 bytes at index_head_dim 128). + self._v41_index_k_data_bytes = manager.index_head_dim // 2 + self._v41_index_k_scale_bytes = manager.index_head_dim // manager.quant_block_size + # [num_local_layers, max_num_tokens]: row of token t's window KV in + # get_buffers(L, SWA) viewed [pages * tokens_per_block, head_dim], at + # local layer layer_offsets[L]. + self.v41_swa_slot_mapping = self.get_empty( + self.cuda_graph_buffers, + (manager.num_local_layers, self.max_num_tokens), + cache_name="v41_swa_slot_mapping", + dtype=torch.int64, + capture_graph=capture_graph, + ) + # Per ratio r, [max_num_tokens], in the packed order of + # compressed_position_ids_cuda[r]: row of the entry in + # get_buffers(kv_source, COMPRESS) viewed [pages * block, head_dim], + # block = tokens_per_block // r. + self.v41_compress_slot_mapping = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name=f"v41_compress_slot_mapping_{compress_ratio}", + dtype=torch.int64, + capture_graph=capture_graph, + ) + for compress_ratio in self._compress_ratios_sorted + } + # Same rows: flat byte offsets of the entry's code bytes and scale word + # in get_indexer_k_cache_buffers(index owner), uint8 + # [pages, block, 1, 68]; per page block * 64 code bytes, then block * 4 + # scale bytes. + self.v41_index_k_slot_data = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name=f"v41_index_k_slot_data_{compress_ratio}", + dtype=torch.int64, + capture_graph=capture_graph, + ) + for compress_ratio in self._compress_ratios_sorted + } + self.v41_index_k_slot_scale = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name=f"v41_index_k_slot_scale_{compress_ratio}", + dtype=torch.int64, + capture_graph=capture_graph, + ) + for compress_ratio in self._compress_ratios_sorted + } + # Context-only prefill index-K inputs; see _prepare_v41_ctx_index_k(). + self._reset_v41_ctx_index_k() + self._init_v41_replay_buffers(capture_graph) + + def _init_v41_replay_buffers(self, capture_graph: bool) -> None: + """Encoder SWA bounded replay fields, filled by _prepare_v41_replay(). + + Without replay rows in the step they hold the plain values: no replay rows, + every row a compute row, compute start = cached tokens, compute cu = + cu_seq_lens, SWA floor 0, no Engram history rows. + """ + num_seqs = self.max_num_sequences + + def buf(name: str, shape: tuple, dtype: torch.dtype) -> torch.Tensor: + return self.get_empty( + self.cuda_graph_buffers, + shape, + cache_name=name, + dtype=dtype, + capture_graph=capture_graph, + ) + + # Per request row: replay rows at the head of its rows this step. + self.v41_num_replay_tokens_cuda = buf( + "v41_num_replay_tokens_cuda", (num_seqs,), torch.int32 + ) + self.v41_num_replay_tokens_cpu = torch.zeros( + (num_seqs,), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() + ) + self.v41_num_replay_tokens_total = 0 + # Per request row: lowest position its window may attend (the replay start). + self.v41_swa_floor_cuda = buf("v41_swa_floor_cuda", (num_seqs,), torch.int32) + # Per context request row: the sparse attention's seq_lens, its kv length + # minus its SWA floor (see _prepare_v41_swa_seq_lens()). + self.v41_swa_seq_lens_cuda = buf("v41_swa_seq_lens_cuda", (num_seqs,), torch.int32) + # Per token row: whether it is a replay row. + self.v41_is_replay_token_cuda = buf( + "v41_is_replay_token_cuda", (self.max_num_tokens,), torch.bool + ) + # Rows that compute main KV / index K / compressor state, packed in row order. + self.v41_compute_token_index_cuda = buf( + "v41_compute_token_index_cuda", (self.max_num_tokens,), torch.int32 + ) + self.v41_num_compute_tokens = 0 + # Per request row: first computed position (cached + replay rows), and the + # cumulative row counts of the packed compute rows. + self.v41_compute_start_pos_cuda = buf( + "v41_compute_start_pos_cuda", (num_seqs,), torch.int32 + ) + self.v41_compute_cu_seq_lens_cuda = buf( + "v41_compute_cu_seq_lens_cuda", (num_seqs + 1,), torch.int32 + ) + # Engram history of the replay starts, as engram_hash input rows: ids, + # context request row, per-context-row cu, and start position. + history = max(1, self.v41_layout.engram_history_len) + self.v41_replay_history_len = 0 + self.v41_replay_history_token_ids_cuda = buf( + "v41_replay_history_token_ids_cuda", (num_seqs * history,), torch.int32 + ) + self.v41_replay_history_req_idx_cuda = buf( + "v41_replay_history_req_idx_cuda", (num_seqs * history,), torch.int32 + ) + self.v41_replay_history_cu_seq_lens_cuda = buf( + "v41_replay_history_cu_seq_lens_cuda", (num_seqs + 1,), torch.int32 + ) + self.v41_replay_history_start_pos_cuda = buf( + "v41_replay_history_start_pos_cuda", (num_seqs,), torch.int32 + ) + + def _prepare_v41_replay(self, num_requests: int) -> None: + """Encoder SWA bounded replay rows of this step, from the cache manager's plans. + + A context request whose chunk starts at its replay start (a prefix-cache hit of + ``reused`` tokens under replay; ``cached_token_lens`` is then ``reused - R``) + has ``R = min(window, reused)`` replay rows at the head of its rows: cached + tokens re-forwarded to rebuild the sliding windows. Filled here: + + - ``v41_num_replay_tokens_cuda/_cpu[b]`` = R (0 for other rows), + ``v41_num_replay_tokens_total`` (host); + - ``v41_is_replay_token_cuda[t]``; + - ``v41_compute_token_index_cuda[:v41_num_compute_tokens]``: the other rows; + - ``v41_compute_start_pos_cuda[b]`` = cached + R and + ``v41_compute_cu_seq_lens_cuda``: the compute rows as a step starting at the + reuse boundary; + - ``v41_swa_floor_cuda[b]`` = the replay start (0 without replay): the + window of a replay row starts no earlier (tech report 3.2.2); + - ``v41_swa_seq_lens_cuda[b]`` = kv length - floor for each context + request (``_prepare_v41_swa_seq_lens()``); + - the Engram history rows of each replay start. + + The compressed stream (``past_kv_lens``, ``new_comp_kv_lens``, positions, the + COMPRESS / index-K write slots) is then computed from the compute start, so no + cached main-KV or index-K entry is rewritten. + """ + manager = self.kv_cache_manager + num_contexts = self.num_contexts + num_tokens = self.num_tokens + device = self.v41_num_replay_tokens_cuda.device + replay = self.v41_num_replay_tokens_cpu + replay[:num_requests] = 0 + cached = self.cached_token_lens_cpu[:num_requests] + history_ids: List[int] = [] + history_req: List[int] = [] + history_lens = [0] * num_contexts + history_start = [0] * num_contexts + if num_contexts > 0 and getattr(manager, "v41_encoder_replay", False): + plans = manager.get_v41_replay_plans( + self.request_ids[:num_contexts], cached[:num_contexts].tolist() + ) + for b, plan in enumerate(plans): + if plan is None: + continue + if int(self.seq_lens[b]) < plan.num_tokens: + raise RuntimeError( + f"request {self.request_ids[b]}: a chunk of {int(self.seq_lens[b])} " + f"rows cannot hold its {plan.num_tokens} replay rows" + ) + replay[b] = plan.num_tokens + ids = plan.history_token_ids + history_ids.extend(ids) + history_req.extend([b] * len(ids)) + history_lens[b] = len(ids) + history_start[b] = plan.start - len(ids) + total = int(replay[:num_requests].sum()) + self.v41_num_replay_tokens_total = total + self.v41_num_compute_tokens = num_tokens - total + replay_cuda = self.v41_num_replay_tokens_cuda[:num_requests] + replay_cuda.copy_(replay[:num_requests], non_blocking=True) + cached_cuda = self.cached_token_lens_cuda[:num_requests] + cu = self.cu_seq_lens_cuda[: num_requests + 1] + if total == 0: + self.v41_swa_floor_cuda[:num_requests].zero_() + self.v41_compute_start_pos_cuda[:num_requests].copy_(cached_cuda) + self.v41_compute_cu_seq_lens_cuda[: num_requests + 1].copy_(cu) + self.v41_is_replay_token_cuda[:num_tokens].zero_() + torch.arange(num_tokens, out=self.v41_compute_token_index_cuda[:num_tokens]) + self._prepare_v41_swa_seq_lens() + else: + cu = cu.to(torch.int32) + self.v41_swa_floor_cuda[:num_requests] = torch.where(replay_cuda > 0, cached_cuda, 0) + self._prepare_v41_swa_seq_lens() + self.v41_compute_start_pos_cuda[:num_requests] = cached_cuda + replay_cuda + compute_cu = torch.nn.functional.pad( + torch.cumsum(self._seq_lens_cuda[:num_requests].to(torch.int32) - replay_cuda, 0), + (1, 0), + ).to(torch.int32) + self.v41_compute_cu_seq_lens_cuda[: num_requests + 1] = compute_cu + rows = torch.arange(num_tokens, dtype=torch.int32, device=device) + req_idx = self.req_idx_per_token[:num_tokens].to(torch.int64) + self.v41_is_replay_token_cuda[:num_tokens] = (rows - cu[req_idx]) < replay_cuda[req_idx] + num_compute = num_tokens - total + if num_compute > 0: + crow = rows[:num_compute] + creq = torch.searchsorted(compute_cu[1:], crow, right=True).to(torch.int64) + self.v41_compute_token_index_cuda[:num_compute] = ( + crow - compute_cu[creq] + cu[creq] + replay_cuda[creq] + ) + num_history = len(history_ids) + self.v41_replay_history_len = num_history + if num_history > 0: + history_cu = [0] + for n in history_lens: + history_cu.append(history_cu[-1] + n) + self.v41_replay_history_token_ids_cuda[:num_history].copy_( + torch.tensor(history_ids, dtype=torch.int32), non_blocking=True + ) + self.v41_replay_history_req_idx_cuda[:num_history].copy_( + torch.tensor(history_req, dtype=torch.int32), non_blocking=True + ) + self.v41_replay_history_cu_seq_lens_cuda[: num_contexts + 1].copy_( + torch.tensor(history_cu, dtype=torch.int32), non_blocking=True + ) + self.v41_replay_history_start_pos_cuda[:num_contexts].copy_( + torch.tensor(history_start, dtype=torch.int32), non_blocking=True + ) + + def _prepare_v41_swa_seq_lens(self) -> None: + """``v41_swa_seq_lens_cuda[b]`` for each context request ``b``: its kv length + minus its SWA floor, the ``seq_lens`` of the context sparse attention call. + + The kernel takes a query's window count from ``seq_lens``, + ``min(window, seq_lens - q_len + j + 1)`` for query ``j`` of a request, and + reads a -1 window column inside that count as an all-zero K/V row, not as + a mask. A replay request's windows start at its floor, its replay start + (the chunk's first position), so its kv length would count up to + ``window - 1`` columns below the floor; ``kv_len - floor`` counts exactly + its floored window, ``min(window, p - floor + 1)`` keys at position ``p``. + With floor 0 (no replay) this is the kv length itself. Generation rows' + windows are never floored and are not covered: they pass their kv + lengths, which speculative decoding may move after prepare(). Refreshed + by on_update_kv_lens() whenever the step has context requests. + """ + n = self.num_contexts + if n: + torch.sub( + self.kv_lens_cuda_runtime[:n], + self.v41_swa_floor_cuda[:n], + out=self.v41_swa_seq_lens_cuda[:n], + ) + + @property + def engram_token_block_table(self) -> Optional[torch.Tensor]: + """[max_num_sequences, max_blocks] page table of the Engram token history pool.""" + layout = self.v41_layout + if layout is None or layout.engram_token_owner is None: + return None + local_layer = self.kv_cache_manager.layer_offsets.get(layout.engram_token_owner) + if local_layer is None: + return None + return self.sliding_block_tables[local_layer, DeepseekV4AttentionType.ENGRAM_TOKEN.value] + def prepare_for_indexer_k_cache(self): """Prepare the shared indexer K-cache decode table for DSA kernels.""" # INDEXER_COMPRESS uses shared page indices, so the generic DSA @@ -475,9 +851,24 @@ def prepare_for_block_tables(self): num_contexts=self.num_contexts, num_seqs=self.num_seqs, ) + if self.v41_layout is not None: + for compress_ratio, table in self.indexer_compress_block_tables.items(): + host_table = self.host_indexer_compress_block_tables[compress_ratio] + self.kv_cache_manager.copy_batch_indexer_compress_block_tables( + host_table, + self.request_ids, + beam_width=self.beam_width, + num_contexts=self.num_contexts, + num_seqs=self.num_seqs, + compress_ratio=compress_ratio, + ) + table[: self.num_seqs].copy_(host_table[: self.num_seqs], non_blocking=True) + + def prepare_for_deepseek_v4_indices(self, token_positions=None) -> torch.Tensor: + """Prepare SWA/compressed local indices and sparse_mla_topk_lens. - def prepare_for_deepseek_v4_indices(self, token_positions=None): - """Prepare SWA/compressed local indices and sparse_mla_topk_lens.""" + Returns the per-token absolute positions used, ``[num_tokens]`` int32. + """ window_size = self.window_size device = self.swa_local_indices_cuda.device @@ -495,6 +886,21 @@ def prepare_for_deepseek_v4_indices(self, token_positions=None): offsets = token_idx - cu_seq_lens[req_idx].to(torch.int32) token_positions = cached_tokens[req_idx].to(torch.int32) + offsets + if self.v41_layout is not None: + num_tokens = token_positions.shape[0] + swa_floor = self.v41_swa_floor_cuda[self.req_idx_per_token[:num_tokens].to(torch.int64)] + self._prepare_deepseek_v41_indices_compiled( + token_positions, + swa_floor, + window_size, + self.sparse_mla_topk, + self.swa_local_indices_cuda, + self.sparse_mla_topk_lens, + self.compressed_visible_lens_cuda, + self._topk_len_ratios, + ) + return token_positions + if self.use_fp8_ds_mla: self.token_positions_cuda[: token_positions.shape[0]] = token_positions @@ -508,6 +914,7 @@ def prepare_for_deepseek_v4_indices(self, token_positions=None): self.sparse_mla_topk_lens, self._compress_ratios_sorted, ) + return token_positions @staticmethod @maybe_compile(dynamic=True, options={"max-autotune": True}) @@ -572,10 +979,60 @@ def _prepare_deepseek_v4_indices_compiled( raise ValueError(f"Unsupported compress_ratio: {compress_ratio}") sparse_mla_topk_lens_bufs[compress_ratio][:num_tokens] = total_count.to(torch.int32) + @staticmethod + @maybe_compile(dynamic=True, options={"max-autotune": True}) + def _prepare_deepseek_v41_indices_compiled( + token_positions: torch.Tensor, + swa_floor: torch.Tensor, + window_size: int, + sparse_mla_topk: int, + swa_local_indices_buf: torch.Tensor, + sparse_mla_topk_lens_bufs: Dict[int, torch.Tensor], + compressed_visible_lens_bufs: Dict[int, torch.Tensor], + layer_ratios: list, + ): + """DeepSeek-V4.1 SWA local indices, visible compressed lengths and topk_lens. + + For a query at absolute position ``p`` (prefill, chunked prefill, + decode and speculative-verify rows alike): + - window: positions ``max(floor, p - window_size + 1) .. p`` in order, + then -1, i.e. the reference's causal prefill window, and in decode the + ring's slots as positions; ``floor`` is 0, or the replay start of an + encoder-replay step's request (the window truncated to the replay + segment, tech report 3.2.2); + - ``compressed_visible_lens[r] = (p + 1) // r``, the reference's + ``compress_lens``; + - ``sparse_mla_topk_lens[r]``: ``window_size`` for window-only layers + (the kernel takes the window count from ``seq_lens``), and + ``window_size + min((p + 1) // r, index_topk)`` for compressed ones, + the number of valid top-k entries the indexer packs at the front. + """ + device = token_positions.device + num_tokens = token_positions.shape[0] + + positions = token_positions.unsqueeze(1) + swa_offsets = torch.arange(window_size, dtype=torch.int32, device=device) + swa_start = torch.maximum(positions - window_size + 1, swa_floor.unsqueeze(1)).clamp(min=0) + swa_indices = swa_start + swa_offsets + swa_indices = torch.where(swa_indices > positions, -1, swa_indices).to(torch.int32) + swa_local_indices_buf[:num_tokens] = swa_indices + + kv_lens = token_positions + 1 + for compress_ratio in layer_ratios: + if compress_ratio == 0: + total_count = torch.full_like(kv_lens, window_size) + else: + visible = kv_lens // compress_ratio + compressed_visible_lens_bufs[compress_ratio][:num_tokens] = visible.to(torch.int32) + total_count = window_size + visible.clamp(max=sparse_mla_topk) + sparse_mla_topk_lens_bufs[compress_ratio][:num_tokens] = total_count.to(torch.int32) + def _build_cache_buffer_data_pointers( self, manager: "DeepseekV4CacheManager", compress_ratios_by_layer: list[int] ) -> tuple[dict[int, int], dict[int, int], dict[int, int]]: """Build sparse cache pointers for a target or draft manager.""" + if _v41_layout_of(manager) is not None: + return self._build_v41_cache_buffer_data_pointers(manager) sparse_mla_base_ptrs = {1: manager.swa_pool_ptr} for ratio, compress_pool_ptr in manager.compress_pool_ptrs.items(): sparse_mla_base_ptrs[ratio] = compress_pool_ptr @@ -588,7 +1045,42 @@ def _build_cache_buffer_data_pointers( } return sparse_mla_base_ptrs, swa_buffer_ptrs, compressed_buffer_ptrs + @staticmethod + def _build_v41_cache_buffer_data_pointers( + manager: "DeepseekV4CacheManager", + ) -> tuple[dict[int, int], dict[int, int], dict[int, int]]: + """DeepSeek-V4.1 pool pointers keyed by layer ratio and by layer. + + ``sparse_mla_base_ptrs``: 0 -> the SWA pool, r -> the COMPRESS base of + the first local KV source of ratio r. ``compressed_buffer_ptrs[L]`` is + the COMPRESS base of ``kv_source(L)`` for every compressed layer, so a + reader layer resolves to its owner's pool instead of a missing entry. + """ + layout = manager.v41_layout + sparse_mla_base_ptrs = {0: manager.swa_pool_ptr} + sparse_mla_base_ptrs.update(manager.compress_pool_ptrs) + swa_buffer_ptrs = {layer_idx: manager.swa_pool_ptr for layer_idx in manager.pp_layers} + owner_ptrs: dict[int, int] = {} + compressed_buffer_ptrs = {} + for layer_idx in manager.pp_layers: + owner = layout.kv_source(layer_idx) + if owner is None: + continue + if owner not in owner_ptrs: + owner_ptrs[owner] = manager.get_buffers( + owner, DeepseekV4AttentionType.COMPRESS + ).data_ptr() + compressed_buffer_ptrs[layer_idx] = owner_ptrs[owner] + return sparse_mla_base_ptrs, swa_buffer_ptrs, compressed_buffer_ptrs + def _init_cache_buffer_data_pointers(self): + if self.v41_layout is not None: + ( + self.sparse_mla_base_ptrs, + self.swa_buffer_ptrs, + self.compressed_buffer_ptrs, + ) = self._build_v41_cache_buffer_data_pointers(self.kv_cache_manager) + return # If MTP is enabled, enlarge the compress ratios by max_draft_tokens - 1. extend_compress_ratios = self.compress_ratios + [self.compress_ratios[-1]] * ( self.max_draft_tokens - 1 @@ -610,6 +1102,11 @@ def _init_draft_sparse_buffers(self): if not isinstance(draft_mgr, DeepseekV4CacheManager): return + if self.v41_layout is not None: + raise NotImplementedError( + "DeepSeek-V4.1 has no separate draft KV cache: its DSpark stages keep their " + "windows in the drafter." + ) # Current DSv4 MTP layers are SWA-only. Fail fast if a future model # introduces a compressed or indexer MTP layer. @@ -622,7 +1119,7 @@ def _init_draft_sparse_buffers(self): draft_block_table_shape = ( draft_mgr.num_local_layers, - len(DEEPSEEK_V4_SLIDING_ATTENTION), + _num_sliding_table_columns(draft_mgr), self.max_num_sequences, draft_mgr.max_blocks_per_seq, ) @@ -757,6 +1254,9 @@ def prepare(self) -> None: # For indices conversion self.prepare_for_indices_conversion() + if self.v41_layout is not None: + self._prepare_v41_replay(num_requests) + has_sparse_layers = DEEPSEEK_V4_SPARSE_RATIO in self.compress_ratio_set # For block offsets @@ -770,7 +1270,7 @@ def prepare(self) -> None: self.prepare_for_spec_decode(kv_lens) # For DeepSeek-V4 indices - self.prepare_for_deepseek_v4_indices() + token_positions = self.prepare_for_deepseek_v4_indices() # Prepare metadata for indexer (only needed when sparse layers exist) if has_sparse_layers: @@ -784,6 +1284,9 @@ def prepare(self) -> None: num_generations = self.num_generations kv_lens_slice = kv_lens[:num_requests] cached_slice = cached_token_lens[:num_requests] + if self.v41_layout is not None: + # Replay rows compute no compressed entry: the stream starts after them. + cached_slice = cached_slice + self.v41_num_replay_tokens_cpu[:num_requests] # Host-side per-ratio ctx compressed-token counts (Python ints), so # _compute_ctx_compressed_position_ids never reads a device scalar @@ -823,6 +1326,8 @@ def prepare(self) -> None: self.cached_token_lens_cuda[:num_requests] + self._seq_lens_cuda[:num_requests] ) cached_tokens_cuda = self.cached_token_lens_cuda[:num_requests] + if self.v41_layout is not None: + cached_tokens_cuda = self.v41_compute_start_pos_cuda[:num_requests] self.prepare_compressed_kv_metadata(kv_lens_cuda, cached_tokens_cuda, ctx_output_sizes) self._compute_compressed_mask( @@ -834,6 +1339,10 @@ def prepare(self) -> None: self._compress_ratios_sorted, ) + if self.v41_layout is not None: + self._prepare_v41_address_book(token_positions) + self._prepare_v41_ctx_index_k(token_positions) + def _sync_gen_tokens_per_seq(self, num_gen_tokens: int) -> int: """Synchronize the compressor's scalar and per-request token counts. @@ -941,7 +1450,11 @@ def on_update_kv_lens(self): ctx_output_sizes = ( self._ctx_output_sizes if getattr(self, "num_chunked_ctx_requests", 0) == 0 else None ) - self.prepare_compressed_kv_metadata(kv_lens, cached_tokens, ctx_output_sizes) + compressed_start = cached_tokens + if self.v41_layout is not None: + # Replay rows compute no compressed entry (see _prepare_v41_replay). + compressed_start = cached_tokens + self.v41_num_replay_tokens_cuda[:batch_size] + self.prepare_compressed_kv_metadata(kv_lens, compressed_start, ctx_output_sizes) self._compute_compressed_mask( self.new_comp_kv_lens_cuda, @@ -963,6 +1476,203 @@ def on_update_kv_lens(self): self.prepare_for_deepseek_v4_indices(token_positions) + if self.v41_layout is not None: + self._prepare_v41_address_book(token_positions) + self._prepare_v41_swa_seq_lens() + + @staticmethod + def _v41_pages( + block_table: torch.Tensor, + req_idx: torch.Tensor, + entry_idx: torch.Tensor, + entries_per_page: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Page id and in-page slot of entry ``entry_idx`` of request row ``req_idx``. + + ``block_table`` is ``[..., rows, max_blocks]``; the page has its leading + dims followed by ``entry_idx``'s. The page is -1 where the entry is + negative, lies past the table, or sits on an unallocated page. The + block index is clamped before the gather, so stale rows of a captured + graph never read out of bounds. Graph-safe: no host synchronization. + """ + max_blocks = block_table.shape[-1] + block = torch.div(entry_idx, entries_per_page, rounding_mode="floor") + in_table = (entry_idx >= 0) & (block < max_blocks) + page = block_table[..., req_idx, block.clamp(0, max_blocks - 1)].to(torch.int64) + page = torch.where(in_table & (page >= 0), page, -1) + return page, torch.remainder(entry_idx, entries_per_page) + + def _v41_index_k_offsets( + self, + block_table: torch.Tensor, + req_idx: torch.Tensor, + entry_idx: torch.Tensor, + entries_per_page: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Flat byte offsets of an index-K entry's code bytes and scale word. + + A page of ``block`` entries holds ``block * data_bytes`` code bytes, + then ``block * scale_bytes`` scale bytes (the DeepGEMM paged layout, + as in the DSA indexer's ``_compute_slot_mappings``). -1 where the page + is -1. + """ + data_bytes = self._v41_index_k_data_bytes + scale_bytes = self._v41_index_k_scale_bytes + page, slot = self._v41_pages(block_table, req_idx, entry_idx, entries_per_page) + page_base = page * (entries_per_page * (data_bytes + scale_bytes)) + valid = page >= 0 + data = torch.where(valid, page_base + slot * data_bytes, -1) + scale = torch.where( + valid, page_base + entries_per_page * data_bytes + slot * scale_bytes, -1 + ) + return data, scale + + def _prepare_v41_address_book(self, token_positions: torch.Tensor) -> None: + """DeepSeek-V4.1 write addresses of this step's KV rows, on device. + + Runs at the end of prepare() and of on_update_kv_lens(), so it is + recomputed inside CUDA-graph capture and replay: it only reads device + tensors and host ints (the token count, ``num_seqs``, + ``num_total_compressed_tokens``), and writes the persistent buffers in + place. + + - ``v41_swa_slot_mapping[l, t]`` for local layer ``l`` and token ``t`` + at position ``p`` of request row ``b``: + ``page * tokens_per_block + p % tokens_per_block`` with + ``page = sliding_block_tables[l, SWA, b, p // tokens_per_block]``. + - Per ratio r, row ``i`` of the packed compressed stream (the order of + ``compressed_position_ids_cuda[r]``), request row ``b`` from + ``cu_new_comp_kv_cuda[r]``, entry ``k = compressed_position_ids // r``, + ``block = tokens_per_block // r``: + ``v41_compress_slot_mapping[r][i] = page * block + k % block`` with + ``page = compress_block_tables[r][b, k // block]``, and + ``v41_index_k_slot_data[r][i]`` / ``v41_index_k_slot_scale[r][i]`` + from ``indexer_compress_block_tables[r]``. + + Rows whose ``compressed_mask_cuda[r]`` is False (reserved generation + slots) or whose page is unallocated get -1. The masked rows' positions + in ``compressed_position_ids_cuda[r]`` are set to 0: they are derived + past the last request's entries and can reach ``max_seq_len``, and a + RoPE table lookup reads every row. + """ + manager = self.kv_cache_manager + tokens_per_block = manager.tokens_per_block + device = token_positions.device + num_tokens = token_positions.shape[0] + + req_idx = self.req_idx_per_token[:num_tokens].to(torch.int64) + positions = token_positions.to(torch.int64) + swa_table = self.sliding_block_tables[:, DeepseekV4AttentionType.SWA.value] + page, slot = self._v41_pages(swa_table, req_idx, positions, tokens_per_block) + self.v41_swa_slot_mapping[:, :num_tokens] = torch.where( + page >= 0, page * tokens_per_block + slot, -1 + ) + + batch_size = self.num_seqs + for compress_ratio in self._compress_ratios_sorted: + total = self.num_total_compressed_tokens[compress_ratio] + if total == 0: + continue + entries_per_page = tokens_per_block // compress_ratio + cu_new_comp = self.cu_new_comp_kv_cuda[compress_ratio][: batch_size + 1] + rows = torch.arange(total, dtype=cu_new_comp.dtype, device=device) + row_req = torch.searchsorted(cu_new_comp, rows, right=True) - 1 + row_req = row_req.clamp(0, batch_size - 1).to(torch.int64) + valid = self.compressed_mask_cuda[compress_ratio][:total] + position_ids = self.compressed_position_ids_cuda[compress_ratio][:total] + entry = torch.where( + valid, + torch.div(position_ids, compress_ratio, rounding_mode="floor").to(torch.int64), + -1, + ) + + page, slot = self._v41_pages( + self.compress_block_tables[compress_ratio], row_req, entry, entries_per_page + ) + self.v41_compress_slot_mapping[compress_ratio][:total] = torch.where( + page >= 0, page * entries_per_page + slot, -1 + ) + data, scale = self._v41_index_k_offsets( + self.indexer_compress_block_tables[compress_ratio], + row_req, + entry, + entries_per_page, + ) + self.v41_index_k_slot_data[compress_ratio][:total] = data + self.v41_index_k_slot_scale[compress_ratio][:total] = scale + position_ids.masked_fill_(~valid, 0) + + def _reset_v41_ctx_index_k(self) -> None: + ratios = self._compress_ratios_sorted + self.v41_max_ctx_q_len = 0 + self.v41_ctx_index_k_len = {compress_ratio: 0 for compress_ratio in ratios} + self.v41_ctx_index_k_gather_data = {compress_ratio: None for compress_ratio in ratios} + self.v41_ctx_index_k_gather_scale = {compress_ratio: None for compress_ratio in ratios} + self.v41_ctx_index_ks = {compress_ratio: None for compress_ratio in ratios} + self.v41_ctx_index_ke = {compress_ratio: None for compress_ratio in ratios} + + def _prepare_v41_ctx_index_k(self, token_positions: torch.Tensor) -> None: + """DeepSeek-V4.1 prefill indexer inputs of the context requests (eager). + + Per ratio r, with ``kv_len_b = cached_b + seq_len_b`` of context + request ``b``: + + - ``v41_ctx_index_k_len[r]`` (host int): ``S_r = sum_b kv_len_b // r``. + - ``v41_ctx_index_k_gather_data[r]`` / ``v41_ctx_index_k_gather_scale[r]``: + int64 ``[S_r]`` byte offsets into the index-K pool of every + compressed entry ``k = 0 .. kv_len_b // r - 1`` of every context + request (cached entries included), request after request in batch + order; -1 on an unallocated page. + - ``v41_ctx_index_ks[r]`` / ``v41_ctx_index_ke[r]``: int32 + ``[num_ctx_tokens]``; context token ``t`` at position ``p`` of + request ``b`` scores rows ``[ks, ke)`` of the gathered array: ``ks`` + = the first row of ``b``, ``ke = ks + (p + 1) // r``. + - ``v41_max_ctx_q_len`` (host int): the longest context chunk. + + Sizes come from host lengths, so nothing reads the device. Every call + replaces the previous step's values: without context requests the host + ints are 0 and the tensors None. on_update_kv_lens() does not touch + them. + """ + self._reset_v41_ctx_index_k() + num_contexts = self.num_contexts + if num_contexts == 0: + return + device = token_positions.device + num_ctx_tokens = self.num_ctx_tokens + seq_lens = self.seq_lens[:num_contexts] + kv_lens = self.cached_token_lens_cpu[:num_contexts] + seq_lens + kv_lens_cuda = ( + self.cached_token_lens_cuda[:num_contexts] + self._seq_lens_cuda[:num_contexts] + ).to(torch.int64) + self.v41_max_ctx_q_len = int(seq_lens.max()) + + ctx_req = self.req_idx_per_token[:num_ctx_tokens].to(torch.int64) + visible_end = token_positions[:num_ctx_tokens].to(torch.int64) + 1 + tokens_per_block = self.kv_cache_manager.tokens_per_block + for compress_ratio in self._compress_ratios_sorted: + total = int((kv_lens // compress_ratio).sum()) + cu_comp = torch.nn.functional.pad( + torch.cumsum(kv_lens_cuda // compress_ratio, dim=0), (1, 0) + ) + rows = torch.arange(total, dtype=torch.int64, device=device) + row_req = torch.searchsorted(cu_comp, rows, right=True) - 1 + entry = rows - cu_comp[row_req] + data, scale = self._v41_index_k_offsets( + self.indexer_compress_block_tables[compress_ratio], + row_req, + entry, + tokens_per_block // compress_ratio, + ) + ks = cu_comp[ctx_req] + self.v41_ctx_index_k_len[compress_ratio] = total + self.v41_ctx_index_k_gather_data[compress_ratio] = data + self.v41_ctx_index_k_gather_scale[compress_ratio] = scale + self.v41_ctx_index_ks[compress_ratio] = ks.to(torch.int32) + self.v41_ctx_index_ke[compress_ratio] = ( + ks + torch.div(visible_end, compress_ratio, rounding_mode="floor") + ).to(torch.int32) + @staticmethod @maybe_compile(dynamic=True, options={"max-autotune": True}) def _compute_per_ratio_kv_lens( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/params.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/params.py index baa60f4f5154..b6d080db3758 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/params.py @@ -28,6 +28,12 @@ class DeepseekV4AttentionType(Enum): COMPRESS = 5 INDEXER_COMPRESS = 6 + # Sliding-window role of DeepSeek-V4.1 only: the Engram token history, one + # int32 compressed token id per position. Appended after the DeepSeek-V4 + # roles and kept out of DEEPSEEK_V4_SLIDING_ATTENTION / DEEPSEEK_V4_ATTENTION, + # so DeepSeek-V4's table widths and disaggregated layer ids do not change. + ENGRAM_TOKEN = 7 + @property def role(self) -> DataRole: return DataRole(f"deepseek_v4_{self.name.lower()}") @@ -48,6 +54,14 @@ def role(self) -> DataRole: DeepseekV4AttentionType.COMPRESS, DeepseekV4AttentionType.INDEXER_COMPRESS, ) +DEEPSEEK_V4_ATTENTION = DEEPSEEK_V4_SLIDING_ATTENTION + DEEPSEEK_V4_NON_SLIDING_ATTENTION + + +def is_sliding_attention_type(attn_type: DeepseekV4AttentionType) -> bool: + return ( + attn_type in DEEPSEEK_V4_SLIDING_ATTENTION + or attn_type == DeepseekV4AttentionType.ENGRAM_TOKEN + ) def is_overlap_compressor(compress_ratio: int) -> bool: @@ -80,6 +94,9 @@ def compress_ratio_has_attention(compress_ratio: int, attn_type: DeepseekV4Atten return is_sparse if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE: return is_sparse + if attn_type == DeepseekV4AttentionType.ENGRAM_TOKEN: + # Owned per layer in the DeepSeek-V4.1 layout, never by compress ratio. + return False raise ValueError(f"Unsupported DeepSeek-V4 attention type: {attn_type}") diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/v41_layout.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/v41_layout.py new file mode 100644 index 000000000000..2ef0c50e3c1e --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/v41_layout.py @@ -0,0 +1,602 @@ +# 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. +"""DeepSeek-V4.1 KV-cache layout: which layer owns which pool and who reads it. + +A DeepSeek-V4.1 model can run on the DeepSeek-V4 cache manager and attention +metadata in a separate layout mode. ``ModelConfig.from_pretrained`` gives a +DeepSeek-V4.1 checkpoint the CSA2 sparse config, which the in-tree DeepSeek-V4.1 +model runs on; a model built for this layout mode installs the ``deepseek_v4`` +sparse config instead (``use_deepseek_v41_layout_mode``). With that config, the +mode is selected by the checkpoint's own config +(``architectures[0] == "DeepseekV41ForCausalLM"`` or a ``deepseek_v41`` +``model_type``), never by a user option, and it reads the per-layer roles from +that config rather than from ``DeepSeekV4SparseAttentionConfig.compress_ratios`` +(whose validator maps ratio 0 to 1, erasing the difference between a +window-only layer and a ratio-1 compressed layer). + +Per backbone layer ``L`` with ``r = compress_ratios[L]``: + +- ``r == 0``: window-only. Owns an SWA pool and nothing else. +- ``r > 0``: attends the window plus the compressed KV of ``kv_source(L)``, the + latest KV-source layer at or before ``L``. Only KV-source layers own + ``COMPRESS`` (main KV) and ``INDEXER_COMPRESS`` (index K) pools; the other + compressed layers are readers and own no compressed pool. KV sources with + ``r > 1`` also own the compressor pending pair (``COMPRESSOR_KV`` / + ``COMPRESSOR_SCORE``); ratio 1 has no pending state. +- ``index_source(L)``: the latest index-source layer at or before ``L``, whose + top-k selection ``L`` consumes. An index source that is not a KV source (a + "Reindex" layer) scores its own queries over the index K of its KV source. +- The Engram token history (``ENGRAM_TOKEN``, one int32 compressed token id + per position) is owned by the first Engram layer and shared by all of them. + +The reader-to-owner maps live here and not in the cache manager's +``(layer, role) -> manager layer`` table, whose disaggregated-transfer encoding +assumes one owner per entry. +""" + +from __future__ import annotations + +import bisect +import math +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Iterable, Optional, Tuple + +from .params import DeepseekV4AttentionType + +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig + +DEEPSEEK_V41_ARCHITECTURES = ("DeepseekV41ForCausalLM",) +DEEPSEEK_V41_MODEL_TYPES = ("deepseek_v41", "deepseek_v41_text") +DEEPSEEK_V41_SUPPORTED_RATIOS = (0, 1, 2) + +# Storage of the SWA and main-KV pools; index K is native MXFP4 in both modes. +# "bf16" stores the reference's fake-quantized values in bf16 containers +# (1024 B per entry). "native" stores the same values as uint8 entries of codes +# followed by their block scales: SWA E4M3 [512] + UE8M0 [16] = 528 B, main KV +# packed E2M1 [256] + E4M3 [32] = 288 B (no global scale). +DEEPSEEK_V41_KV_STORAGE_ENV = "TRTLLM_DSV41_KV_STORAGE" +DEEPSEEK_V41_KV_STORAGE_BF16 = "bf16" +DEEPSEEK_V41_KV_STORAGE_NATIVE = "native" +DEEPSEEK_V41_KV_STORAGES = (DEEPSEEK_V41_KV_STORAGE_BF16, DEEPSEEK_V41_KV_STORAGE_NATIVE) +# Storage used when TRTLLM_DSV41_KV_STORAGE is unset. +DEEPSEEK_V41_DEFAULT_KV_STORAGE = DEEPSEEK_V41_KV_STORAGE_NATIVE + +# Encoder SWA bounded replay (tech report 3.2.2): the prefix cache keeps only global KV +# (main KV + index K); a hit of ``reused`` tokens re-forwards the last +# ``min(window, reused)`` cached tokens as replay rows, which rebuild the sliding windows. +DEEPSEEK_V41_ENCODER_REPLAY_ENV = "TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY" +# Replay mode used when TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY is unset. The cache +# manager resolves an unset switch off where replay is not supported (disaggregated +# serving, a KV connector, a separately reused draft pool, a token +# budget of at most one window, or a backend without the rebuild flag); an explicit +# 1 there raises instead. +DEEPSEEK_V41_DEFAULT_ENCODER_REPLAY = True +# Decoder SWA bounded replay (tech report 3.2.2): the layers after the last KV-source +# layer run, in a request's final context chunk, over its last ``window`` prompt +# positions only, so their sliding-window KV holds that tail and nothing below it. The +# model applies it; the cache manager resolves the switch against the engine (unset: +# on, attention TP or DP, except off when block reuse is on while encoder replay is +# off, since a prefix-cache hit inside a tail would then keep window +# KV no forward of the request wrote). +DEEPSEEK_V41_DECODER_REPLAY_ENV = "TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY" +_TRUE_STRINGS = ("1", "true", "on", "yes") +_FALSE_STRINGS = ("0", "false", "off", "no") + +# Roles a DeepSeek-V4.1 layout allocates. INDEXER_COMPRESSOR_* (V4's ratio-4 +# overlap indexer compressor) never exists in V4.1. +DEEPSEEK_V41_SLIDING_ATTENTION = ( + DeepseekV4AttentionType.SWA, + DeepseekV4AttentionType.COMPRESSOR_KV, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + DeepseekV4AttentionType.ENGRAM_TOKEN, +) +DEEPSEEK_V41_NON_SLIDING_ATTENTION = ( + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, +) +DEEPSEEK_V41_ATTENTION = DEEPSEEK_V41_SLIDING_ATTENTION + DEEPSEEK_V41_NON_SLIDING_ATTENTION +# Width of the per-layer sliding block table in V4.1 mode, indexed by +# ``attn_type.value`` exactly as in V4. Columns of roles V4.1 never allocates +# stay BAD_PAGE_INDEX. +DEEPSEEK_V41_SLIDING_TABLE_COLUMNS = max(t.value for t in DEEPSEEK_V41_SLIDING_ATTENTION) + 1 + + +def _cfg_get(config: Any, name: str, default: Any = None) -> Any: + if config is None: + return default + if isinstance(config, dict): + return config.get(name, default) + return getattr(config, name, default) + + +def _text_config(config: Any) -> Any: + text_config = _cfg_get(config, "text_config") + return text_config if text_config is not None else config + + +def is_deepseek_v41_config(config: Any) -> bool: + """Whether ``config`` (composite or text sub-config) is a DeepSeek-V4.1 checkpoint.""" + if config is None: + return False + architectures = _cfg_get(config, "architectures") or () + if len(architectures) > 0 and architectures[0] in DEEPSEEK_V41_ARCHITECTURES: + return True + if _cfg_get(config, "model_type") in DEEPSEEK_V41_MODEL_TYPES: + return True + return _cfg_get(_cfg_get(config, "text_config"), "model_type") in DEEPSEEK_V41_MODEL_TYPES + + +def get_deepseek_v41_kv_storage() -> str: + storage = os.environ.get(DEEPSEEK_V41_KV_STORAGE_ENV, DEEPSEEK_V41_DEFAULT_KV_STORAGE) + storage = storage.strip().lower() or DEEPSEEK_V41_DEFAULT_KV_STORAGE + if storage not in DEEPSEEK_V41_KV_STORAGES: + raise ValueError( + f"{DEEPSEEK_V41_KV_STORAGE_ENV}={storage!r} is not one of {DEEPSEEK_V41_KV_STORAGES}." + ) + return storage + + +def _get_switch_setting(name: str) -> Optional[bool]: + """The explicit value of boolean switch ``name``, or None when unset.""" + value = os.environ.get(name, "").strip().lower() + if not value: + return None + if value in _TRUE_STRINGS: + return True + if value in _FALSE_STRINGS: + return False + raise ValueError(f"{name}={value!r} is not one of {_TRUE_STRINGS + _FALSE_STRINGS}.") + + +def get_deepseek_v41_encoder_replay_setting() -> Optional[bool]: + """The explicit value of ``TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY``, or None when unset.""" + return _get_switch_setting(DEEPSEEK_V41_ENCODER_REPLAY_ENV) + + +def get_deepseek_v41_decoder_replay_setting() -> Optional[bool]: + """The explicit value of ``TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY``, or None when unset.""" + return _get_switch_setting(DEEPSEEK_V41_DECODER_REPLAY_ENV) + + +def get_deepseek_v41_encoder_replay() -> bool: + """The switch's value, ``DEEPSEEK_V41_DEFAULT_ENCODER_REPLAY`` when unset (before the + cache manager resolves an unset switch against the engine configuration).""" + setting = get_deepseek_v41_encoder_replay_setting() + return DEEPSEEK_V41_DEFAULT_ENCODER_REPLAY if setting is None else setting + + +def deepseek_v41_native_entry_split( + attn_type: DeepseekV4AttentionType, head_dim: int +) -> Tuple[int, int]: + """(code bytes, scale bytes) of one native SWA or main-KV entry. + + An entry is one contiguous uint8 row: the codes, then their block scales. + SWA: E4M3 codes, one UE8M0 byte per 32 values. Main KV: E2M1 codes packed + two per byte (even element in the low nibble), one E4M3 byte per 16 values. + """ + if attn_type == DeepseekV4AttentionType.SWA: + return head_dim, head_dim // 32 + if attn_type == DeepseekV4AttentionType.COMPRESS: + return head_dim // 2, head_dim // 16 + raise ValueError(f"DeepSeek-V4.1 stores {attn_type.name} in one format only.") + + +def deepseek_v41_token_bytes( + attn_type: DeepseekV4AttentionType, + head_dim: int, + index_head_dim: int, + kv_storage: str = DEEPSEEK_V41_KV_STORAGE_BF16, +) -> int: + """Bytes one cache entry of ``attn_type`` occupies in the V4.1 layout.""" + if attn_type in (DeepseekV4AttentionType.SWA, DeepseekV4AttentionType.COMPRESS): + if kv_storage == DEEPSEEK_V41_KV_STORAGE_NATIVE: + return sum(deepseek_v41_native_entry_split(attn_type, head_dim)) + return head_dim * 2 + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + # MXFP4: packed E2M1 codes plus one UE8M0 byte per 32 values, in every storage mode. + return index_head_dim // 2 + index_head_dim // 32 + if attn_type in ( + DeepseekV4AttentionType.COMPRESSOR_KV, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + ): + return head_dim * 4 + if attn_type == DeepseekV4AttentionType.ENGRAM_TOKEN: + return 4 + raise ValueError(f"DeepSeek-V4.1 has no cache role {attn_type.name}.") + + +@dataclass(frozen=True) +class DeepseekV41Layout: + """Per-layer cache roles of a DeepSeek-V4.1 backbone.""" + + # One entry per backbone layer (DSpark stage entries of the checkpoint's + # list are dropped): 0 = window only, 1 or 2 = window plus compressed KV. + compress_ratios: Tuple[int, ...] + kv_source_layer_ids: Tuple[int, ...] + index_source_layer_ids: Tuple[int, ...] + candidate_source_layer_id: Optional[int] + engram_layer_ids: Tuple[int, ...] + window_size: int + head_dim: int + index_head_dim: int + index_topk: int + # Positions of Engram history an n-gram hash reads before a run start. + engram_history_len: int + kv_storage: str = DEEPSEEK_V41_DEFAULT_KV_STORAGE + # Encoder SWA bounded replay: windowed roles stay out of the prefix cache and a + # reuse hit re-forwards the last ``window_size`` cached tokens. + encoder_replay: bool = DEEPSEEK_V41_DEFAULT_ENCODER_REPLAY + + def __post_init__(self) -> None: + num_layers = len(self.compress_ratios) + bad = sorted(set(self.compress_ratios) - set(DEEPSEEK_V41_SUPPORTED_RATIOS)) + if bad: + raise ValueError( + f"DeepSeek-V4.1 compress_ratios must be in {DEEPSEEK_V41_SUPPORTED_RATIOS}, got {bad}." + ) + for name, ids in ( + ("kv_source_layer_ids", self.kv_source_layer_ids), + ("index_source_layer_ids", self.index_source_layer_ids), + ("engram_layer_ids", self.engram_layer_ids), + ): + if list(ids) != sorted(set(ids)) or any(not 0 <= i < num_layers for i in ids): + raise ValueError( + f"DeepSeek-V4.1 {name} must be sorted, unique backbone layer ids in " + f"[0, {num_layers}), got {list(ids)}." + ) + for layer in self.kv_source_layer_ids: + if self.compress_ratios[layer] == 0: + raise ValueError(f"KV-source layer {layer} has compress ratio 0.") + for layer in self.index_source_layer_ids: + if self.compress_ratios[layer] == 0: + raise ValueError(f"Index-source layer {layer} has compress ratio 0.") + for layer, ratio in enumerate(self.compress_ratios): + if ratio == 0: + continue + kv_source = self.kv_source(layer) + if kv_source is None or self.compress_ratios[kv_source] != ratio: + raise ValueError( + f"Layer {layer} (ratio {ratio}) has no KV-source layer of the same ratio " + f"at or before it (kv_source_layer_ids={list(self.kv_source_layer_ids)})." + ) + index_source = self.index_source(layer) + if index_source is None or self.kv_source(index_source) != kv_source: + raise ValueError( + f"Layer {layer} reads the compressed KV of layer {kv_source} but its index " + f"source {index_source} selects over a different pool." + ) + if self.candidate_source_layer_id is not None and ( + self.candidate_source_layer_id not in self.kv_source_layer_ids + or self.candidate_source_layer_id not in self.index_source_layer_ids + ): + raise ValueError( + f"candidate_source_layer_id {self.candidate_source_layer_id} must be both a " + "KV source and an index source." + ) + if self.kv_storage not in DEEPSEEK_V41_KV_STORAGES: + raise ValueError(f"Unknown DeepSeek-V4.1 KV storage {self.kv_storage!r}.") + + @classmethod + def from_config( + cls, + config: Any, + kv_storage: Optional[str] = None, + encoder_replay: Optional[bool] = None, + ) -> "DeepseekV41Layout": + """Build the layout from a DeepSeek-V4.1 composite config or its text sub-config.""" + text = _text_config(config) + + def required(name: str) -> Any: + value = _cfg_get(text, name) + if value is None: + raise ValueError(f"DeepSeek-V4.1 config is missing {name!r}.") + return value + + num_layers = int(required("num_hidden_layers")) + ratios = [int(r) for r in required("compress_ratios")] + if len(ratios) < num_layers: + raise ValueError( + f"compress_ratios has {len(ratios)} entries for {num_layers} backbone layers." + ) + head_dim = _cfg_get(text, "head_dim") + if head_dim is None: + head_dim = int(required("kv_lora_rank")) + int(required("qk_rope_head_dim")) + candidate = _cfg_get(text, "candidate_source_layer_id") + return cls( + compress_ratios=tuple(ratios[:num_layers]), + kv_source_layer_ids=tuple(int(i) for i in required("kv_source_layer_ids")), + index_source_layer_ids=tuple(int(i) for i in required("index_source_layer_ids")), + candidate_source_layer_id=None if candidate is None else int(candidate), + engram_layer_ids=tuple(int(i) for i in (_cfg_get(text, "engram_layer_ids") or ())), + window_size=int(required("sliding_window")), + head_dim=int(head_dim), + index_head_dim=int(_cfg_get(text, "index_head_dim", 128)), + index_topk=int(_cfg_get(text, "index_topk", 512)), + engram_history_len=max(0, int(_cfg_get(text, "engram_max_ngram_size", 4)) - 1), + kv_storage=get_deepseek_v41_kv_storage() if kv_storage is None else kv_storage, + encoder_replay=( + get_deepseek_v41_encoder_replay() if encoder_replay is None else encoder_replay + ), + ) + + @property + def num_layers(self) -> int: + return len(self.compress_ratios) + + @property + def compressed_ratios(self) -> Tuple[int, ...]: + """Ratios that have a compressed stream (a COMPRESS pool per KV source).""" + return tuple(sorted({self.compress_ratios[i] for i in self.kv_source_layer_ids})) + + @property + def layer_ratios(self) -> Tuple[int, ...]: + """Every ratio a backbone layer has, window-only (0) included.""" + return tuple(sorted(set(self.compress_ratios))) + + @property + def replay_window(self) -> int: + """Most cached tokens a reuse hit re-forwards (``n_win``, the sliding window).""" + return self.window_size + + @property + def reuse_alignment(self) -> int: + """A reuse boundary under encoder replay is a multiple of every compress ratio, so + no compressed group straddles it and no compressor state is needed there.""" + return math.lcm(*self.compressed_ratios) if self.compressed_ratios else 1 + + @property + def engram_token_owner(self) -> Optional[int]: + return self.engram_layer_ids[0] if self.engram_layer_ids else None + + @staticmethod + def _latest_at_or_before(ids: Tuple[int, ...], layer: int) -> Optional[int]: + pos = bisect.bisect_right(ids, layer) + return ids[pos - 1] if pos > 0 else None + + def kv_source(self, layer: int) -> Optional[int]: + """Owner of the compressed KV and index K ``layer`` reads; None for window-only layers.""" + if self.compress_ratios[layer] == 0: + return None + return self._latest_at_or_before(self.kv_source_layer_ids, layer) + + def index_source(self, layer: int) -> Optional[int]: + """Layer whose top-k selection ``layer`` consumes; None for window-only layers.""" + if self.compress_ratios[layer] == 0: + return None + return self._latest_at_or_before(self.index_source_layer_ids, layer) + + def is_kv_source(self, layer: int) -> bool: + return layer in self.kv_source_layer_ids + + def has_attention(self, layer: int, attn_type: DeepseekV4AttentionType) -> bool: + """Whether ``layer`` owns a pool of ``attn_type`` (owner-aware, not ratio-based).""" + if not 0 <= layer < self.num_layers: + return False + if attn_type == DeepseekV4AttentionType.SWA: + return True + if attn_type in ( + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ): + return self.is_kv_source(layer) + if attn_type in ( + DeepseekV4AttentionType.COMPRESSOR_KV, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + ): + return self.is_kv_source(layer) and self.compress_ratios[layer] > 1 + if attn_type == DeepseekV4AttentionType.ENGRAM_TOKEN: + return layer == self.engram_token_owner + return False + + def owner(self, layer: int, attn_type: DeepseekV4AttentionType) -> Optional[int]: + """The layer whose pool of ``attn_type`` ``layer`` uses, or None if it uses none.""" + if attn_type == DeepseekV4AttentionType.SWA: + return layer + if attn_type in ( + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ): + return self.kv_source(layer) + if attn_type == DeepseekV4AttentionType.ENGRAM_TOKEN: + return self.engram_token_owner if layer in self.engram_layer_ids else None + return layer if self.has_attention(layer, attn_type) else None + + def attention_types(self, layer: int) -> Tuple[DeepseekV4AttentionType, ...]: + return tuple(t for t in DEEPSEEK_V41_ATTENTION if self.has_attention(layer, t)) + + def window(self, attn_type: DeepseekV4AttentionType, max_draft_len: int = 0) -> Optional[int]: + """Sliding-window size in tokens of a windowed role; None for persistent roles.""" + if attn_type == DeepseekV4AttentionType.SWA: + base = self.window_size + elif attn_type in ( + DeepseekV4AttentionType.COMPRESSOR_KV, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + ): + # One pending group of the largest ratio, indexed by absolute position. + base = max(self.compressed_ratios, default=1) + elif attn_type == DeepseekV4AttentionType.ENGRAM_TOKEN: + # The manager keeps positions [h + 1 - window, h] live, h being the + # first position of the step, and the hash of h reads the + # engram_history_len positions before it. A replay step starts + # replay_window positions before the reuse boundary h, and its + # first rows read the history before that. + base = self.engram_history_len + 1 + if self.encoder_replay: + base += self.replay_window + else: + return None + return base + max_draft_len + + def stage_dependency_error(self, layers: Iterable[int]) -> Optional[str]: + """Why a pipeline stage holding ``layers`` cannot run on its own pools, or None. + + Every cross-layer read of a V4.1 backbone must stay on the stage: a compressed + layer reads the main KV and index K of its KV source and the top-k of its index + source; an index source of the candidate source's ratio after it applies + that layer's candidate-block mask; an Engram layer hashes over the token history owned by + the first Engram layer. A pool, or a per-step product, of another stage is not + reachable from this one. + """ + local = set(layers) + cand = self.candidate_source_layer_id + for layer in sorted(local): + if self.compress_ratios[layer]: + for what, source in ( + ("KV source", self.kv_source(layer)), + ("index source", self.index_source(layer)), + ): + if source not in local: + return ( + f"layer {layer} reads its {what} layer {source}, held by another stage" + ) + if ( + cand is not None + and layer > cand + and layer in self.index_source_layer_ids + and self.compress_ratios[layer] == self.compress_ratios[cand] + and cand not in local + ): + return ( + f"layer {layer} applies the candidate-block mask of layer {cand}, held by " + "another stage" + ) + if layer in self.engram_layer_ids and self.engram_token_owner not in local: + return ( + f"Engram layer {layer} reads the token history of layer " + f"{self.engram_token_owner}, held by another stage" + ) + return None + + def pipeline_boundaries(self) -> Tuple[int, ...]: + """Layers ``c`` at which a two-stage split ``[0, c) | [c, num_layers)`` is valid.""" + return tuple( + c + for c in range(1, self.num_layers) + if self.stage_dependency_error(range(c)) is None + and self.stage_dependency_error(range(c, self.num_layers)) is None + ) + + def persistent_bytes_per_token(self, layers: Iterable[int]) -> int: + """Global (non-windowed) cache bytes per token of ``layers``: main KV plus index K.""" + total = 0 + for layer in layers: + if not self.is_kv_source(layer): + continue + ratio = self.compress_ratios[layer] + for attn_type in DEEPSEEK_V41_NON_SLIDING_ATTENTION: + total += ( + deepseek_v41_token_bytes( + attn_type, self.head_dim, self.index_head_dim, self.kv_storage + ) + // ratio + ) + return total + + +def get_deepseek_v41_layout(config: Any) -> Optional[DeepseekV41Layout]: + """The V4.1 layout of ``config``, or None when it is not a DeepSeek-V4.1 checkpoint.""" + if not is_deepseek_v41_config(config): + return None + return DeepseekV41Layout.from_config(config) + + +def deepseek_v41_sparse_attention_config( + pretrained_config: Any, +) -> "DeepSeekV4SparseAttentionConfig": + """The ``deepseek_v4`` sparse config that runs a DeepSeek-V4.1 checkpoint in the layout mode. + + The ``deepseek_v4`` algorithm selects ``DeepseekV4CacheManager`` and + ``DeepseekV4TrtllmAttentionMetadata``, which enter the V4.1 layout mode from + ``pretrained_config``. The per-layer layout is not expressible in this config: + ratio 0 (window only) and ratio 1 (window plus the global KV of the ratio-1 + source layer) are different roles, and ``DeepSeekV4SparseAttentionConfig`` + normalizes 0 to 1. The manager and the metadata therefore read the layout from + ``pretrained_config``; ``compress_ratios`` here is a placeholder of the + checkpoint's length in which every layer has ratio 1. The window and the indexer + geometry come from the checkpoint; the other indexer settings take the values + ``ModelConfig.from_pretrained`` gives a DeepSeek-V4 checkpoint loaded without a + user sparse config, with the index K stored as MXFP4. + + Args: + pretrained_config: DeepSeek-V4.1 composite config or its text sub-config. + """ + from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig + + text = _text_config(pretrained_config) + checkpoint_ratios = _cfg_get(text, "compress_ratios") + if not checkpoint_ratios: + raise ValueError("DeepSeek-V4.1 requires per-layer `compress_ratios` in the checkpoint.") + window_size = _cfg_get(text, "sliding_window") + if window_size is None: + raise ValueError("DeepSeek-V4.1 config is missing 'sliding_window'.") + return DeepSeekV4SparseAttentionConfig( + compress_ratios=[1] * len(checkpoint_ratios), + window_size=int(window_size), + index_n_heads=_cfg_get(text, "index_n_heads"), + index_head_dim=int(_cfg_get(text, "index_head_dim", 128)), + index_topk=int(_cfg_get(text, "index_topk", 512)), + indexer_max_chunk_size=None, + skip_indexer_for_short_seqs=True, + use_cute_dsl_topk=False, + use_cute_dsl_paged_mqa_logits=False, + q_split_threshold=8192, + indexer_rope_interleave=False, + enable_heuristic_topk=False, + use_self_sampling_topk=True, + use_gvr_emission=False, + indexer_k_dtype="fp4", + ) + + +def use_deepseek_v41_layout_mode(model_config: "ModelConfig") -> None: + """Make the engine run ``model_config``'s DeepSeek-V4.1 checkpoint in the layout mode. + + Replaces ``model_config.sparse_attention_config`` with + ``deepseek_v41_sparse_attention_config(model_config.pretrained_config)``. The + engine reads the sparse config off the built model's ``model_config`` to pick + the KV cache manager and the attention metadata, so a model calls this from its + constructor. ``ModelConfig.from_pretrained`` has given the checkpoint the CSA2 + config by then, the only one it accepts for DeepSeek-V4.1; a CSA2 config that + carries user-set fields is refused rather than silently dropped. + + Args: + model_config: The (frozen) ``ModelConfig`` of a DeepSeek-V4.1 checkpoint. + """ + pretrained_config = model_config.pretrained_config + if not is_deepseek_v41_config(pretrained_config): + raise ValueError("use_deepseek_v41_layout_mode() needs a DeepSeek-V4.1 checkpoint config.") + current = model_config.sparse_attention_config + if current is not None: + algorithm = getattr(current, "algorithm", None) + if algorithm == "deepseek_v4": + return + user_fields = set(getattr(current, "model_fields_set", ())) + if algorithm != "csa2" or user_fields: + raise ValueError( + "The DeepSeek-V4.1 layout mode runs on the deepseek_v4 sparse runtime and " + "takes its settings from the checkpoint; got sparse_attention_config " + f"{algorithm!r} with user-set fields {sorted(user_fields)}. Leave " + "sparse_attention_config unset." + ) + config = deepseek_v41_sparse_attention_config(pretrained_config) + frozen = model_config._frozen + model_config._frozen = False + try: + model_config.sparse_attention_config = config + finally: + model_config._frozen = frozen diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 6fac5ce27ff1..2dba83e25826 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -362,6 +362,27 @@ def autotune(tune_mode: bool = True, autotuner._dist = prev_dist +@contextlib.contextmanager +def fallback_tactic(): + """Make every ``AutoTuner.choose_one`` call inside the scope return the fallback. + + Inside the scope ``choose_one`` returns ``runners[0]`` with tactic ``-1`` + without profiling and without reading or writing the profiling cache, so + the choice does not depend on what this process tuned or loaded from a + cache file (``autotune(cache_path=...)``). ``autotune(tune_mode=False)`` + only stops tuning: a cached tactic is still taken. Use this for a + computation replicated across ranks that must launch the same kernels on + every rank. + """ + autotuner = AutoTuner.get() + old = autotuner.force_fallback_tactic + autotuner.force_fallback_tactic = True + try: + yield + finally: + autotuner.force_fallback_tactic = old + + @dataclass class AutoTunerStatistics: """Statistics collected by the AutoTuner. @@ -941,6 +962,8 @@ def __init__(self, warmup=2, repeat=10, stream_delay_micro_secs=1000): self._primed_cached_tactics: set[tuple[str, str, str, str]] = set() self.is_tuning_mode = False self.skip_dynamic_tuning_buckets = False + # Set by fallback_tactic(): choose_one returns the fallback, no cache lookup. + self.force_fallback_tactic = False # Timing backend: globaltimer kernel vs cuda events. # TLLM_PROFILING_TIMER env var overrides auto-detection: @@ -1139,6 +1162,9 @@ def choose_one( 'kwargs': kwargs, }) + if self.force_fallback_tactic: + return (runners[0], -1) + input_shapes = tuple(self._get_input_sizes(inputs)) # For ops opted out of the persistent cache (e.g. JIT-driven diff --git a/tensorrt_llm/_torch/custom_ops/e4m3_quantize_pow2_dynamic.py b/tensorrt_llm/_torch/custom_ops/e4m3_quantize_pow2_dynamic.py new file mode 100644 index 000000000000..9c2f9ce6f3e9 --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/e4m3_quantize_pow2_dynamic.py @@ -0,0 +1,132 @@ +# 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. +"""Dynamic per-tensor E4M3 quantization with a power-of-two scale: +``trtllm::e4m3_quantize_pow2_dynamic``. + +Over the whole tensor ``x`` (bf16 / fp16 / fp32, contiguous):: + + amax = max(max|x|, 1e-4) # fp32 + s = 2^ceil(log2(amax * (1 / 448))) # from the fp32 bits, exact + out = e4m3_rne(clamp(x * (1 / s), -448, 448)) # exact scaling by a power of two + scale_out[0] = s * scale_mult # one fp32 multiply + +so ``out * s`` is ``x`` rounded to E4M3's 3-bit mantissa (no value saturates: +``|x| / s <= 448``), and ``scale_out`` can be handed on directly as a fused +dequantization factor (e.g. an attention's ``bmm1_scale`` with +``scale_mult = softmax_scale * kv_scale``). The scale is the same recipe as the +MXFP8 activation quantizer's UE8M0 scale, taken over the whole tensor. + +Three launches: a memset of a one-float amax scratch, a grid of block-amax +programs with an atomic max, and a grid of quantizing programs (program 0 also +writes ``scale_out``). Graph-capturable; no host synchronization. +""" + +import torch +import triton +import triton.language as tl + +_BLOCK = 4096 +_AMAX_FLOOR = 1e-4 +_IN_DTYPES = (torch.bfloat16, torch.float16, torch.float32) + + +@triton.jit +def _amax_kernel(x_ptr, amax_ptr, n, BLOCK: tl.constexpr): + off = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(x_ptr + off, mask=off < n, other=0.0).to(tl.float32) + tl.atomic_max(amax_ptr, tl.max(tl.abs(x), axis=0)) + + +@triton.jit +def _pow2_scale(amax_ptr, floor): + a = tl.maximum(tl.load(amax_ptr), floor) * (1.0 / 448.0) + bits = a.to(tl.int32, bitcast=True) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) # biased exponent of the ceil + return (e << 23).to(tl.float32, bitcast=True) + + +@triton.jit +def _quant_kernel(x_ptr, out_ptr, amax_ptr, scale_ptr, n, scale_mult, floor, BLOCK: tl.constexpr): + s = _pow2_scale(amax_ptr, floor) + inv = ((254 << 23) - s.to(tl.int32, bitcast=True)).to(tl.float32, bitcast=True) # exact 1/s + pid = tl.program_id(0) + off = pid.to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(x_ptr + off, mask=off < n, other=0.0).to(tl.float32) + v = x * inv + v = tl.where(v != v, v, tl.minimum(tl.maximum(v, -448.0), 448.0)) # NaN stays NaN + tl.store(out_ptr + off, v.to(tl.float8e4nv).to(tl.uint8, bitcast=True), mask=off < n) + if pid == 0: + tl.store(scale_ptr, s * scale_mult) + + +@torch.library.custom_op("trtllm::e4m3_quantize_pow2_dynamic", mutates_args=("out", "scale_out")) +def e4m3_quantize_pow2_dynamic( + x: torch.Tensor, + out: torch.Tensor, + scale_out: torch.Tensor, + scale_mult: float, +) -> None: + """Quantize ``x`` to E4M3 into ``out`` with one power-of-two scale ``s`` derived from + ``max|x|`` on the device, and write ``s * scale_mult`` to ``scale_out[0]`` (fp32).""" + if not x.is_cuda or x.dtype not in _IN_DTYPES or not x.is_contiguous(): + raise ValueError( + f"x must be a contiguous bf16/fp16/fp32 CUDA tensor, got {x.dtype} on {x.device}" + ) + if out.device != x.device or out.dtype != torch.float8_e4m3fn or not out.is_contiguous(): + raise ValueError(f"out must be a contiguous float8_e4m3fn tensor on {x.device}") + if out.numel() != x.numel(): + raise ValueError(f"out has {out.numel()} elements, x has {x.numel()}") + if scale_out.device != x.device or scale_out.dtype != torch.float32 or scale_out.numel() < 1: + raise ValueError( + f"scale_out must be a float32 CUDA tensor with >= 1 element, got {scale_out.dtype}" + ) + n = x.numel() + amax = torch.zeros(1, dtype=torch.float32, device=x.device) + if n: + grid = (triton.cdiv(n, _BLOCK),) + _amax_kernel[grid](x, amax, n, BLOCK=_BLOCK, num_warps=8) + _quant_kernel[grid]( + x, + out.view(torch.uint8), + amax, + scale_out, + n, + float(scale_mult), + _AMAX_FLOOR, + BLOCK=_BLOCK, + num_warps=8, + ) + else: + _quant_kernel[(1,)]( + x, + out.view(torch.uint8), + amax, + scale_out, + 0, + float(scale_mult), + _AMAX_FLOOR, + BLOCK=_BLOCK, + num_warps=8, + ) + + +@e4m3_quantize_pow2_dynamic.register_fake +def _e4m3_quantize_pow2_dynamic_fake( + x: torch.Tensor, + out: torch.Tensor, + scale_out: torch.Tensor, + scale_mult: float, +) -> None: + return None diff --git a/tensorrt_llm/_torch/custom_ops/mxfp8_dequantize.py b/tensorrt_llm/_torch/custom_ops/mxfp8_dequantize.py new file mode 100644 index 000000000000..7e5c735aa1ee --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/mxfp8_dequantize.py @@ -0,0 +1,259 @@ +# 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. +"""MXFP8 row gather + dequantization to bf16: ``trtllm::mxfp8_dequantize``. + +The source is a pool of native MXFP8 rows: ``codes`` holds E4M3 codes +``[slots, W]`` (or paged ``[pages, page_size, W]``) and ``scales`` one UE8M0 +byte per 32-element block ``[..., W // 32]`` with the same slot geometry. +For output row ``r``:: + + s = src_slots[r] (r when src_slots is None) + d = dst_slots[r] (r when dst_slots is None) + out[d, i] = bf16(e4m3(codes[s, i]) * 2^(scales[s, i // 32] - 127)) + +A source slot that is negative or ``>= slots`` yields a row of zeros; a +destination slot that is negative or ``>= out slots`` is skipped. Slot ``s`` +of a paged tensor is ``(s // page_size, s % page_size)``; leading strides are +free, so interleaved ``[.., 528]`` byte rows and page-footer layouts are +passed as two strided views. One Triton program per output row; no +autotuning. +""" + +from typing import Optional + +import torch +import triton +import triton.language as tl + +MXFP8_BLOCK = 32 + +_CODE_DTYPES = (torch.float8_e4m3fn, torch.uint8) +_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) +_SLOT_DTYPES = (torch.int32, torch.int64) + + +@triton.jit +def _mxfp8_dequantize_kernel( + codes_ptr, + scales_ptr, + out_ptr, + src_ptr, + dst_ptr, + src_stride, + dst_stride, + codes_page_stride, + codes_slot_stride, + scales_page_stride, + scales_slot_stride, + src_page_size, + src_num_slots, + out_page_stride, + out_slot_stride, + out_page_size, + out_num_slots, + W: tl.constexpr, + NUM_BLOCKS_P2: tl.constexpr, + HAS_SRC: tl.constexpr, + HAS_DST: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + if HAS_DST: + dst = tl.load(dst_ptr + row * dst_stride).to(tl.int64) + else: + dst = row + if (dst >= 0) & (dst < out_num_slots): + if HAS_SRC: + src = tl.load(src_ptr + row * src_stride).to(tl.int64) + else: + src = row + valid = (src >= 0) & (src < src_num_slots) + src = tl.where(valid, src, 0) + + blk = tl.arange(0, NUM_BLOCKS_P2) + col = blk[:, None] * 32 + tl.arange(0, 32)[None, :] + mask = (col < W) & valid + s_page = src // src_page_size + s_in = src % src_page_size + + c = tl.load( + codes_ptr + s_page * codes_page_stride + s_in * codes_slot_stride + col, + mask=mask, + other=0, + ) + c = c.to(tl.float8e4nv, bitcast=True) + b = tl.load( + scales_ptr + s_page * scales_page_stride + s_in * scales_slot_stride + blk, + mask=(blk * 32 < W) & valid, + other=127, + ).to(tl.int32) + # UE8M0 byte b is 2^(b - 127): the fp32 exponent field itself for b in + # 1..254, the fp32 subnormal 2^-127 for b = 0, NaN for b = 255. + bits = tl.where(b == 0, 0x00400000, tl.where(b == 255, 0x7FC00000, b << 23)) + scale = bits.to(tl.float32, bitcast=True) + # An e4m3 value times a power of two is exact in fp32; the one rounding + # is the RNE conversion to bf16 (exact too unless the product falls in + # bf16's subnormal range). + y = (c.to(tl.float32) * scale[:, None]).to(tl.bfloat16) + + d_page = dst // out_page_size + d_in = dst % out_page_size + tl.store(out_ptr + d_page * out_page_stride + d_in * out_slot_stride + col, y, mask=col < W) + + +def _slot_geometry(t: torch.Tensor, name: str) -> tuple[int, int, int, int]: + """(page_size, num_slots, page_stride, slot_stride) of a [slots, C] or [pages, page_size, C] tensor.""" + if t.dim() == 2: + return 1, t.shape[0], t.stride(0), 0 + if t.dim() == 3: + return t.shape[1], t.shape[0] * t.shape[1], t.stride(0), t.stride(1) + raise ValueError( + f"{name} must be [slots, C] or [pages, page_size, C], got shape {tuple(t.shape)}" + ) + + +def _check_slots(slots: torch.Tensor, name: str, device: torch.device) -> None: + if slots.device != device: + raise ValueError(f"{name} must be on {device}, got {slots.device}") + if slots.dtype not in _SLOT_DTYPES: + raise ValueError(f"{name} must be int32 or int64, got {slots.dtype}") + if slots.dim() != 1: + raise ValueError(f"{name} must be 1-D, got shape {tuple(slots.shape)}") + + +@torch.library.custom_op("trtllm::mxfp8_dequantize", mutates_args=("out",)) +def mxfp8_dequantize( + codes: torch.Tensor, + scales: torch.Tensor, + out: torch.Tensor, + src_slots: Optional[torch.Tensor] = None, + dst_slots: Optional[torch.Tensor] = None, +) -> None: + """Gather MXFP8 rows (E4M3 codes x UE8M0 per 32) from a pool and write them to ``out`` as bf16. + + ``codes``: float8_e4m3fn / uint8 ``[slots, W]`` or ``[pages, page_size, W]``, + ``W % 32 == 0``, unit last-dim stride. ``scales``: uint8 / float8_e8m0fnu + ``codes.shape[:-1] + (W // 32,)``, unit last-dim stride. ``out``: bf16 + ``[out_slots, W]`` or ``[pages, page_size, W]``, unit last-dim stride. + ``src_slots`` / ``dst_slots``: int32/int64 ``[rows]`` (same length when both + are given). Row ``r`` reads slot ``src_slots[r]`` (``r`` when None; zeros + when negative or ``>= slots``) and writes slot ``dst_slots[r]`` of ``out`` + (``r`` when None; skipped when negative or ``>= out_slots``). With both + None, every source slot is dequantized to the same slot of ``out``. + """ + if not codes.is_cuda: + raise ValueError(f"codes must be a CUDA tensor, got device {codes.device}") + device = codes.device + if codes.dtype not in _CODE_DTYPES: + raise ValueError(f"codes must be float8_e4m3fn or uint8, got {codes.dtype}") + width = codes.shape[-1] if codes.dim() > 0 else 0 + if width <= 0 or width % MXFP8_BLOCK != 0: + raise ValueError( + f"codes' last dim must be a positive multiple of {MXFP8_BLOCK}, got {width}" + ) + if codes.stride(-1) != 1: + raise ValueError( + f"codes must have a unit stride along its last dim, got strides {codes.stride()}" + ) + src_page_size, src_num_slots, codes_page_stride, codes_slot_stride = _slot_geometry( + codes, "codes" + ) + + if scales.device != device: + raise ValueError(f"scales must be on {device}, got {scales.device}") + if scales.dtype not in _SCALE_DTYPES: + raise ValueError(f"scales must be uint8 or float8_e8m0fnu, got {scales.dtype}") + expected_scales = tuple(codes.shape[:-1]) + (width // MXFP8_BLOCK,) + if tuple(scales.shape) != expected_scales or scales.stride(-1) != 1: + raise ValueError( + f"scales must be {expected_scales} with unit last-dim stride, " + f"got shape {tuple(scales.shape)} strides {scales.stride()}" + ) + _, _, scales_page_stride, scales_slot_stride = _slot_geometry(scales, "scales") + + if out.device != device: + raise ValueError(f"out must be on {device}, got {out.device}") + if out.dtype != torch.bfloat16: + raise ValueError(f"out must be bfloat16, got {out.dtype}") + if out.dim() == 0 or out.shape[-1] != width or out.stride(-1) != 1: + raise ValueError( + f"out's last dim must be {width} with unit stride, got shape {tuple(out.shape)} strides {out.stride()}" + ) + out_page_size, out_num_slots, out_page_stride, out_slot_stride = _slot_geometry(out, "out") + + if src_slots is not None: + _check_slots(src_slots, "src_slots", device) + if dst_slots is not None: + _check_slots(dst_slots, "dst_slots", device) + if src_slots is not None and dst_slots is not None and src_slots.shape[0] != dst_slots.shape[0]: + raise ValueError( + f"src_slots and dst_slots must have the same length, got {src_slots.shape[0]} and {dst_slots.shape[0]}" + ) + if src_slots is not None: + num_rows = src_slots.shape[0] + elif dst_slots is not None: + num_rows = dst_slots.shape[0] + else: + num_rows = src_num_slots + if dst_slots is None and out_num_slots < num_rows: + raise ValueError( + f"without dst_slots out needs >= {num_rows} slots (one per row), got {out_num_slots}" + ) + + if num_rows == 0: + return + codes_arg = codes.view(torch.uint8) if codes.dtype != torch.uint8 else codes + scales_arg = scales.view(torch.uint8) if scales.dtype != torch.uint8 else scales + # Never dereferenced when the matching HAS_* constexpr is False. + src_arg, src_stride = ( + (src_slots, src_slots.stride(0)) if src_slots is not None else (scales_arg, 0) + ) + dst_arg, dst_stride = ( + (dst_slots, dst_slots.stride(0)) if dst_slots is not None else (scales_arg, 0) + ) + _mxfp8_dequantize_kernel[(num_rows,)]( + codes_arg, + scales_arg, + out, + src_arg, + dst_arg, + src_stride, + dst_stride, + codes_page_stride, + codes_slot_stride, + scales_page_stride, + scales_slot_stride, + src_page_size, + src_num_slots, + out_page_stride, + out_slot_stride, + out_page_size, + out_num_slots, + W=width, + NUM_BLOCKS_P2=triton.next_power_of_2(width // MXFP8_BLOCK), + HAS_SRC=src_slots is not None, + HAS_DST=dst_slots is not None, + num_warps=4, + ) + + +@mxfp8_dequantize.register_fake +def _mxfp8_dequantize_fake( + codes: torch.Tensor, + scales: torch.Tensor, + out: torch.Tensor, + src_slots: Optional[torch.Tensor] = None, + dst_slots: Optional[torch.Tensor] = None, +) -> None: + return None diff --git a/tensorrt_llm/_torch/custom_ops/mxfp8_fake_quantize.py b/tensorrt_llm/_torch/custom_ops/mxfp8_fake_quantize.py new file mode 100644 index 000000000000..c84db38e9f49 --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/mxfp8_fake_quantize.py @@ -0,0 +1,246 @@ +# 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. +"""MXFP8 fake quantization with an optional slot scatter: ``trtllm::mxfp8_fake_quantize``. + +Per row of ``x`` and per 32-element block along its last dim:: + + amax = max(max|x_block|, 1e-4) # fp32 + e = ceil(log2(fp32(amax * fp32(1/448)))) # from the IEEE bits + scale = 2^e # UE8M0 byte = e + 127 + q = e4m3_rne(clamp(x / scale, -448, 448)) + +and the row is written to slot ``slot_mapping[row]`` of ``out`` (row ``row`` +when no mapping is given; a negative or out-of-range slot is skipped), either +as the dequantized values ``bf16(q * scale)`` (a bf16 ``out``) or as the raw +e4m3 codes plus one UE8M0 byte per block (an fp8/uint8 ``out`` together with +``out_scale``). The dequantized values are exact in bf16, so both forms hold +the same value set. + +``out`` is ``[slots, W]`` or paged ``[pages, page_size, W]`` with arbitrary +leading strides; slot ``s`` of a paged ``out`` is ``(s // page_size, s % page_size)``. +One Triton program per row of ``x``; no autotuning. +""" + +from typing import Optional + +import torch +import triton +import triton.language as tl + +MXFP8_BLOCK = 32 +_E4M3_MAX = 448.0 +_AMAX_FLOOR = 1e-4 + +_NATIVE_DATA_DTYPES = (torch.float8_e4m3fn, torch.uint8) +_NATIVE_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) +_SLOT_DTYPES = (torch.int32, torch.int64) + + +@triton.jit +def _mxfp8_fake_quantize_kernel( + x_ptr, + out_ptr, + scale_ptr, + slot_ptr, + x_row_stride, + slot_stride, + out_page_stride, + out_slot_stride, + scale_page_stride, + scale_slot_stride, + page_size, + num_slots, + W: tl.constexpr, + NUM_BLOCKS_P2: tl.constexpr, + HAS_SLOT: tl.constexpr, + NATIVE: tl.constexpr, + AMAX_FLOOR: tl.constexpr, + E4M3_MAX: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + if HAS_SLOT: + slot = tl.load(slot_ptr + row * slot_stride).to(tl.int64) + else: + slot = row + if (slot >= 0) & (slot < num_slots): + page = slot // page_size + in_page = slot % page_size + + blk = tl.arange(0, NUM_BLOCKS_P2) + col = blk[:, None] * 32 + tl.arange(0, 32)[None, :] + mask = col < W + x = tl.load(x_ptr + row * x_row_stride + col, mask=mask, other=0.0).to(tl.float32) + + amax = tl.maximum(tl.max(tl.abs(x), axis=1), AMAX_FLOOR) + # 2^ceil(log2(amax / 448)) from the IEEE bits of the fp32 product with the + # fp32 reciprocal: the exponent, plus one when any mantissa bit is set. + bits = (amax * (1.0 / E4M3_MAX)).to(tl.int32, bitcast=True) + e = ((bits >> 23) & 0xFF) - 127 + ((bits & 0x7FFFFF) != 0).to(tl.int32) + scale = ((e + 127) << 23).to(tl.float32, bitcast=True) + # The scale is a power of two, so multiplying by its exact reciprocal is + # the IEEE quotient x / scale. + inv_scale = ((127 - e) << 23).to(tl.float32, bitcast=True) + q = tl.clamp(x * inv_scale[:, None], -E4M3_MAX, E4M3_MAX).to(tl.float8e4nv) + + out_off = page * out_page_stride + in_page * out_slot_stride + col + if NATIVE: + tl.store(out_ptr + out_off, q, mask=mask) + scale_off = page * scale_page_stride + in_page * scale_slot_stride + blk + tl.store(scale_ptr + scale_off, (e + 127).to(tl.uint8), mask=blk * 32 < W) + else: + y = (q.to(tl.float32) * scale[:, None]).to(tl.bfloat16) + tl.store(out_ptr + out_off, y, mask=mask) + + +def _slot_geometry(t: torch.Tensor, name: str) -> tuple[int, int, int, int]: + """(page_size, num_slots, page_stride, slot_stride) of a [slots, C] or [pages, page_size, C] tensor.""" + if t.dim() == 2: + return 1, t.shape[0], t.stride(0), 0 + if t.dim() == 3: + return t.shape[1], t.shape[0] * t.shape[1], t.stride(0), t.stride(1) + raise ValueError( + f"{name} must be [slots, C] or [pages, page_size, C], got shape {tuple(t.shape)}" + ) + + +@torch.library.custom_op("trtllm::mxfp8_fake_quantize", mutates_args=("out", "out_scale")) +def mxfp8_fake_quantize( + x: torch.Tensor, + out: torch.Tensor, + slot_mapping: Optional[torch.Tensor] = None, + out_scale: Optional[torch.Tensor] = None, +) -> None: + """Fake-quantize each row of ``x`` to MXFP8 (E4M3, UE8M0 per 32, amax floor 1e-4) into ``out``. + + ``x``: bf16 ``[..., W]``, ``W % 32 == 0``, unit last-dim stride, leading dims + viewable as one row dim. ``out``: bf16 (dequantized values) or + float8_e4m3fn / uint8 (codes, with ``out_scale`` uint8 / float8_e8m0fnu + ``[..., W // 32]`` of the same leading shape), ``[slots, W]`` or + ``[pages, page_size, W]``. ``slot_mapping``: int32/int64 ``[rows]``; row + ``r`` goes to slot ``slot_mapping[r]``, skipped when negative or + ``>= slots``. Without it row ``r`` goes to slot ``r``. + """ + if not x.is_cuda: + raise ValueError(f"x must be a CUDA tensor, got device {x.device}") + if x.dtype != torch.bfloat16: + raise ValueError(f"x must be bfloat16, got {x.dtype}") + if x.dim() < 2: + raise ValueError(f"x must have rank >= 2, got shape {tuple(x.shape)}") + width = x.shape[-1] + if width <= 0 or width % MXFP8_BLOCK != 0: + raise ValueError(f"x's last dim must be a positive multiple of {MXFP8_BLOCK}, got {width}") + if x.stride(-1) != 1: + raise ValueError(f"x must have a unit stride along its last dim, got strides {x.stride()}") + try: + x2 = x.view(-1, width) + except RuntimeError as err: + raise ValueError( + f"x's leading dims must be viewable as one row dim, got strides {x.stride()}" + ) from err + num_rows = x2.shape[0] + + if out.device != x.device: + raise ValueError(f"out must be on {x.device}, got {out.device}") + if out.shape[-1] != width or out.stride(-1) != 1: + raise ValueError( + f"out's last dim must be {width} with unit stride, got shape {tuple(out.shape)} strides {out.stride()}" + ) + page_size, num_slots, out_page_stride, out_slot_stride = _slot_geometry(out, "out") + + native = out.dtype != torch.bfloat16 + if native: + if out.dtype not in _NATIVE_DATA_DTYPES: + raise ValueError(f"out must be bfloat16, float8_e4m3fn or uint8, got {out.dtype}") + if out_scale is None: + raise ValueError( + "a float8_e4m3fn/uint8 out (native codes) needs out_scale for the UE8M0 bytes" + ) + if out_scale.device != x.device: + raise ValueError(f"out_scale must be on {x.device}, got {out_scale.device}") + if out_scale.dtype not in _NATIVE_SCALE_DTYPES: + raise ValueError(f"out_scale must be uint8 or float8_e8m0fnu, got {out_scale.dtype}") + if ( + tuple(out_scale.shape) != tuple(out.shape[:-1]) + (width // MXFP8_BLOCK,) + or out_scale.stride(-1) != 1 + ): + raise ValueError( + f"out_scale must be {tuple(out.shape[:-1]) + (width // MXFP8_BLOCK,)} with unit last-dim stride, " + f"got shape {tuple(out_scale.shape)} strides {out_scale.stride()}" + ) + _, _, scale_page_stride, scale_slot_stride = _slot_geometry(out_scale, "out_scale") + out_arg = out.view(torch.float8_e4m3fn) if out.dtype == torch.uint8 else out + scale_arg = out_scale.view(torch.uint8) if out_scale.dtype != torch.uint8 else out_scale + else: + if out_scale is not None: + raise ValueError( + "out_scale is only written with a float8_e4m3fn/uint8 out; a bfloat16 out takes None" + ) + scale_page_stride = scale_slot_stride = 0 + out_arg = out + scale_arg = ( + out # never dereferenced: NATIVE is a constexpr that compiles the scale store out + ) + + if slot_mapping is None: + if num_slots < num_rows: + raise ValueError( + f"without slot_mapping out needs >= {num_rows} slots (one per row), got {num_slots}" + ) + slot_arg, slot_stride = x2, 0 # never dereferenced: HAS_SLOT is a constexpr + else: + if slot_mapping.device != x.device: + raise ValueError(f"slot_mapping must be on {x.device}, got {slot_mapping.device}") + if slot_mapping.dtype not in _SLOT_DTYPES: + raise ValueError(f"slot_mapping must be int32 or int64, got {slot_mapping.dtype}") + if slot_mapping.dim() != 1 or slot_mapping.shape[0] != num_rows: + raise ValueError( + f"slot_mapping must be [{num_rows}] (one slot per row), got {tuple(slot_mapping.shape)}" + ) + slot_arg, slot_stride = slot_mapping, slot_mapping.stride(0) + + if num_rows == 0: + return + num_blocks = width // MXFP8_BLOCK + _mxfp8_fake_quantize_kernel[(num_rows,)]( + x2, + out_arg, + scale_arg, + slot_arg, + x2.stride(0), + slot_stride, + out_page_stride, + out_slot_stride, + scale_page_stride, + scale_slot_stride, + page_size, + num_slots, + W=width, + NUM_BLOCKS_P2=triton.next_power_of_2(num_blocks), + HAS_SLOT=slot_mapping is not None, + NATIVE=native, + AMAX_FLOOR=_AMAX_FLOOR, + E4M3_MAX=_E4M3_MAX, + num_warps=4, + ) + + +@mxfp8_fake_quantize.register_fake +def _mxfp8_fake_quantize_fake( + x: torch.Tensor, + out: torch.Tensor, + slot_mapping: Optional[torch.Tensor] = None, + out_scale: Optional[torch.Tensor] = None, +) -> None: + return None diff --git a/tensorrt_llm/_torch/custom_ops/mxfp8_to_e4m3.py b/tensorrt_llm/_torch/custom_ops/mxfp8_to_e4m3.py new file mode 100644 index 000000000000..10bee384cdd1 --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/mxfp8_to_e4m3.py @@ -0,0 +1,278 @@ +# 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. +"""MXFP8 row gather + transcode to per-tensor-scaled E4M3: ``trtllm::mxfp8_to_e4m3``. + +The source is a pool of native MXFP8 rows: ``codes`` holds E4M3 codes +``[slots, W]`` (or paged ``[pages, page_size, W]``) and ``scales`` one UE8M0 +byte per 32-element block ``[..., W // 32]`` with the same slot geometry. +For output row ``r``:: + + s = src_slots[r] (r when src_slots is None) + d = dst_slots[r] (r when dst_slots is None) + v = e4m3(codes[s, i]) * 2^(scales[s, i // 32] - 127) (exact in fp32) + out[d, i] = e4m3_rne(clamp(v * inv_scale, -448, 448)) + +``inv_scale`` is one fp32 factor for the whole output (a per-tensor FP8 +scale ``1 / inv_scale``). With ``inv_scale`` a power of two, ``v * inv_scale`` +is exact, and the transcode is exact for every value whose magnitude lands in +E4M3's normal range ``[2^-6, 448]``: an MXFP8 value is an E4M3 code times a +power of two, so it keeps its 3-bit mantissa. Values below ``2^-6`` round to +E4M3 subnormals; values above 448 saturate. + +A source slot that is negative or ``>= slots`` yields a row of zeros; a +destination slot that is negative or ``>= out slots`` is skipped. Slot ``s`` +of a paged tensor is ``(s // page_size, s % page_size)``; leading strides are +free, so interleaved ``[.., 528]`` byte rows are passed as two strided views. +One Triton program per output row; no autotuning. +""" + +from typing import Optional + +import torch +import triton +import triton.language as tl + +MXFP8_BLOCK = 32 +_E4M3_MAX = 448.0 + +_CODE_DTYPES = (torch.float8_e4m3fn, torch.uint8) +_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) +_OUT_DTYPES = (torch.float8_e4m3fn,) +_SLOT_DTYPES = (torch.int32, torch.int64) + + +@triton.jit +def _mxfp8_to_e4m3_kernel( + codes_ptr, + scales_ptr, + out_ptr, + src_ptr, + dst_ptr, + inv_scale, + src_stride, + dst_stride, + codes_page_stride, + codes_slot_stride, + scales_page_stride, + scales_slot_stride, + src_page_size, + src_num_slots, + out_page_stride, + out_slot_stride, + out_page_size, + out_num_slots, + W: tl.constexpr, + NUM_BLOCKS_P2: tl.constexpr, + HAS_SRC: tl.constexpr, + HAS_DST: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + if HAS_DST: + dst = tl.load(dst_ptr + row * dst_stride).to(tl.int64) + else: + dst = row + if (dst >= 0) & (dst < out_num_slots): + if HAS_SRC: + src = tl.load(src_ptr + row * src_stride).to(tl.int64) + else: + src = row + valid = (src >= 0) & (src < src_num_slots) + src = tl.where(valid, src, 0) + + blk = tl.arange(0, NUM_BLOCKS_P2) + col = blk[:, None] * 32 + tl.arange(0, 32)[None, :] + mask = (col < W) & valid + s_page = src // src_page_size + s_in = src % src_page_size + + c = tl.load( + codes_ptr + s_page * codes_page_stride + s_in * codes_slot_stride + col, + mask=mask, + other=0, + ) + c = c.to(tl.float8e4nv, bitcast=True) + b = tl.load( + scales_ptr + s_page * scales_page_stride + s_in * scales_slot_stride + blk, + mask=(blk * 32 < W) & valid, + other=127, + ).to(tl.int32) + # UE8M0 byte b is 2^(b - 127): the fp32 exponent field itself for b in + # 1..254, the fp32 subnormal 2^-127 for b = 0, NaN for b = 255. + bits = tl.where(b == 0, 0x00400000, tl.where(b == 255, 0x7FC00000, b << 23)) + scale = bits.to(tl.float32, bitcast=True) + v = c.to(tl.float32) * scale[:, None] * inv_scale + v = tl.where(v != v, v, tl.minimum(tl.maximum(v, -448.0), 448.0)) # NaN stays NaN + y = v.to(tl.float8e4nv) + + d_page = dst // out_page_size + d_in = dst % out_page_size + tl.store( + out_ptr + d_page * out_page_stride + d_in * out_slot_stride + col, + y.to(tl.uint8, bitcast=True), + mask=col < W, + ) + + +def _slot_geometry(t: torch.Tensor, name: str) -> tuple[int, int, int, int]: + """(page_size, num_slots, page_stride, slot_stride) of a [slots, C] or [pages, page_size, C] tensor.""" + if t.dim() == 2: + return 1, t.shape[0], t.stride(0), 0 + if t.dim() == 3: + return t.shape[1], t.shape[0] * t.shape[1], t.stride(0), t.stride(1) + raise ValueError( + f"{name} must be [slots, C] or [pages, page_size, C], got shape {tuple(t.shape)}" + ) + + +def _check_slots(slots: torch.Tensor, name: str, device: torch.device) -> None: + if slots.device != device: + raise ValueError(f"{name} must be on {device}, got {slots.device}") + if slots.dtype not in _SLOT_DTYPES: + raise ValueError(f"{name} must be int32 or int64, got {slots.dtype}") + if slots.dim() != 1: + raise ValueError(f"{name} must be 1-D, got shape {tuple(slots.shape)}") + + +@torch.library.custom_op("trtllm::mxfp8_to_e4m3", mutates_args=("out",)) +def mxfp8_to_e4m3( + codes: torch.Tensor, + scales: torch.Tensor, + out: torch.Tensor, + inv_scale: float, + src_slots: Optional[torch.Tensor] = None, + dst_slots: Optional[torch.Tensor] = None, +) -> None: + """Gather MXFP8 rows (E4M3 codes x UE8M0 per 32) from a pool and write them to ``out`` as E4M3 + with one per-tensor scale: ``out = e4m3_rne(clamp(value * inv_scale, +-448))``. + + ``codes``: float8_e4m3fn / uint8 ``[slots, W]`` or ``[pages, page_size, W]``, + ``W % 32 == 0``, unit last-dim stride. ``scales``: uint8 / float8_e8m0fnu + ``codes.shape[:-1] + (W // 32,)``, unit last-dim stride. ``out``: float8_e4m3fn + ``[out_slots, W]`` or ``[pages, page_size, W]``, unit last-dim stride. + ``src_slots`` / ``dst_slots``: int32/int64 ``[rows]``. Row ``r`` reads slot + ``src_slots[r]`` (``r`` when None; zeros when negative or ``>= slots``) and + writes slot ``dst_slots[r]`` of ``out`` (``r`` when None; skipped when + negative or ``>= out_slots``). + """ + if not codes.is_cuda: + raise ValueError(f"codes must be a CUDA tensor, got device {codes.device}") + device = codes.device + if codes.dtype not in _CODE_DTYPES: + raise ValueError(f"codes must be float8_e4m3fn or uint8, got {codes.dtype}") + width = codes.shape[-1] if codes.dim() > 0 else 0 + if width <= 0 or width % MXFP8_BLOCK != 0: + raise ValueError( + f"codes' last dim must be a positive multiple of {MXFP8_BLOCK}, got {width}" + ) + if codes.stride(-1) != 1: + raise ValueError( + f"codes must have a unit stride along its last dim, got strides {codes.stride()}" + ) + src_page_size, src_num_slots, codes_page_stride, codes_slot_stride = _slot_geometry( + codes, "codes" + ) + + if scales.device != device: + raise ValueError(f"scales must be on {device}, got {scales.device}") + if scales.dtype not in _SCALE_DTYPES: + raise ValueError(f"scales must be uint8 or float8_e8m0fnu, got {scales.dtype}") + expected_scales = tuple(codes.shape[:-1]) + (width // MXFP8_BLOCK,) + if tuple(scales.shape) != expected_scales or scales.stride(-1) != 1: + raise ValueError( + f"scales must be {expected_scales} with unit last-dim stride, " + f"got shape {tuple(scales.shape)} strides {scales.stride()}" + ) + _, _, scales_page_stride, scales_slot_stride = _slot_geometry(scales, "scales") + + if out.device != device: + raise ValueError(f"out must be on {device}, got {out.device}") + if out.dtype not in _OUT_DTYPES: + raise ValueError(f"out must be float8_e4m3fn, got {out.dtype}") + if out.dim() == 0 or out.shape[-1] != width or out.stride(-1) != 1: + raise ValueError( + f"out's last dim must be {width} with unit stride, got shape {tuple(out.shape)} strides {out.stride()}" + ) + out_page_size, out_num_slots, out_page_stride, out_slot_stride = _slot_geometry(out, "out") + if not (inv_scale > 0.0): + raise ValueError(f"inv_scale must be a positive float, got {inv_scale}") + + if src_slots is not None: + _check_slots(src_slots, "src_slots", device) + if dst_slots is not None: + _check_slots(dst_slots, "dst_slots", device) + if src_slots is not None and dst_slots is not None and src_slots.shape[0] != dst_slots.shape[0]: + raise ValueError( + f"src_slots and dst_slots must have the same length, got {src_slots.shape[0]} and {dst_slots.shape[0]}" + ) + if src_slots is not None: + num_rows = src_slots.shape[0] + elif dst_slots is not None: + num_rows = dst_slots.shape[0] + else: + num_rows = src_num_slots + if dst_slots is None and out_num_slots < num_rows: + raise ValueError( + f"without dst_slots out needs >= {num_rows} slots (one per row), got {out_num_slots}" + ) + + if num_rows == 0: + return + codes_arg = codes.view(torch.uint8) if codes.dtype != torch.uint8 else codes + scales_arg = scales.view(torch.uint8) if scales.dtype != torch.uint8 else scales + out_arg = out.view(torch.uint8) + # Never dereferenced when the matching HAS_* constexpr is False. + src_arg, src_stride = ( + (src_slots, src_slots.stride(0)) if src_slots is not None else (scales_arg, 0) + ) + dst_arg, dst_stride = ( + (dst_slots, dst_slots.stride(0)) if dst_slots is not None else (scales_arg, 0) + ) + _mxfp8_to_e4m3_kernel[(num_rows,)]( + codes_arg, + scales_arg, + out_arg, + src_arg, + dst_arg, + float(inv_scale), + src_stride, + dst_stride, + codes_page_stride, + codes_slot_stride, + scales_page_stride, + scales_slot_stride, + src_page_size, + src_num_slots, + out_page_stride, + out_slot_stride, + out_page_size, + out_num_slots, + W=width, + NUM_BLOCKS_P2=triton.next_power_of_2(width // MXFP8_BLOCK), + HAS_SRC=src_slots is not None, + HAS_DST=dst_slots is not None, + num_warps=4, + ) + + +@mxfp8_to_e4m3.register_fake +def _mxfp8_to_e4m3_fake( + codes: torch.Tensor, + scales: torch.Tensor, + out: torch.Tensor, + inv_scale: float, + src_slots: Optional[torch.Tensor] = None, + dst_slots: Optional[torch.Tensor] = None, +) -> None: + return None diff --git a/tensorrt_llm/_torch/custom_ops/nvfp4_fake_quantize.py b/tensorrt_llm/_torch/custom_ops/nvfp4_fake_quantize.py new file mode 100644 index 000000000000..a5f1a2b06fd9 --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/nvfp4_fake_quantize.py @@ -0,0 +1,292 @@ +# 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. +"""NVFP4 (E2M1 x E4M3 per 16, no global scale) fake quantization with an optional slot +scatter: ``trtllm::nvfp4_fake_quantize``. + +Per row of ``x`` and per 16-element block along its last dim:: + + amax = max(max|x_block|, 6 * 2^-9) # fp32, exact for bf16 input + scale = e4m3_rne(min(amax / 6, 448)) # IEEE fp32 division; E4M3 byte + y = x / scale # IEEE fp32 division + code = e2m1_rne(clamp(y, -6, 6)) # ties to the even code; sign bit of y kept + +and the row is written to slot ``slot_mapping[row]`` of ``out`` (row ``row`` when no +mapping is given; a negative or out-of-range slot is skipped), either as the +dequantized values ``bf16(E2M1[code] * scale)`` (a bf16 ``out``; exact in bf16) or as +the packed codes, two per byte with the even element in the low nibble, plus one E4M3 +byte per block (a uint8 ``out`` together with ``out_scale``). A negative value whose +magnitude rounds to zero keeps its sign: code ``0x8``, dequantized ``-0.0``. + +``out`` is ``[slots, C]`` or paged ``[pages, page_size, C]`` with arbitrary leading +strides; slot ``s`` of a paged ``out`` is ``(s // page_size, s % page_size)``. The +native pair may be two views into one ``[.., W/2 + W/16]`` byte row (288 B at W=512). +One Triton program per row of ``x``; no autotuning. +""" + +from typing import Optional + +import torch +import triton +import triton.language as tl + +NVFP4_BLOCK = 16 +_E4M3_MAX = 448.0 +_AMAX_FLOOR = 6.0 * 2.0**-9 + +_NATIVE_DATA_DTYPES = (torch.uint8,) +_NATIVE_SCALE_DTYPES = (torch.float8_e4m3fn, torch.uint8) +_SLOT_DTYPES = (torch.int32, torch.int64) + + +@triton.jit +def _e2m1_code(y): + """fp32 -> (E2M1 magnitude code 0..7, sign bit) with round-to-nearest-even. + + The magnitude code is the number of midpoints between neighbouring E2M1 + magnitudes (0, 0.5, 1, 1.5, 2, 3, 4, 6) that |y| lies above; a value exactly on a + midpoint goes to the neighbour with the even code: up at 0.75, 1.75, 3.5, down at + 0.25, 1.25, 2.5, 5. Values beyond 5 saturate to code 7 (6.0). + """ + a = tl.abs(y) + code = (a > 0.25).to(tl.int32) + code += (a >= 0.75).to(tl.int32) + code += (a > 1.25).to(tl.int32) + code += (a >= 1.75).to(tl.int32) + code += (a > 2.5).to(tl.int32) + code += (a >= 3.5).to(tl.int32) + code += (a > 5.0).to(tl.int32) + sign = (y.to(tl.int32, bitcast=True) >> 31) & 1 + return code, sign + + +@triton.jit +def _e2m1_dequant(code, sign, scale): + """E2M1[code] * scale with the sign bit applied by bit (keeps -0.0), fp32 (exact).""" + mag = tl.where( + code < 4, code.to(tl.float32) * 0.5, tl.where(code == 7, 6.0, (code - 2).to(tl.float32)) + ) + v = mag * scale + return (v.to(tl.int32, bitcast=True) | (sign << 31)).to(tl.float32, bitcast=True) + + +@triton.jit +def _nvfp4_fake_quantize_kernel( + x_ptr, + out_ptr, + scale_ptr, + slot_ptr, + x_row_stride, + slot_stride, + out_page_stride, + out_slot_stride, + scale_page_stride, + scale_slot_stride, + page_size, + num_slots, + W: tl.constexpr, + NUM_BLOCKS_P2: tl.constexpr, + HAS_SLOT: tl.constexpr, + NATIVE: tl.constexpr, + AMAX_FLOOR: tl.constexpr, + E4M3_MAX: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + if HAS_SLOT: + slot = tl.load(slot_ptr + row * slot_stride).to(tl.int64) + else: + slot = row + if (slot >= 0) & (slot < num_slots): + page = slot // page_size + in_page = slot % page_size + + blk = tl.arange(0, NUM_BLOCKS_P2) + pair = tl.arange(0, 8) + col_even = blk[:, None] * 16 + pair[None, :] * 2 # [blocks, 8] + mask = col_even < W + x_row = x_ptr + row * x_row_stride + xe = tl.load(x_row + col_even, mask=mask, other=0.0).to(tl.float32) + xo = tl.load(x_row + col_even + 1, mask=mask, other=0.0).to(tl.float32) + + amax = tl.maximum(tl.max(tl.abs(xe), axis=1), tl.max(tl.abs(xo), axis=1)) + amax = tl.maximum(amax, AMAX_FLOOR) + six = tl.full(amax.shape, 6.0, tl.float32) + scale_f8 = tl.minimum(tl.math.div_rn(amax, six), E4M3_MAX).to( + tl.float8e4nv, fp_downcast_rounding="rtne" + ) + scale = scale_f8.to(tl.float32)[:, None] # [blocks, 1] + scale_b = tl.broadcast_to(scale, xe.shape) + + ce, se = _e2m1_code(tl.math.div_rn(xe, scale_b)) + co, so = _e2m1_code(tl.math.div_rn(xo, scale_b)) + + if NATIVE: + byte = (ce | (se << 3)) | ((co | (so << 3)) << 4) + out_off = ( + page * out_page_stride + + in_page * out_slot_stride + + blk[:, None] * 8 + + pair[None, :] + ) + tl.store(out_ptr + out_off, byte.to(tl.uint8), mask=mask) + scale_off = page * scale_page_stride + in_page * scale_slot_stride + blk + tl.store(scale_ptr + scale_off, scale_f8, mask=blk * 16 < W) + else: + out_row = out_ptr + page * out_page_stride + in_page * out_slot_stride + tl.store(out_row + col_even, _e2m1_dequant(ce, se, scale_b).to(tl.bfloat16), mask=mask) + tl.store( + out_row + col_even + 1, _e2m1_dequant(co, so, scale_b).to(tl.bfloat16), mask=mask + ) + + +def _slot_geometry(t: torch.Tensor, name: str) -> tuple[int, int, int, int]: + """(page_size, num_slots, page_stride, slot_stride) of a [slots, C] or [pages, page_size, C] tensor.""" + if t.dim() == 2: + return 1, t.shape[0], t.stride(0), 0 + if t.dim() == 3: + return t.shape[1], t.shape[0] * t.shape[1], t.stride(0), t.stride(1) + raise ValueError( + f"{name} must be [slots, C] or [pages, page_size, C], got shape {tuple(t.shape)}" + ) + + +@torch.library.custom_op("trtllm::nvfp4_fake_quantize", mutates_args=("out", "out_scale")) +def nvfp4_fake_quantize( + x: torch.Tensor, + out: torch.Tensor, + slot_mapping: Optional[torch.Tensor] = None, + out_scale: Optional[torch.Tensor] = None, +) -> None: + """Fake-quantize each row of ``x`` to NVFP4 (E2M1, E4M3 per 16, amax floor 6*2^-9, no global scale). + + ``x``: bf16 ``[..., W]``, ``W % 16 == 0``, unit last-dim stride, leading dims + viewable as one row dim. ``out``: bf16 ``[.., W]`` (dequantized values) or uint8 + ``[.., W // 2]`` (packed codes, even element in the low nibble, with ``out_scale`` + float8_e4m3fn / uint8 ``[.., W // 16]`` of the same leading shape), ``[slots, C]`` + or ``[pages, page_size, C]``. ``slot_mapping``: int32/int64 ``[rows]``; row ``r`` + goes to slot ``slot_mapping[r]``, skipped when negative or ``>= slots``. Without it + row ``r`` goes to slot ``r``. + """ + if not x.is_cuda: + raise ValueError(f"x must be a CUDA tensor, got device {x.device}") + if x.dtype != torch.bfloat16: + raise ValueError(f"x must be bfloat16, got {x.dtype}") + if x.dim() < 2: + raise ValueError(f"x must have rank >= 2, got shape {tuple(x.shape)}") + width = x.shape[-1] + if width <= 0 or width % NVFP4_BLOCK != 0: + raise ValueError(f"x's last dim must be a positive multiple of {NVFP4_BLOCK}, got {width}") + if x.stride(-1) != 1: + raise ValueError(f"x must have a unit stride along its last dim, got strides {x.stride()}") + try: + x2 = x.view(-1, width) + except RuntimeError as err: + raise ValueError( + f"x's leading dims must be viewable as one row dim, got strides {x.stride()}" + ) from err + num_rows = x2.shape[0] + + if out.device != x.device: + raise ValueError(f"out must be on {x.device}, got {out.device}") + if out.dtype != torch.bfloat16 and out.dtype not in _NATIVE_DATA_DTYPES: + raise ValueError(f"out must be bfloat16 or uint8, got {out.dtype}") + native = out.dtype != torch.bfloat16 + out_width = width // 2 if native else width + if out.shape[-1] != out_width or out.stride(-1) != 1: + raise ValueError( + f"out's last dim must be {out_width} ({'packed E2M1 bytes' if native else 'bf16 values'}) with unit " + f"stride, got shape {tuple(out.shape)} strides {out.stride()}" + ) + page_size, num_slots, out_page_stride, out_slot_stride = _slot_geometry(out, "out") + + if native: + if out_scale is None: + raise ValueError( + "a uint8 out (packed E2M1 codes) needs out_scale for the E4M3 block scales" + ) + if out_scale.device != x.device: + raise ValueError(f"out_scale must be on {x.device}, got {out_scale.device}") + if out_scale.dtype not in _NATIVE_SCALE_DTYPES: + raise ValueError(f"out_scale must be float8_e4m3fn or uint8, got {out_scale.dtype}") + want = tuple(out.shape[:-1]) + (width // NVFP4_BLOCK,) + if tuple(out_scale.shape) != want or out_scale.stride(-1) != 1: + raise ValueError( + f"out_scale must be {want} with unit last-dim stride, " + f"got shape {tuple(out_scale.shape)} strides {out_scale.stride()}" + ) + _, _, scale_page_stride, scale_slot_stride = _slot_geometry(out_scale, "out_scale") + scale_arg = ( + out_scale.view(torch.float8_e4m3fn) if out_scale.dtype == torch.uint8 else out_scale + ) + else: + if out_scale is not None: + raise ValueError( + "out_scale is only written with a uint8 out; a bfloat16 out takes None" + ) + scale_page_stride = scale_slot_stride = 0 + scale_arg = ( + out # never dereferenced: NATIVE is a constexpr that compiles the scale store out + ) + + if slot_mapping is None: + if num_slots < num_rows: + raise ValueError( + f"without slot_mapping out needs >= {num_rows} slots (one per row), got {num_slots}" + ) + slot_arg, slot_stride = x2, 0 # never dereferenced: HAS_SLOT is a constexpr + else: + if slot_mapping.device != x.device: + raise ValueError(f"slot_mapping must be on {x.device}, got {slot_mapping.device}") + if slot_mapping.dtype not in _SLOT_DTYPES: + raise ValueError(f"slot_mapping must be int32 or int64, got {slot_mapping.dtype}") + if slot_mapping.dim() != 1 or slot_mapping.shape[0] != num_rows: + raise ValueError( + f"slot_mapping must be [{num_rows}] (one slot per row), got {tuple(slot_mapping.shape)}" + ) + slot_arg, slot_stride = slot_mapping, slot_mapping.stride(0) + + if num_rows == 0: + return + num_blocks = width // NVFP4_BLOCK + _nvfp4_fake_quantize_kernel[(num_rows,)]( + x2, + out, + scale_arg, + slot_arg, + x2.stride(0), + slot_stride, + out_page_stride, + out_slot_stride, + scale_page_stride, + scale_slot_stride, + page_size, + num_slots, + W=width, + NUM_BLOCKS_P2=triton.next_power_of_2(num_blocks), + HAS_SLOT=slot_mapping is not None, + NATIVE=native, + AMAX_FLOOR=_AMAX_FLOOR, + E4M3_MAX=_E4M3_MAX, + num_warps=4, + ) + + +@nvfp4_fake_quantize.register_fake +def _nvfp4_fake_quantize_fake( + x: torch.Tensor, + out: torch.Tensor, + slot_mapping: Optional[torch.Tensor] = None, + out_scale: Optional[torch.Tensor] = None, +) -> None: + return None diff --git a/tensorrt_llm/_torch/custom_ops/nvfp4_kv_to_bf16.py b/tensorrt_llm/_torch/custom_ops/nvfp4_kv_to_bf16.py new file mode 100644 index 000000000000..6ae4393dcedd --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/nvfp4_kv_to_bf16.py @@ -0,0 +1,204 @@ +# 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. +"""Paged NVFP4 KV gather + dequantize to bf16: ``trtllm::nvfp4_kv_to_bf16``. + +The source is a paged pool of NVFP4 rows (DeepSeek-V4.1's 288-byte compressed-KV +row): ``kv_cache`` holds ``W/2`` bytes of packed E2M1 codes per row (even element +in the low nibble) and ``kv_cache_sf`` one E4M3 scale byte per 16 elements. For +every request ``b`` and every position ``t < min(seq_lens[b], L)``:: + + page = block_tables[b, t // page_size] + out[b, t, 0, i] = bf16((E2M1[code_i] * e4m3(sf_{i // 16})) * global_scale) if 0 <= page < pages + +Positions at or past ``seq_lens[b]`` (or ``L``) and rows whose page id is out of +range are left unwritten. The product ``E2M1 x E4M3`` has at most 6 significant +bits and lies in [2^-10, 2688], so with ``global_scale`` 1 every value is exact +in bf16. This is the value set, the write set and the operand layout of +FlashInfer's ``nvfp4_kv_dequantize_paged`` as the catalog's +``nvfp4_kv_dequantize`` calls it, in one launch instead of two. + +``kv_cache`` / ``kv_cache_sf`` are ``[pages, page_size, 1, C]`` with a unit +last-dim stride and free leading strides, so the two halves of one 288-byte +row are passed as two views. ``out`` is a contiguous bf16 ``[B, L, 1, W]``. +Grid ``(ceil(L / 8), B)``, eight rows per program; no autotuning. +""" + +import torch +import triton +import triton.language as tl + +NVFP4_BLOCK = 16 +_ROWS = 8 + +_CODE_DTYPES = (torch.uint8,) +_SF_DTYPES = (torch.uint8, torch.float8_e4m3fn) + + +@triton.jit +def _e2m1_value(n): + """4-bit E2M1 code (int32, sign in bit 3) -> fp32 value; -0.0 for code 8.""" + m = n & 7 + mag = tl.where( + m < 4, + m.to(tl.float32) * 0.5, + tl.where(m == 4, 2.0, tl.where(m == 5, 3.0, tl.where(m == 6, 4.0, 6.0))), + ) + return tl.where((n & 8) != 0, -mag, mag) + + +@triton.jit +def _nvfp4_kv_to_bf16_kernel( + codes_ptr, + sf_ptr, + bt_ptr, + seq_ptr, + gs_ptr, + out_ptr, + codes_page_stride, + codes_slot_stride, + sf_page_stride, + sf_slot_stride, + bt_stride, + num_pages, + page_size, + L, + W: tl.constexpr, + ROWS: tl.constexpr, +): + b = tl.program_id(1).to(tl.int64) + t0 = tl.program_id(0).to(tl.int64) * ROWS + n = tl.minimum(tl.load(seq_ptr + b).to(tl.int64), L) + if t0 < n: + gs = tl.load(gs_ptr) + blk = tl.arange(0, W // 16) + j = tl.arange(0, 16) + col = blk[:, None] * 16 + j[None, :] # element index + byte = col // 2 + hi = (col % 2) == 1 + for r in tl.static_range(ROWS): + t = t0 + r + if t < n: + page = tl.load(bt_ptr + b * bt_stride + t // page_size).to(tl.int64) + if (page >= 0) & (page < num_pages): + slot = t % page_size + c = tl.load( + codes_ptr + page * codes_page_stride + slot * codes_slot_stride + byte + ).to(tl.int32) + nib = tl.where(hi, c >> 4, c) & 15 + s = tl.load(sf_ptr + page * sf_page_stride + slot * sf_slot_stride + blk) + s = s.to(tl.float8e4nv, bitcast=True).to(tl.float32) + v = (_e2m1_value(nib) * s[:, None]) * gs + tl.store(out_ptr + (b * L + t) * W + col, v.to(tl.bfloat16)) + + +@torch.library.custom_op("trtllm::nvfp4_kv_to_bf16", mutates_args=("out",)) +def nvfp4_kv_to_bf16( + kv_cache: torch.Tensor, + kv_cache_sf: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + global_scale: torch.Tensor, + out: torch.Tensor, +) -> None: + """Gather paged NVFP4 rows (packed E2M1 x E4M3 per 16) into ``out`` ([B, L, 1, W] bf16, + contiguous) as ``bf16((E2M1 * scale) * global_scale)`` for every ``t < min(seq_lens[b], L)`` + whose page id is in range; everything else of ``out`` is left unwritten.""" + if not kv_cache.is_cuda: + raise ValueError(f"kv_cache must be a CUDA tensor, got device {kv_cache.device}") + device = kv_cache.device + if kv_cache.dtype not in _CODE_DTYPES: + raise ValueError(f"kv_cache must be uint8 (packed E2M1), got {kv_cache.dtype}") + if kv_cache_sf.dtype not in _SF_DTYPES: + raise ValueError(f"kv_cache_sf must be uint8 or float8_e4m3fn, got {kv_cache_sf.dtype}") + if ( + kv_cache.dim() != 4 + or kv_cache_sf.dim() != 4 + or kv_cache.shape[2] != 1 + or kv_cache_sf.shape[2] != 1 + ): + raise ValueError( + "kv_cache / kv_cache_sf must be [pages, page_size, 1, C], got " + f"{tuple(kv_cache.shape)} / {tuple(kv_cache_sf.shape)}" + ) + if tuple(kv_cache.shape[:2]) != tuple(kv_cache_sf.shape[:2]): + raise ValueError( + f"page geometry differs: {tuple(kv_cache.shape)} vs {tuple(kv_cache_sf.shape)}" + ) + if kv_cache.stride(-1) != 1 or kv_cache_sf.stride(-1) != 1: + raise ValueError("kv_cache and kv_cache_sf need a unit last-dim stride") + width = kv_cache.shape[-1] * 2 + if width % NVFP4_BLOCK != 0 or kv_cache_sf.shape[-1] * NVFP4_BLOCK != width: + raise ValueError( + f"expected {width // NVFP4_BLOCK} scale bytes per row, got {kv_cache_sf.shape[-1]}" + ) + if width & (width - 1): + raise ValueError(f"row width must be a power of two, got {width}") + if out.device != device or out.dtype != torch.bfloat16: + raise ValueError(f"out must be bfloat16 on {device}, got {out.dtype} on {out.device}") + if out.dim() != 4 or out.shape[2] != 1 or out.shape[3] != width or not out.is_contiguous(): + raise ValueError(f"out must be a contiguous [B, L, 1, {width}], got {tuple(out.shape)}") + if block_tables.device != device or block_tables.dtype not in (torch.int32, torch.int64): + raise ValueError(f"block_tables must be int32/int64 on {device}") + if block_tables.dim() != 2 or block_tables.stride(-1) != 1: + raise ValueError( + f"block_tables must be [B, max_pages] with a unit last-dim stride, got {block_tables.shape}" + ) + if seq_lens.device != device or seq_lens.dtype != torch.int32 or seq_lens.dim() != 1: + raise ValueError("seq_lens must be a 1-D int32 CUDA tensor") + if ( + global_scale.device != device + or global_scale.dtype != torch.float32 + or global_scale.numel() != 1 + ): + raise ValueError("global_scale must be a float32 CUDA tensor of one element") + B, L = out.shape[0], out.shape[1] + if seq_lens.shape[0] != B or block_tables.shape[0] != B: + raise ValueError( + f"seq_lens {tuple(seq_lens.shape)} / block_tables {tuple(block_tables.shape)} vs B={B}" + ) + if B == 0 or L == 0: + return + sf_arg = kv_cache_sf.view(torch.uint8) if kv_cache_sf.dtype != torch.uint8 else kv_cache_sf + _nvfp4_kv_to_bf16_kernel[(triton.cdiv(L, _ROWS), B)]( + kv_cache, + sf_arg, + block_tables, + seq_lens, + global_scale, + out, + kv_cache.stride(0), + kv_cache.stride(1), + sf_arg.stride(0), + sf_arg.stride(1), + block_tables.stride(0), + kv_cache.shape[0], + kv_cache.shape[1], + L, + W=width, + ROWS=_ROWS, + num_warps=4, + ) + + +@nvfp4_kv_to_bf16.register_fake +def _nvfp4_kv_to_bf16_fake( + kv_cache: torch.Tensor, + kv_cache_sf: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + global_scale: torch.Tensor, + out: torch.Tensor, +) -> None: + return None diff --git a/tensorrt_llm/_torch/custom_ops/nvfp4_kv_to_e4m3.py b/tensorrt_llm/_torch/custom_ops/nvfp4_kv_to_e4m3.py new file mode 100644 index 000000000000..194be6baa5db --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/nvfp4_kv_to_e4m3.py @@ -0,0 +1,205 @@ +# 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. +"""Paged NVFP4 KV gather + transcode to per-tensor-scaled E4M3: ``trtllm::nvfp4_kv_to_e4m3``. + +The source is a paged pool of NVFP4 rows without a global scale (DeepSeek-V4.1's +288-byte compressed-KV row): ``kv_cache`` holds ``W/2`` bytes of packed E2M1 +codes per row (even element in the low nibble) and ``kv_cache_sf`` one E4M3 +scale byte per 16 elements. For every request ``b`` and every position +``t < min(seq_lens[b], L)``:: + + page = block_tables[b, t // page_size] + v_i = E2M1[code_i] * e4m3(sf_{i // 16}) (exact in fp32) + out[b, t, 0, i] = e4m3_rne(clamp(v_i * inv_scale, -448, 448)) if 0 <= page < pages + +Positions at or past ``seq_lens[b]`` (or ``L``) and rows whose page id is out of +range are left unwritten. ``inv_scale`` is one fp32 factor for the whole +output (a per-tensor FP8 scale ``1 / inv_scale``). The transcode is lossy in +general: ``E2M1 x E4M3`` has up to 6 significant bits and E4M3 keeps 4. + +``kv_cache`` / ``kv_cache_sf`` are ``[pages, page_size, 1, C]`` with a unit +last-dim stride and free leading strides, so the two halves of one 288-byte +row are passed as two views. ``out`` is a contiguous float8_e4m3fn +``[B, L, 1, W]``. Grid ``(ceil(L / 8), B)``, eight rows per program; no +autotuning. +""" + +import torch +import triton +import triton.language as tl + +NVFP4_BLOCK = 16 +_ROWS = 8 + +_CODE_DTYPES = (torch.uint8,) +_SF_DTYPES = (torch.uint8, torch.float8_e4m3fn) + + +@triton.jit +def _e2m1_value(n): + """4-bit E2M1 code (int32, sign in bit 3) -> fp32 value; -0.0 for code 8.""" + m = n & 7 + mag = tl.where( + m < 4, + m.to(tl.float32) * 0.5, + tl.where(m == 4, 2.0, tl.where(m == 5, 3.0, tl.where(m == 6, 4.0, 6.0))), + ) + return tl.where((n & 8) != 0, -mag, mag) + + +@triton.jit +def _nvfp4_kv_to_e4m3_kernel( + codes_ptr, + sf_ptr, + bt_ptr, + seq_ptr, + out_ptr, + inv_scale, + codes_page_stride, + codes_slot_stride, + sf_page_stride, + sf_slot_stride, + bt_stride, + num_pages, + page_size, + L, + W: tl.constexpr, + ROWS: tl.constexpr, +): + b = tl.program_id(1).to(tl.int64) + t0 = tl.program_id(0).to(tl.int64) * ROWS + n = tl.minimum(tl.load(seq_ptr + b).to(tl.int64), L) + if t0 < n: + blk = tl.arange(0, W // 16) + j = tl.arange(0, 16) + col = blk[:, None] * 16 + j[None, :] # element index + byte = col // 2 + hi = (col % 2) == 1 + for r in tl.static_range(ROWS): + t = t0 + r + if t < n: + page = tl.load(bt_ptr + b * bt_stride + t // page_size).to(tl.int64) + if (page >= 0) & (page < num_pages): + slot = t % page_size + c = tl.load( + codes_ptr + page * codes_page_stride + slot * codes_slot_stride + byte + ).to(tl.int32) + nib = tl.where(hi, c >> 4, c) & 15 + s = tl.load(sf_ptr + page * sf_page_stride + slot * sf_slot_stride + blk) + s = s.to(tl.float8e4nv, bitcast=True).to(tl.float32) + v = _e2m1_value(nib) * s[:, None] * inv_scale + v = tl.where( + v != v, v, tl.minimum(tl.maximum(v, -448.0), 448.0) + ) # NaN stays NaN + y = v.to(tl.float8e4nv).to(tl.uint8, bitcast=True) + tl.store(out_ptr + (b * L + t) * W + col, y) + + +@torch.library.custom_op("trtllm::nvfp4_kv_to_e4m3", mutates_args=("out",)) +def nvfp4_kv_to_e4m3( + kv_cache: torch.Tensor, + kv_cache_sf: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + inv_scale: float, +) -> None: + """Gather paged NVFP4 rows (packed E2M1 x E4M3 per 16, no global scale) into ``out`` + ([B, L, 1, W] float8_e4m3fn, contiguous) as ``e4m3_rne(clamp(value * inv_scale, +-448))`` + for every ``t < min(seq_lens[b], L)`` whose page id is in range; everything else of + ``out`` is left unwritten.""" + if not kv_cache.is_cuda: + raise ValueError(f"kv_cache must be a CUDA tensor, got device {kv_cache.device}") + device = kv_cache.device + if kv_cache.dtype not in _CODE_DTYPES: + raise ValueError(f"kv_cache must be uint8 (packed E2M1), got {kv_cache.dtype}") + if kv_cache_sf.dtype not in _SF_DTYPES: + raise ValueError(f"kv_cache_sf must be uint8 or float8_e4m3fn, got {kv_cache_sf.dtype}") + if ( + kv_cache.dim() != 4 + or kv_cache_sf.dim() != 4 + or kv_cache.shape[2] != 1 + or kv_cache_sf.shape[2] != 1 + ): + raise ValueError( + "kv_cache / kv_cache_sf must be [pages, page_size, 1, C], got " + f"{tuple(kv_cache.shape)} / {tuple(kv_cache_sf.shape)}" + ) + if tuple(kv_cache.shape[:2]) != tuple(kv_cache_sf.shape[:2]): + raise ValueError( + f"page geometry differs: {tuple(kv_cache.shape)} vs {tuple(kv_cache_sf.shape)}" + ) + if kv_cache.stride(-1) != 1 or kv_cache_sf.stride(-1) != 1: + raise ValueError("kv_cache and kv_cache_sf need a unit last-dim stride") + width = kv_cache.shape[-1] * 2 + if width % NVFP4_BLOCK != 0 or kv_cache_sf.shape[-1] * NVFP4_BLOCK != width: + raise ValueError( + f"expected {width // NVFP4_BLOCK} scale bytes per row, got {kv_cache_sf.shape[-1]}" + ) + if width & (width - 1): + raise ValueError(f"row width must be a power of two, got {width}") + if out.device != device or out.dtype != torch.float8_e4m3fn: + raise ValueError(f"out must be float8_e4m3fn on {device}, got {out.dtype} on {out.device}") + if out.dim() != 4 or out.shape[2] != 1 or out.shape[3] != width or not out.is_contiguous(): + raise ValueError(f"out must be a contiguous [B, L, 1, {width}], got {tuple(out.shape)}") + if block_tables.device != device or block_tables.dtype not in (torch.int32, torch.int64): + raise ValueError(f"block_tables must be int32/int64 on {device}") + if block_tables.dim() != 2 or block_tables.stride(-1) != 1: + raise ValueError( + f"block_tables must be [B, max_pages] with a unit last-dim stride, got {block_tables.shape}" + ) + if seq_lens.device != device or seq_lens.dtype != torch.int32 or seq_lens.dim() != 1: + raise ValueError("seq_lens must be a 1-D int32 CUDA tensor") + B, L = out.shape[0], out.shape[1] + if seq_lens.shape[0] != B or block_tables.shape[0] != B: + raise ValueError( + f"seq_lens {tuple(seq_lens.shape)} / block_tables {tuple(block_tables.shape)} vs B={B}" + ) + if not (inv_scale > 0.0): + raise ValueError(f"inv_scale must be a positive float, got {inv_scale}") + if B == 0 or L == 0: + return + sf_arg = kv_cache_sf.view(torch.uint8) if kv_cache_sf.dtype != torch.uint8 else kv_cache_sf + _nvfp4_kv_to_e4m3_kernel[(triton.cdiv(L, _ROWS), B)]( + kv_cache, + sf_arg, + block_tables, + seq_lens, + out.view(torch.uint8), + float(inv_scale), + kv_cache.stride(0), + kv_cache.stride(1), + sf_arg.stride(0), + sf_arg.stride(1), + block_tables.stride(0), + kv_cache.shape[0], + kv_cache.shape[1], + L, + W=width, + ROWS=_ROWS, + num_warps=4, + ) + + +@nvfp4_kv_to_e4m3.register_fake +def _nvfp4_kv_to_e4m3_fake( + kv_cache: torch.Tensor, + kv_cache_sf: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + inv_scale: float, +) -> None: + return None diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 346739d27479..c305f1bb829d 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -494,9 +494,13 @@ def _compute_global_layer_ids(manager, lg_idx: int) -> List[int]: # a rank with only ratio=128 layers won't have INDEXER_* types), so using # max(local values) would produce different num_attn_types across ranks, # causing the same (model_layer, attn_type) to map to different global IDs. - first_key = next(iter(manager._layer_attn_to_layer_id.keys())) - attn_type_class = type(first_key[1]) - num_attn_types = max(e.value for e in attn_type_class) + 1 + # A manager whose role type has members its layout never allocates states + # the width itself (``disagg_num_attn_types``, the same on every rank). + num_attn_types = getattr(manager, "disagg_num_attn_types", None) + if num_attn_types is None: + first_key = next(iter(manager._layer_attn_to_layer_id.keys())) + attn_type_class = type(first_key[1]) + num_attn_types = max(e.value for e in attn_type_class) + 1 return [inverse[lid][0] * num_attn_types + inverse[lid][1] for lid in local_layer_ids] diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py index c697aed2af76..7073907a1b65 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py @@ -65,20 +65,41 @@ def _uses_context_only_lazy_loading(self, model: Any) -> bool: and getattr(getattr(model, "model", None), "disagg_context_only", False) is True) + def _uses_model_requested_lazy_loading(self, model: Any) -> bool: + """Whether the model asks for lazy safetensors handles. + + A model class that sets ``requires_lazy_safetensors = True`` streams the + rank-local slices it owns out of the mmapped handles in its + ``load_weights``, so its checkpoint is never materialized in host RAM. + """ + return (type(self) is HfCheckpointLoader + and type(self.weight_loader) is HfWeightLoader + and getattr(model, "requires_lazy_safetensors", False) is True) + + def _lazy_loading_reason(self, model: Any) -> Optional[str]: + """Why this load must hand out lazy safetensors slices, or None.""" + if self._uses_context_only_lazy_loading(model): + return ("context-only model loading requires lazy safetensors " + "without full-checkpoint read-ahead") + if self._uses_model_requested_lazy_loading(model): + return ("the model streams rank-local slices out of lazy " + "safetensors handles") + return None + def load_weights(self, checkpoint_dir: str, mapping: Mapping, **kwargs) -> dict[str, Any]: - """Keep omitted decoder tensors lazy on a context-only worker.""" - if self._uses_context_only_lazy_loading(kwargs.get("model")): - # The model remapper drops unowned keys before touching their - # slices. Eager loading or read-ahead would read those bytes anyway. + """Keep the checkpoint lazy on a context-only worker, or when the model + asks for lazy safetensors handles.""" + lazy_reason = self._lazy_loading_reason(kwargs.get("model")) + if lazy_reason is not None: + # The model reads only the slices it owns. Eager loading or + # read-ahead would read every byte anyway. loader = self.weight_loader loader._reset_checkpoint_io_status() status = loader._last_checkpoint_io_status status.selected = "native" if loader.checkpoint_io_policy != "native": - status.fallback_reason = ( - "context-only model loading requires lazy safetensors " - "without full-checkpoint read-ahead") + status.fallback_reason = lazy_reason weights = loader._load_lazy_safetensors( checkpoint_dir, kwargs.get("use_consolidated", False)) status.effective = "native" @@ -92,8 +113,8 @@ def open_weight_session(self, checkpoint_dir: str, mapping: Mapping, """Delegate the optimized session only for the built-in HF path.""" if (type(self) is HfCheckpointLoader and type(self.weight_loader) is HfWeightLoader - and self.weight_loader.checkpoint_io_policy != "native" and - not self._uses_context_only_lazy_loading(kwargs.get("model"))): + and self.weight_loader.checkpoint_io_policy != "native" + and self._lazy_loading_reason(kwargs.get("model")) is None): with self.weight_loader.open_weight_session(checkpoint_dir, mapping=mapping, **kwargs) as weights: diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 6b415f7b7a7a..09070329c04f 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1419,6 +1419,28 @@ def _encode_dummy_inputs(self) -> Optional[torch.Tensor]: finally: self._dummy_encoder_inputs = [] + def _profile_inline_mm_encoder(self) -> Optional[torch.Tensor]: + """Run the largest encoder work of a model that encodes inside its forward. + + Such a model runs its multimodal encoder within the language model's + context step rather than through ``encode_multimodal_inputs``, so the + text-only dummy requests never reach it. It exposes + ``profile_inline_mm_encoder(max_num_tokens)``, which runs the largest + encoder work one step can carry and returns the storage that step keeps + while its language-model layers run. The measured peak then covers the + encoder like the separately encoded models above. + """ + model = self._model_engine.model + profile = getattr(model, "profile_inline_mm_encoder", None) + if profile is None or self._llm_args.disable_mm_encoder: + return None + if isinstance( + self._profiling_stage_data, + dict) and not self._profiling_stage_data.get("enable_mm_reqs"): + return None + with torch.inference_mode(): + return profile(max_num_tokens=self._max_num_tokens) + def _get_multimodal_encoder_memory_reserve(self, profiled_output_bytes: int = 0 ) -> int: @@ -1668,6 +1690,8 @@ def configure_kv_cache_capacity(self, # equivalent runtime-owned embeddings resident while the text-only # LLM dummy fills max_num_tokens. encoder_profile_output = self._encode_dummy_inputs() + if encoder_profile_output is None: + encoder_profile_output = self._profile_inline_mm_encoder() if encoder_profile_output is not None: profiled_output_bytes = (encoder_profile_output.numel() * encoder_profile_output.element_size()) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index f09883c04bcf..64efade9a4f4 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -3101,6 +3101,8 @@ def _apply_steady_gen_fast_prepare( attn_metadata.beam_width = 1 attn_metadata.request_ids = cache["request_ids"] attn_metadata.prompt_lens = cache["prompt_lens"] + attn_metadata.context_last_chunk = [] + attn_metadata.context_return_logits = [] attn_metadata.num_contexts = 0 attn_metadata.num_chunked_ctx_requests = 0 attn_metadata.kv_cache_params = KVCacheParams( @@ -3272,6 +3274,8 @@ def _prepare_tp_inputs( sequence_lengths = [] # per sequence prompt_lengths = [] # per sequence request_ids = [] # per request + context_last_chunk = [] # per context request + context_return_logits = [] # per context request gather_ids = [] position_ids = [] # per sequence num_cached_tokens_per_seq = [] # per sequence @@ -3306,6 +3310,8 @@ def _prepare_tp_inputs( for request in scheduled_requests.context_requests: request_ids.append(request.py_request_id) + context_last_chunk.append(bool(request.is_last_context_chunk)) + context_return_logits.append(bool(getattr(request, "py_return_context_logits", False))) draft_lens.append(0) begin_compute = request.context_current_position end_compute = begin_compute + request.context_chunk_size @@ -4137,6 +4143,8 @@ def previous_seq_slots_device(): attn_metadata.request_ids = request_ids attn_metadata.prompt_lens = prompt_lengths + attn_metadata.context_last_chunk = context_last_chunk + attn_metadata.context_return_logits = context_return_logits attn_metadata.num_contexts = scheduled_requests.num_context_requests remote_tail_modes = { getattr(request, "py_csa2_remote_tail_mode", None) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 78be3c9add38..30967204dd49 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6218,6 +6218,9 @@ def _validate_request(self, request: LlmRequest): # Perform sampler-specific validation self.sampler.validate_request(request) + # Perform resource-manager-specific validation + self.resource_manager.validate_request(request) + # Check if request has enough budget self._validate_request_budget(request) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 4a186c864a54..0eee069f05b2 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -67,6 +67,8 @@ f"SM{sm_version}" for sm_version in _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS) FP4_MLA_TOKENS_PER_BLOCK = 128 +# DeepseekV4CacheManager in DeepSeek-V4.1 layout mode requires this page size. +_DEEPSEEK_V41_TOKENS_PER_BLOCK = 128 class _ExecutorMemoryMonitor: @@ -877,6 +879,23 @@ def allocation_scope(current_stage: ExecutorMemoryType): logger.info( f"Change tokens_per_block to: {tokens_per_block} for using FlashMLA" ) + elif getattr(sparse_attention_config, "algorithm", + None) == "deepseek_v4": + from ..attention.backends.sparse.deepseek_v4.v41_layout import \ + is_deepseek_v41_config + if (is_deepseek_v41_config(config) + and tokens_per_block != _DEEPSEEK_V41_TOKENS_PER_BLOCK): + # DeepseekV4CacheManager's V4.1 layout addresses every pool + # in 128-token pages and rejects any other page size, so the + # default of 32 would fail at cache-manager construction. A + # DeepSeek-V4.1 checkpoint reaches this sparse config only + # through use_deepseek_v41_layout_mode(); the in-tree model + # runs on CSA2 and keeps its page size. + tokens_per_block = _DEEPSEEK_V41_TOKENS_PER_BLOCK + kv_cache_config.tokens_per_block = tokens_per_block + logger.info( + f"Change tokens_per_block to: {tokens_per_block} for " + "DeepSeek-V4.1") sm_version = get_sm_version() if (kv_cache_config.enable_block_reuse and sm_version diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index b17d3429cdb7..2ae014b24506 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -230,6 +230,15 @@ def kv_block_budget_applies(self) -> bool: """ return False + def validate_request(self, request: LlmRequest) -> None: + """Raise ``ValueError`` for a request this manager's engine cannot serve. + + Called at admission (``PyExecutor._validate_request``), so only the + offending request fails rather than the executor step that would run + it. The default accepts every request. + """ + return None + def add_dummy_requests(self, request_ids: List[int]): pass @@ -3134,6 +3143,12 @@ def get_request_kv_block_budgets( ) return budgets + def validate_request(self, request: LlmRequest) -> None: + """Let every resource manager refuse the request at admission.""" + for resource_manager in self.resource_managers.values(): + if hasattr(resource_manager, "validate_request"): + resource_manager.validate_request(request) + def get_unchecked_kv_block_budget_pools(self) -> List[ResourceManagerType]: """KV pools present but excluded from the admission check, i.e. those whose ``get_request_kv_block_budget`` returns ``None`` for every diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 7c2c2a1b0e30..0666cbda9007 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -1195,13 +1195,20 @@ def create_input_processor( f"checkpoint_format={checkpoint_format}; skipping HF config load.") if config is not None: - try: - model_cls, _ = get_model_architecture(config) - input_processor_cls = INPUT_PROCESSOR_REGISTRY._input_processors_cls_by_model_type \ - .get(model_cls) - except RuntimeError: # unregistered model - logger.info("Unregistered model, using DefaultInputProcessor") - input_processor_cls = None + # A modeling_v2 target family (TRTLLM_MODELING_V2) that reads its own + # multimodal inputs brings its processor; otherwise the class registered + # under the checkpoint's architecture decides. + from tensorrt_llm._torch._experimental.modeling_v2 import \ + modeling_v2_input_processor + input_processor_cls = modeling_v2_input_processor(config) + if input_processor_cls is None: + try: + model_cls, _ = get_model_architecture(config) + input_processor_cls = INPUT_PROCESSOR_REGISTRY._input_processors_cls_by_model_type \ + .get(model_cls) + except RuntimeError: # unregistered model + logger.info("Unregistered model, using DefaultInputProcessor") + input_processor_cls = None if input_processor_cls is not None: if (issubclass(input_processor_cls, BaseMultimodalInputProcessor) and input_processor_cls.supports_tokenization_cache): diff --git a/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py b/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py index fba84a22ef9c..b43342278fdc 100644 --- a/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py +++ b/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py @@ -159,7 +159,9 @@ def _find_last_user_index(messages: list[dict[str, Any]]) -> int: def _merge_tool_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: merged: list[dict[str, Any]] = [] for message in messages: - message = copy.deepcopy(message) + # A message parsed by the OpenAI server carries its pending media loads (coroutines, + # which cannot be deep-copied) under "media"; the renderer never reads them. + message = copy.deepcopy({key: value for key, value in message.items() if key != "media"}) role = message.get("role") if role == "tool": diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index e7424fdcfbab..b67355bc0846 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -138,6 +138,27 @@ deepseek-ai/DeepSeek-V4.1-Flash: accuracy: 93.3283 threshold: 91.0 num_samples: 1319 +nvidia/DeepSeek-V4.1-Flash-NVFP4: + # Read by the modeling_v2 target's gates (test_modeling_v2_deepseek_v41.py). + # Self-measured on an independent engine: the model card publishes no GSM8K + # number. vLLM 0.1.dev20904+g179dd0fa9 (vllm/vllm-openai:deepseekv41-flash-0909, + # arm64 sha256:d84a123255b822fc22508635218000187221794f59c0694c33b0650d1e377d58), + # 4x GB300 TP4, NVFP4 experts on the FLASHINFER_TRTLLM MoE backend, no + # speculative decoding, 2026-09-23. Protocol = trtllm-eval gsm8k with NO chat + # template: 5-shot plain completion (task default), random_seed=0 (dataset + # shuffle + few-shot seed as LmEvalEvaluator does), greedy, + # max_output_length=512, max_input_length=4096, full 1319 samples, + # lm_eval==0.4.10. strict-match 92.343, flexible-extract 92.267; the value is + # their mean, which is what trtllm-eval reports without a scores filter. + # The chat-template protocol was measured too and rejected: its two filters + # diverge (81.73 vs 95.30). The DSpark row (this checkpoint's mtp.* stages + # are its DSpark drafter) carries the same number because speculative + # verification is output-preserving. + - quant_algo: MIXED_PRECISION + accuracy: 92.30 + - quant_algo: MIXED_PRECISION + spec_dec_algo: DSpark + accuracy: 92.30 deepseek-ai/DeepSeek-V4-Flash-Base: # Base (pretrained, non-instruct) checkpoint, so GSM8K lands well below the # instruct DeepSeek-V4-Flash above. GSM8K measurements from diff --git a/tests/integration/defs/accuracy/references/mmmu.yaml b/tests/integration/defs/accuracy/references/mmmu.yaml index 36f1a0aeaa99..d6a3f3b2178b 100644 --- a/tests/integration/defs/accuracy/references/mmmu.yaml +++ b/tests/integration/defs/accuracy/references/mmmu.yaml @@ -79,3 +79,22 @@ moonshotai/Kimi-K3: zai-org/GLM-5.3-Flash: - quant_algo: FP8_BLOCK_SCALES accuracy: 77.22 +# DeepSeek-V4.1-Flash NVFP4, read by the modeling_v2 target's gates +# (test_modeling_v2_deepseek_v41.py). Self-measured on an independent engine +# (the model card publishes MMMU-Pro with thinking at 262k tokens, a different +# benchmark). vLLM 0.1.dev20904+g179dd0fa9 +# (vllm/vllm-openai:deepseekv41-flash-0909, arm64 +# sha256:d84a123255b822fc22508635218000187221794f59c0694c33b0650d1e377d58), +# 4x GB300 TP4, no speculative decoding, 2026-09-23. Protocol = trtllm-eval +# mmmu (mmmu_val, 900 samples, 0-shot, greedy, max_output_length=512, +# random_seed=0, lm_eval==0.4.10) with thinking OFF: the DeepSeek-V4.1 chat +# rendering, one <｜deepseek_image｜> per image placed BEFORE the question text +# (lm-eval's "" markers stripped), blocks joined by a blank line. acc +# 66.889 +/- 1.511. The DSpark row (the checkpoint's mtp.* drafter) carries the +# same number because speculative verification is output-preserving. +nvidia/DeepSeek-V4.1-Flash-NVFP4: + - quant_algo: MIXED_PRECISION + accuracy: 66.89 + - quant_algo: MIXED_PRECISION + spec_dec_algo: DSpark + accuracy: 66.89 diff --git a/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py b/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py index 1b28045609dd..421a56d85894 100644 --- a/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py +++ b/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py @@ -34,6 +34,7 @@ import pytest from tensorrt_llm import LLM +from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV from tensorrt_llm._utils import get_sm_version from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, MTPDecodingConfig @@ -111,8 +112,14 @@ def test_gsm8k_identity_vs_mtp3(self, mocker, monkeypatch): task = GSM8K(self.MODEL_NAME) # One dict for both legs: the comparison only means anything if the two - # engines were built under the same mode. - modeling_v2 = modeling_v2_llm_args("require", monkeypatch) + # engines were built under the same mode. It also asks each engine to + # check its step contract on its first forward -- the contract is about + # metadata that engine prepared, so the only place it can be checked + # honestly is a run that prepared it; a synthesized metadata object + # would be checking the synthesizer. + modeling_v2 = modeling_v2_llm_args( + "require", monkeypatch, **{MODELING_V2_VALIDATE_ENV: "1"} + ) with LLM(self.MODEL_PATH, **modeling_v2, **self.DEP4) as llm: identity = task.evaluate(llm) @@ -165,7 +172,13 @@ def test_mtp3_acceptance(self, mode, mocker, monkeypatch): # makes the two comparable rather than less so. Only this leg needs it: # the target builds that path itself and never consults the factory. # It rides along with the switch because it is read on the ranks too. - deep_ep = {"TRTLLM_CAN_USE_DEEP_EP": "0"} if mode == "off" else {} + # Ask the engine to check its step contract on its first forward. The + # contract is about metadata that engine prepared, so the only place it + # can be checked honestly is a run that prepared it; a synthesized + # metadata object would be checking the synthesizer. + rank_env = {MODELING_V2_VALIDATE_ENV: "1"} + if mode == "off": + rank_env["TRTLLM_CAN_USE_DEEP_EP"] = "0" mocker.patch.dict(GSM8K.EVALUATE_KWARGS, _SCORES_FILTER) @@ -175,7 +188,7 @@ def test_mtp3_acceptance(self, mode, mocker, monkeypatch): kv_cache_config=self.MTP3_KV, cuda_graph_config=CudaGraphConfig(), enable_iter_perf_stats=True, - **modeling_v2_llm_args(mode, monkeypatch, **deep_ep), + **modeling_v2_llm_args(mode, monkeypatch, **rank_env), **self.DEP4, ) as llm: task = GSM8K(self.MODEL_NAME) diff --git a/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v41.py b/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v41.py new file mode 100644 index 000000000000..59cad9456dec --- /dev/null +++ b/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v41.py @@ -0,0 +1,203 @@ +# 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. +"""Whole-model gates for the deepseek_v41 modeling_v2 target. + +One file per model family, beside the other accuracy suites rather than inside +test_llm_api_pytorch.py: these gate a parallel implementation, and reading them +next to the built-in DeepSeek-V4.1 tests (``TestDeepSeekV41*`` there) would +invite treating one as a variant of the other. + +The target is ``ModelingV2DeepseekV41FlashNvfp4Sm103Unified``: one modeling +codebase that folds three 4-rank members at runtime from the mapping -- tep4 +(attention TP4 + MoE EP4), dep4 (attention DP4 + MoE EP4) and tp4 (attention +TP4 + MoE TP4). Each case gates against the anchor in +``references/{gsm8k,mmmu}.yaml``, keyed by this checkpoint's name. + +Every case runs the target's default feature set: the checkpoint's DSpark +drafter (its ``mtp.*`` stages) at ``max_draft_len`` 5, CUDA graphs, chunked +prefill, the overlap scheduler and KV block reuse, with the Engram tables in +host memory and the KV pools in native (FP8/FP4) storage. The encoder and +decoder SWA bounded replays are left to the cache manager, which turns them on +for every member (MMMU's reuse-off engine has no prefix hit to replay). MMMU +turns block reuse off, because multimodal requests with reuse on are recorded +as crashing and the anchor was measured with prefix caching off. + +Every test here runs under ``TRTLLM_MODELING_V2=require``, which the case puts +in place itself through ``modeling_v2_env``: the switch travels as an +``LLM(...)`` field, so it reaches every rank whether the LLM API spawned the +ranks or ``trtllm-llmapi-launch`` started them. Under ``"auto"`` a configuration +that missed the target's criteria would fall back to the built-in +implementation instead of failing, and a pass would then be reported as the +target's. Each case also asks the target to check its step contract on its +first forward, as the other modeling_v2 gates do. The feature switches are left +to the launch; a case only checks that each one is at its default, which is +what ``_require_default_features`` does. On a host where the LLM API cannot +spawn its workers, launch the file under +``mpirun -n 4 --oversubscribe trtllm-llmapi-launch pytest ...``. + +The protocol is the one the anchors were measured under: ``trtllm-eval +gsm8k`` 5-shot plain completion with no chat template, and ``trtllm-eval +mmmu`` 0-shot with the V4.1 non-thinking chat rendering; ``random_seed`` 0, +greedy, 512 output tokens, ``max_seq_len`` 16384 and ``max_batch_size`` 32. +""" + +import os + +import pytest + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.engram_host_table import ( # noqa: E501 + HOST_OFFLOAD_ENV, +) +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.v41_layout import ( + DEEPSEEK_V41_DECODER_REPLAY_ENV, + DEEPSEEK_V41_ENCODER_REPLAY_ENV, + DEEPSEEK_V41_KV_STORAGE_ENV, + DEEPSEEK_V41_KV_STORAGE_NATIVE, +) +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.llmapi import CudaGraphConfig, DSparkDecodingConfig, KvCacheConfig + +from ..conftest import llm_models_root +from .accuracy_core import GSM8K, MMMU, LlmapiAccuracyTestHarness +from .modeling_v2_env import modeling_v2_llm_args + +_MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4.1-Flash-NVFP4" + +# The target asserts its own SM at construction: certification is per GPU +# architecture, and a receipt from another one says nothing here. +skip_not_sm103 = pytest.mark.skipif( + get_sm_version() != 103, + reason="the deepseek_v41 modeling_v2 target is certified on sm_103 only", +) +skip_no_checkpoint = pytest.mark.skipif( + not os.path.isdir(_MODEL_PATH), reason=f"checkpoint not staged at {_MODEL_PATH}" +) + +# Values each feature switch may hold for a case to measure the default +# feature set. Unset is always the default. The two replays accept nothing +# else: unset lets the cache manager resolve them per engine, while an explicit +# value turns one off or forces it on. +_DEFAULT_FEATURE_VALUES = { + HOST_OFFLOAD_ENV: ("1", "true", "yes", "on"), + DEEPSEEK_V41_KV_STORAGE_ENV: (DEEPSEEK_V41_KV_STORAGE_NATIVE,), + DEEPSEEK_V41_ENCODER_REPLAY_ENV: (), + DEEPSEEK_V41_DECODER_REPLAY_ENV: (), +} + + +def _require_default_features() -> None: + """Skip unless the ranks were started with every feature at its default. + + Not a failure: the ranks read these switches as the launch set them, so a + launch that overrode one measures a different engine than the one these + anchors gate, and has nothing to say about them. + """ + for name, accepted in _DEFAULT_FEATURE_VALUES.items(): + value = os.environ.get(name) + if value is not None and value.strip().lower() not in accepted: + pytest.skip(f"{name}={value!r}; these gates measure the target's default features") + + +@skip_not_sm103 +@skip_no_checkpoint +@pytest.mark.skip_less_device(4) +class TestModelingV2DeepseekV41FlashNvfp4Sm103Unified(LlmapiAccuracyTestHarness): + """deepseek-v4.1-flash-nvfp4 / sm_103 / unified (tep4, dep4, tp4).""" + + MODEL_NAME = "nvidia/DeepSeek-V4.1-Flash-NVFP4" + MODEL_PATH = _MODEL_PATH + + # The three members of the unified target. Routing derives the member from + # the mapping, and the target asserts it against the one the engine built. + # tp4 leaves the MoE split unset, as its gate did: it resolves to MoE TP4 + # (moe_tp_size 4, moe_ep_size 1). + MEMBERS = { + "tep4": dict(tensor_parallel_size=4, moe_expert_parallel_size=4), + "dep4": dict(tensor_parallel_size=4, moe_expert_parallel_size=4, enable_attention_dp=True), + "tp4": dict(tensor_parallel_size=4), + } + + # max_position_embeddings is 1048576; left to default, the engine would + # size its KV cache for a million-token sequence. + MAX_SEQ_LEN = 16384 + MAX_BATCH_SIZE = 32 + + # The anchors' protocol: greedy, 512 output tokens. The input caps are the + # tasks' own (GSM8K 4096, MMMU 8192), and the MMMU stop string is the one + # `trtllm-eval mmmu` passes. + MAX_OUTPUT_LEN = 512 + MMMU_STOP = "<|endoftext|>" + + # The tokenizer ships no chat template; the DeepSeek-V4 tokenizer renders + # the V4.1 non-thinking chat prompt, byte-identical to the checkpoint's + # encoding/encoding.py on all 900 MMMU prompts. + MMMU_TOKENIZER = "deepseek_v4" + MMMU_EVALUATOR_KWARGS = dict(chat_template_kwargs={"thinking": False}) + + def _llm( + self, member: str, monkeypatch: pytest.MonkeyPatch, *, multimodal: bool = False + ) -> LLM: + return LLM( + self.MODEL_PATH, + max_seq_len=self.MAX_SEQ_LEN, + max_batch_size=self.MAX_BATCH_SIZE, + max_num_tokens=8192, + enable_chunked_prefill=True, + disable_overlap_scheduler=False, + cuda_graph_config=CudaGraphConfig( + enable_padding=True, max_batch_size=self.MAX_BATCH_SIZE + ), + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.5, enable_block_reuse=not multimodal + ), + # The checkpoint's mtp.* stages are its DSpark block drafter, loaded + # from the target checkpoint itself. + speculative_config=DSparkDecodingConfig(max_draft_len=5), + custom_tokenizer=self.MMMU_TOKENIZER if multimodal else None, + **self.MEMBERS[member], + **modeling_v2_llm_args("require", monkeypatch, **{MODELING_V2_VALIDATE_ENV: "1"}), + ) + + @pytest.mark.parametrize("member", ["tep4", "dep4", "tp4"]) + def test_gsm8k(self, member, monkeypatch): + _require_default_features() + with self._llm(member, monkeypatch) as llm: + task = GSM8K(self.MODEL_NAME) + # No scores filter: the anchor is the mean of strict-match and + # flexible-extract, which is what `trtllm-eval gsm8k` reports and + # what this task averages by default. + task.evaluate( + llm, + sampling_params=SamplingParams( + max_tokens=self.MAX_OUTPUT_LEN, truncate_prompt_tokens=GSM8K.MAX_INPUT_LEN + ), + ) + + @pytest.mark.parametrize("member", ["tep4"]) + def test_mmmu(self, member, monkeypatch): + _require_default_features() + with self._llm(member, monkeypatch, multimodal=True) as llm: + task = MMMU(self.MODEL_NAME) + task.evaluate( + llm, + sampling_params=SamplingParams( + max_tokens=self.MAX_OUTPUT_LEN, + truncate_prompt_tokens=MMMU.MAX_INPUT_LEN, + stop=self.MMMU_STOP, + ), + extra_evaluator_kwargs=self.MMMU_EVALUATOR_KWARGS, + ) diff --git a/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v41_replay.py b/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v41_replay.py new file mode 100644 index 000000000000..6f53e24f902f --- /dev/null +++ b/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v41_replay.py @@ -0,0 +1,142 @@ +# 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. +"""Bitwise gates of DeepSeek-V4.1's encoder SWA bounded replay on the modeling_v2 target. + +On a prefix-cache hit of ``reused`` tokens, encoder replay re-forwards the last +``R = min(128, reused)`` cached tokens (they rebuild the sliding windows the prefix +cache does not keep) and charges them to the step's token budget. So the hit's +context chunks start R positions earlier than the same hit's would without +replay. Chunked prefill is not invariant to where chunks start: cuBLAS picks its +kernel by row count, so a row's bits depend on the row count of its step. A hit +with replay is therefore not bitwise equal to the same hit without it once the +prompt spans more than one chunk. What replay does guarantee, and what these +cases gate: + +* a hit of at most 128 tokens replays from position 0, so it runs exactly the + cold run's chunks and its tokens and logprobs are bitwise equal to the cold + run's. The producer's first step has as many rows as the cold run's first + chunk, so the cached rows were computed at the same row count; +* two producers of one prefix give bitwise-equal hits; +* the same cold request twice is bitwise equal. + +The engine is multi-chunk (``max_num_tokens`` 2048 against a ~4.6K-token prompt), +greedy, without speculation (so generation logprobs are returned), with both +replay switches unset: the cache manager resolves them on at tep4 and dep4. At +dep4 every request is pinned to attention-DP rank 0, since a prefix hit needs the +producer's KV on the same rank. ``TRTLLM_MODELING_V2=require`` and the step +contract check travel on the ``LLM`` (``modeling_v2_env``), as in the accuracy +file; on a host where the LLM API cannot spawn its workers, launch under +``mpirun -n 4 --oversubscribe trtllm-llmapi-launch pytest ...``. +""" + +import os + +import pytest + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.v41_layout import ( + DEEPSEEK_V41_DECODER_REPLAY_ENV, + DEEPSEEK_V41_ENCODER_REPLAY_ENV, +) +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, SchedulingParams + +from .modeling_v2_env import modeling_v2_llm_args +from .test_modeling_v2_deepseek_v41 import _MODEL_PATH, skip_no_checkpoint, skip_not_sm103 + +_MEMBERS = { + "tep4": dict(tensor_parallel_size=4, moe_expert_parallel_size=4), + "dep4": dict(tensor_parallel_size=4, moe_expert_parallel_size=4, enable_attention_dp=True), +} +_BUDGET = 2048 +_HIT = 120 # a hit of at most one window: replay from position 0 +_TAIL = 200 + + +def _prompt(tokenizer): + """A ~4.6K-token passkey haystack: three chunks at a 2048-token budget.""" + filler = ( + "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again. " + ) + text = ( + filler * 120 + + "The pass key is 73520. Remember it. 73520 is the pass key. " + + filler * 120 + + "What is the pass key? The pass key is" + ) + ids = tokenizer.encode(text, add_special_tokens=False) + assert len(ids) > 2 * _BUDGET, len(ids) + return ids + + +def _tail(n: int, seed: int): + return [1000 + (seed * 7919 + i * 104729) % 59000 for i in range(n)] + + +@skip_not_sm103 +@skip_no_checkpoint +@pytest.mark.skip_less_device(4) +@pytest.mark.parametrize("member", ["tep4", "dep4"]) +def test_encoder_replay_bitwise(member, monkeypatch): + for name in (DEEPSEEK_V41_ENCODER_REPLAY_ENV, DEEPSEEK_V41_DECODER_REPLAY_ENV): + if os.environ.get(name) is not None: + pytest.skip( + f"{name} is set; these cases gate the replays as the cache manager resolves them" + ) + with LLM( + _MODEL_PATH, + max_seq_len=16384, + max_batch_size=8, + max_num_tokens=_BUDGET, + enable_chunked_prefill=True, + cuda_graph_config=CudaGraphConfig(), + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.5, enable_block_reuse=True), + **_MEMBERS[member], + **modeling_v2_llm_args("require", monkeypatch, **{MODELING_V2_VALIDATE_ENV: "1"}), + ) as llm: + sp = SamplingParams(max_tokens=16, temperature=0.0, logprobs=5) + pin = {} + if member == "dep4": + pin["scheduling_params"] = [ + SchedulingParams(attention_dp_rank=0, attention_dp_relax=False) + ] + + def run(ids, salt): + out = llm.generate([ids], sp, cache_salt=[salt], **pin)[0] + o = out.outputs[0] + logprobs = [{t: v.logprob for t, v in step.items()} for step in o.logprobs] + return list(o.token_ids), logprobs, out.cached_tokens + + ids = _prompt(llm.tokenizer) + cold = run(ids, "cold") + assert cold[:2] == run(ids, "cold-again")[:2], "the same cold request twice differs" + + # A hit of _HIT tokens. The producer's first step has _BUDGET rows, as the cold + # run's first chunk does, so its cached rows were computed at the same row count. + run(ids[:_HIT] + _tail(_BUDGET - _HIT, 1), "short") + short = run(ids, "short") + assert short[:2] == cold[:2], ( + "a hit replayed from position 0 runs the cold run's chunks and must match it bitwise" + ) + + # Two producers of the prompt's first N - 60 tokens (a hit inside its last + # window), with different tails of one length: bitwise-equal hits. + k = len(ids) - 60 + run(ids[:k] + _tail(_TAIL, 2), "x") + cross = run(ids, "x") + run(ids[:k] + _tail(_TAIL, 3), "y") + cross2 = run(ids, "y") + assert cross[2] > 0 and cross2[2] > 0, (cross[2], cross2[2]) + assert cross[:2] == cross2[:2], "two producers of one prefix gave different hits" diff --git a/tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py b/tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py index 03708a9037db..148210a4ff86 100644 --- a/tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py +++ b/tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py @@ -29,6 +29,7 @@ import pytest from tensorrt_llm import LLM +from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV from tensorrt_llm._utils import get_sm_version from ..conftest import llm_models_root @@ -71,6 +72,17 @@ def test_gsm8k(self, mocker, monkeypatch): mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", 8192) mocker.patch.dict(GSM8K.EVALUATE_KWARGS, {"scores_filter": "exact_match,flexible-extract"}) - with LLM(self.MODEL_PATH, **modeling_v2_llm_args("require", monkeypatch)) as llm: + # The target checks its step contract on the first forward and only + # when asked. Asking here rather than in a unit test is deliberate: the + # contract is about metadata this engine prepared -- the field set, the + # int32 position ids, the two KV pools this checkpoint's alternating + # window sizes produce -- and a synthesized metadata object would be + # checking the synthesizer. One forward of the run below pays for it. + # It rides with the switch: the check runs where the metadata does, + # which is the rank rather than this process. + with LLM( + self.MODEL_PATH, + **modeling_v2_llm_args("require", monkeypatch, **{MODELING_V2_VALIDATE_ENV: "1"}), + ) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm, extra_evaluator_kwargs=self.extra_evaluator_kwargs) diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 5ec4a2aabce6..357d96ac3cf7 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -52,6 +52,7 @@ l0_b300: - unittest/_torch/thop/serial - unittest/_torch/executor # 250s - unittest/_torch/disaggregation + - unittest/disaggregated/test_deepseek_v41_kv_transfer.py # ------------- modules (non-MoE) --------------- - unittest/_torch/modules/test_fused_add_rms_norm_quant.py - unittest/_torch/modules/test_fused_activation_quant.py @@ -132,8 +133,9 @@ l0_b300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_chunked_prefill[quant_dtype=fp8-kv_cache_reuse=True-fp8kv=True-overlap_scheduler=True] # Cover nvbugs 6084445 - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] - # TIMEOUT measured: 348 cases in 5m34s on one GB300, so 30 is 5x headroom. - - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm TIMEOUT (30) + # TIMEOUT: ~47 min estimated on one GB300 -- 5m34s for the 348 base cases plus ~37 min that pr-a's + # receipts measured for the shared files carrying DeepSeek-V4.1 cells -- so 90 is ~2x headroom. + - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm --ignore=unittest/_torch/modeling_v2/models/deepseek_v41 TIMEOUT (90) - accuracy/test_modeling_v2_gpt_oss.py::TestModelingV2GptOss120bSm103Tp1::test_gsm8k - condition: ranges: @@ -150,6 +152,11 @@ l0_b300: stage: post_merge backend: pytorch tests: + # TIMEOUT from pr-a's whole-file receipts on one GB300: attention 144 min, quantization 80 min, + # the rest 35 min; each is ~1.7-1.9x headroom. + - unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention -k "not test_capture_set_replays_identical_across_processes_and_gpus and not test_bitwise_identical_across_calls_processes_and_gpus" TIMEOUT (240) + - unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization TIMEOUT (150) + - unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified --ignore=unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention --ignore=unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization --ignore=unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm -k "not test_vmm_table_created_by_a_one_gpu_process and not test_real_host_table_four_ranks_read_the_same_mapping" TIMEOUT (60) # ---- moved to post-merge (MoE CI optimization) ---- - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "CUTLASS" - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index c25b2148d14d..422c15800145 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -17,9 +17,21 @@ l0_gb300_multi_gpus: tests: # One ordered sequence in one job: the cases share module-global rank state. - unittest/_torch/modeling_v2/comm + - unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm + - unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_flashinfer_top_k_ragged_deterministic.py::test_capture_set_replays_identical_across_processes_and_gpus + - unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_flashinfer_top_k_ragged_deterministic.py::test_bitwise_identical_across_calls_processes_and_gpus + - unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/embedding/test_modeling_v2_engram_gather.py::test_vmm_table_created_by_a_one_gpu_process + - unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/embedding/test_modeling_v2_engram_gather.py::test_real_host_table_four_ranks_read_the_same_mapping - accuracy/test_modeling_v2_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_gsm8k_identity_vs_mtp3 - accuracy/test_modeling_v2_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[modeling_v2] - accuracy/test_modeling_v2_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[stock] + # TIMEOUT measured on 4x GB300: 4.5-8.5 min per case, engine boot included. + - accuracy/test_modeling_v2_deepseek_v41.py::TestModelingV2DeepseekV41FlashNvfp4Sm103Unified::test_gsm8k[tep4] TIMEOUT (90) + - accuracy/test_modeling_v2_deepseek_v41.py::TestModelingV2DeepseekV41FlashNvfp4Sm103Unified::test_gsm8k[dep4] TIMEOUT (90) + - accuracy/test_modeling_v2_deepseek_v41.py::TestModelingV2DeepseekV41FlashNvfp4Sm103Unified::test_gsm8k[tp4] TIMEOUT (90) + - accuracy/test_modeling_v2_deepseek_v41.py::TestModelingV2DeepseekV41FlashNvfp4Sm103Unified::test_mmmu[tep4] TIMEOUT (90) + - accuracy/test_modeling_v2_deepseek_v41_replay.py::test_encoder_replay_bitwise[tep4] TIMEOUT (90) + - accuracy/test_modeling_v2_deepseek_v41_replay.py::test_encoder_replay_bitwise[dep4] TIMEOUT (90) # Covers tests/unittest/_torch/attention/. Two sub-trees moved in here from elsewhere # under tests/unittest/_torch/ and have no entry of their own on any list, so this entry # is what picks them up: diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v41_decoder_replay.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v41_decoder_replay.py new file mode 100644 index 000000000000..376a5a3e14ab --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v41_decoder_replay.py @@ -0,0 +1,284 @@ +# 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. +"""DeepSeek-V4.1 decoder SWA bounded replay: the runtime half. + +The model runs decoder replay; the V4.1 cache manager resolves where it may run. +Its window KV of the layers after the last KV-source layer holds each prompt's last +``window`` positions only, so with block reuse on it needs encoder replay, which +re-forwards a hit's tail rows. Under decoder replay the manager also refuses, at +admission, the context logits and prompt logprobs those layers never compute. It admits +every prompt length the engine admits, and so it needs no second check when a request +paused for recompute comes back with its generated tokens in its prompt. +""" + +import os +from unittest.mock import MagicMock + +import pytest +from utils.util import skip_pre_blackwell + +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + cache_manager as v4_cache_manager, +) +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.cache_manager import ( + _v41_decoder_replay_refusal, + _v41_encoder_replay_refusal, +) +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.v41_layout import ( + DEEPSEEK_V41_DECODER_REPLAY_ENV, + DEEPSEEK_V41_ENCODER_REPLAY_ENV, + get_deepseek_v41_decoder_replay_setting, + get_deepseek_v41_layout, +) +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm.bindings import SamplingConfig +from tensorrt_llm.mapping import Mapping + +from .test_deepseek_v41_encoder_replay import make_manager, prompt +from .test_deepseek_v41_layout import WINDOW, make_v41_config + +# candidate_topk_blocks x candidate_block_size of the checkpoint (2048 x 8): past this +# many visible positions layer 20's candidate blocks bind. +CANDIDATE_POSITIONS = 16384 + + +def test_switch_and_refusal(): + """The switch parses like the encoder one; with block reuse on, decoder replay is + refused exactly when encoder replay is off, naming why encoder replay is off.""" + saved = os.environ.pop(DEEPSEEK_V41_DECODER_REPLAY_ENV, None) + try: + assert get_deepseek_v41_decoder_replay_setting() is None + for value, expected in (("1", True), ("on", True), ("0", False), ("No", False)): + os.environ[DEEPSEEK_V41_DECODER_REPLAY_ENV] = value + assert get_deepseek_v41_decoder_replay_setting() is expected + os.environ[DEEPSEEK_V41_DECODER_REPLAY_ENV] = "maybe" + with pytest.raises(ValueError, match=DEEPSEEK_V41_DECODER_REPLAY_ENV): + get_deepseek_v41_decoder_replay_setting() + finally: + os.environ.pop(DEEPSEEK_V41_DECODER_REPLAY_ENV, None) + if saved is not None: + os.environ[DEEPSEEK_V41_DECODER_REPLAY_ENV] = saved + + layout = get_deepseek_v41_layout(make_v41_config()) + assert _v41_decoder_replay_refusal(layout, True, None) is None + assert _v41_decoder_replay_refusal(layout, False, "a KV cache connector") is None + refusal = _v41_decoder_replay_refusal(layout, True, "a KV cache connector") + assert refusal is not None and "block reuse" in refusal and "a KV cache connector" in refusal + assert f"last {WINDOW} positions" in refusal + # A KV connector turns encoder replay off (a manager with a stand-in connector is not + # built); the manager passes its short reason on to the decoder refusal. + _, reason = _v41_encoder_replay_refusal( + layout, Mapping(), 8192, {"kv_connector_manager": object()} + ) + assert reason.split(": ", 1)[0] == "a KV cache connector" + + +# (encoder switch, block reuse, manager kwargs, encoder replay resolved, why encoder +# replay is off as the decoder line and error name it -- None where decoder replay may run) +RESOLUTION_CASES = { + "tp": (None, True, {}, True, None), + "encoder_off_reuse_off": ("0", False, {}, False, None), + "encoder_off": ("0", True, {}, False, f"{DEEPSEEK_V41_ENCODER_REPLAY_ENV}=0"), + "disagg": (None, True, {"is_disagg": True}, False, "disaggregated serving"), + "disagg_reuse_off": (None, False, {"is_disagg": True}, False, None), + "small_token_budget": ( + None, + True, + {"max_num_tokens": WINDOW}, + False, + f"max_num_tokens <= {WINDOW} (got {WINDOW})", + ), + # Attention DP: each rank's manager replays its own requests' hits, and the model + # bounds its own requests; both replays run there as under attention TP. + "attention_dp": ( + None, + True, + {"mapping": Mapping(world_size=1, rank=0, tp_size=1, enable_attention_dp=True)}, + True, + None, + ), +} + + +@skip_pre_blackwell +@pytest.mark.parametrize("case", list(RESOLUTION_CASES)) +def test_resolves_with_encoder_replay(case, capfd, monkeypatch): + """Ruling R21 for decoder replay: with block reuse on it needs encoder replay. Unset, + it resolves on (attention TP or DP) unless reuse is on and encoder replay is off, then + off with one line naming why; an explicit 1 there raises before any pool exists. An + explicit 0 is always off.""" + encoder_switch, reuse, kwargs, encoder_on, why = RESOLUTION_CASES[case] + monkeypatch.setattr(v4_cache_manager, "_V41_DECODER_REPLAY_ANNOUNCED", set()) + refused = why is not None + + def build(decoder_switch): + if decoder_switch is None: + monkeypatch.delenv(DEEPSEEK_V41_DECODER_REPLAY_ENV, raising=False) + else: + monkeypatch.setenv(DEEPSEEK_V41_DECODER_REPLAY_ENV, decoder_switch) + manager, _ = make_manager(encoder_on, switch=encoder_switch, block_reuse=reuse, **kwargs) + manager.shutdown() + return manager + + capfd.readouterr() + manager = build(None) + assert (manager.v41_encoder_replay_off_reason is None) is encoder_on + assert manager.v41_decoder_replay is not refused, case + lines = [ + line + for line in capfd.readouterr().out.splitlines() + if DEEPSEEK_V41_DECODER_REPLAY_ENV in line + ] + if refused: + assert len(lines) == 1, lines + assert "resolves off" in lines[0] and why in lines[0] and "block reuse" in lines[0] + # Once per process: the engine builds a KV-size estimation manager first. + assert build(None).v41_decoder_replay is False + assert DEEPSEEK_V41_DECODER_REPLAY_ENV not in capfd.readouterr().out + else: + assert lines == [], lines + + if refused: + with pytest.raises(ValueError, match=DEEPSEEK_V41_DECODER_REPLAY_ENV) as error: + build("1") + assert why in str(error.value) + else: + assert build("1").v41_decoder_replay is True + assert build("0").v41_decoder_replay is False + assert DEEPSEEK_V41_DECODER_REPLAY_ENV not in capfd.readouterr().out + + +def admission_request( + request_id: int, *, context_logits: bool, generation_only: bool = False, length: int = 300 +): + request = LlmRequest( + request_id=request_id, + max_new_tokens=8, + input_tokens=prompt(request_id, length), + sampling_config=SamplingConfig(), + is_streaming=False, + return_context_logits=context_logits, + ) + if generation_only: + request.py_llm_request_type = LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY + return request + + +def admit(resource_manager: ResourceManager, request: LlmRequest) -> None: + """PyExecutor's own admission check, with everything but the resource managers stubbed.""" + executor = MagicMock() + executor.kv_cache_transceiver = None + executor.max_beam_width = 1 + executor.resource_manager = resource_manager + PyExecutor._validate_request(executor, request) + + +# (encoder switch, decoder switch, block reuse, decoder replay resolved) +ADMISSION_CASES = { + "replay_on": (None, None, True, True), + "switch_off": (None, "0", True, False), + "resolved_off": ("0", None, True, False), +} + + +@skip_pre_blackwell +@pytest.mark.parametrize("case", list(ADMISSION_CASES)) +def test_admission_refuses_context_logits_under_replay(case, monkeypatch): + """Under decoder replay a request for context logits or prompt logprobs (which the + executor serves through context logits) fails at admission with a message naming the + switch, instead of raising inside the forward, which fails the whole step. Other + requests, a generation-only request (its context ran on another engine) and every + request without decoder replay are admitted.""" + encoder_switch, decoder_switch, reuse, decoder_on = ADMISSION_CASES[case] + if decoder_switch is None: + monkeypatch.delenv(DEEPSEEK_V41_DECODER_REPLAY_ENV, raising=False) + else: + monkeypatch.setenv(DEEPSEEK_V41_DECODER_REPLAY_ENV, decoder_switch) + manager, _ = make_manager(encoder_switch is None, switch=encoder_switch, block_reuse=reuse) + try: + assert manager.v41_decoder_replay is decoder_on + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: manager}) + admit(resource_manager, admission_request(1, context_logits=False)) + admit(resource_manager, admission_request(2, context_logits=True, generation_only=True)) + refused = admission_request(3, context_logits=True) + if decoder_on: + with pytest.raises( + ValueError, match="return_context_logits / prompt_logprobs" + ) as error: + admit(resource_manager, refused) + assert f"{DEEPSEEK_V41_DECODER_REPLAY_ENV}=0" in str(error.value) + assert f"last {WINDOW} positions" in str(error.value) + else: + admit(resource_manager, refused) + finally: + manager.shutdown() + + +@skip_pre_blackwell +@pytest.mark.parametrize("decoder_switch", [None, "0"], ids=["replay_on", "switch_off"]) +def test_admission_admits_long_and_recomputed_prompts(decoder_switch, monkeypatch): + """Decoder replay serves every prompt length the engine admits, past the candidate + blocks' 2048 x 8 positions too, so admission refuses none by length. A request paused + for recompute comes back with its generated tokens in its prompt + (``PyExecutor._pause_recompute_request``); admitting it again gives the same answer.""" + if decoder_switch is None: + monkeypatch.delenv(DEEPSEEK_V41_DECODER_REPLAY_ENV, raising=False) + else: + monkeypatch.setenv(DEEPSEEK_V41_DECODER_REPLAY_ENV, decoder_switch) + max_seq_len = 2 * CANDIDATE_POSITIONS + manager, _ = make_manager(True, switch=None, max_seq_len=max_seq_len, max_batch_size=1) + try: + assert manager.v41_decoder_replay is (decoder_switch is None) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: manager}) + for n, length in enumerate((CANDIDATE_POSITIONS, CANDIDATE_POSITIONS + 1, max_seq_len - 8)): + admit(resource_manager, admission_request(n + 1, context_logits=False, length=length)) + + prompt_len, generated = CANDIDATE_POSITIONS - 384, 1000 + request = LlmRequest( + request_id=9, + max_new_tokens=4000, + input_tokens=prompt(9, prompt_len), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + admit(resource_manager, request) + for t in range(generated): + request.add_new_token(t % 1000, 0) + PyExecutor._pause_recompute_request(MagicMock(), request) + assert request.prompt_len == prompt_len + generated > CANDIDATE_POSITIONS + admit(resource_manager, request) + finally: + manager.shutdown() + + +def test_resource_manager_validate_request_skips_managers_without_the_hook(): + """Resource managers that predate the hook, or unset slots, are skipped.""" + seen = [] + + class Hooked: + def validate_request(self, request): + seen.append(request) + + resource_manager = ResourceManager( + { + ResourceManagerType.SEQ_SLOT_MANAGER: object(), + ResourceManagerType.SPEC_RESOURCE_MANAGER: None, + ResourceManagerType.KV_CACHE_MANAGER: Hooked(), + } + ) + resource_manager.validate_request("request") + assert seen == ["request"] diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v41_encoder_replay.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v41_encoder_replay.py new file mode 100644 index 000000000000..3a1a350f2dc0 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v41_encoder_replay.py @@ -0,0 +1,613 @@ +# 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. +"""DeepSeek-V4.1 encoder SWA bounded replay: the runtime half. + +The prefix cache keeps main KV and index K only; a hit of ``reused`` tokens (a +multiple of the compress ratio) schedules the last ``min(window, reused)`` cached +tokens as replay rows at the head of the first chunk, charged to the token budget, +and the attention metadata tells the target which rows they are. +""" + +import os +from contextlib import contextmanager +from typing import Iterator, List, Optional, Sequence + +import pytest +import torch +from utils.util import skip_pre_blackwell + +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + DeepseekV4CacheManager, + DeepseekV4TrtllmAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + cache_manager as v4_cache_manager, +) +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.cache_manager import ( + _v41_encoder_replay_refusal, +) +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.v41_layout import ( + DEEPSEEK_V41_DEFAULT_ENCODER_REPLAY, + DEEPSEEK_V41_ENCODER_REPLAY_ENV, + get_deepseek_v41_encoder_replay, + get_deepseek_v41_layout, +) +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.bindings import DataType, SamplingConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import ( + CapacitySchedulerPolicy, + DeepSeekV4SparseAttentionConfig, + KvCacheConfig, +) +from tensorrt_llm.mapping import Mapping + +from .test_deepseek_v41_layout import ( + HEAD_DIM, + INDEX_HEAD_DIM, + TOKENS_PER_BLOCK, + WINDOW, + T, + kv_storage_env, + make_v41_config, + reference_page_indices, +) + +ENGRAM_HISTORY = 3 # engram_max_ngram_size 4 + + +@contextmanager +def replay_env(value: Optional[str]) -> Iterator[None]: + saved = os.environ.get(DEEPSEEK_V41_ENCODER_REPLAY_ENV) + try: + if value is None: + os.environ.pop(DEEPSEEK_V41_ENCODER_REPLAY_ENV, None) + else: + os.environ[DEEPSEEK_V41_ENCODER_REPLAY_ENV] = value + yield + finally: + if saved is None: + os.environ.pop(DEEPSEEK_V41_ENCODER_REPLAY_ENV, None) + else: + os.environ[DEEPSEEK_V41_ENCODER_REPLAY_ENV] = saved + + +def make_manager( + replay: bool, + *, + block_reuse: bool = True, + max_batch_size: int = 8, + max_seq_len: int = 4096, + max_num_tokens: int = 8192, + mapping: Optional[Mapping] = None, + switch: Optional[str] = "", + **extra, +) -> tuple: + """A V4.1-layout manager whose replay mode resolves to ``replay``. ``switch`` is the + environment value it is built under: "" = "1" / "0" from ``replay``, None = unset.""" + config = make_v41_config() + sparse_config = DeepSeekV4SparseAttentionConfig( + index_head_dim=INDEX_HEAD_DIM, + window_size=WINDOW, + compress_ratios=[1] * len(config.compress_ratios), + ) + kv_cache_config = KvCacheConfig( + enable_block_reuse=block_reuse, + max_tokens=max_seq_len * max_batch_size, + event_buffer_max_size=0, + ) + env = ("1" if replay else "0") if switch == "" else switch + with replay_env(env), kv_storage_env("bf16"): + manager = DeepseekV4CacheManager( + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=config.num_hidden_layers, + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + max_input_len=max_seq_len, + mapping=mapping or Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=max_num_tokens, + sparse_attn_config=sparse_config, + pretrained_config=config, + **extra, + ) + assert manager.v41_encoder_replay == replay + return manager, sparse_config + + +def make_request(request_id: int, tokens: Sequence[int]) -> LlmRequest: + return LlmRequest( + request_id=request_id, + max_new_tokens=8, + input_tokens=list(tokens), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + + +def prompt(seed: int, length: int) -> List[int]: + g = torch.Generator().manual_seed(seed) + return torch.randint(10, 120000, (length,), generator=g).tolist() + + +def run_context(manager: DeepseekV4CacheManager, request: LlmRequest) -> None: + """Schedule the whole context, run it, and commit it to the prefix cache.""" + assert manager.prepare_context(request) + assert manager.resize_context(request, request.context_chunk_size) + request.context_current_position = request.prompt_len + batch = ScheduledRequests() + batch.context_requests_last_chunk = [request] + manager.update_context_resources(batch) + + +def pool_stats(manager: DeepseekV4CacheManager, layer: int, attn_type): + impl = manager.impl + lc = impl.get_layer_group_id(manager._layer_attn_to_layer_id[layer, attn_type]) + pool_group = impl.get_life_cycle_pool_group_indices()[lc] + return impl.get_storage_statistics()[pool_group] + + +# --------------------------------------------------------------------------- +# Switch and layout (host only). +# --------------------------------------------------------------------------- + + +def test_switch_and_layout(): + with replay_env(None): + assert get_deepseek_v41_encoder_replay() is DEEPSEEK_V41_DEFAULT_ENCODER_REPLAY + for value, expected in ( + ("1", True), + ("true", True), + ("ON", True), + ("0", False), + ("off", False), + ): + with replay_env(value): + assert get_deepseek_v41_encoder_replay() is expected + with replay_env("maybe"), pytest.raises(ValueError, match=DEEPSEEK_V41_ENCODER_REPLAY_ENV): + get_deepseek_v41_encoder_replay() + + for replay in (False, True): + with replay_env("1" if replay else "0"): + layout = get_deepseek_v41_layout(make_v41_config()) + assert layout.encoder_replay is replay + assert layout.replay_window == WINDOW + assert layout.reuse_alignment == 2 # compress ratios {1, 2} + for draft in (0, 5): + assert layout.window(T.SWA, draft) == WINDOW + draft + assert layout.window(T.COMPRESSOR_KV, draft) == 2 + draft + # A replay step starts `window` positions before the reuse boundary the + # manager keeps live, and its first rows hash the 3 positions before it. + assert layout.window(T.ENGRAM_TOKEN, draft) == ( + ENGRAM_HISTORY + 1 + (WINDOW if replay else 0) + draft + ) + assert layout.window(T.COMPRESS, draft) is None + + +@skip_pre_blackwell +@pytest.mark.parametrize( + "case", + ["disagg", "small_token_budget", "joint_draft_reuse"], +) +def test_unsupported_combinations_fail_fast(case): + kwargs = {} + match = DEEPSEEK_V41_ENCODER_REPLAY_ENV + if case == "disagg": + kwargs["is_disagg"] = True + match = "disaggregated" + elif case == "small_token_budget": + kwargs["max_num_tokens"] = WINDOW + match = "max_num_tokens" + else: + kwargs["joint_kv_cache_reuse"] = True + match = "draft KV pool" + with pytest.raises(ValueError, match=match): + make_manager(True, **kwargs) + if case == "small_token_budget": + # Refused only because of replay. + manager, _ = make_manager(False, **kwargs) + manager.shutdown() + + +@skip_pre_blackwell +@pytest.mark.parametrize( + "case", ["tp", "attention_dp", "disagg", "connector", "small_token_budget"] +) +def test_unset_switch_resolves_per_configuration(case, capfd, monkeypatch): + """Ruling R21: unset, replay is on exactly where it is supported (aggregated, no + connector; attention TP or DP: under attention DP each rank's manager replays the + hits of its own requests) and off elsewhere with one line naming why; an explicit + 1 on an unsupported configuration raises and an explicit 0 is always off.""" + kwargs = {} + reason = None + if case == "attention_dp": + kwargs["mapping"] = Mapping(world_size=1, rank=0, tp_size=1, enable_attention_dp=True) + elif case == "disagg": + kwargs["is_disagg"] = True + reason = "disaggregated" + elif case == "connector": + kwargs["kv_connector_manager"] = object() + reason = "KV cache connector" + layout = get_deepseek_v41_layout(make_v41_config()) + refusal = _v41_encoder_replay_refusal(layout, Mapping(), 8192, kwargs) + assert refusal is not None and reason in refusal[1], refusal + # The explicit switch raises before any pool exists; a manager with a stand-in + # connector is not built. + with pytest.raises(ValueError, match=reason): + make_manager(True, switch="1", **kwargs) + return + elif case == "small_token_budget": + kwargs["max_num_tokens"] = WINDOW + reason = "max_num_tokens" + supported = case in ("tp", "attention_dp") + assert DEEPSEEK_V41_DEFAULT_ENCODER_REPLAY is True + # The line is once per process: start from none printed, whatever ran before. + monkeypatch.setattr(v4_cache_manager, "_V41_ENCODER_REPLAY_ANNOUNCED", set()) + capfd.readouterr() + manager, _ = make_manager(supported, switch=None, **kwargs) + manager.shutdown() + out = capfd.readouterr().out + if supported: + assert DEEPSEEK_V41_ENCODER_REPLAY_ENV not in out, out + else: + lines = [line for line in out.splitlines() if DEEPSEEK_V41_ENCODER_REPLAY_ENV in line] + assert len(lines) == 1 and "resolves off" in lines[0] and reason in lines[0], out + # Once per process: a second manager (the engine builds an estimation one + # first) does not repeat it. + manager, _ = make_manager(False, switch=None, **kwargs) + manager.shutdown() + assert DEEPSEEK_V41_ENCODER_REPLAY_ENV not in capfd.readouterr().out + with pytest.raises(ValueError, match=reason): + make_manager(True, switch="1", **kwargs) + manager, _ = make_manager(False, switch="0", **kwargs) + manager.shutdown() + + +# --------------------------------------------------------------------------- +# The prefix cache and the reuse boundary (GPU). +# --------------------------------------------------------------------------- + + +@skip_pre_blackwell +@pytest.mark.parametrize("replay", [False, True], ids=["replay_off", "replay_on"]) +def test_prefix_cache_keeps_global_kv_only(replay: bool): + """With replay on, a finished request leaves main KV and index K in the prefix + cache and no window page (SWA, compressor state, Engram history).""" + manager, _ = make_manager(replay) + try: + request = make_request(0, prompt(0, 1000)) + run_context(manager, request) + manager.free_resources(request) + cached = { + role: pool_stats(manager, layer, role).evictable + for layer, role in ( + (0, T.SWA), + (2, T.COMPRESSOR_KV), + (1, T.ENGRAM_TOKEN), + (2, T.COMPRESS), + (20, T.COMPRESS), + ) + } + assert cached[T.COMPRESS] > 0, cached + for role in (T.SWA, T.COMPRESSOR_KV, T.ENGRAM_TOKEN): + assert (cached[role] == 0) == replay, (role, cached) + finally: + manager.shutdown() + + +# (shared prefix length with the cached request, expected reuse boundary) +HIT_CASES = { + "mid_block_even": (700, 700), + "mid_block_odd": (701, 700), + "block_aligned": (896, 896), + "below_window": (101, 100), + "whole_prompt": (1000, 1000), +} + + +@skip_pre_blackwell +@pytest.mark.parametrize("name", list(HIT_CASES)) +def test_hit_schedules_replay_rows(name: str): + shared, reused = HIT_CASES[name] + manager, _ = make_manager(True) + requests = [] + try: + first = make_request(0, prompt(1, 1000)) + requests.append(first) + run_context(manager, first) # kept alive: its window pages stay private + tokens = first.get_tokens(0)[:shared] + prompt(2, 300) + second = make_request(1, tokens) + requests.append(second) + assert manager.prepare_context(second) + + num_replay = min(WINDOW, reused) + start = reused - num_replay + kv_cache = manager.kv_cache_map[second.py_request_id] + assert kv_cache.num_committed_tokens == reused # the cache's claim, even + assert reused % 2 == 0 + # The cursor sits at the replay start: the replay rows head the first chunk + # and count as context tokens of it (the scheduler's budget). + assert second.context_current_position == start + assert second.prepopulated_prompt_len == start + assert second.is_first_context_chunk + assert second.context_remaining_length == len(tokens) - start + plan = manager.get_v41_replay_plan(second.py_request_id, start) + assert plan is not None + assert (plan.start, plan.num_tokens, plan.reused) == (start, num_replay, reused) + assert list(plan.history_token_ids) == tokens[max(0, start - ENGRAM_HISTORY) : start] + assert manager.get_v41_replay_plan(second.py_request_id, start + 1) is None + + # A chunk that cannot hold every replay row is refused. + assert not manager.resize_context(second, num_replay - 1) + assert manager.resize_context(second, second.context_chunk_size) + + request_ids = [first.py_request_id, second.py_request_id] + # Window pages of the replay rows (and the Engram history before them) are + # allocated and private: never the cached request's pages. + window_starts = ( + (0, T.SWA, start), + (39, T.SWA, start), + (1, T.ENGRAM_TOKEN, max(0, start - 3)), + ) + for layer, role, lo in window_starts: + pages = reference_page_indices(manager, request_ids, 2, layer, role, pages_mode(role)) + for ordinal in range(lo // TOKENS_PER_BLOCK, (reused - 1) // TOKENS_PER_BLOCK + 1): + assert int(pages[1, ordinal]) >= 0, (role, ordinal) + assert int(pages[1, ordinal]) != int(pages[0, ordinal]), (role, ordinal) + # Main KV and index K of the fully cached blocks are the cached request's. + for owner in (2, 20): + for role in (T.COMPRESS, T.INDEXER_COMPRESS): + pages = reference_page_indices( + manager, request_ids, 2, owner, role, pages_mode(role) + ) + for ordinal in range(reused // TOKENS_PER_BLOCK): + assert int(pages[1, ordinal]) == int(pages[0, ordinal]), (owner, role, ordinal) + finally: + for request in requests: + manager.free_resources(request) + manager.shutdown() + + +def pages_mode(role): + from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode + + if role in (T.COMPRESS, T.INDEXER_COMPRESS): + return PageIndexMode.SHARED + return PageIndexMode.PER_LAYER + + +@skip_pre_blackwell +def test_replay_off_keeps_the_claim(): + """Replay off: a hit settles the cursor at the claim and schedules no replay rows.""" + manager, _ = make_manager(False) + requests = [] + try: + first = make_request(0, prompt(1, 1000)) + requests.append(first) + run_context(manager, first) + second = make_request(1, first.get_tokens(0) + prompt(2, 50)) + requests.append(second) + assert manager.prepare_context(second) + claimed = manager.kv_cache_map[second.py_request_id].num_committed_tokens + assert second.context_current_position == claimed + assert manager.get_v41_replay_plan(second.py_request_id, claimed) is None + finally: + for request in requests: + manager.free_resources(request) + manager.shutdown() + + +@skip_pre_blackwell +@pytest.mark.parametrize("chunked", [False, True], ids=["whole_context", "chunked"]) +def test_replay_rows_are_charged_to_the_token_budget(chunked: bool): + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + from tensorrt_llm.llmapi.llm_args import ContextChunkingPolicy + + manager, _ = make_manager(True) + requests = [] + try: + first = make_request(0, prompt(1, 1000)) + requests.append(first) + run_context(manager, first) + manager.free_resources(first) + requests.remove(first) + tokens = first.get_tokens(0)[:700] + prompt(2, 300) # 1000 tokens, 700 cached + start = 700 - WINDOW + needed = len(tokens) - start # suffix 300 + replay 128 + + def schedule(budget: int, request_id: int): + request = make_request(request_id, tokens) + requests.append(request) + ctx_chunk_config = ( + (ContextChunkingPolicy.FIRST_COME_FIRST_SERVED, TOKENS_PER_BLOCK) + if chunked + else None + ) + scheduler = KVCacheV2Scheduler( + max_batch_size=8, + max_num_tokens=budget, + kv_cache_manager=manager, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ctx_chunk_config=ctx_chunk_config, + ) + return request, scheduler.schedule_request([request], set()) + + if chunked: + # The first chunk starts at the replay start and holds the replay rows. + request, out = schedule(2 * TOKENS_PER_BLOCK, 10) + assert [r.py_request_id for r in out.context_requests] == [10] + assert request.context_current_position == start + assert request.context_chunk_size == 2 * TOKENS_PER_BLOCK + assert manager.get_v41_replay_plan(10, start) is not None + else: + request, out = schedule(needed, 10) + assert [r.py_request_id for r in out.context_requests] == [10] + assert request.context_current_position == start + assert request.context_chunk_size == needed + # One token short: the replay rows are charged, so the context no + # longer fits, and the failed admission drops the claim and the plan. + request, out = schedule(needed - 1, 11) + assert out.context_requests == [] + assert request.context_current_position == 0 + assert 11 not in manager.kv_cache_map + assert manager.get_v41_replay_plan(11, start) is None + finally: + for request in requests: + manager.free_resources(request) + manager.shutdown() + + +# --------------------------------------------------------------------------- +# What the target sees (GPU). +# --------------------------------------------------------------------------- + + +def build_metadata(manager, sparse_config, requests, num_contexts, cached, new_lens): + metadata = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor(new_lens, dtype=torch.int), + request_ids=[r.py_request_id for r in requests], + max_num_requests=manager.max_batch_size, + num_contexts=num_contexts, + prompt_lens=[c + n for c, n in zip(cached, new_lens)], + max_num_tokens=8192, + kv_cache_manager=manager, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + sparse_attention_config=sparse_config, + ) + metadata.prepare() + torch.cuda.synchronize() + return metadata + + +@skip_pre_blackwell +@pytest.mark.parametrize("replay", [False, True], ids=["replay_off", "replay_on"]) +def test_metadata_exposes_replay_rows(replay: bool): + """A batch of [hit context, cold context, generation]: the hit's first rows are + replay rows; the compressed stream and its write slots start at the reuse + boundary; the replay rows' window starts at the replay start.""" + manager, sparse_config = make_manager(replay) + requests = [] + try: + first = make_request(0, prompt(1, 1000)) + run_context(manager, first) + manager.free_resources(first) + + hit = make_request(1, first.get_tokens(0)[:701] + prompt(2, 299)) # 1000 tokens + cold = make_request(2, prompt(3, 257)) + gen = make_request(3, prompt(4, 130)) + for request in (hit, cold, gen): + requests.append(request) + assert manager.prepare_context(request) + assert manager.resize_context(request, request.context_chunk_size) + claimed = manager.kv_cache_map[hit.py_request_id].num_committed_tokens + reused = 700 if replay else claimed + num_replay = WINDOW if replay else 0 + start = reused - num_replay + assert hit.context_current_position == start + cached = [start, 0, 129] + new_lens = [1000 - start, 257, 1] + metadata = build_metadata(manager, sparse_config, requests, 2, cached, new_lens) + num_tokens = sum(new_lens) + rows_of = [range(0, new_lens[0]), range(new_lens[0], new_lens[0] + 257), [num_tokens - 1]] + + assert metadata.v41_num_replay_tokens_cpu[:3].tolist() == [num_replay, 0, 0] + assert metadata.v41_num_replay_tokens_cuda[:3].tolist() == [num_replay, 0, 0] + assert metadata.v41_num_replay_tokens_total == num_replay + is_replay = metadata.v41_is_replay_token_cuda[:num_tokens].tolist() + assert is_replay == [t < num_replay for t in range(num_tokens)] + assert metadata.v41_num_compute_tokens == num_tokens - num_replay + compute = metadata.v41_compute_token_index_cuda[: num_tokens - num_replay].tolist() + assert compute == list(range(num_replay, num_tokens)) + compute_start = [start + num_replay, 0, 129] + assert metadata.v41_compute_start_pos_cuda[:3].tolist() == compute_start + compute_lens = [new_lens[0] - num_replay, 257, 1] + assert metadata.v41_compute_cu_seq_lens_cuda[:4].tolist() == [ + 0, + compute_lens[0], + compute_lens[0] + 257, + compute_lens[0] + 258, + ] + + # Compressed stream: new entries from the compute start only, so no cached + # main-KV / index-K entry of the replay rows is written. + for r in (1, 2): + kv_lens = [c + n for c, n in zip(cached, new_lens)] + new_comp = [k // r - s // r for k, s in zip(kv_lens, compute_start)] + assert metadata.new_comp_kv_lens_cuda[r][:3].tolist() == new_comp + assert metadata.past_kv_lens_cuda[r][:3].tolist() == [s // r for s in compute_start] + assert metadata.num_total_compressed_tokens[r] == sum(new_comp) + positions = metadata.compressed_position_ids_cuda[r][: new_comp[0]].tolist() + assert positions == [(compute_start[0] // r + k) * r for k in range(new_comp[0])] + slots = metadata.v41_compress_slot_mapping[r][: new_comp[0]].cpu() + table = metadata.compress_block_tables[r][0].cpu() + per_page = TOKENS_PER_BLOCK // r + for k, slot in zip(range(compute_start[0] // r, kv_lens[0] // r), slots.tolist()): + assert slot == int(table[k // per_page]) * per_page + k % per_page + + # Window of each query: truncated at the replay start for the hit's rows. + swa = metadata.swa_local_indices_cuda[:num_tokens].cpu() + floors = [start if replay else 0, 0, 0] + assert metadata.v41_swa_floor_cuda[:3].tolist() == floors + # The context attention's seq_lens: kv length minus the floor. + kv_lens = [c + n for c, n in zip(cached, new_lens)] + seq_lens = metadata.v41_swa_seq_lens_cuda[:2].tolist() + assert seq_lens == [kv_lens[0] - floors[0], kv_lens[1]] + for b, rows in enumerate(rows_of): + for j, t in enumerate(rows): + p = cached[b] + j + want = list(range(max(floors[b], p - WINDOW + 1), p + 1)) + row = swa[t].tolist() + assert row[: len(want)] == want, (b, p) + assert all(i == -1 for i in row[len(want) :]), (b, p) + if b < 2: + # The kernel's window count covers exactly the window: no -1. + assert min(WINDOW, seq_lens[b] - new_lens[b] + j + 1) == len(want), (b, p) + # Every row writes its window entry, replay rows included. + assert int((metadata.v41_swa_slot_mapping[0, :num_tokens] < 0).sum()) == 0 + + # The Engram history before the replay start, as engram_hash input rows. + if replay: + ids = hit.get_tokens(0)[start - ENGRAM_HISTORY : start] + assert metadata.v41_replay_history_len == ENGRAM_HISTORY + assert metadata.v41_replay_history_token_ids_cuda[:3].tolist() == ids + assert metadata.v41_replay_history_req_idx_cuda[:3].tolist() == [0, 0, 0] + assert metadata.v41_replay_history_cu_seq_lens_cuda[:3].tolist() == [0, 3, 3] + assert metadata.v41_replay_history_start_pos_cuda[:2].tolist() == [start - 3, 0] + table = metadata.engram_token_block_table[0].cpu() + for p in range(start - ENGRAM_HISTORY, reused): + assert int(table[p // TOKENS_PER_BLOCK]) >= 0, p + else: + assert metadata.v41_replay_history_len == 0 + + # A later chunk of the same request carries no replay rows. + later = build_metadata( + manager, sparse_config, requests[:1], 1, [start + 256], [1000 - start - 256] + ) + assert later.v41_num_replay_tokens_total == 0 + assert later.v41_compute_start_pos_cuda[:1].tolist() == [start + 256] + assert later.v41_swa_seq_lens_cuda[:1].tolist() == [1000] + finally: + for request in requests: + manager.free_resources(request) + manager.shutdown() diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v41_layout.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v41_layout.py new file mode 100644 index 000000000000..5b0ca645f770 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v41_layout.py @@ -0,0 +1,1735 @@ +# 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. +"""DeepSeek-V4.1 layout mode of DeepseekV4CacheManager and DeepseekV4TrtllmAttentionMetadata.""" + +import os +from contextlib import contextmanager +from types import SimpleNamespace +from typing import Iterator, List, Optional, Sequence, Tuple + +import pytest +import torch +from utils.util import skip_pre_blackwell + +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + DeepseekV4CacheManager, + DeepseekV4TrtllmAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DEEPSEEK_V4_SLIDING_ATTENTION, + DeepseekV4AttentionType, + compress_ratio_has_attention, +) +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.v41_layout import ( + DEEPSEEK_V41_DEFAULT_KV_STORAGE, + DEEPSEEK_V41_KV_STORAGE_BF16, + DEEPSEEK_V41_KV_STORAGE_ENV, + DEEPSEEK_V41_KV_STORAGE_NATIVE, + DeepseekV41Layout, + deepseek_v41_native_entry_split, + deepseek_v41_sparse_attention_config, + deepseek_v41_token_bytes, + get_deepseek_v41_kv_storage, + get_deepseek_v41_layout, + is_deepseek_v41_config, + use_deepseek_v41_layout_mode, +) +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm.bindings import DataType, SamplingConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import ( + CSA2SparseAttentionConfig, + DeepSeekV4SparseAttentionConfig, + KvCacheConfig, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX, AttnLifeCycle, PageIndexMode + +T = DeepseekV4AttentionType +TOKENS_PER_BLOCK = 128 +WINDOW = 128 +INDEX_TOPK = 512 +HEAD_DIM = 512 +INDEX_HEAD_DIM = 128 + +# The checkpoint's text_config fields the layout reads (config.json of +# DeepSeek-V4.1-Flash-NVFP4), 40 backbone layers plus 3 DSpark stages. +V41_TEXT_FIELDS = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + candidate_topk_blocks=2048, + candidate_block_size=8, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=WINDOW, + head_dim=HEAD_DIM, + index_head_dim=INDEX_HEAD_DIM, + index_topk=INDEX_TOPK, + kv_lora_rank=448, + qk_rope_head_dim=64, +) + + +def make_v41_config(**overrides) -> SimpleNamespace: + """A composite config as ModelConfig sees it: text fields mirrored to the top level.""" + text = dict(V41_TEXT_FIELDS, **overrides) + top = {k: v for k, v in text.items() if k != "model_type"} + return SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **top, + ) + + +@contextmanager +def kv_storage_env(kv_storage: Optional[str]) -> Iterator[None]: + """Set TRTLLM_DSV41_KV_STORAGE for the block (None: unset); the previous value is restored.""" + saved = os.environ.get(DEEPSEEK_V41_KV_STORAGE_ENV) + try: + if kv_storage is None: + os.environ.pop(DEEPSEEK_V41_KV_STORAGE_ENV, None) + else: + os.environ[DEEPSEEK_V41_KV_STORAGE_ENV] = kv_storage + yield + finally: + if saved is None: + os.environ.pop(DEEPSEEK_V41_KV_STORAGE_ENV, None) + else: + os.environ[DEEPSEEK_V41_KV_STORAGE_ENV] = saved + + +# Bytes per entry: SWA / main KV in bf16 containers or native, index K MXFP4 in both. +ENTRY_BYTES = { + DEEPSEEK_V41_KV_STORAGE_BF16: {"SWA": 1024, "COMPRESS": 1024, "INDEXER_COMPRESS": 68}, + DEEPSEEK_V41_KV_STORAGE_NATIVE: {"SWA": 528, "COMPRESS": 288, "INDEXER_COMPRESS": 68}, +} +# Persistent (global) KV bytes per token, R16: 3 * (C + 68) / 2 + (C + 68). +PERSISTENT_BYTES_PER_TOKEN = { + DEEPSEEK_V41_KV_STORAGE_BF16: 2730, + DEEPSEEK_V41_KV_STORAGE_NATIVE: 890, +} + + +# --------------------------------------------------------------------------- +# Reference: the checkpoint's inference/model.py index rules. +# --------------------------------------------------------------------------- + + +def ref_get_window_topk_idxs(window_size: int, bsz: int, seqlen: int, start_pos: int): + """model.py:409-426 (without the lru_cache).""" + if start_pos == 0: + end = torch.arange(seqlen).unsqueeze(1) + idxs = (end - window_size + 1).clamp(0) + torch.arange(min(seqlen, window_size)) + idxs = torch.where(idxs > end, -1, idxs) + else: + oldest = start_pos % window_size + 1 + idxs = torch.cat([torch.arange(oldest, window_size), torch.arange(oldest)]) + idxs = torch.where(idxs > start_pos, -1, idxs) + return idxs.int().unsqueeze(0).expand(bsz, -1, -1).contiguous() + + +def ref_window_positions(position: int) -> List[int]: + """Window positions a query at ``position`` attends, per the reference. + + A decode step (one query at start_pos = position) sees the ring; its slots + are mapped back to the positions they hold. A query inside a prompt sees + row ``position`` of the causal prefill rule (start_pos == 0); chunked + prefill and verify rows are queries inside the prompt at their absolute + positions. + """ + if position == 0: + row = ref_get_window_topk_idxs(WINDOW, 1, 1, 0)[0, 0].tolist() + return sorted(i for i in row if i >= 0) + ring = ref_get_window_topk_idxs(WINDOW, 1, 1, position)[0, 0].tolist() + decode = sorted(position - ((position - s) % WINDOW) for s in ring if s >= 0) + prefill_row = ref_get_window_topk_idxs(WINDOW, 1, position + 1, 0)[0, position].tolist() + prefill = sorted(i for i in prefill_row if i >= 0) + assert decode == prefill, (position, decode, prefill) + return prefill + + +def ref_valid_compressed(position: int, end_pos: int, ratio: int) -> int: + """Valid entries of the reference indexer top-k row (model.py:563-580). + + ``compress_lens`` = (i + 1) // ratio in prefill, end_pos // ratio in + decode; ``topk = min(index_topk, end_pos // ratio)``; entries at or past + compress_lens become -1, and sort after the valid ones. + """ + topk = min(INDEX_TOPK, end_pos // ratio) + return min(topk, (position + 1) // ratio) + + +# --------------------------------------------------------------------------- +# Layout (host only). +# --------------------------------------------------------------------------- + + +def test_layout_detection_and_role_table(): + config = make_v41_config() + assert is_deepseek_v41_config(config) + assert is_deepseek_v41_config(config.text_config) + assert not is_deepseek_v41_config( + SimpleNamespace(architectures=["DeepseekV4ForCausalLM"], model_type="deepseek_v4") + ) + assert get_deepseek_v41_layout(SimpleNamespace(architectures=["DeepseekV4ForCausalLM"])) is None + + layout = get_deepseek_v41_layout(config) + assert layout.num_layers == 40 + assert layout.compressed_ratios == (1, 2) + assert layout.layer_ratios == (0, 1, 2) + # plan.md section 2 "Layer roles". + for layer in (0, 1): + assert layout.kv_source(layer) is None and layout.index_source(layer) is None + assert layout.attention_types(layer) == ( + (T.SWA, T.ENGRAM_TOKEN) if layer == 1 else (T.SWA,) + ) + for layer in range(2, 20): + assert layout.kv_source(layer) == max(s for s in (2, 8, 14) if s <= layer) + assert layout.index_source(layer) == layout.kv_source(layer) + for layer in range(20, 40): + assert layout.kv_source(layer) == 20 + assert layout.index_source(layer) == max(s for s in (20, 24, 28, 32, 36) if s <= layer) + owners = { + attn_type: [layer for layer in range(40) if layout.has_attention(layer, attn_type)] + for attn_type in T + } + assert owners[T.SWA] == list(range(40)) + assert owners[T.COMPRESS] == [2, 8, 14, 20] + assert owners[T.INDEXER_COMPRESS] == [2, 8, 14, 20] + assert owners[T.COMPRESSOR_KV] == [2, 8, 14] + assert owners[T.COMPRESSOR_SCORE] == [2, 8, 14] + assert owners[T.ENGRAM_TOKEN] == [1] + assert owners[T.INDEXER_COMPRESSOR_KV] == [] + assert owners[T.INDEXER_COMPRESSOR_SCORE] == [] + assert layout.owner(14, T.ENGRAM_TOKEN) == 1 + assert layout.owner(27, T.COMPRESS) == 20 + assert layout.owner(27, T.COMPRESSOR_KV) is None + # Windows: 128 + D, 2 + D, 4 + D (history 3 plus the step's own position; + # plus the 128-token replay span under encoder SWA bounded replay); + # persistent roles have none. + assert layout.window(T.SWA, 5) == 133 + assert layout.window(T.COMPRESSOR_KV, 5) == 7 + assert layout.window(T.ENGRAM_TOKEN, 5) == 9 + (WINDOW if layout.encoder_replay else 0) + assert layout.window(T.COMPRESS, 5) is None + # R16: 3 * (C + 68) / 2 + (C + 68), C = 1024 (bf16 containers) or 288 (native). + for storage, expected in PERSISTENT_BYTES_PER_TOKEN.items(): + pinned = DeepseekV41Layout.from_config(config, kv_storage=storage) + assert pinned.kv_storage == storage + assert pinned.persistent_bytes_per_token(range(40)) == expected, storage + + +@pytest.mark.parametrize("max_draft_len", [0, 5]) +def test_windows_keep_every_position_the_step_reads(max_draft_len: int): + """Each windowed role's window, read through KVCacheManagerV2's live-window rule, + keeps every position a step starting at history ``c`` reads. + + KVCacheManagerV2 treats ``[c + 1 - window, c]`` as live (the window counts the + in-flight token at position ``c``); blocks before it are stale and are neither + kept for the request nor mapped on a prefix hit. What the V4.1 forward reads + before position ``c``: + + - SWA: ``c - 127 .. c - 1`` (the attention window of token ``c``); + - compressor pending pair: ``c - 1`` when ``c`` is odd (ratio 2); + - Engram token history: ``c - 3 .. c - 1`` (the 4-gram ending at ``c``). + + Measured before the Engram window counted the in-flight token: a prefix hit at + ``c = 4610`` (``c % 128 == 2``) left the page of ``c - 3`` unmapped. + """ + layout = get_deepseek_v41_layout(make_v41_config()) + reads = { + T.SWA: lambda c: range(max(0, c - (layout.window_size - 1)), c), + T.COMPRESSOR_KV: lambda c: range(c - 1, c) if c % 2 else range(0), + T.COMPRESSOR_SCORE: lambda c: range(c - 1, c) if c % 2 else range(0), + T.ENGRAM_TOKEN: lambda c: range(max(0, c - layout.engram_history_len), c), + } + for attn_type, needed in reads.items(): + lc = AttnLifeCycle.make(layout.window(attn_type, max_draft_len), 0, TOKENS_PER_BLOCK) + for c in range(1, 4 * TOKENS_PER_BLOCK + 3): + stale = lc.get_stale_range(c, TOKENS_PER_BLOCK) + for p in needed(c): + assert not (stale.beg <= p // TOKENS_PER_BLOCK < stale.end), (attn_type, c, p) + + +def test_kv_storage_switch(): + with kv_storage_env(None): + assert get_deepseek_v41_kv_storage() == DEEPSEEK_V41_DEFAULT_KV_STORAGE + assert get_deepseek_v41_layout(make_v41_config()).kv_storage == ( + DEEPSEEK_V41_DEFAULT_KV_STORAGE + ) + for value, expected in ( + ("bf16", DEEPSEEK_V41_KV_STORAGE_BF16), + (" Native ", DEEPSEEK_V41_KV_STORAGE_NATIVE), + ("", DEEPSEEK_V41_DEFAULT_KV_STORAGE), + ): + with kv_storage_env(value): + assert get_deepseek_v41_kv_storage() == expected, value + assert get_deepseek_v41_layout(make_v41_config()).kv_storage == expected, value + with kv_storage_env("fp8"), pytest.raises(ValueError, match="is not one of"): + get_deepseek_v41_layout(make_v41_config()) + + # Native entries: codes, then block scales. + assert deepseek_v41_native_entry_split(T.SWA, HEAD_DIM) == (512, 16) + assert deepseek_v41_native_entry_split(T.COMPRESS, HEAD_DIM) == (256, 32) + with pytest.raises(ValueError): + deepseek_v41_native_entry_split(T.INDEXER_COMPRESS, HEAD_DIM) + for storage, table in ENTRY_BYTES.items(): + for name, expected in table.items(): + assert ( + deepseek_v41_token_bytes(T[name], HEAD_DIM, INDEX_HEAD_DIM, storage) == expected + ), (storage, name) + # The windowed state roles do not depend on the storage mode. + assert deepseek_v41_token_bytes(T.COMPRESSOR_KV, HEAD_DIM, INDEX_HEAD_DIM, storage) == 2048 + assert deepseek_v41_token_bytes(T.ENGRAM_TOKEN, HEAD_DIM, INDEX_HEAD_DIM, storage) == 4 + + +@pytest.mark.parametrize( + "overrides, message", + [ + (dict(compress_ratios=[0, 0, 4] + [1] * 40), "must be in"), + (dict(kv_source_layer_ids=[8, 14, 20]), "no KV-source layer"), + (dict(index_source_layer_ids=[2, 8, 14, 24]), "index"), + (dict(kv_source_layer_ids=[1, 8, 14, 20]), "ratio 0"), + ], +) +def test_layout_rejects_inconsistent_roles(overrides, message): + with pytest.raises(ValueError, match=message): + get_deepseek_v41_layout(make_v41_config(**overrides)) + + +def test_v4_role_predicates_unchanged(): + assert T.ENGRAM_TOKEN not in DEEPSEEK_V4_SLIDING_ATTENTION + assert len(DEEPSEEK_V4_SLIDING_ATTENTION) == 5 + for ratio in (1, 4, 128): + assert not compress_ratio_has_attention(ratio, T.ENGRAM_TOKEN) + # Disaggregated synthetic layer ids keep DeepSeek-V4's encoding width (7 roles); + # only the DeepSeek-V4.1 layout mode counts ENGRAM_TOKEN. + v4_manager = DeepseekV4CacheManager.__new__(DeepseekV4CacheManager) + assert v4_manager.v41_layout is None + assert v4_manager.disagg_num_attn_types == 7 + v41_manager = DeepseekV4CacheManager.__new__(DeepseekV4CacheManager) + v41_manager.v41_layout = get_deepseek_v41_layout(make_v41_config()) + assert v41_manager.disagg_num_attn_types == T.ENGRAM_TOKEN.value + 1 + + +def test_layout_mode_sparse_config(): + """The deepseek_v4 sparse config of the layout mode, and its installation on a + ModelConfig that from_pretrained gave the CSA2 config.""" + from tensorrt_llm._torch.model_config import ModelConfig + + config = make_v41_config() + sparse = deepseek_v41_sparse_attention_config(config) + assert sparse.algorithm == "deepseek_v4" + # A placeholder: the manager and the metadata read the layout from the checkpoint. + assert sparse.compress_ratios == [1] * len(V41_TEXT_FIELDS["compress_ratios"]) + assert sparse.window_size == WINDOW + assert (sparse.index_head_dim, sparse.index_topk) == (INDEX_HEAD_DIM, INDEX_TOPK) + assert sparse.indexer_k_dtype == "fp4" + with pytest.raises(ValueError, match="compress_ratios"): + deepseek_v41_sparse_attention_config(make_v41_config(compress_ratios=[])) + + model_config = ModelConfig(pretrained_config=config) + assert isinstance(model_config.sparse_attention_config, CSA2SparseAttentionConfig) + model_config._frozen = True + use_deepseek_v41_layout_mode(model_config) + assert model_config._frozen + installed = model_config.sparse_attention_config + assert isinstance(installed, DeepSeekV4SparseAttentionConfig) + assert installed == sparse + use_deepseek_v41_layout_mode(model_config) # idempotent + assert model_config.sparse_attention_config is installed + + # A CSA2 config the user tuned would be dropped silently: refused instead. + model_config = ModelConfig( + pretrained_config=config, + sparse_attention_config=CSA2SparseAttentionConfig(use_fp8_staging=False), + ) + with pytest.raises(ValueError, match="use_fp8_staging"): + use_deepseek_v41_layout_mode(model_config) + # Not a DeepSeek-V4.1 checkpoint. + with pytest.raises(ValueError, match="DeepSeek-V4.1"): + use_deepseek_v41_layout_mode( + SimpleNamespace( + pretrained_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]), + sparse_attention_config=None, + ) + ) + + +@pytest.mark.parametrize("pool_ratio", [None, [1.0]]) +def test_v41_cache_config_keeps_the_long_decode_constraint(pool_ratio): + """While the pool split is derived, the layout mode's pools must hold one request at + max_seq_len beside minimal decode requests filling the resident batch, as DeepSeek-V4's + do. Without it, a 1,048,576-token, 256-request manager got 9,400 SWA pages instead of + 10,920 and could not hold the 119th request of a step.""" + from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + BatchDesc, + GpuCacheTierConfig, + KVCacheDesc, + KVCacheManagerConfig, + ) + + with kv_storage_env(DEEPSEEK_V41_KV_STORAGE_NATIVE): + layout = get_deepseek_v41_layout(make_v41_config()) + manager = object.__new__(DeepseekV4CacheManager) + manager.v41_layout = layout + manager.pp_layers = list(range(layout.num_layers)) + manager._max_draft_len = 0 + manager.head_dim = HEAD_DIM + manager.index_head_dim = INDEX_HEAD_DIM + manager.tokens_per_block = TOKENS_PER_BLOCK + manager.compressed_block_sizes = [ + TOKENS_PER_BLOCK // ratio if ratio else TOKENS_PER_BLOCK for ratio in layout.compress_ratios + ] + manager.max_seq_len = 1 << 20 + manager.max_batch_size = 256 + manager.is_estimating_kv_cache = False + manager.max_cuda_graph_batch_size = None + manager.max_attention_window_vec = [None] * layout.num_layers + manager.max_draft_len = 0 + manager.num_extra_kv_tokens = 0 + + context_constraint = BatchDesc([KVCacheDesc(capacity=16384, history_length=0)]) + base_config = KVCacheManagerConfig( + tokens_per_block=TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=1 << 30)], + layers=[], + constraints=[context_constraint] if pool_ratio is None else [], + initial_pool_ratio=pool_ratio, + ) + config = manager._build_cache_config(base_config) + + assert len(config.layers) == manager._num_manager_layers + if pool_ratio is None: + assert config.constraints == [ + context_constraint, + BatchDesc( + [KVCacheDesc(capacity=1 << 20, history_length=(1 << 20) - 1)] + + [KVCacheDesc(capacity=1, history_length=0)] * 255 + ), + ] + assert base_config.constraints == [context_constraint] + else: + assert config.initial_pool_ratio == pool_ratio + assert config.constraints == [] + + +def test_pipeline_stage_split(): + """A stage runs on its own pools only; for this checkpoint that leaves boundaries at 1 and 20.""" + layout = get_deepseek_v41_layout(make_v41_config()) + assert layout.pipeline_boundaries() == (1, 20) + # PP2's even split and PP3 as 1 | 19 | 20 are valid. + for stage in (range(0, 20), range(20, 40), range(0, 1), range(1, 20)): + assert layout.stage_dependency_error(stage) is None, stage + cases = { + range(10, 20): "layer 10 reads its KV source layer 8", + range(21, 40): "layer 21 reads its KV source layer 20", + range(30, 40): "layer 30 reads its KV source layer 20", + range(14, 40): "Engram layer 14 reads the token history of layer 1", + } + for stage, message in cases.items(): + err = layout.stage_dependency_error(stage) + assert err is not None and err.startswith(message), (stage, err) + # A re-index layer needs the candidate source's mask (a synthetic role table + # where layer 24 has a KV source of its own but still applies layer 20's mask). + other = get_deepseek_v41_layout(make_v41_config(kv_source_layer_ids=[2, 8, 14, 20, 24])) + err = other.stage_dependency_error(range(24, 40)) + assert err is not None and "candidate-block mask of layer 20" in err, err + + +# --------------------------------------------------------------------------- +# Cache manager and metadata (GPU). +# --------------------------------------------------------------------------- + + +def create_v41_manager( + max_batch_size: int = 8, + max_seq_len: int = 4096, + spec_config: Optional[object] = None, + enable_swa_scratch_reuse: bool = True, + config: Optional[SimpleNamespace] = None, + kv_storage: Optional[str] = DEEPSEEK_V41_KV_STORAGE_BF16, + mapping: Optional[Mapping] = None, +) -> Tuple[DeepseekV4CacheManager, DeepSeekV4SparseAttentionConfig]: + """A V4.1 manager; ``kv_storage`` pins TRTLLM_DSV41_KV_STORAGE while it is built (None: unset).""" + with kv_storage_env(kv_storage): + return _create_v41_manager( + max_batch_size, max_seq_len, spec_config, enable_swa_scratch_reuse, config, mapping + ) + + +def _create_v41_manager( + max_batch_size: int, + max_seq_len: int, + spec_config: Optional[object], + enable_swa_scratch_reuse: bool, + config: Optional[SimpleNamespace], + mapping: Optional[Mapping] = None, +) -> Tuple[DeepseekV4CacheManager, DeepSeekV4SparseAttentionConfig]: + config = make_v41_config() if config is None else config + # What ModelConfig synthesizes for V4.1: a placeholder all-1 ratio list. + sparse_config = DeepSeekV4SparseAttentionConfig( + index_head_dim=INDEX_HEAD_DIM, + window_size=WINDOW, + compress_ratios=[1] * len(config.compress_ratios), + ) + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + max_tokens=max_seq_len * max_batch_size, + event_buffer_max_size=0, + enable_swa_scratch_reuse=enable_swa_scratch_reuse, + ) + manager = DeepseekV4CacheManager( + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=config.num_hidden_layers, + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + max_input_len=max_seq_len, + mapping=mapping + if mapping is not None + else Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=8192, + sparse_attn_config=sparse_config, + pretrained_config=config, + spec_config=spec_config, + ) + return manager, sparse_config + + +def make_request(request_id: int, prompt_len: int) -> LlmRequest: + return LlmRequest( + request_id=request_id, + max_new_tokens=64, + input_tokens=list(range(prompt_len)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + + +def allocate(manager: DeepseekV4CacheManager, lengths: Sequence[int]) -> List[LlmRequest]: + requests = [] + for request_id, length in enumerate(lengths): + request = make_request(request_id, length) + assert manager.prepare_context(request) + assert manager.resize_context(request, request.context_chunk_size) + requests.append(request) + return requests + + +def reference_page_indices( + manager: DeepseekV4CacheManager, + request_ids: List[int], + num_contexts: int, + layer_idx: int, + attn_type: DeepseekV4AttentionType, + page_index_mode: PageIndexMode, +) -> torch.Tensor: + """Page table rows through the V2 manager's own Python page-index converter.""" + layer_id = manager._layer_attn_to_layer_id[layer_idx, attn_type] + pool_id = manager.layer_to_pool_mapping_dict[layer_id] + converter = manager.impl.get_page_index_converter(layer_id, attn_type.role) + copy_idx = manager.index_mapper.get_copy_index(request_ids, num_contexts, 1) + expected = torch.full( + (len(request_ids), manager.max_blocks_per_seq), BAD_PAGE_INDEX, dtype=torch.int32 + ) + for row, request_id in enumerate(request_ids): + base = manager.host_kv_cache_block_offsets[pool_id, int(copy_idx[row]), 0].tolist() + scratch = None + if row < num_contexts and page_index_mode == PageIndexMode.PER_LAYER: + scratch = manager.kv_cache_map[request_id].get_scratch_desc(pool_id) + converted = converter(base, page_index_mode, scratch) + expected[row, : len(converted)] = torch.tensor(converted, dtype=torch.int32) + return expected + + +@skip_pre_blackwell +def test_manager_allocates_owner_pools_only(): + manager, _ = create_v41_manager() + try: + layout = manager.v41_layout + assert layout is not None + keys = set(manager._layer_attn_to_layer_id) + expected = { + (layer, attn_type) + for layer in range(40) + for attn_type in T + if layout.has_attention(layer, attn_type) + } + assert keys == expected + # Readers own no compressed pool; the V4 per-ratio predicates are not used. + assert (3, T.COMPRESS) not in keys and (24, T.INDEXER_COMPRESS) not in keys + assert not any(t in (T.INDEXER_COMPRESSOR_KV, T.INDEXER_COMPRESSOR_SCORE) for _, t in keys) + assert manager.get_disagg_ignored_roles() == frozenset() + assert manager.num_sliding_attention_columns == T.ENGRAM_TOKEN.value + 1 + + # Page geometry and element types of every role. + geometry = { + (0, T.SWA): (TOKENS_PER_BLOCK, HEAD_DIM, torch.bfloat16), + (39, T.SWA): (TOKENS_PER_BLOCK, HEAD_DIM, torch.bfloat16), + (2, T.COMPRESS): (64, HEAD_DIM, torch.bfloat16), + (8, T.COMPRESS): (64, HEAD_DIM, torch.bfloat16), + (20, T.COMPRESS): (128, HEAD_DIM, torch.bfloat16), + (2, T.INDEXER_COMPRESS): (64, 68, torch.uint8), + (20, T.INDEXER_COMPRESS): (128, 68, torch.uint8), + (14, T.COMPRESSOR_KV): (TOKENS_PER_BLOCK, HEAD_DIM, torch.float32), + (14, T.COMPRESSOR_SCORE): (TOKENS_PER_BLOCK, HEAD_DIM, torch.float32), + (1, T.ENGRAM_TOKEN): (TOKENS_PER_BLOCK, 1, torch.int32), + } + for (layer, attn_type), (entries, width, dtype) in geometry.items(): + buffer = manager.get_buffers(layer, attn_type) + assert buffer.shape[1:] == (entries, width), (layer, attn_type, buffer.shape) + assert buffer.dtype == dtype, (layer, attn_type, buffer.dtype) + assert manager.get_indexer_k_cache_buffers(20).shape[1:] == (128, 1, 68) + assert manager.compressed_block_sizes[2] == 64 + assert manager.compressed_block_sizes[33] == 128 + + # Distinct owners of one ratio have distinct pool bases. + bases = {manager.get_buffers(owner, T.COMPRESS).data_ptr() for owner in (2, 8, 14, 20)} + assert len(bases) == 4 + assert manager.compress_pool_ptrs[2] == manager.get_buffers(2, T.COMPRESS).data_ptr() + assert manager.compress_pool_ptrs[1] == manager.get_buffers(20, T.COMPRESS).data_ptr() + + # Windowed vs persistent lifecycles. + windows = { + layer_config.layer_id: layer_config.sliding_window_size + for layer_config in manager.kv_cache_manager_py_config.layers + } + for (layer, attn_type), layer_id in manager._layer_attn_to_layer_id.items(): + assert windows[layer_id] == layout.window(attn_type, 0), (layer, attn_type) + + # Core global KV: main KV bf16 + native MXFP4 index K on the four owners. + assert manager.get_cache_bytes_per_token() == 2730 + pool_bytes_per_token = 0 + for owner in (2, 8, 14, 20): + for attn_type in (T.COMPRESS, T.INDEXER_COMPRESS): + buffer = manager.get_buffers(owner, attn_type) + page_bytes = buffer.shape[1] * buffer.shape[2] * buffer.element_size() + pool_bytes_per_token += page_bytes // TOKENS_PER_BLOCK + assert pool_bytes_per_token == 2730 + model_config = SimpleNamespace( + pretrained_config=make_v41_config(), + sparse_attention_config=None, + spec_config=None, + ) + with kv_storage_env(DEEPSEEK_V41_KV_STORAGE_BF16): + per_token, per_batch = DeepseekV4CacheManager.get_cache_size_per_token( + model_config, Mapping(), tokens_per_block=TOKENS_PER_BLOCK, max_batch_size=2 + ) + assert per_token == 2730 + # Windows per request, each up to two pages (a window straddles a page + # boundary): SWA on 40 layers, the compressor pair on 3, Engram on 1 + # (three pages under encoder replay, whose Engram window spans the replay). + engram_pages = 3 if get_deepseek_v41_layout(make_v41_config()).encoder_replay else 2 + assert per_batch == 2 * TOKENS_PER_BLOCK * ( + 2 * (40 * 1024 + 3 * 2 * 2048) + engram_pages * 4 + ) + finally: + manager.shutdown() + + +@skip_pre_blackwell +def test_manager_rejects_unsupported_v41_settings(): + with pytest.raises(ValueError, match="is not one of"): + create_v41_manager(kv_storage="fp8") + with pytest.raises(ValueError, match="must equal"): + create_v41_manager(config=make_v41_config(sliding_window=64)) + + +@skip_pre_blackwell +def test_manager_native_pools(): + """TRTLLM_DSV41_KV_STORAGE=native: uint8 SWA / main-KV entries, 890 B/token.""" + manager, _ = create_v41_manager(kv_storage=DEEPSEEK_V41_KV_STORAGE_NATIVE) + try: + layout = manager.v41_layout + assert layout.kv_storage == DEEPSEEK_V41_KV_STORAGE_NATIVE + # The same owner-only roles as bf16 storage; only the entry format changes. + keys = set(manager._layer_attn_to_layer_id) + assert keys == { + (layer, attn_type) + for layer in range(40) + for attn_type in T + if layout.has_attention(layer, attn_type) + } + + geometry = { + **{(layer, T.SWA): (TOKENS_PER_BLOCK, 528, torch.uint8) for layer in range(40)}, + (2, T.COMPRESS): (64, 288, torch.uint8), + (8, T.COMPRESS): (64, 288, torch.uint8), + (14, T.COMPRESS): (64, 288, torch.uint8), + (20, T.COMPRESS): (128, 288, torch.uint8), + (2, T.INDEXER_COMPRESS): (64, 68, torch.uint8), + (20, T.INDEXER_COMPRESS): (128, 68, torch.uint8), + (14, T.COMPRESSOR_KV): (TOKENS_PER_BLOCK, HEAD_DIM, torch.float32), + (14, T.COMPRESSOR_SCORE): (TOKENS_PER_BLOCK, HEAD_DIM, torch.float32), + (1, T.ENGRAM_TOKEN): (TOKENS_PER_BLOCK, 1, torch.int32), + } + for (layer, attn_type), (entries, width, dtype) in geometry.items(): + buffer = manager.get_buffers(layer, attn_type) + assert buffer.shape[1:] == (entries, width), (layer, attn_type, buffer.shape) + assert buffer.dtype == dtype, (layer, attn_type, buffer.dtype) + assert buffer.is_contiguous(), (layer, attn_type) + assert manager.get_indexer_k_cache_buffers(20).shape[1:] == (128, 1, 68) + # Every SWA view starts at the shared SWA base, as in bf16 storage. + assert {manager.get_buffers(layer, T.SWA).data_ptr() for layer in range(40)} == { + manager.swa_pool_ptr + } + + # (codes, scales) views: codes first, then scales, within each entry. + for layer, attn_type, codes_width, scales_width in ( + (0, T.SWA, 512, 16), + (39, T.SWA, 512, 16), + (8, T.COMPRESS, 256, 32), + (20, T.COMPRESS, 256, 32), + ): + buffer = manager.get_buffers(layer, attn_type) + codes, scales = manager.get_v41_native_kv_views(layer, attn_type) + entry = codes_width + scales_width + assert codes.dtype == scales.dtype == torch.uint8 + assert codes.shape == buffer.shape[:2] + (codes_width,) + assert scales.shape == buffer.shape[:2] + (scales_width,) + assert codes.stride() == scales.stride() == (buffer.shape[1] * entry, entry, 1) + assert codes.data_ptr() == buffer.data_ptr() + assert scales.data_ptr() == buffer.data_ptr() + codes_width + for layer, attn_type in ((2, T.INDEXER_COMPRESS), (14, T.COMPRESSOR_KV)): + with pytest.raises(RuntimeError, match="not a DeepSeek-V4.1 native pool"): + manager.get_v41_native_kv_views(layer, attn_type) + + # Distinct owners of one ratio have distinct pool bases. + bases = {manager.get_buffers(owner, T.COMPRESS).data_ptr() for owner in (2, 8, 14, 20)} + assert len(bases) == 4 + assert manager.compress_pool_ptrs[2] == manager.get_buffers(2, T.COMPRESS).data_ptr() + assert manager.compress_pool_ptrs[1] == manager.get_buffers(20, T.COMPRESS).data_ptr() + + # Windows are unchanged by the storage mode. + windows = { + layer_config.layer_id: layer_config.sliding_window_size + for layer_config in manager.kv_cache_manager_py_config.layers + } + for (layer, attn_type), layer_id in manager._layer_attn_to_layer_id.items(): + assert windows[layer_id] == layout.window(attn_type, 0), (layer, attn_type) + + # 890 B/token three ways: the manager's figure, the pools' page geometry, + # and the engine's static sizing hook. + assert manager.get_cache_bytes_per_token() == 890 + pool_bytes_per_token = 0 + for owner in (2, 8, 14, 20): + for attn_type in (T.COMPRESS, T.INDEXER_COMPRESS): + buffer = manager.get_buffers(owner, attn_type) + page_bytes = buffer.shape[1] * buffer.shape[2] * buffer.element_size() + pool_bytes_per_token += page_bytes // TOKENS_PER_BLOCK + assert pool_bytes_per_token == 890 + model_config = SimpleNamespace( + pretrained_config=make_v41_config(), + sparse_attention_config=None, + spec_config=None, + ) + with kv_storage_env(DEEPSEEK_V41_KV_STORAGE_NATIVE): + per_token, per_batch = DeepseekV4CacheManager.get_cache_size_per_token( + model_config, Mapping(), tokens_per_block=TOKENS_PER_BLOCK, max_batch_size=2 + ) + assert per_token == 890 + # SWA windows at 528 B/entry; compressor pair and Engram unchanged. + engram_pages = 3 if get_deepseek_v41_layout(make_v41_config()).encoder_replay else 2 + assert per_batch == 2 * TOKENS_PER_BLOCK * ( + 2 * (40 * 528 + 3 * 2 * 2048) + engram_pages * 4 + ) + + # Pool sizing follows the native pools: for the same max_tokens the + # quota is smaller and the persistent pools still hold max_tokens. + bf16_manager, _ = create_v41_manager(kv_storage=DEEPSEEK_V41_KV_STORAGE_BF16) + try: + assert bf16_manager.get_cache_bytes_per_token() == 2730 + bf16_quota = int(bf16_manager.impl.get_quota(0)) + finally: + bf16_manager.shutdown() + assert int(manager.impl.get_quota(0)) < bf16_quota + max_tokens = 4096 * 8 # create_v41_manager: max_seq_len * max_batch_size + for owner in (2, 20): + buffer = manager.get_buffers(owner, T.COMPRESS) + assert buffer.shape[0] * buffer.shape[1] * layout.compress_ratios[owner] >= max_tokens + finally: + manager.shutdown() + + +@skip_pre_blackwell +@pytest.mark.parametrize( + "kv_storage", [DEEPSEEK_V41_KV_STORAGE_BF16, DEEPSEEK_V41_KV_STORAGE_NATIVE] +) +@pytest.mark.parametrize("scratch_reuse", [False, True]) +def test_manager_page_tables_per_ratio(scratch_reuse: bool, kv_storage: str): + # max_batch_size sizes the windowed pools for 8 resident requests; three + # concurrent fresh contexts do not fit V2's pool split at max_batch_size=3 + # with SWA scratch reuse (one context plus two decodes is its typical step). + manager, _ = create_v41_manager(enable_swa_scratch_reuse=scratch_reuse, kv_storage=kv_storage) + requests = [] + try: + prompt_len = 2 * TOKENS_PER_BLOCK + 1 + requests = allocate(manager, [prompt_len] * 3) + request_ids = [r.py_request_id for r in requests] + num_contexts = 3 + + # Per-ratio COMPRESS and INDEXER_COMPRESS tables, shared by that ratio's owners. + for ratio, owners in ((2, (2, 8, 14)), (1, (20,))): + for attn_type in (T.COMPRESS, T.INDEXER_COMPRESS): + actual = torch.empty( + len(request_ids), manager.max_blocks_per_seq, dtype=torch.int32 + ) + if attn_type == T.COMPRESS: + manager.copy_batch_compress_block_tables( + actual, + request_ids, + compress_ratio=ratio, + beam_width=1, + num_contexts=num_contexts, + num_seqs=len(request_ids), + ) + else: + manager.copy_batch_indexer_compress_block_tables( + actual, + request_ids, + beam_width=1, + num_contexts=num_contexts, + num_seqs=len(request_ids), + compress_ratio=ratio, + ) + for owner in owners: + expected = reference_page_indices( + manager, request_ids, num_contexts, owner, attn_type, PageIndexMode.SHARED + ) + torch.testing.assert_close(actual, expected) + assert (actual[:, :3] != BAD_PAGE_INDEX).all() + + # Sliding tables: SWA on every layer, compressor pair, Engram history. + shape = ( + manager.num_local_layers, + manager.num_sliding_attention_columns, + len(request_ids), + manager.max_blocks_per_seq, + ) + actual = torch.empty(shape, dtype=torch.int32, device="cuda") + with torch.cuda.stream(manager._stream): + manager.compute_sliding_block_tables(request_ids, num_contexts=num_contexts) + manager.copy_batch_sliding_block_tables( + actual, request_ids, num_contexts=num_contexts, num_seqs=len(request_ids) + ) + manager._stream.synchronize() + actual = actual.cpu() + for (layer, attn_type), _ in manager._layer_attn_to_layer_id.items(): + if attn_type in (T.COMPRESS, T.INDEXER_COMPRESS): + continue + expected = reference_page_indices( + manager, request_ids, num_contexts, layer, attn_type, PageIndexMode.PER_LAYER + ) + torch.testing.assert_close( + actual[manager.layer_offsets[layer], attn_type.value], expected + ) + # Columns of roles V4.1 never allocates stay empty. + for attn_type in (T.INDEXER_COMPRESSOR_KV, T.INDEXER_COMPRESSOR_SCORE): + assert (actual[:, attn_type.value] == BAD_PAGE_INDEX).all() + + # Owners of one ratio share page ids but not storage. + table = torch.empty(len(request_ids), manager.max_blocks_per_seq, dtype=torch.int32) + manager.copy_batch_compress_block_tables( + table, request_ids, compress_ratio=2, beam_width=1, num_contexts=3, num_seqs=3 + ) + pages = table[:, :3].flatten().tolist() + for value, owner in enumerate((2, 8, 14), start=1): + manager.get_buffers(owner, T.COMPRESS)[pages] = float(value) + for value, owner in enumerate((2, 8, 14), start=1): + assert (manager.get_buffers(owner, T.COMPRESS)[pages] == float(value)).all() + finally: + for request in requests: + manager.free_resources(request) + manager.shutdown() + + +@skip_pre_blackwell +@pytest.mark.parametrize("pp_rank", [0, 1]) +def test_manager_pp2_stage_pools(pp_rank: int): + """PP2's even split: each stage holds its own layers' pools, and the page table of + a ratio with no KV source on the stage is empty rather than an error.""" + mapping = Mapping(world_size=2, rank=pp_rank, tp_size=1, pp_size=2, gpus_per_node=2) + manager, sparse_config = create_v41_manager(mapping=mapping) + requests = [] + try: + layout = manager.v41_layout + layers = list(range(20)) if pp_rank == 0 else list(range(20, 40)) + assert list(manager.pp_layers) == layers + expected = { + (layer, attn_type) + for layer in layers + for attn_type in T + if layout.has_attention(layer, attn_type) + } + assert set(manager._layer_attn_to_layer_id) == expected + owners = {2: (2, 8, 14)} if pp_rank == 0 else {1: (20,)} + absent = 1 if pp_rank == 0 else 2 + assert set(manager.compress_pool_ptrs) == set(owners) + # Persistent bytes per token of the stage: 3 ratio-2 owners, or 1 ratio-1 owner. + assert manager.get_cache_bytes_per_token() == (3 * 1092 // 2 if pp_rank == 0 else 1092) + + prompt_len = 2 * TOKENS_PER_BLOCK + 1 + requests = allocate(manager, [prompt_len] * 3) + request_ids = [r.py_request_id for r in requests] + for ratio in (1, 2): + for attn_type in (T.COMPRESS, T.INDEXER_COMPRESS): + actual = torch.empty( + len(request_ids), manager.max_blocks_per_seq, dtype=torch.int32 + ) + if attn_type == T.COMPRESS: + manager.copy_batch_compress_block_tables( + actual, + request_ids, + compress_ratio=ratio, + beam_width=1, + num_contexts=3, + num_seqs=3, + ) + else: + manager.copy_batch_indexer_compress_block_tables( + actual, + request_ids, + beam_width=1, + num_contexts=3, + num_seqs=3, + compress_ratio=ratio, + ) + if ratio == absent: + assert (actual == BAD_PAGE_INDEX).all(), (ratio, attn_type) + continue + for owner in owners[ratio]: + ref = reference_page_indices( + manager, request_ids, 3, owner, attn_type, PageIndexMode.SHARED + ) + torch.testing.assert_close(actual, ref) + # The metadata builds on the stage (the absent ratio's slots stay -1). + metadata = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor([prompt_len] * 3, dtype=torch.int), + request_ids=request_ids, + max_num_requests=manager.max_batch_size, + num_contexts=3, + prompt_lens=[prompt_len] * 3, + max_num_tokens=8192, + kv_cache_manager=manager, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0, 0, 0]), + mapping=mapping, + sparse_attention_config=sparse_config, + ) + metadata.prepare() + torch.cuda.synchronize() + total = metadata.num_total_compressed_tokens[absent] + assert total > 0 + assert (metadata.v41_compress_slot_mapping[absent][:total] == -1).all() + present = 2 if pp_rank == 0 else 1 + total = metadata.num_total_compressed_tokens[present] + assert (metadata.v41_compress_slot_mapping[present][:total] >= 0).all() + assert (metadata.engram_token_block_table is None) == (pp_rank == 1) + finally: + for request in requests: + manager.free_resources(request) + manager.shutdown() + + +@skip_pre_blackwell +def test_manager_rejects_a_stage_split_across_a_kv_source(): + mapping = Mapping( + world_size=2, rank=1, tp_size=1, pp_size=2, gpus_per_node=2, pp_partition=[10, 30] + ) + with pytest.raises( + ValueError, match="layer 10 reads its KV source layer 8.*layers \\[1, 20\\]" + ): + create_v41_manager(mapping=mapping) + + +def collect_block_tables( + manager: DeepseekV4CacheManager, request_ids: List[int], num_contexts: int +) -> dict: + """Per-ratio COMPRESS / INDEXER_COMPRESS tables and all sliding tables of one batch.""" + num_seqs = len(request_ids) + tables = {} + for ratio in (1, 2): + table = torch.empty(num_seqs, manager.max_blocks_per_seq, dtype=torch.int32) + manager.copy_batch_compress_block_tables( + table, request_ids, ratio, beam_width=1, num_contexts=num_contexts, num_seqs=num_seqs + ) + tables[T.COMPRESS, ratio] = table.clone() + manager.copy_batch_indexer_compress_block_tables( + table, + request_ids, + beam_width=1, + num_contexts=num_contexts, + num_seqs=num_seqs, + compress_ratio=ratio, + ) + tables[T.INDEXER_COMPRESS, ratio] = table.clone() + sliding = torch.empty( + manager.num_local_layers, + manager.num_sliding_attention_columns, + num_seqs, + manager.max_blocks_per_seq, + dtype=torch.int32, + device="cuda", + ) + with torch.cuda.stream(manager._stream): + manager.compute_sliding_block_tables(request_ids, num_contexts=num_contexts) + manager.copy_batch_sliding_block_tables( + sliding, request_ids, num_contexts=num_contexts, num_seqs=num_seqs + ) + manager._stream.synchronize() + tables["sliding"] = sliding.cpu() + return tables + + +@skip_pre_blackwell +def test_block_tables_independent_of_kv_storage(): + """Native storage changes entry bytes only: the same batch gets the same page ids.""" + tables = {} + for kv_storage in (DEEPSEEK_V41_KV_STORAGE_BF16, DEEPSEEK_V41_KV_STORAGE_NATIVE): + manager, _ = create_v41_manager(kv_storage=kv_storage) + requests = [] + try: + requests = allocate(manager, [257, 700, 130]) + tables[kv_storage] = collect_block_tables( + manager, [r.py_request_id for r in requests], num_contexts=3 + ) + finally: + for request in requests: + manager.free_resources(request) + manager.shutdown() + bf16, native = tables[DEEPSEEK_V41_KV_STORAGE_BF16], tables[DEEPSEEK_V41_KV_STORAGE_NATIVE] + assert bf16.keys() == native.keys() + for key in bf16: + torch.testing.assert_close(native[key], bf16[key], rtol=0, atol=0, msg=str(key)) + # 130 tokens at ratio 2 are 65 entries: two 64-entry pages for every request. + assert (native[T.COMPRESS, 2][:, :2] != BAD_PAGE_INDEX).all() + + +@skip_pre_blackwell +@pytest.mark.parametrize( + "kv_storage", [DEEPSEEK_V41_KV_STORAGE_BF16, DEEPSEEK_V41_KV_STORAGE_NATIVE] +) +def test_disagg_page_table_enumerates_v41_roles(kv_storage: str): + manager, _ = create_v41_manager(max_batch_size=1, max_seq_len=512, kv_storage=kv_storage) + try: + page_table = build_page_table_from_manager(manager) + roles = set() + # role -> per-page byte sizes the transfer views carry for it. + view_page_bytes = {} + for layer_group in page_table.layer_groups: + for pool_view in layer_group.pool_views: + roles.update(str(role) for role in pool_view.pool_role) + sizes = {int(entry[2]) for entry in pool_view.buffer_entries} + assert len(sizes) == 1, (pool_view.pool_role, sizes) + for role in pool_view.pool_role: + view_page_bytes.setdefault(str(role), set()).update(sizes) + for attn_type in ( + T.SWA, + T.COMPRESS, + T.INDEXER_COMPRESS, + T.COMPRESSOR_KV, + T.COMPRESSOR_SCORE, + T.ENGRAM_TOKEN, + ): + assert str(attn_type.role) in roles, (attn_type, roles) + # The role-generic page-table walk sees the pools' own page sizes. + entry = ENTRY_BYTES[kv_storage] + assert view_page_bytes == { + str(T.SWA.role): {TOKENS_PER_BLOCK * entry["SWA"]}, + str(T.COMPRESS.role): {64 * entry["COMPRESS"], 128 * entry["COMPRESS"]}, + str(T.INDEXER_COMPRESS.role): {64 * 68, 128 * 68}, + str(T.COMPRESSOR_KV.role): {TOKENS_PER_BLOCK * 2048}, + str(T.COMPRESSOR_SCORE.role): {TOKENS_PER_BLOCK * 2048}, + str(T.ENGRAM_TOKEN.role): {TOKENS_PER_BLOCK * 4}, + } + finally: + manager.shutdown() + + +# Each scenario: (prompt_len used for allocation, cached tokens, new tokens) per +# request, context requests first. +PREFILL = [(1, 0, 1), (130, 0, 130), (700, 0, 700)] +CHUNKED = [(1000, 512, 300), (1500, 1024, 476), (257, 0, 257)] +DECODE = [(2, 1, 1), (128, 127, 1), (129, 128, 1), (2049, 2048, 1)] +VERIFY = [(40, 34, 6), (300, 294, 6), (1100, 1094, 6)] + + +def build_metadata( + manager, sparse_config, requests, num_contexts, cached, new_lens, max_draft_len: int = 0 +): + metadata = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor(new_lens, dtype=torch.int), + request_ids=[r.py_request_id for r in requests], + max_num_requests=manager.max_batch_size, + num_contexts=num_contexts, + prompt_lens=[c + n for c, n in zip(cached, new_lens)], + max_num_tokens=8192, + kv_cache_manager=manager, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + sparse_attention_config=sparse_config, + ) + if max_draft_len > 0: + # What the engine does before prepare() when drafting. + metadata.update_spec_dec_param( + batch_size=len(requests), + is_spec_decoding_enabled=False, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, + max_draft_len=max_draft_len, + max_total_draft_tokens=max_draft_len, + ) + metadata.prepare() + torch.cuda.synchronize() + return metadata + + +@skip_pre_blackwell +@pytest.mark.parametrize( + "scenario, contexts", + [(PREFILL, True), (CHUNKED, True), (DECODE, False), (VERIFY, False)], + ids=["prefill", "chunked_prefill", "decode", "dspark_verify"], +) +def test_metadata_indices_match_reference(scenario, contexts): + manager, sparse_config = create_v41_manager() + requests = [] + try: + requests = allocate(manager, [alloc for alloc, _, _ in scenario]) + cached = [c for _, c, _ in scenario] + new_lens = [n for _, _, n in scenario] + num_contexts = len(scenario) if contexts else 0 + metadata = build_metadata(manager, sparse_config, requests, num_contexts, cached, new_lens) + assert metadata.is_deepseek_v41 + assert metadata.compress_ratio_set == {1, 2} + assert set(metadata.sparse_mla_topk_lens) == {0, 1, 2} + + num_tokens = sum(new_lens) + swa = metadata.swa_local_indices_cuda[:num_tokens].cpu() + topk_lens = {r: metadata.sparse_mla_topk_lens[r][:num_tokens].cpu() for r in (0, 1, 2)} + visible = {r: metadata.compressed_visible_lens_cuda[r][:num_tokens].cpu() for r in (1, 2)} + req_idx = metadata.req_idx_per_token[:num_tokens].cpu() + + token = 0 + for b, (c, n) in enumerate(zip(cached, new_lens)): + end_pos = c + n + for j in range(n): + p = c + j + assert int(req_idx[token]) == b + row = swa[token].tolist() + want = ref_window_positions(p) + # Valid window rows packed at the front in position order, -1 after. + assert row[: len(want)] == want, (p, row[:8]) + assert all(i == -1 for i in row[len(want) :]), p + assert int(topk_lens[0][token]) == WINDOW + for r in (1, 2): + assert int(visible[r][token]) == (p + 1) // r + assert int(topk_lens[r][token]) == WINDOW + ref_valid_compressed(p, end_pos, r) + token += 1 + + # Per-ratio compressed lengths, new entries and group-first positions. + for r in (1, 2): + kv_lens = [c + n for c, n in zip(cached, new_lens)] + comp = metadata.compressed_kv_lens_cuda[r][: len(scenario)].cpu().tolist() + new_comp = metadata.new_comp_kv_lens_cuda[r][: len(scenario)].cpu().tolist() + assert comp == [k // r for k in kv_lens] + assert new_comp == [k // r - c // r for k, c in zip(kv_lens, cached)] + cu = metadata.cu_new_comp_kv_cuda[r][: len(scenario) + 1].cpu().tolist() + assert cu == [0] + list(torch.tensor(new_comp).cumsum(0).tolist()) + positions = metadata.compressed_position_ids_cuda[r].cpu().tolist() + for b, c in enumerate(cached): + for k in range(new_comp[b]): + # RoPE position of a compressed entry: its group's first token. + assert positions[cu[b] + k] == (c // r + k) * r + if contexts: + assert metadata.max_ctx_compressed_tokens[r] == max(new_comp) + + # Page tables per ratio, and the Engram history table. + request_ids = [r.py_request_id for r in requests] + for r, owner in ((2, 2), (1, 20)): + expected = reference_page_indices( + manager, request_ids, num_contexts, owner, T.COMPRESS, PageIndexMode.SHARED + ) + torch.testing.assert_close( + metadata.compress_block_tables[r][: len(scenario)].cpu(), expected + ) + expected = reference_page_indices( + manager, request_ids, num_contexts, owner, T.INDEXER_COMPRESS, PageIndexMode.SHARED + ) + torch.testing.assert_close( + metadata.indexer_compress_block_tables[r][: len(scenario)].cpu(), expected + ) + expected = reference_page_indices( + manager, request_ids, num_contexts, 1, T.ENGRAM_TOKEN, PageIndexMode.PER_LAYER + ) + torch.testing.assert_close( + metadata.engram_token_block_table[: len(scenario)].cpu(), expected + ) + for r in (1, 2): + assert metadata.max_compressed_kv_lens[r] == metadata.max_seq_len // r + assert metadata.max_compressed_kv_lens[r] >= int(visible[r].max()) + + # Pool pointers resolve reader layers to their owners (no KeyError). + for layer in range(40): + owner = manager.v41_layout.kv_source(layer) + if owner is None: + assert layer not in metadata.compressed_buffer_ptrs + else: + assert ( + metadata.compressed_buffer_ptrs[layer] + == manager.get_buffers(owner, T.COMPRESS).data_ptr() + ) + assert metadata.swa_buffer_ptrs[layer] == manager.swa_pool_ptr + assert metadata.sparse_mla_base_ptrs[0] == manager.swa_pool_ptr + assert metadata.sparse_mla_base_ptrs[2] == manager.compress_pool_ptrs[2] + assert metadata.sparse_mla_base_ptrs[1] == manager.compress_pool_ptrs[1] + finally: + for request in requests: + manager.free_resources(request) + manager.shutdown() + + +# --------------------------------------------------------------------------- +# Address book: write slots and prefill index-K inputs (GPU). +# --------------------------------------------------------------------------- + +INDEX_K_DATA_BYTES = INDEX_HEAD_DIM // 2 +INDEX_K_SCALE_BYTES = INDEX_HEAD_DIM // 32 +INDEX_K_ENTRY_BYTES = INDEX_K_DATA_BYTES + INDEX_K_SCALE_BYTES +# The index-K / main-KV owner whose pool a ratio's pages address. +RATIO_OWNER = {2: 2, 1: 20} + +# (scenario, number of context requests); contexts first. Prefill mixes odd and +# even lengths and spans up to six pages; chunked prefill has odd cached counts +# (a ratio-2 group straddles the chunk boundary); the decode batch ends near +# max_seq_len so its reserved ratio-2 slots derive positions at max_seq_len. +ADDRESS_BOOK_SCENARIOS = { + "prefill": ([(1, 0, 1), (130, 0, 130), (257, 0, 257), (700, 0, 700), (512, 0, 512)], 5), + "chunked_prefill": ( + [(1000, 513, 300), (1500, 1024, 476), (257, 0, 257), (900, 255, 301)], + 4, + ), + "decode": ([(2, 1, 1), (128, 127, 1), (129, 128, 1), (4095, 4094, 1)], 0), + "mixed": ([(300, 0, 300), (901, 513, 388), (129, 128, 1), (2049, 2048, 1), (64, 63, 1)], 2), + "dspark_verify": (VERIFY, 0), +} + + +def ref_page_slot(table: List[List[int]], b: int, entry: int, per_page: int) -> Tuple[int, int]: + return table[b][entry // per_page], entry % per_page + + +def ref_index_k_offsets(page: int, slot: int, per_page: int) -> Tuple[int, int]: + """Page-blocked MXFP4 index K: per page ``per_page`` x 64 code bytes, then x 4 scale bytes.""" + base = page * per_page * INDEX_K_ENTRY_BYTES + return ( + base + slot * INDEX_K_DATA_BYTES, + base + per_page * INDEX_K_DATA_BYTES + slot * INDEX_K_SCALE_BYTES, + ) + + +def ref_compressed_rows(cached, new_lens, num_contexts, ratio) -> List[Optional[Tuple[int, int]]]: + """(request, entry) of every packed compressed row; None for a reserved slot. + + Context entries, then generation entries, in request order; then the + generation requests' unused reserved slots, ceil(next_n / ratio) each. + """ + rows = [ + (b, k) + for b, (c, n) in enumerate(zip(cached, new_lens)) + for k in range(c // ratio, (c + n) // ratio) + ] + num_ctx_rows = sum(1 for b, _ in rows if b < num_contexts) + gen_lens = new_lens[num_contexts:] + reserved = len(gen_lens) * -(-gen_lens[0] // ratio) if gen_lens else 0 + total = num_ctx_rows + reserved + assert len(rows) <= total + return rows + [None] * (total - len(rows)) + + +def read_address_book(metadata, num_tokens: int) -> dict: + book = {"swa": metadata.v41_swa_slot_mapping[:, :num_tokens].cpu()} + for r in (1, 2): + total = metadata.num_total_compressed_tokens[r] + book[("positions", r)] = metadata.compressed_position_ids_cuda[r][:total].cpu() + book[("compress", r)] = metadata.v41_compress_slot_mapping[r][:total].cpu() + book[("data", r)] = metadata.v41_index_k_slot_data[r][:total].cpu() + book[("scale", r)] = metadata.v41_index_k_slot_scale[r][:total].cpu() + return book + + +def check_address_book(metadata, manager, cached, new_lens, num_contexts) -> None: + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.kernels import ( + dsv4_local_to_global_indices, + ) + + tpb = manager.tokens_per_block + num_tokens = sum(new_lens) + tokens = [(b, c + j) for b, (c, n) in enumerate(zip(cached, new_lens)) for j in range(n)] + book = read_address_book(metadata, num_tokens) + + # (A) SWA write rows, every local layer. + swa = book["swa"] + assert metadata.v41_swa_slot_mapping.dtype == torch.int64 + assert metadata.v41_swa_slot_mapping.shape == (manager.num_local_layers, 8192) + sliding = metadata.sliding_block_tables[:, T.SWA.value, : len(cached)].cpu().tolist() + swa_rows = manager.get_buffers(0, T.SWA).shape[0] * tpb + for local_layer in range(manager.num_local_layers): + expected = [] + for b, p in tokens: + page, slot = ref_page_slot(sliding[local_layer], b, p, tpb) + expected.append(page * tpb + slot if page >= 0 else -1) + got = swa[local_layer].tolist() + assert got == expected, local_layer + # Every token of the step has its own row in the pool. + assert min(got) >= 0 and max(got) < swa_rows + assert len(set(got)) == num_tokens + # The attention reader maps each token's own window position to the same row. + req_id = metadata.req_idx_per_token[:num_tokens] + own_column = torch.tensor([min(p, WINDOW - 1) for _, p in tokens], device="cuda") + for layer in range(40): + local_layer = manager.layer_offsets[layer] + read_rows = dsv4_local_to_global_indices( + req_id, + metadata.sliding_block_tables[local_layer, T.SWA.value], + metadata.swa_local_indices_cuda[:num_tokens], + tpb, + 0, + None, + None, + 0, + 0, + False, + ) + own = read_rows.gather(1, own_column.unsqueeze(1)).squeeze(1).cpu().to(torch.int64) + torch.testing.assert_close(own, swa[local_layer], rtol=0, atol=0) + + # (B) Compressed stream write rows and index-K offsets, per ratio. + for r in (1, 2): + per_page = tpb // r + rows = ref_compressed_rows(cached, new_lens, num_contexts, r) + assert metadata.num_total_compressed_tokens[r] == len(rows) + comp_table = metadata.compress_block_tables[r][: len(cached)].cpu().tolist() + index_table = metadata.indexer_compress_block_tables[r][: len(cached)].cpu().tolist() + mask = metadata.compressed_mask_cuda[r][: len(rows)].cpu().tolist() + positions = book[("positions", r)].tolist() + slots = book[("compress", r)].tolist() + data = book[("data", r)].tolist() + scale = book[("scale", r)].tolist() + for buffers in ( + metadata.v41_compress_slot_mapping, + metadata.v41_index_k_slot_data, + metadata.v41_index_k_slot_scale, + ): + assert buffers[r].dtype == torch.int64 and buffers[r].shape == (8192,) + for i, row in enumerate(rows): + if row is None: + assert not mask[i], (r, i) + assert (slots[i], data[i], scale[i], positions[i]) == (-1, -1, -1, 0), (r, i) + continue + b, k = row + assert mask[i] and positions[i] == k * r, (r, i) + page, slot = ref_page_slot(comp_table, b, k, per_page) + assert page >= 0 and slots[i] == page * per_page + slot, (r, i) + page, slot = ref_page_slot(index_table, b, k, per_page) + assert page >= 0 and (data[i], scale[i]) == ref_index_k_offsets(page, slot, per_page) + # Positions feed an unchecked RoPE table lookup. + assert all(0 <= p < metadata.max_seq_len for p in positions), r + valid = [i for i, row in enumerate(rows) if row is not None] + valid_slots = [slots[i] for i in valid] + assert len(set(valid_slots)) == len(valid) + owner = RATIO_OWNER[r] + assert ( + max(valid_slots, default=0) < manager.get_buffers(owner, T.COMPRESS).shape[0] * per_page + ) + index_bytes = manager.get_indexer_k_cache_buffers(owner).numel() + assert max((data[i] for i in valid), default=0) + INDEX_K_DATA_BYTES <= index_bytes + assert max((scale[i] for i in valid), default=0) + INDEX_K_SCALE_BYTES <= index_bytes + if valid: + # The compressed-segment reader maps (request, entry) to the same row. + read_rows = dsv4_local_to_global_indices( + torch.tensor([rows[i][0] for i in valid], dtype=torch.int32, device="cuda"), + metadata.sliding_block_tables[0, T.SWA.value], + torch.full((len(valid), 1), -1, dtype=torch.int32, device="cuda"), + tpb, + 0, + metadata.compress_block_tables[r], + torch.tensor([[rows[i][1]] for i in valid], dtype=torch.int32, device="cuda"), + per_page, + 0, + False, + ) + assert read_rows[:, 1].cpu().tolist() == valid_slots + + # (C) Prefill index-K gather offsets and per-token row ranges. + if num_contexts == 0: + assert metadata.v41_max_ctx_q_len == 0 + for r in (1, 2): + assert metadata.v41_ctx_index_k_len[r] == 0 + assert metadata.v41_ctx_index_k_gather_data[r] is None + assert metadata.v41_ctx_index_k_gather_scale[r] is None + assert metadata.v41_ctx_index_ks[r] is None + assert metadata.v41_ctx_index_ke[r] is None + return + assert metadata.v41_max_ctx_q_len == max(new_lens[:num_contexts]) + ctx_tokens = [(b, p) for b, p in tokens if b < num_contexts] + for r in (1, 2): + per_page = tpb // r + index_table = metadata.indexer_compress_block_tables[r][:num_contexts].cpu().tolist() + entries = [ + (b, k) for b in range(num_contexts) for k in range((cached[b] + new_lens[b]) // r) + ] + comp_lens = [(cached[b] + new_lens[b]) // r for b in range(num_contexts)] + first_row = [sum(comp_lens[:b]) for b in range(num_contexts)] + assert metadata.v41_ctx_index_k_len[r] == len(entries) + gather_data = metadata.v41_ctx_index_k_gather_data[r] + gather_scale = metadata.v41_ctx_index_k_gather_scale[r] + assert gather_data.dtype == gather_scale.dtype == torch.int64 + assert gather_data.shape == gather_scale.shape == (len(entries),) + expected = [] + for b, k in entries: + page, slot = ref_page_slot(index_table, b, k, per_page) + assert page >= 0 + expected.append(ref_index_k_offsets(page, slot, per_page)) + got = list(zip(gather_data.cpu().tolist(), gather_scale.cpu().tolist())) + assert got == expected, r + # The entries this step writes are gathered from where they are written. + rows = ref_compressed_rows(cached, new_lens, num_contexts, r) + write = dict(zip(rows, zip(book[("data", r)].tolist(), book[("scale", r)].tolist()))) + for j, (b, k) in enumerate(entries): + if (b, k) in write: + assert write[(b, k)] == got[j] + ks = metadata.v41_ctx_index_ks[r] + ke = metadata.v41_ctx_index_ke[r] + assert ks.dtype == ke.dtype == torch.int32 + assert ks.shape == ke.shape == (len(ctx_tokens),) + assert ks.cpu().tolist() == [first_row[b] for b, _ in ctx_tokens] + assert ke.cpu().tolist() == [first_row[b] + (p + 1) // r for b, p in ctx_tokens] + + +def poison_address_book(metadata) -> None: + metadata.v41_swa_slot_mapping.fill_(-7) + for r in (1, 2): + metadata.v41_compress_slot_mapping[r].fill_(-7) + metadata.v41_index_k_slot_data[r].fill_(-7) + metadata.v41_index_k_slot_scale[r].fill_(-7) + + +def assert_same_book(actual: dict, expected: dict) -> None: + assert actual.keys() == expected.keys() + for key in expected: + torch.testing.assert_close(actual[key], expected[key], rtol=0, atol=0, msg=str(key)) + + +def replay_address_book_in_cuda_graph(metadata, positions: List[int]) -> None: + """The address-book rebuild captures, and replay reads the live page tables.""" + num_tokens = len(positions) + token_positions = torch.tensor(positions, dtype=torch.int32, device="cuda") + before = read_address_book(metadata, num_tokens) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + metadata._prepare_v41_address_book(token_positions) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + metadata._prepare_v41_address_book(token_positions) + poison_address_book(metadata) + graph.replay() + torch.cuda.synchronize() + assert_same_book(read_address_book(metadata, num_tokens), before) + + # Move every allocated page one page up; replay follows the tables. + tables = [ + metadata.sliding_block_tables, + *metadata.compress_block_tables.values(), + *metadata.indexer_compress_block_tables.values(), + ] + saved = [table.clone() for table in tables] + try: + for table in tables: + table.add_((table >= 0).to(table.dtype)) + graph.replay() + torch.cuda.synchronize() + moved = read_address_book(metadata, num_tokens) + finally: + for table, original in zip(tables, saved): + table.copy_(original) + step = {"swa": TOKENS_PER_BLOCK} + for r in (1, 2): + per_page = TOKENS_PER_BLOCK // r + step[("positions", r)] = 0 + step[("compress", r)] = per_page + step[("data", r)] = per_page * INDEX_K_ENTRY_BYTES + step[("scale", r)] = per_page * INDEX_K_ENTRY_BYTES + expected = { + key: value if step[key] == 0 else torch.where(value >= 0, value + step[key], value) + for key, value in before.items() + } + assert_same_book(moved, expected) + + +@skip_pre_blackwell +@pytest.mark.parametrize("name", list(ADDRESS_BOOK_SCENARIOS)) +def test_metadata_address_book_matches_reference(name): + scenario, num_contexts = ADDRESS_BOOK_SCENARIOS[name] + manager, sparse_config = create_v41_manager() + requests = [] + try: + requests = allocate(manager, [alloc for alloc, _, _ in scenario]) + cached = [c for _, c, _ in scenario] + new_lens = [n for _, _, n in scenario] + max_draft_len = new_lens[-1] - 1 if num_contexts < len(scenario) else 0 + metadata = build_metadata( + manager, sparse_config, requests, num_contexts, cached, new_lens, max_draft_len + ) + check_address_book(metadata, manager, cached, new_lens, num_contexts) + num_tokens = sum(new_lens) + positions = [c + j for c, n in zip(cached, new_lens) for j in range(n)] + + if num_contexts > 0: + # The prefill index-K inputs are sized from host lengths: no sync. + gathered = {r: metadata.v41_ctx_index_k_gather_data[r].clone() for r in (1, 2)} + token_positions = torch.tensor(positions, dtype=torch.int32, device="cuda") + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + metadata._prepare_v41_ctx_index_k(token_positions) + finally: + torch.cuda.set_sync_debug_mode("default") + for r in (1, 2): + assert torch.equal(metadata.v41_ctx_index_k_gather_data[r], gathered[r]) + + # on_update_kv_lens() (inside a captured step) rebuilds the address + # book on device: poison it, and the same values come back. + before = read_address_book(metadata, num_tokens) + ctx_fields = { + field: dict(getattr(metadata, field)) + for field in ( + "v41_ctx_index_k_len", + "v41_ctx_index_k_gather_data", + "v41_ctx_index_k_gather_scale", + "v41_ctx_index_ks", + "v41_ctx_index_ke", + ) + } + poison_address_book(metadata) + metadata.on_update_kv_lens() + torch.cuda.synchronize() + assert_same_book(read_address_book(metadata, num_tokens), before) + for field, values in ctx_fields.items(): + assert getattr(metadata, field) == values, field + check_address_book(metadata, manager, cached, new_lens, num_contexts) + + if num_contexts == 0: + replay_address_book_in_cuda_graph(metadata, positions) + finally: + for request in requests: + manager.free_resources(request) + manager.shutdown() + + +# --------------------------------------------------------------------------- +# Native storage: native writers -> manager pools -> dequantizing readers (GPU). +# --------------------------------------------------------------------------- + + +def multiscale_rows(rows: int, block: int) -> torch.Tensor: + """bf16 [rows, 512], one magnitude per quantization block, row 0's first block all zero.""" + x = torch.randn(rows, HEAD_DIM // block, block) + x = x * torch.exp2(torch.randint(-12, 9, (rows, HEAD_DIM // block, 1)).float()) + x = x.reshape(rows, HEAD_DIM).to(torch.bfloat16) + if rows > 0: + x[0, :block] = 0 + return x.cuda() + + +def assert_bits_equal(actual: torch.Tensor, expected: torch.Tensor, msg: str) -> None: + assert actual.shape == expected.shape, (msg, actual.shape, expected.shape) + mismatches = int((actual.view(torch.int16) != expected.view(torch.int16)).sum()) + assert mismatches == 0, (msg, mismatches) + + +@skip_pre_blackwell +@pytest.mark.parametrize("name", ["prefill", "mixed", "dspark_verify"]) +def test_native_pools_round_trip_with_native_kernels(name): + """Native writers fill the manager's pools at the address book's rows; the + dequantizing readers return the bf16-container value set bit for bit.""" + from tensorrt_llm._torch.custom_ops.mxfp8_dequantize import mxfp8_dequantize + from tensorrt_llm._torch.custom_ops.mxfp8_fake_quantize import mxfp8_fake_quantize + from tensorrt_llm._torch.custom_ops.nvfp4_fake_quantize import nvfp4_fake_quantize + from tensorrt_llm._torch.custom_ops.nvfp4_kv_to_bf16 import ( + nvfp4_kv_to_bf16 as nvfp4_kv_dequantize, + ) + + scenario, num_contexts = ADDRESS_BOOK_SCENARIOS[name] + manager, sparse_config = create_v41_manager(kv_storage=DEEPSEEK_V41_KV_STORAGE_NATIVE) + requests = [] + try: + requests = allocate(manager, [alloc for alloc, _, _ in scenario]) + cached = [c for _, c, _ in scenario] + new_lens = [n for _, _, n in scenario] + max_draft_len = new_lens[-1] - 1 if num_contexts < len(scenario) else 0 + metadata = build_metadata( + manager, sparse_config, requests, num_contexts, cached, new_lens, max_draft_len + ) + num_tokens = sum(new_lens) + torch.manual_seed(0) + + # SWA: window KV of several layers into the one shared SWA pool, E4M3 x + # UE8M0/32. Layer by layer, as the forward runs: with SWA scratch reuse a + # context's out-of-window rows are scratch pages shared by all layers. + full_swa = max( + (manager.get_buffers(layer, T.SWA) for layer in range(40)), key=lambda b: b.shape[0] + ) + swa_before = full_swa.clone() + written = [] + for layer in (0, 1, 20, 39): + codes, scales = manager.get_v41_native_kv_views(layer, T.SWA) + x = multiscale_rows(num_tokens, 32) + ref = torch.empty_like(x) + mxfp8_fake_quantize(x, ref) # the bf16-container value set + slots = metadata.v41_swa_slot_mapping[manager.layer_offsets[layer]][:num_tokens] + assert int(slots.min()) >= 0 and len(set(slots.tolist())) == num_tokens + mxfp8_fake_quantize(x, codes, slots, scales) + written.append(slots) + # Compact gather of the step's rows. + compact = torch.full_like(ref, float("nan")) + mxfp8_dequantize(codes, scales, compact, src_slots=slots) + assert_bits_equal(compact, ref, f"SWA gather, layer {layer}") + # Shadow pool with the native pool's geometry: rows keep their numbering. + shadow = torch.full(codes.shape, float("nan"), dtype=torch.bfloat16, device="cuda") + mxfp8_dequantize(codes, scales, shadow, src_slots=slots, dst_slots=slots) + assert_bits_equal(shadow.view(-1, HEAD_DIM)[slots], ref, f"SWA shadow, layer {layer}") + del shadow + # Nothing outside the step's rows changed. + untouched = torch.ones( + full_swa.shape[0] * TOKENS_PER_BLOCK, dtype=torch.bool, device="cuda" + ) + untouched[torch.cat(written)] = False + assert torch.equal(full_swa.view(-1, 528)[untouched], swa_before.view(-1, 528)[untouched]) + del swa_before + + # Main KV: every owner of a ratio at the same rows, E2M1 x E4M3/16, no global scale. + global_scale = torch.ones(1, dtype=torch.float32, device="cuda") + for r, owners in ((2, (2, 8, 14)), (1, (20,))): + per_page = TOKENS_PER_BLOCK // r + n = metadata.num_total_compressed_tokens[r] + slots = metadata.v41_compress_slot_mapping[r][:n] + valid = slots >= 0 + valid_slots = slots[valid] + packed_rows = ref_compressed_rows(cached, new_lens, num_contexts, r) + assert len(packed_rows) == n + before = { + owner: ( + manager.get_buffers(owner, T.COMPRESS).clone(), + manager.get_indexer_k_cache_buffers(owner).clone(), + ) + for owner in owners + } + refs = {} + for owner in owners: + codes, scales = manager.get_v41_native_kv_views(owner, T.COMPRESS) + latent = multiscale_rows(n, 16) + ref = torch.empty_like(latent) + nvfp4_fake_quantize(latent, ref) # the bf16-container value set + nvfp4_fake_quantize(latent, codes, slots, scales) + refs[owner] = ref + torch.cuda.synchronize() + bases = {owner: manager.get_buffers(owner, T.COMPRESS).data_ptr() for owner in owners} + for owner in owners: + buffer = manager.get_buffers(owner, T.COMPRESS) + codes, scales = manager.get_v41_native_kv_views(owner, T.COMPRESS) + ref = refs[owner] + # Only the step's entries changed. The owners of one ratio share + # a coalesced page slot (their bases are whole pages apart), so + # an owner's view also spans its neighbours' entries: every + # changed row must be one some owner wrote. + old_buffer, old_index_k = before[owner] + changed = (buffer.view(-1, 288) != old_buffer.view(-1, 288)).any(dim=1) + changed = set(torch.nonzero(changed).flatten().tolist()) + written = set() + for other, base in bases.items(): + offset = base - bases[owner] + assert offset % (per_page * 288) == 0, (owner, other, offset) + written.update(s + offset // 288 for s in valid_slots.tolist()) + assert changed <= written, sorted(changed - written)[:8] + assert len(changed) > 0 or valid_slots.numel() == 0 + # The index K beside the main KV is untouched. + assert torch.equal(manager.get_indexer_k_cache_buffers(owner), old_index_k) + if valid_slots.numel() > 0: + # Row gather (decode / verify top-k form): one entry per page. + rows = buffer.view(-1, 288) + out = torch.full( + (valid_slots.numel(), 1, 1, HEAD_DIM), + float("nan"), + dtype=torch.bfloat16, + device="cuda", + ) + nvfp4_kv_dequantize( + rows[:, :256][:, None, None], + rows[:, 256:][:, None, None], + valid_slots.to(torch.int32).view(-1, 1).contiguous(), + torch.ones(valid_slots.numel(), dtype=torch.int32, device="cuda"), + global_scale, + out, + ) + assert_bits_equal( + out[:, 0, 0], ref[valid], f"COMPRESS gather r={r} owner {owner}" + ) + # Whole context through the per-ratio page table, for requests + # whose every entry this step wrote. + fresh = [b for b in range(num_contexts) if cached[b] == 0 and new_lens[b] >= r] + if fresh: + lens = [new_lens[b] // r for b in fresh] + table = metadata.compress_block_tables[r][torch.tensor(fresh, device="cuda")] + out = torch.full( + (len(fresh), max(lens), 1, HEAD_DIM), + float("nan"), + dtype=torch.bfloat16, + device="cuda", + ) + nvfp4_kv_dequantize( + codes.unsqueeze(2), + scales.unsqueeze(2), + table.to(torch.int32).contiguous(), + torch.tensor(lens, dtype=torch.int32, device="cuda"), + global_scale, + out, + ) + row_of = {row: i for i, row in enumerate(packed_rows) if row is not None} + for i, b in enumerate(fresh): + expected = ref[[row_of[(b, k)] for k in range(lens[i])]] + assert_bits_equal( + out[i, : lens[i], 0], + expected, + f"COMPRESS context r={r} owner {owner} b={b}", + ) + finally: + for request in requests: + manager.free_resources(request) + manager.shutdown() diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index ef9f3a58e574..c43c6bbcbd48 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -205,11 +205,12 @@ def _make_remote_tail_executor_and_request( scheduler=types.SimpleNamespace(max_context_length=8192), max_num_tokens=8192, resource_manager=types.SimpleNamespace( + validate_request=lambda _request: None, resource_managers={ ResourceManagerType.KV_CACHE_MANAGER: types.SimpleNamespace( enable_block_reuse=False ) - } + }, ), ) request = types.SimpleNamespace( diff --git a/tests/unittest/_torch/modeling_v2/_validating.py b/tests/unittest/_torch/modeling_v2/_validating.py new file mode 100644 index 000000000000..a163e4e67f7b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/_validating.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The only caller of ``OpWrapper.is_valid``. + +The guard lives beside the op, in the product tree, because that is where a +target author reads it. What *runs* it lives here, because a served engine has +nothing to gain from it: everything ``is_valid`` checks is a property of how a +target calls the op, fixed when the target was written, and a guard of ours +firing in production is our bug taking down a working deployment. + +So the wrapper's ``raw_call`` goes straight to the kernel, and this interposes +``is_valid`` for the duration of a test. Patching is on the subclass rather +than the instance: ``raw_call`` is looked up on the type, and every entry has +its own class, so one entry's guard never reaches another's. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Iterator + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper + + +@contextmanager +def validating(*wrappers: OpWrapper) -> Iterator[None]: + """Run each wrapper's ``is_valid`` before its op call, inside this block. + + Driving the certified cells through here is what proves the guard admits + what it is supposed to admit. A guard that rejected a shipped target's own + inputs would otherwise sit unnoticed until someone turned it on. + """ + saved = [] + for w in wrappers: + cls = type(w) + assert isinstance(w, OpWrapper), f"{w!r} is not a catalog entry" + original = cls.raw_call + saved.append((cls, original)) + + def guarded(self, *args, _original=original, **kwargs): + from tensorrt_llm._torch._experimental.modeling_v2.catalog import _op + + # A key's generation stays 0 forever if it is only ever bound + # before the first forward -- a construction-time constant can + # never be flagged. A key bound mid-forward is flagged the + # moment a later forward advances the counter without rebinding + # it: its stamp is then behind the engine's current generation. + stale = sorted( + key + for key, gen in self._generation.items() + if gen >= 1 and gen < _op._STEP_GENERATION + ) + assert not stale, ( + f"{type(self).__name__} bound {stale} during a forward and did " + f"not rebind {'it' if len(stale) == 1 else 'them'} this one: " + "this call is using the previous step's metadata -- wrong " + "output, no error, outside validation" + ) + self.is_valid(*args, **kwargs) + return _original(self, *args, **kwargs) + + cls.raw_call = guarded + try: + yield + finally: + for cls, original in reversed(saved): + cls.raw_call = original + + +@pytest.fixture +def certifying(): + """Arm one entry for the rest of the test: drive its guard, and refuse to + finish unless its own reference and gate were the ones used. + + ``validating`` interposes ``is_valid`` and nothing more, which is right for + a guard test but leaves the certification itself unchecked: an entry can + carry a mirror that nothing calls and a band that nothing applies, and the + type system is satisfied while both claims are dead. That is not + hypothetical -- ``mla_rope_append_paged_kv_assign_q`` ships a 38-line + ``reference`` its own test never reaches, comparing against a local copy + instead. + + A fixture rather than a context manager because the comparison does not + happen where the call does. In every entry that uses ``validating`` today + the op call sits inside the block and ``reference``/``compare`` are + statements after it -- and in the paged-cache entries the block is buried + inside a helper, tens of lines from the comparison. A block that spanned + both would have to span the whole test. + + What this proves is that the entry's own members were the ones exercised, + not that the exercise was meaningful: ``op.compare(x, x)`` also counts. + That is the same bound the other static gates in this tree state about + themselves, and it is worth what it costs -- a dead mirror is invisible + without it. + """ + armed: list[tuple[type, str, object, bool]] = [] + seen: dict[str, int] = {} + + def arm(*wrappers: OpWrapper) -> None: + for w in wrappers: + assert isinstance(w, OpWrapper), f"{w!r} is not a catalog entry" + cls = type(w) + for name in ("raw_call", "reference", "compare"): + original = getattr(cls, name) + # Whether the subclass owns it decides how to put it back: + # `compare` and `is_valid` have base defaults, and restoring an + # inherited method with setattr would copy it onto the subclass + # permanently. + owned = name in cls.__dict__ + armed.append((cls, name, original, owned)) + + def counted(self, *args, _o=original, _n=name, **kwargs): + seen[_n] = seen.get(_n, 0) + 1 + if _n == "raw_call": + self.is_valid(*args, **kwargs) + return _o(self, *args, **kwargs) + + setattr(cls, name, counted) + + yield arm + + for cls, name, original, owned in reversed(armed): + if owned: + setattr(cls, name, original) + else: + delattr(cls, name) + + if not armed: + return + missing = [n for n in ("raw_call", "reference", "compare") if not seen.get(n)] + assert not missing, ( + f"the entry was armed but {', '.join(missing)} never ran: a cell driven " + f"without the entry's own reference and gate certifies nothing it claims" + ) diff --git a/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_flashinfer_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_flashinfer_silu_and_mul.py index 54338b915cb5..e61dfdfe9e28 100644 --- a/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_flashinfer_silu_and_mul.py +++ b/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_flashinfer_silu_and_mul.py @@ -1,81 +1,365 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the flashinfer_silu_and_mul catalog entry.""" +"""GPU test for the flashinfer_silu_and_mul catalog entry. +The entry owns what is certified. This file turns a cell's spec into tensors +and drives the inputs the guard exists to refuse. +""" + +import pytest import torch import torch.nn.functional as F from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( - flashinfer_silu_and_mul, + flashinfer_silu_and_mul as op, ) +__extra_import_path__ = [".."] +from _validating import validating # noqa: E402 — needs the path declared above + assert torch.cuda.is_available(), "flashinfer_silu_and_mul requires a CUDA device" -def _ref_silu_and_mul(x: torch.Tensor) -> torch.Tensor: - """fp32-accumulated reference: silu(x[..., :d]) * x[..., d:].""" - gate, up = x.float().chunk(2, dim=-1) - return (F.silu(gate) * up).to(x.dtype) +def _build(spec, seed): + """A cell's input; `data="vit"` is the vision-tower data of `_vit_input`, else plain randn.""" + if spec.get("data") == "vit": + rows, width = spec["shape"] + assert width == 2 * _VIT_D and spec["dtype"] is torch.bfloat16, spec + return _vit_input(rows, seed) + g = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(spec["shape"], generator=g, device="cuda").to(spec["dtype"]) -def _check(x: torch.Tensor) -> None: - out = flashinfer_silu_and_mul(x) - ref = _ref_silu_and_mul(x) - assert out.shape == x.shape[:-1] + (x.shape[-1] // 2,) - assert out.dtype == x.dtype - torch.testing.assert_close(out, ref) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:40] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + x = _build(cell.spec, seed=abs(hash(cell.why)) % 2**31) + with validating(op): + out = op(x) + assert out.shape == (*x.shape[:-1], x.shape[-1] // 2) + assert out.data_ptr() != x.data_ptr(), "the op must return a new tensor" + op.compare(out, op.reference(x)) + if cell.spec.get("data") == "vit": + # The vision-tower gate on top: 1 ulp everywhere but the fast-math flushes. + _assert_flushes_only_in_planted_rows(_assert_vit_gates(out, x, cell.why[:40])) -def test_bf16_2d() -> None: - torch.manual_seed(0) - # decode-like (few tokens) and prefill-like (many tokens) shapes; - # last dim is 2 * intermediate_size (gate half then up half) - for num_tokens, two_d in [(1, 8192), (4, 28672), (2048, 8192)]: - x = torch.randn(num_tokens, two_d, dtype=torch.bfloat16, device="cuda") - _check(x) +@pytest.mark.parametrize( + "width,match", + [ + (4095, "must be even"), + # 24 bf16 elements is 48 bytes per row, which the op's own row check + # accepts; the half is 24 bytes, so the second half starts mid-vector. + (24, "16-byte vectors"), + # The only width that reaches the third assert: for the half to be a + # whole number of 16-byte vectors *and* smaller than one, it has to be + # zero. Worth a case anyway -- it is the difference between a guard + # that is unreachable and one that is merely narrow. + (0, "at least one 16-byte vector"), + ], +) +def test_guard_refuses_widths_the_kernel_faults_on(width, match) -> None: + """A misaligned half kills the CUDA context rather than raising. + That is why these are `is_valid` and not left to the op: the failure is + unrecoverable and its message names neither the op nor the shape. + """ + x = torch.randn(4, width, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match=match): + with validating(op): + op(x) -def test_bf16_3d() -> None: - torch.manual_seed(1) - x = torch.randn(4, 16, 2048, dtype=torch.bfloat16, device="cuda") - _check(x) +# ─── The DeepSeek-V4.1 vision tower's MLP call: every row count 1..17037 ────── +# +# The multimodal target's ViT block runs `flashinfer_silu_and_mul(cublas_mm(h2, w1_t))`: +# a contiguous bf16 [P, 5632] input whose first 2816 columns are the gate and the last +# 2816 the up projection (the checkpoint's w1 is [gate; up]), no clamp, one call per +# tower call. The tower groups a request's images into calls of at most 17037 patches +# (<= 9216 per image, <= 16 images), so P, the row count, is anything in 1..17037. It +# runs eagerly (the tower is never CUDA-graph captured). Rows only size the grid: one +# CTA per row, `min(d / 8, 1024)` threads, the row offset computed in int64. -def test_bf16_edge_sizes() -> None: - # 16 is the smallest legal last dim (d = 8 = one vector); - # 16400 gives d = 8200 > 8192 = blockDim * vec_size, exercising the - # scalar remainder loop after the vectorized loop - torch.manual_seed(2) - for two_d in [16, 16400]: - x = torch.randn(16, two_d, dtype=torch.bfloat16, device="cuda") - _check(x) +_VIT_D = 2816 +_VIT_N_MAX = 17037 +# Tower-call row counts measured on the prototype (one image: 1517, 8772, 9189, the +# per-image maximum 9216, the mid-size images 2340 / 2352 and a padded 2433; several +# images: 3 x 1560, 3774 + 1551, 8787 + 8250 = 17037, and a 6-image request split into +# 15657 + 3120), then values either side of the 352-thread CTA's and 128-row edges. +_VIT_ROWS = (1517, 2340, 2352, 2433, 3120, 4680, 5325, 8772, 8787, 9189, 9216, 15657, 17037) +_VIT_EDGE_ROWS = (1, 2, 3, 127, 128, 129, 351, 352, 353, 8191, 8192, 8193, 17036) -def test_a_misaligned_half_is_rejected_before_dispatch() -> None: - """A final dim of 24 passes the op's own check and faults the kernel. +def _vit_input(n: int, seed: int) -> torch.Tensor: + """bf16 [n, 2 * 2816]: N(0, 3) gates and ups, with edge rows planted at fixed places. - The op validates the *row* (`shape[-1] * itemsize % 16 == 0`, which 24 - satisfies at 48 bytes), but the kernel vectorizes over each half - separately, and at 24 the second half starts 24 bytes in -- mid-vector. - The launch then dies with `CUDA misaligned address`, which poisons the - context for everything after it rather than raising something a caller - could catch. The wrapper's precondition is what keeps that off the device, - so this asserts it raises *without* calling the op. + Row 0 all zeros; row 1 gates of magnitude 40..300 (both signs), far past any + SwiGLU clamp; row 2 ups of magnitude 1e3..3e4; row 3 gates of magnitude 1e-30..1e-20; + the last row gates of -80..-1000 (silu underflows to -0). They sit inside every + prefix the sweep takes (rows 0..3) or at its top (the last row). """ - for width in (24, 40, 56): # d * 2 bytes = 24, 40, 56 -- none a multiple of 16 - x = torch.randn(4, width, dtype=torch.bfloat16, device="cuda") - try: - flashinfer_silu_and_mul(x) - except AssertionError: - continue - raise AssertionError(f"shape[-1]={width} should have been rejected before dispatch") - - -def test_an_odd_width_is_rejected() -> None: - """There is no half to split at an odd width; the op would truncate.""" - x = torch.randn(4, 33, dtype=torch.bfloat16, device="cuda") + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(n, 2 * _VIT_D, generator=g, device="cuda") * 3.0 + sign = torch.randint(0, 2, (_VIT_D,), generator=g, device="cuda").float() * 2 - 1 + + def mag(lo: float, hi: float) -> torch.Tensor: + return lo + (hi - lo) * torch.rand(_VIT_D, generator=g, device="cuda") + + if n > 0: + x[0] = 0.0 + if n > 1: + x[1, :_VIT_D] = sign * mag(40.0, 300.0) + if n > 2: + x[2, _VIT_D:] = sign * mag(1e3, 3e4) + if n > 3: + x[3, :_VIT_D] = sign * torch.exp(mag(-69.0, -46.0)) # 1e-30 .. 1e-20 + if n > 4: + x[n - 1, :_VIT_D] = -mag(80.0, 1000.0) + return x.to(torch.bfloat16) + + +def _ordered_bits(t: torch.Tensor) -> torch.Tensor: + """bf16 bit patterns mapped to integers that are monotonic in value (-0 and +0 coincide).""" + b = t.view(torch.int16).int() + return torch.where(b < 0, -(b & 0x7FFF), b) + + +def _ulp_distance(out: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: + """Per-element distance in bf16 ulps (int32).""" + return (_ordered_bits(out) - _ordered_bits(ref)).abs() + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor): + """(max distance in bf16 ulps, fraction of elements bit-identical).""" + ulps = int(_ulp_distance(out, ref).max().item()) + exact = (out.view(torch.int16) == ref.view(torch.int16)).double().mean().item() + return ulps, exact + + +# The kernel is built with fast math. Two consequences, pinned over every finite bf16 +# gate in test_vit_width_every_finite_bf16_gate: (a) the silu division `x / (1 + e^-x)` +# is the fast reciprocal division, which returns 0 once the denominator passes 2^126, +# i.e. for gates in (-88.72, -87.33] (bf16: -88.5, -88, -87.5), where the exact result +# is still nonzero (|ref| <= 8.8e-37 * |up|); from -88.72 down, e^-x overflows and the +# fp32 reference is 0 as well. (b) fp32 subnormals are flushed: where silu(gate) or the +# product is below 2^-126 in magnitude, the kernel returns 0. +_FTZ = 2.0**-126 +_FDIV_ZERO = (-88.72, -87.33) + + +def _flush_set(x: torch.Tensor) -> torch.Tensor: + """Elements whose kernel result may be 0 where the fp32 reference is a tiny nonzero.""" + xf = x.float() + d = x.shape[-1] // 2 + g, u = xf[..., :d], xf[..., d:] + s = F.silu(g) + return (s.abs() < _FTZ) | ((s * u).abs() < _FTZ) | ((g > _FDIV_ZERO[0]) & (g <= _FDIV_ZERO[1])) + + +def _assert_vit_gates(out: torch.Tensor, x: torch.Tensor, where: str) -> torch.Tensor: + """The certification gates at the ViT call: default bf16 assert_close against the + fp32 reference (silu and product in fp32, one rounding); every element within 1 bf16 + ulp of it except the fast-math flushes above, which must be exactly zero; and >= 0.9999 + of the other elements bit-identical (__expf is the only other difference). Returns the + flushed elements, so a caller can say where they are allowed to be.""" + ref = op.reference(x) + assert out.shape == ref.shape and out.dtype == torch.bfloat16 + op.compare(out, ref) + dist = _ulp_distance(out, ref) + flushed = _flush_set(x) & (out == 0) & (dist > 1) + ok = (dist <= 1) | flushed + n_bad = int((~ok).sum().item()) + exact = (out.view(torch.int16) == ref.view(torch.int16))[~flushed].double().mean().item() + worst = int(dist[~flushed].max().item()) if dist.numel() else 0 + print( + f"[vit {where}] rows={x.shape[0]}: max {worst} ulp outside {int(flushed.sum().item())} " + f"fast-math flushes to 0, bit-exact {exact:.6f}" + ) + assert n_bad == 0 and exact >= 0.9999, (where, n_bad, worst, exact) + return flushed + + +def _assert_flushes_only_in_planted_rows(flushed: torch.Tensor) -> None: + """On _vit_input data the flush set is only reachable through the planted rows of + large negative gates: row 1 (-300..-40) and the last row (-1000..-80).""" + n = flushed.shape[0] + rows = set(flushed.any(1).nonzero().flatten().tolist()) + assert rows <= {1, n - 1}, sorted(rows)[:8] + + +def test_v41_vit_width_n_max_against_the_reference() -> None: + """The 17037-row call (the tower's largest) against the fp32 reference, edge rows included. + + No clamp is applied: row 1's gates reach 300 and its outputs are 300-scale, where a + SwiGLU limit of 10 (the text model's shared-expert form) would cap them. + """ + x = _vit_input(_VIT_N_MAX, seed=2816) + out = op(x) + _assert_flushes_only_in_planted_rows(_assert_vit_gates(out, x, "n_max")) + assert out.is_contiguous() and out.data_ptr() != x.data_ptr() + assert torch.equal(out[0], torch.zeros_like(out[0])) # silu(0) * 0 + big = out[1].float().abs().max().item() + assert big > 100.0, f"row 1 (gates 40..300) peaked at {big}: something clamped it" + assert bool(torch.isfinite(out.float()).all()) + + +def test_v41_vit_width_wrong_references_are_rejected() -> None: + """Controls for the gates above: swapped [up; gate] halves and a SwiGLU clamp at 10 both fail them.""" + x = _vit_input(4096, seed=2817) + out = op(x) + xf = x.float() + gate, up = xf[:, :_VIT_D], xf[:, _VIT_D:] + swapped = (F.silu(up) * gate).to(torch.bfloat16) + clamped = (F.silu(gate.clamp(max=10.0)) * up.clamp(-10.0, 10.0)).to(torch.bfloat16) + for name, wrong in (("halves swapped", swapped), ("clamp 10", clamped)): + ulps, exact = _ulp_stats(out, wrong) + print(f"[control {name}] max {ulps} ulp, bit-exact {exact:.6f}") + assert ulps > 1 and (exact < 0.9999 or ulps > 1000), (name, ulps, exact) + # The clamp shows only where a value is past the limit: rows 1 and 2 carry those. + ulps_rows12, _ = _ulp_stats(out[1:3], clamped[1:3]) + assert ulps_rows12 > 1000, ulps_rows12 + + +def test_v41_vit_width_every_row_count_is_a_prefix_of_n_max() -> None: + """Every P = 1..17037: the op on the first P rows returns exactly the + first P rows of the 17037-row result, which is held to the reference above. + + Each output is NaN-poisoned before it is freed, so a row the kernel skipped cannot + inherit correct bits from the previous (one row shorter) call's recycled block; the + fraction of calls whose output landed on the previous call's block is reported. Then + the same kernel (the op body's flashinfer call, with a caller-owned `out`) writes into + a NaN buffer one row longer than P, ascending: row P must still be NaN after the call. + """ + from flashinfer.activation import silu_and_mul as _flashinfer_silu_and_mul + + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + x = _vit_input(_VIT_N_MAX, seed=2818) + full = op(x) + _assert_vit_gates(full, x, "sweep n_max") + full_bits = full.view(torch.int16) + bad = torch.zeros((), dtype=torch.int64, device="cuda") + reused, prev_ptr = 0, None + for r in range(1, _VIT_N_MAX + 1): + out = op(x[:r]) + assert out.shape == (r, _VIT_D) + reused += int(out.data_ptr() == prev_ptr) + bad += (out.view(torch.int16) != full_bits[:r]).sum() + out.fill_(float("nan")) + prev_ptr = out.data_ptr() + del out + assert bad.item() == 0, f"{bad.item()} elements differ from the 17037-row result's rows" + print( + f"[vit sweep] P=1..{_VIT_N_MAX}: every prefix bitwise; {reused} of {_VIT_N_MAX - 1} " + f"outputs landed on the previous call's NaN-poisoned block" + ) + + buf = torch.full((_VIT_N_MAX + 1, _VIT_D), float("nan"), dtype=torch.bfloat16, device="cuda") + tail_written = torch.zeros((), dtype=torch.int64, device="cuda") + pdl = get_env_enable_pdl() + for r in range(1, _VIT_N_MAX + 1): + _flashinfer_silu_and_mul(x[:r], out=buf[:r], enable_pdl=pdl) + tail_written += (~torch.isnan(buf[r])).sum() + assert tail_written.item() == 0, f"{tail_written.item()} elements written past the last row" + assert torch.equal(buf[:_VIT_N_MAX].view(torch.int16), full_bits) + + +def test_v41_vit_width_row_sweep_harness_sees_a_one_row_shift() -> None: + """Control for the sweep: the same bitwise check rejects a read one row off at every r tried.""" + x = _vit_input(_VIT_N_MAX, seed=2818) + full = op(x) + for r in (1, 2, 4, 351, 352, 353, 8192, _VIT_N_MAX - 1): + shifted = op(x[1 : r + 1]) + n_bad = int((shifted.view(torch.int16) != full.view(torch.int16)[:r]).sum().item()) + print(f"[control row shift] r={r}: {n_bad} / {r * _VIT_D} elements differ") + assert n_bad > r * _VIT_D // 2 + + +def test_v41_vit_width_one_kernel_and_launch_shape_over_the_rows() -> None: + """One kernel symbol at every probed P, block fixed at 352 threads, grid (P, 1, 1). + + Read from the profiler's own trace of each launch; the flashinfer module is JIT-built + once per activation, never per shape. + """ + import json + import os + import tempfile + + from torch.profiler import ProfilerActivity, profile + + rows = sorted(set(_VIT_ROWS) | set(_VIT_EDGE_ROWS) | set(range(1, 10))) + x = _vit_input(_VIT_N_MAX, seed=2819) + op(x[:1]) # build / load the module outside the trace + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for r in rows: + op(x[:r]) + torch.cuda.synchronize() + fd, path = tempfile.mkstemp(suffix=".json") + os.close(fd) try: - flashinfer_silu_and_mul(x) - except AssertionError: - return - raise AssertionError("an odd shape[-1] should have been rejected") + prof.export_chrome_trace(path) + with open(path) as f: + trace = json.load(f) + finally: + os.remove(path) + kernels = [e for e in trace["traceEvents"] if e.get("cat") == "kernel"] + assert len(kernels) == len(rows), ( + len(kernels), + len(rows), + sorted({k["name"] for k in kernels}), + ) + names = {k["name"] for k in kernels} + assert len(names) == 1 and "act_and_mul_kernel" in next(iter(names)), names + grids = [tuple(k["args"]["grid"]) for k in sorted(kernels, key=lambda k: k["ts"])] + blocks = {tuple(k["args"]["block"]) for k in kernels} + assert blocks == {(352, 1, 1)}, blocks + assert grids == [(r, 1, 1) for r in rows], list(zip(rows, grids))[:8] + print( + f"[vit probe] {len(rows)} row counts 1..{max(rows)}: one kernel {next(iter(names))[:80]}, " + f"block (352,1,1), grid (P,1,1)" + ) + + +@pytest.mark.parametrize("rows", _VIT_ROWS + _VIT_EDGE_ROWS) +def test_v41_vit_width_fresh_inputs_at_the_tower_row_counts(rows: int) -> None: + """Each measured tower row count on its own fresh input, against the reference.""" + x = _vit_input(rows, seed=3000 + rows) + _assert_flushes_only_in_planted_rows(_assert_vit_gates(op(x), x, f"fresh P={rows}")) + + +def test_v41_vit_width_every_finite_bf16_gate() -> None: + """Every finite bf16 gate value against ups 1, 1.5, -3, 1e4 and 1e-4: within 1 ulp of + the fp32 reference everywhere except the pinned fast-math flush set (_flush_set), where + the result is 0 or within 1 ulp; the flush set's reference magnitudes stay below + 1e-36 * max(1, |up|), and it is exactly where the kernel disagrees with the reference.""" + bits = torch.arange(-32768, 32768, dtype=torch.int32).to(torch.int16) + g = bits.view(torch.bfloat16) + g = g[torch.isfinite(g.float())].cuda() + n = g.numel() + rows = (n + _VIT_D - 1) // _VIT_D + gate = torch.zeros(rows * _VIT_D, dtype=torch.bfloat16, device="cuda") + gate[:n] = g + for up in (1.0, 1.5, -3.0, 1e4, 1e-4): + x = torch.empty(rows, 2 * _VIT_D, dtype=torch.bfloat16, device="cuda") + x[:, :_VIT_D] = gate.view(rows, _VIT_D) + x[:, _VIT_D:] = up + out = op(x) + ref = op.reference(x) + dist = _ulp_distance(out, ref).flatten()[:n] + fl = _flush_set(x).flatten()[:n] + o, r = out.flatten()[:n], ref.flatten()[:n] + wrong = dist > 1 + zero_in_fl = int((fl & (o == 0) & (r != 0)).sum().item()) + print( + f"[vit gates] up={up:g}: {n} finite gates; {int(wrong.sum().item())} > 1 ulp, all in the flush set " + f"({int(fl.sum().item())} elements, {zero_in_fl} returned 0 over a nonzero reference); " + f"max {int(dist[~fl].max().item())} ulp outside it" + ) + assert not bool((wrong & ~fl).any()), g[(wrong & ~fl)][:8].tolist() + assert bool(((o[fl] == 0) | (dist[fl] <= 1)).all()) + assert r[fl].float().abs().max().item() <= 1e-36 * max(1.0, abs(up)) + # The division flush: gates -88.5 / -88 / -87.5 return 0 where the reference is a + # nonzero normal number (at |up| >= 1; at 1e-4 the reference itself underflows). + for v in (-88.5, -88.0, -87.5) if abs(up) >= 1 else (): + k = int((g.float() == v).nonzero()[0]) + assert o[k].item() == 0 and r[k].item() != 0, (up, v, o[k].item(), r[k].item()) diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py deleted file mode 100644 index b71a64912ef6..000000000000 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py +++ /dev/null @@ -1,254 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the fused_qk_norm_rope catalog entry.""" - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( - fused_qk_norm_rope, -) - -assert torch.cuda.is_available(), "fused_qk_norm_rope requires a CUDA device" - -# The kernel rotates the bf16-rounded (x1, x2) pair, so the per-element -# error is absolute in the pair magnitude (up to ~5 after RMS norm with -# weights near 1), not relative to each output element: 0.02 ~= 5 ulp of -# bf16 (2^-8) at magnitude 5. Observed max abs err across all cases is -# 0.016; the default bf16 atol of 1e-5 is unreachable for near-zero -# outputs produced by rotating large pairs. -ATOL = 0.02 -RTOL = 1.6e-2 # torch default for bf16 - - -def _inv_freq(rotary_dim: int, base: float, factor: float, low: float, high: float) -> torch.Tensor: - """YaRN-blended inverse frequencies; plain RoPE when factor == 1.""" - half = rotary_dim // 2 - j = torch.arange(half, dtype=torch.float32, device="cuda") - pos_freqs = base ** (2.0 * j / rotary_dim) - inv_extrapolation = 1.0 / pos_freqs - inv_interpolation = 1.0 / (factor * pos_freqs) - if high == low: - high = high + 0.001 # kernel guards the ramp singularity the same way - ramp = ((j - low) / (high - low)).clamp(0.0, 1.0) - extrapolation_factor = 1.0 - ramp - return ( - inv_interpolation * (1.0 - extrapolation_factor) + inv_extrapolation * extrapolation_factor - ) - - -def _ref( - qkv: torch.Tensor, - num_heads_q: int, - num_heads_k: int, - num_heads_v: int, - head_dim: int, - rotary_dim: int, - eps: float, - q_weight: torch.Tensor, - k_weight: torch.Tensor, - base: float, - is_neox: bool, - position_ids: torch.Tensor, - factor: float = 1.0, - low: float = 0.0, - high: float = 0.0, - attention_factor: float = 1.0, - is_qk_norm: bool = True, - use_gemma: bool = False, - use_mrope: bool = False, - mrope_section1: int = 0, - mrope_section2: int = 0, -) -> torch.Tensor: - """fp32 reference: per-head RMS norm on q/k heads, then RoPE; v untouched.""" - num_heads = num_heads_q + num_heads_k + num_heads_v - x = qkv.float().view(-1, num_heads, head_dim).clone() - half = rotary_dim // 2 - inv_freq = _inv_freq(rotary_dim, base, factor, low, high) - if use_mrope: - angle = position_ids.float()[:, :, None] * inv_freq # [3, T, half] - cos, sin = angle.cos(), angle.sin() - - def pick(c: torch.Tensor) -> torch.Tensor: - out = c[0].clone() - out[:, 1 : mrope_section1 * 3 : 3] = c[1][:, 1 : mrope_section1 * 3 : 3] - out[:, 2 : mrope_section2 * 3 : 3] = c[2][:, 2 : mrope_section2 * 3 : 3] - return out - - cos, sin = pick(cos), pick(sin) - else: - angle = position_ids.float()[:, None] * inv_freq # [T, half] - cos, sin = angle.cos(), angle.sin() - cos = cos[:, None, :] * attention_factor - sin = sin[:, None, :] * attention_factor - for start, count, weight in ( - (0, num_heads_q, q_weight.float()), - (num_heads_q, num_heads_k, k_weight.float()), - ): - h = x[:, start : start + count, :] - if is_qk_norm: - h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + eps) - h = h * ((1.0 + weight) if use_gemma else weight) - r = h[..., :rotary_dim] - if is_neox: - x1, x2 = r[..., :half], r[..., half:] - rotated = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) - else: - x1, x2 = r[..., ::2], r[..., 1::2] - rotated = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2) - x[:, start : start + count, :] = torch.cat([rotated, h[..., rotary_dim:]], -1) - return x.view(qkv.shape[0], -1) - - -def _check( - num_tokens: int, - num_heads_q: int, - num_heads_kv: int, - head_dim: int, - rotary_dim: int, - eps: float, - base: float, - is_neox: bool, - factor: float = 1.0, - low: float = 0.0, - high: float = 0.0, - attention_factor: float = 1.0, - is_qk_norm: bool = True, - use_gemma: bool = False, - use_mrope: bool = False, - mrope_section1: int = 0, - mrope_section2: int = 0, -) -> None: - width = (num_heads_q + 2 * num_heads_kv) * head_dim - qkv = torch.randn(num_tokens, width, dtype=torch.bfloat16, device="cuda") - q_w = torch.randn(head_dim, dtype=torch.bfloat16, device="cuda") * 0.1 + 1.0 - k_w = torch.randn(head_dim, dtype=torch.bfloat16, device="cuda") * 0.1 + 1.0 - shape = (3, num_tokens) if use_mrope else (num_tokens,) - pos = torch.randint(0, 32768, shape, dtype=torch.int32, device="cuda") - expected = _ref( - qkv, num_heads_q, num_heads_kv, num_heads_kv, head_dim, rotary_dim, - eps, q_w, k_w, base, is_neox, pos, factor, low, high, attention_factor, - is_qk_norm, use_gemma, use_mrope, mrope_section1, mrope_section2, - ).to(torch.bfloat16) # fmt: skip - fused_qk_norm_rope( - qkv, num_heads_q, num_heads_kv, num_heads_kv, head_dim, rotary_dim, - eps, q_w, k_w, base, is_neox, pos, factor, low, high, attention_factor, - is_qk_norm, use_gemma, use_mrope, mrope_section1, mrope_section2, - ) # fmt: skip - torch.testing.assert_close(qkv, expected, rtol=RTOL, atol=ATOL) - # v heads must pass through bit-exactly - v = qkv.view(num_tokens, -1, head_dim)[:, num_heads_q + num_heads_kv :, :] - ev = expected.view(num_tokens, -1, head_dim)[:, num_heads_q + num_heads_kv :, :] - assert torch.equal(v, ev) - - -def test_bf16_neox_head_dims() -> None: - torch.manual_seed(0) - # decode-like shapes over every supported head_dim - for head_dim in (64, 128, 256): - _check( - 2, 8, 2, head_dim=head_dim, rotary_dim=head_dim, eps=1e-6, - base=10000.0, is_neox=True, - ) # fmt: skip - - -def test_bf16_neox_prefill() -> None: - torch.manual_seed(1) - # prefill-like: many tokens, Qwen3-32B-like head layout - _check( - 8192, 32, 8, head_dim=128, rotary_dim=128, eps=1e-6, - base=1000000.0, is_neox=True, - ) # fmt: skip - - -def test_bf16_interleaved() -> None: - torch.manual_seed(2) - for num_tokens in (1, 65): - _check( - num_tokens, 8, 2, head_dim=128, rotary_dim=128, eps=1e-6, - base=10000.0, is_neox=False, - ) # fmt: skip - - -def test_bf16_gemma_norm() -> None: - torch.manual_seed(3) - _check( - 17, 4, 4, head_dim=64, rotary_dim=64, eps=1e-6, - base=10000.0, is_neox=True, use_gemma=True, - ) # fmt: skip - - -def test_bf16_rope_only() -> None: - torch.manual_seed(4) - # is_qk_norm=False skips the norm entirely; weights are ignored - _check( - 17, 8, 2, head_dim=128, rotary_dim=128, eps=1e-6, - base=10000.0, is_neox=True, is_qk_norm=False, - ) # fmt: skip - - -def test_bf16_partial_rotary() -> None: - torch.manual_seed(5) - # norm covers the full head_dim; rope covers only the first rotary_dim - for is_neox in (True, False): - _check( - 17, 8, 2, head_dim=128, rotary_dim=64, eps=1e-6, - base=10000.0, is_neox=is_neox, - ) # fmt: skip - - -def test_bf16_yarn() -> None: - torch.manual_seed(6) - _check( - 33, 8, 2, head_dim=128, rotary_dim=128, eps=1e-6, - base=1000000.0, is_neox=True, - factor=4.0, low=4.0, high=20.0, attention_factor=1.2, - ) # fmt: skip - - -def test_bf16_mrope_interleaved() -> None: - torch.manual_seed(7) - _check( - 19, 8, 2, head_dim=128, rotary_dim=128, eps=1e-6, - base=10000.0, is_neox=True, - use_mrope=True, mrope_section1=16, mrope_section2=24, - ) # fmt: skip - - -def test_rejects_out_of_contract() -> None: - torch.manual_seed(8) - pos = torch.arange(4, dtype=torch.int32, device="cuda") - w = torch.ones(128, dtype=torch.bfloat16, device="cuda") - - def call(qkv, head_dim=128, q_w=w, position_ids=pos): - fused_qk_norm_rope( - qkv, 8, 2, 2, head_dim, head_dim, 1e-6, q_w, w, - 10000.0, True, position_ids, - ) # fmt: skip - - # non-bf16 qkv - for dtype in (torch.float16, torch.float32): - try: - call(torch.randn(4, 12 * 128, dtype=dtype, device="cuda")) - raise AssertionError(f"{dtype} qkv unexpectedly accepted") - except RuntimeError: - pass - # unsupported head_dim - try: - w96 = torch.ones(96, dtype=torch.bfloat16, device="cuda") - call( - torch.randn(4, 12 * 96, dtype=torch.bfloat16, device="cuda"), - head_dim=96, - q_w=w96, - ) - raise AssertionError("head_dim=96 unexpectedly accepted") - except RuntimeError: - pass - # non-int32 position_ids - try: - call( - torch.randn(4, 12 * 128, dtype=torch.bfloat16, device="cuda"), - position_ids=pos.to(torch.int64), - ) - raise AssertionError("int64 position_ids unexpectedly accepted") - except RuntimeError: - pass diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py index cb0a1f090d4c..93e81e3a8b39 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py @@ -164,16 +164,33 @@ page. Wrong batch states (the flag left off, context_lengths carrying the full KV length or 0, a replaced in-window cached page) are gated far outside the tolerance band. +10. No-KV-cache (pool-free) mode, the DeepSeek-V4.1 vision tower's surface: + every pool argument None, 16/16/64 bf16, mask_type 0, variable-length + images (1 to 9216 tokens, up to 16 per call) with and without + cu_q_seqlens. Gated against a per-image fp32 reference, with a causal-mask + rival and a cross-image-leak rival each required far outside the band; + mask_type 1 shown causal per image; the batch state pinned (given + cu_q_seqlens overrides context_lengths; a context_lengths / sequence_length + mismatch is silently wrong; update_kv_cache=False is rejected; page size, + max lengths and the input type are inert; rows past T untouched); and + CUDA-graph capture/replay -- fixed to the captured split without + max_context_q_len_override, following the split rewritten into the device + length tensors with it, which is the target's two-slot (image, pad) form. """ import math from typing import Dict, List, NamedTuple, Optional, Tuple +import pytest import torch import torch.nn.functional as F from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( - thop_attention, + MLA_CONTEXT, + STANDARD, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + thop_attention as op, ) from tensorrt_llm._torch.attention.backends.interface import RopeParams from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType, KVCacheManager @@ -181,6 +198,9 @@ from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping +__extra_import_path__ = [".."] +from _validating import validating # noqa: E402 — needs the path declared above + assert torch.cuda.is_available(), "thop_attention requires a CUDA device" # The output is a softmax-weighted combination of unit-scale bf16 v rows; @@ -408,93 +428,94 @@ def call_op( dtype=qkv.dtype, device=qkv.device, ) - thop_attention( - q=qkv, - k=None, # packed QKV rides inside q - v=None, - output=output, - output_sf=None, - workspace_=self.workspace, - sequence_length=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), - host_past_key_value_lengths=torch.tensor(kv_lens, dtype=torch.int32), - host_total_kv_lens=torch.tensor([total_ctx_kv, total_gen_kv], dtype=torch.int32), - context_lengths=torch.tensor(ctx_lens, dtype=torch.int32, device="cuda"), - host_context_lengths=torch.tensor(ctx_lens, dtype=torch.int32), - host_request_types=torch.tensor(req_types, dtype=torch.int32), - max_context_q_len_override=None, - kv_cache_block_offsets=self.block_offsets, - host_kv_cache_pool_pointers=self.pool_pointers, - host_kv_cache_pool_mapping=self.pool_mapping, - cache_indirection=None, - kv_scale_orig_quant=self.kv_scale_orig_quant, - kv_scale_quant_orig=self.kv_scale_quant_orig, - out_scale=None, - rotary_inv_freq=None, - rotary_cos_sin=None, - latent_cache=None, - q_pe=None, - block_ids_per_seq=None, - attention_sinks=attention_sinks, - is_fused_qkv=True, - update_kv_cache=True, - predicted_tokens_per_seq=1, - local_layer_idx=0, - num_heads=self.num_heads, - num_kv_heads=self.num_kv_heads, - head_size=self.head_dim, - tokens_per_block=self.tokens_per_block, - max_num_requests=self.max_batch, - max_context_length=self.max_seq_len, - max_seq_len=( - self.max_seq_len if max_seq_len_override is None else max_seq_len_override - ), - attention_window_size=self.attention_window_size, - beam_width=1, - mask_type=mask_type, - quant_mode=self.quant_mode, - q_scaling=1.0, - position_embedding_type=0, # no in-kernel RoPE - rope_dim=0, - rope_base=10000.0, - rope_scale_type=0, - rope_scale=1.0, - rope_short_m_scale=1.0, - rope_long_m_scale=1.0, - rope_max_positions=1024, - rope_original_max_positions=1024, - use_paged_context_fmha=use_paged_context_fmha, - attention_input_type=0, # mixed - is_mla_enable=False, - chunked_prefill_buffer_batch_size=1, - q_lora_rank=None, - kv_lora_rank=None, - qk_nope_head_dim=None, - qk_rope_head_dim=None, - v_head_dim=None, - rope_append=None, - mrope_rotary_cos_sin=None, - mrope_position_deltas=None, - helix_position_offsets=None, - helix_is_inactive_rank=None, - attention_chunk_size=None, - softmax_stats_tensor=None, - is_spec_decoding_enabled=False, - use_spec_decoding=False, - is_spec_dec_tree=False, - spec_decoding_generation_lengths=None, - spec_decoding_position_offsets_for_cpp=None, - spec_decoding_packed_mask=None, - spec_decoding_bl_tree_mask_offset=None, - spec_decoding_bl_tree_mask=None, - spec_bl_tree_first_sparse_mask_offset_kv=None, - sparse_kv_indices=None, - sparse_kv_offsets=None, - sparse_attn_indices=None, - sparse_attn_offsets=None, - sparse_attn_indices_block_size=0, - num_contexts=num_contexts, - num_ctx_tokens=num_ctx_tokens, - ) + with validating(op): + op( + q=qkv, + k=None, # packed QKV rides inside q + v=None, + output=output, + output_sf=None, + workspace_=self.workspace, + sequence_length=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), + host_past_key_value_lengths=torch.tensor(kv_lens, dtype=torch.int32), + host_total_kv_lens=torch.tensor([total_ctx_kv, total_gen_kv], dtype=torch.int32), + context_lengths=torch.tensor(ctx_lens, dtype=torch.int32, device="cuda"), + host_context_lengths=torch.tensor(ctx_lens, dtype=torch.int32), + host_request_types=torch.tensor(req_types, dtype=torch.int32), + max_context_q_len_override=None, + kv_cache_block_offsets=self.block_offsets, + host_kv_cache_pool_pointers=self.pool_pointers, + host_kv_cache_pool_mapping=self.pool_mapping, + cache_indirection=None, + kv_scale_orig_quant=self.kv_scale_orig_quant, + kv_scale_quant_orig=self.kv_scale_quant_orig, + out_scale=None, + rotary_inv_freq=None, + rotary_cos_sin=None, + latent_cache=None, + q_pe=None, + block_ids_per_seq=None, + attention_sinks=attention_sinks, + is_fused_qkv=True, + update_kv_cache=True, + predicted_tokens_per_seq=1, + local_layer_idx=0, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + tokens_per_block=self.tokens_per_block, + max_num_requests=self.max_batch, + max_context_length=self.max_seq_len, + max_seq_len=( + self.max_seq_len if max_seq_len_override is None else max_seq_len_override + ), + attention_window_size=self.attention_window_size, + beam_width=1, + mask_type=mask_type, + quant_mode=self.quant_mode, + q_scaling=1.0, + position_embedding_type=0, # no in-kernel RoPE + rope_dim=0, + rope_base=10000.0, + rope_scale_type=0, + rope_scale=1.0, + rope_short_m_scale=1.0, + rope_long_m_scale=1.0, + rope_max_positions=1024, + rope_original_max_positions=1024, + use_paged_context_fmha=use_paged_context_fmha, + attention_input_type=0, # mixed + is_mla_enable=False, + chunked_prefill_buffer_batch_size=1, + q_lora_rank=None, + kv_lora_rank=None, + qk_nope_head_dim=None, + qk_rope_head_dim=None, + v_head_dim=None, + rope_append=None, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + attention_chunk_size=None, + softmax_stats_tensor=None, + is_spec_decoding_enabled=False, + use_spec_decoding=False, + is_spec_dec_tree=False, + spec_decoding_generation_lengths=None, + spec_decoding_position_offsets_for_cpp=None, + spec_decoding_packed_mask=None, + spec_decoding_bl_tree_mask_offset=None, + spec_decoding_bl_tree_mask=None, + spec_bl_tree_first_sparse_mask_offset_kv=None, + sparse_kv_indices=None, + sparse_kv_offsets=None, + sparse_attn_indices=None, + sparse_attn_offsets=None, + sparse_attn_indices_block_size=0, + num_contexts=num_contexts, + num_ctx_tokens=num_ctx_tokens, + ) torch.cuda.synchronize() # Record the K/V slices just appended to the cache, per request. @@ -688,14 +709,8 @@ def test_rejects_non_divisible_head_counts() -> None: raising, having computed only the first (6 // 4) * 4 = 4 head columns and left the other two all-zero (the decode path does raise, from xqaDispatcher: 'numQHeads should be multiple of numKVHeads').""" - torch.manual_seed(50) - env = _PagedAttnEnv(num_heads=6, num_kv_heads=4, head_dim=128) - env.add_request(0, 16) - try: - env.call_op(env.random_qkv(16), [16], 1, [0], MASK_CAUSAL) - except AssertionError: - return - raise AssertionError("6q/4kv geometry was not rejected by the wrapper") + with pytest.raises(AssertionError, match="must be a multiple of"): + op.is_valid(num_heads=6, num_kv_heads=4) # ─── Standard configuration, multi-layer shared pool ─────────────────── @@ -879,96 +894,97 @@ def call_op( dtype=qkv.dtype, device=qkv.device, ) - thop_attention( - q=qkv, - k=None, # packed QKV rides inside q - v=None, - output=output, - output_sf=None, - workspace_=self.workspace, - sequence_length=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), - host_past_key_value_lengths=torch.tensor(kv_lens, dtype=torch.int32), - host_total_kv_lens=torch.tensor( - [sum(kv_lens[:num_contexts]), sum(kv_lens[num_contexts:])], - dtype=torch.int32, - ), - context_lengths=torch.tensor(ctx_lens, dtype=torch.int32, device="cuda"), - host_context_lengths=torch.tensor(ctx_lens, dtype=torch.int32), - host_request_types=torch.tensor(req_types, dtype=torch.int32), - max_context_q_len_override=None, - kv_cache_block_offsets=block_offsets, - host_kv_cache_pool_pointers=pool_pointers, - host_kv_cache_pool_mapping=pool_mapping, - cache_indirection=None, - kv_scale_orig_quant=self.kv_scale_orig_quant, - kv_scale_quant_orig=self.kv_scale_quant_orig, - out_scale=None, - rotary_inv_freq=None, - rotary_cos_sin=None, - latent_cache=None, - q_pe=None, - block_ids_per_seq=None, - attention_sinks=None, - is_fused_qkv=True, - update_kv_cache=True, - predicted_tokens_per_seq=1, - local_layer_idx=local_layer_idx, - num_heads=self.num_heads, - num_kv_heads=self.num_kv_heads, - head_size=self.head_dim, - tokens_per_block=self.tokens_per_block, - max_num_requests=self.max_batch, - max_context_length=self.max_seq_len, - max_seq_len=self.max_seq_len, - attention_window_size=( - self.max_seq_len if attention_window_size is None else attention_window_size - ), - beam_width=1, - mask_type=MASK_CAUSAL, - quant_mode=self.quant_mode, - q_scaling=1.0, - position_embedding_type=0, # no in-kernel RoPE - rope_dim=0, - rope_base=10000.0, - rope_scale_type=0, - rope_scale=1.0, - rope_short_m_scale=1.0, - rope_long_m_scale=1.0, - rope_max_positions=1024, - rope_original_max_positions=1024, - use_paged_context_fmha=use_paged_context_fmha, - attention_input_type=0, # mixed - is_mla_enable=False, - chunked_prefill_buffer_batch_size=1, - q_lora_rank=None, - kv_lora_rank=None, - qk_nope_head_dim=None, - qk_rope_head_dim=None, - v_head_dim=None, - rope_append=None, - mrope_rotary_cos_sin=None, - mrope_position_deltas=None, - helix_position_offsets=None, - helix_is_inactive_rank=None, - attention_chunk_size=None, - softmax_stats_tensor=None, - is_spec_decoding_enabled=False, - use_spec_decoding=False, - is_spec_dec_tree=False, - spec_decoding_generation_lengths=None, - spec_decoding_position_offsets_for_cpp=None, - spec_decoding_packed_mask=None, - spec_decoding_bl_tree_mask_offset=None, - spec_decoding_bl_tree_mask=None, - spec_bl_tree_first_sparse_mask_offset_kv=None, - sparse_kv_indices=None, - sparse_kv_offsets=None, - sparse_attn_indices=None, - sparse_attn_offsets=None, - sparse_attn_indices_block_size=0, - num_contexts=num_contexts, - num_ctx_tokens=num_ctx_tokens, - ) + with validating(op): + op( + q=qkv, + k=None, # packed QKV rides inside q + v=None, + output=output, + output_sf=None, + workspace_=self.workspace, + sequence_length=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), + host_past_key_value_lengths=torch.tensor(kv_lens, dtype=torch.int32), + host_total_kv_lens=torch.tensor( + [sum(kv_lens[:num_contexts]), sum(kv_lens[num_contexts:])], + dtype=torch.int32, + ), + context_lengths=torch.tensor(ctx_lens, dtype=torch.int32, device="cuda"), + host_context_lengths=torch.tensor(ctx_lens, dtype=torch.int32), + host_request_types=torch.tensor(req_types, dtype=torch.int32), + max_context_q_len_override=None, + kv_cache_block_offsets=block_offsets, + host_kv_cache_pool_pointers=pool_pointers, + host_kv_cache_pool_mapping=pool_mapping, + cache_indirection=None, + kv_scale_orig_quant=self.kv_scale_orig_quant, + kv_scale_quant_orig=self.kv_scale_quant_orig, + out_scale=None, + rotary_inv_freq=None, + rotary_cos_sin=None, + latent_cache=None, + q_pe=None, + block_ids_per_seq=None, + attention_sinks=None, + is_fused_qkv=True, + update_kv_cache=True, + predicted_tokens_per_seq=1, + local_layer_idx=local_layer_idx, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + tokens_per_block=self.tokens_per_block, + max_num_requests=self.max_batch, + max_context_length=self.max_seq_len, + max_seq_len=self.max_seq_len, + attention_window_size=( + self.max_seq_len if attention_window_size is None else attention_window_size + ), + beam_width=1, + mask_type=MASK_CAUSAL, + quant_mode=self.quant_mode, + q_scaling=1.0, + position_embedding_type=0, # no in-kernel RoPE + rope_dim=0, + rope_base=10000.0, + rope_scale_type=0, + rope_scale=1.0, + rope_short_m_scale=1.0, + rope_long_m_scale=1.0, + rope_max_positions=1024, + rope_original_max_positions=1024, + use_paged_context_fmha=use_paged_context_fmha, + attention_input_type=0, # mixed + is_mla_enable=False, + chunked_prefill_buffer_batch_size=1, + q_lora_rank=None, + kv_lora_rank=None, + qk_nope_head_dim=None, + qk_rope_head_dim=None, + v_head_dim=None, + rope_append=None, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + attention_chunk_size=None, + softmax_stats_tensor=None, + is_spec_decoding_enabled=False, + use_spec_decoding=False, + is_spec_dec_tree=False, + spec_decoding_generation_lengths=None, + spec_decoding_position_offsets_for_cpp=None, + spec_decoding_packed_mask=None, + spec_decoding_bl_tree_mask_offset=None, + spec_decoding_bl_tree_mask=None, + spec_bl_tree_first_sparse_mask_offset_kv=None, + sparse_kv_indices=None, + sparse_kv_offsets=None, + sparse_attn_indices=None, + sparse_attn_offsets=None, + sparse_attn_indices_block_size=0, + num_contexts=num_contexts, + num_ctx_tokens=num_ctx_tokens, + ) torch.cuda.synchronize() if record: @@ -1565,12 +1581,25 @@ def test_sinks_reject_bad_dtype_and_layout() -> None: "too many elements": torch.cat([good, other]), "empty": torch.empty(0, dtype=torch.float32, device="cuda"), } + # One good call before the guard assertions. Not decoration: this file has a + # pre-existing order dependence inside the sinks group. Removing *any* of + # its members makes test_bf16_paged_context_gpt_oss_sinks_and_window die + # with an illegal instruction and take twenty later tests down as setup + # errors. Measured on the unmodified file: dropping two sinks tests gives 20 + # errors, dropping two tests from outside the group gives none. Asserting + # the guard directly, as below, would otherwise remove this test's only op + # call and trip it. + env.call_op(qkv, [16], 1, [0], MASK_CAUSAL, attention_sinks=good, record=False) + + # The guard is asserted against `is_valid` rather than by driving a whole + # paged call and expecting it to abort partway. A guard is a claim about + # arguments; firing it from inside `call_op` leaves the environment + # half-advanced -- output allocated, cache bookkeeping moved on, no op call + # behind either. for label, sinks in bad_layouts.items(): - try: - env.call_op(qkv, [16], 1, [0], MASK_CAUSAL, attention_sinks=sinks, record=False) - except AssertionError: - continue - raise AssertionError(f"{label} attention_sinks was not rejected by the wrapper") + with pytest.raises(AssertionError, match="attention_sinks must be contiguous"): + op.is_valid(num_heads=SINK_HQ, num_kv_heads=SINK_HKV, attention_sinks=sinks) + del label # ─── Cyclic sliding window (attention_window_size < max_seq_len) + sinks ─ @@ -2987,103 +3016,104 @@ def _call_op( ) -> None: ns = len(kv_lens) req_types = [0 if i < num_contexts else 1 for i in range(ns)] - thop_attention( - q=q, - k=k, - v=v, - output=output, - output_sf=None, - workspace_=self.workspace, - sequence_length=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), - host_past_key_value_lengths=torch.tensor( - kv_lens if host_past_lens is None else host_past_lens, - dtype=torch.int32, - ), - host_total_kv_lens=torch.tensor( - [sum(kv_lens[:num_contexts]), sum(kv_lens[num_contexts:])], - dtype=torch.int32, - ), - context_lengths=torch.tensor(ctx_lens, dtype=torch.int32, device="cuda"), - host_context_lengths=torch.tensor(ctx_lens, dtype=torch.int32), - host_request_types=torch.tensor(req_types, dtype=torch.int32), - max_context_q_len_override=None, - kv_cache_block_offsets=self.block_offsets, - host_kv_cache_pool_pointers=self.pool_pointers, - host_kv_cache_pool_mapping=self.pool_mapping, - cache_indirection=None, - kv_scale_orig_quant=self.kv_scale_orig_quant, - kv_scale_quant_orig=self.kv_scale_quant_orig, - out_scale=None, - rotary_inv_freq=self.rotary_inv_freq, - rotary_cos_sin=self.rotary_cos_sin, - latent_cache=latent_cache, - q_pe=q_pe, - block_ids_per_seq=None, - attention_sinks=None, - is_fused_qkv=k is None, - update_kv_cache=True, - predicted_tokens_per_seq=predicted_tokens_per_seq, - local_layer_idx=0, - num_heads=self.num_heads, - num_kv_heads=num_kv_heads, - head_size=head_size, - tokens_per_block=self.tokens_per_block, - max_num_requests=self.max_batch, - max_context_length=self.max_seq_len, - max_seq_len=self.max_seq_len, - attention_window_size=self.max_seq_len, - beam_width=1, - mask_type=mask_type, - quant_mode=self.quant_mode, - q_scaling=self.q_scaling, - position_embedding_type=POSITION_EMBEDDING_TYPE_YARN, - rope_dim=QK_ROPE_HEAD_DIM, - rope_base=self.rope_scalars.rope_base, - rope_scale_type=self.rope_scalars.rope_scale_type, - rope_scale=self.rope_scalars.rope_scale, - rope_short_m_scale=self.rope_scalars.rope_short_m_scale, - rope_long_m_scale=self.rope_scalars.rope_long_m_scale, - rope_max_positions=self.rope_scalars.rope_max_positions, - rope_original_max_positions=self.rope_scalars.rope_original_max_positions, - use_paged_context_fmha=False, - attention_input_type=attention_input_type, - is_mla_enable=True, - chunked_prefill_buffer_batch_size=1, - q_lora_rank=self.q_lora_rank, - kv_lora_rank=KV_LORA_RANK, - qk_nope_head_dim=QK_NOPE_HEAD_DIM, - qk_rope_head_dim=QK_ROPE_HEAD_DIM, - v_head_dim=v_head_dim, - rope_append=True, - mrope_rotary_cos_sin=None, - mrope_position_deltas=None, - helix_position_offsets=None, - helix_is_inactive_rank=None, - attention_chunk_size=None, - softmax_stats_tensor=softmax_stats_tensor, - is_spec_decoding_enabled=False, - use_spec_decoding=False, - is_spec_dec_tree=False, - spec_decoding_generation_lengths=None, - spec_decoding_position_offsets_for_cpp=None, - spec_decoding_packed_mask=None, - spec_decoding_bl_tree_mask_offset=None, - spec_decoding_bl_tree_mask=None, - spec_bl_tree_first_sparse_mask_offset_kv=None, - sparse_kv_indices=None, - sparse_kv_offsets=None, - sparse_attn_indices=None, - sparse_attn_offsets=None, - sparse_attn_indices_block_size=0, - cu_q_seqlens=cu_q_seqlens, - cu_kv_seqlens=cu_kv_seqlens, - fmha_scheduler_counter=fmha_scheduler_counter, - mla_bmm1_scale=mla_bmm1_scale, - mla_bmm2_scale=mla_bmm2_scale, - quant_q_buffer=quant_q_buffer, - num_contexts=num_contexts, - num_ctx_tokens=num_ctx_tokens, - ) + with validating(op): + op( + q=q, + k=k, + v=v, + output=output, + output_sf=None, + workspace_=self.workspace, + sequence_length=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), + host_past_key_value_lengths=torch.tensor( + kv_lens if host_past_lens is None else host_past_lens, + dtype=torch.int32, + ), + host_total_kv_lens=torch.tensor( + [sum(kv_lens[:num_contexts]), sum(kv_lens[num_contexts:])], + dtype=torch.int32, + ), + context_lengths=torch.tensor(ctx_lens, dtype=torch.int32, device="cuda"), + host_context_lengths=torch.tensor(ctx_lens, dtype=torch.int32), + host_request_types=torch.tensor(req_types, dtype=torch.int32), + max_context_q_len_override=None, + kv_cache_block_offsets=self.block_offsets, + host_kv_cache_pool_pointers=self.pool_pointers, + host_kv_cache_pool_mapping=self.pool_mapping, + cache_indirection=None, + kv_scale_orig_quant=self.kv_scale_orig_quant, + kv_scale_quant_orig=self.kv_scale_quant_orig, + out_scale=None, + rotary_inv_freq=self.rotary_inv_freq, + rotary_cos_sin=self.rotary_cos_sin, + latent_cache=latent_cache, + q_pe=q_pe, + block_ids_per_seq=None, + attention_sinks=None, + is_fused_qkv=k is None, + update_kv_cache=True, + predicted_tokens_per_seq=predicted_tokens_per_seq, + local_layer_idx=0, + num_heads=self.num_heads, + num_kv_heads=num_kv_heads, + head_size=head_size, + tokens_per_block=self.tokens_per_block, + max_num_requests=self.max_batch, + max_context_length=self.max_seq_len, + max_seq_len=self.max_seq_len, + attention_window_size=self.max_seq_len, + beam_width=1, + mask_type=mask_type, + quant_mode=self.quant_mode, + q_scaling=self.q_scaling, + position_embedding_type=POSITION_EMBEDDING_TYPE_YARN, + rope_dim=QK_ROPE_HEAD_DIM, + rope_base=self.rope_scalars.rope_base, + rope_scale_type=self.rope_scalars.rope_scale_type, + rope_scale=self.rope_scalars.rope_scale, + rope_short_m_scale=self.rope_scalars.rope_short_m_scale, + rope_long_m_scale=self.rope_scalars.rope_long_m_scale, + rope_max_positions=self.rope_scalars.rope_max_positions, + rope_original_max_positions=self.rope_scalars.rope_original_max_positions, + use_paged_context_fmha=False, + attention_input_type=attention_input_type, + is_mla_enable=True, + chunked_prefill_buffer_batch_size=1, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=KV_LORA_RANK, + qk_nope_head_dim=QK_NOPE_HEAD_DIM, + qk_rope_head_dim=QK_ROPE_HEAD_DIM, + v_head_dim=v_head_dim, + rope_append=True, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + attention_chunk_size=None, + softmax_stats_tensor=softmax_stats_tensor, + is_spec_decoding_enabled=False, + use_spec_decoding=False, + is_spec_dec_tree=False, + spec_decoding_generation_lengths=None, + spec_decoding_position_offsets_for_cpp=None, + spec_decoding_packed_mask=None, + spec_decoding_bl_tree_mask_offset=None, + spec_decoding_bl_tree_mask=None, + spec_bl_tree_first_sparse_mask_offset_kv=None, + sparse_kv_indices=None, + sparse_kv_offsets=None, + sparse_attn_indices=None, + sparse_attn_offsets=None, + sparse_attn_indices_block_size=0, + cu_q_seqlens=cu_q_seqlens, + cu_kv_seqlens=cu_kv_seqlens, + fmha_scheduler_counter=fmha_scheduler_counter, + mla_bmm1_scale=mla_bmm1_scale, + mla_bmm2_scale=mla_bmm2_scale, + quant_q_buffer=quant_q_buffer, + num_contexts=num_contexts, + num_ctx_tokens=num_ctx_tokens, + ) torch.cuda.synchronize() def call_context( @@ -5265,3 +5295,1049 @@ def test_mla_rejects_null_q_lora_rank() -> None: assert "bad optional access" in str(exc), f"unexpected message: {exc}" else: raise AssertionError("q_lora_rank=None was accepted on the MLA path") + + +# ─── No-KV-cache (pool-free) mode: the DeepSeek-V4.1 vision tower ───── +# +# Every pool argument None (kv_cache_block_offsets, host_kv_cache_pool_pointers, +# host_kv_cache_pool_mapping, cache_indirection; tokens_per_block None too): +# the op then runs a plain packed-QKV context FMHA with no append, which is how +# a ViT tower is served. mask_type 0 (padding) is full bidirectional attention +# within each sequence ("image"); sequences never see each other. +# +# Tolerance: the standard ATOL / RTOL above, unchanged. Measured on sm_103 over +# the batches below at q scales 1 and 2 (logit std ~1 and ~2): at most 46% of +# the allowance with these seeds (37/1296/500/2048 at q scale 2); 64% was the +# worst over three other seeds of 16 x 576 at q scale 2, where FlashInfer's +# fmha_varlen and fa2 ragged kernels used the same fraction on the same inputs. +# A sharper softmax (q scale 3) reached 1.03x on one seed of 16 x 576 -- the +# probabilities are a bf16 MMA operand, whose rounding error scales with +# sum(p * |v|) rather than with |out| -- so that regime is reported in the +# contract and not gated here. +# +# Discrimination: every correct result is also required to sit far outside two +# wrong hypotheses built from the same inputs -- a causal mask within each image, +# and attention leaking across images (one merged sequence). Measured on sm_103 +# for the cases below: 57-531x (causal) and 48-694x (leak) outside, against a +# 20x floor. + +VIT_HEADS = 16 +VIT_HEAD_DIM = 64 +VIT_MAX_SEQ_LEN = 16384 # measured inert in this mode (64, 1000, 1296, 65536 identical) +NO_POOL_MIN_SEPARATION = 20.0 + + +def _vit_cu(lens: List[int]) -> torch.Tensor: + cu = [0] + for n in lens: + cu.append(cu[-1] + n) + return torch.tensor(cu, dtype=torch.int32, device="cuda") + + +def _vit_qkv(lens: List[int], q_scale: float, seed: int) -> torch.Tensor: + """Packed [T, 3 * H * D] bf16 q|k|v rows; q scaled to set the softmax sharpness.""" + torch.manual_seed(seed) + qkv = torch.randn(sum(lens), 3 * VIT_HEADS * VIT_HEAD_DIM, dtype=torch.bfloat16, device="cuda") + q_w = VIT_HEADS * VIT_HEAD_DIM + qkv[:, :q_w] = (qkv[:, :q_w].float() * q_scale).to(torch.bfloat16) + return qkv + + +def _no_pool_call( + qkv: torch.Tensor, + lens: List[int], + mask_type: int = MASK_PADDING, + cu_q_seqlens: Optional[torch.Tensor] = None, + cu_kv_seqlens: Optional[torch.Tensor] = None, + ctx_lens: Optional[List[int]] = None, + update_kv_cache: bool = True, + tokens_per_block: Optional[int] = None, + max_seq_len: int = VIT_MAX_SEQ_LEN, + attention_input_type: int = 0, + output: Optional[torch.Tensor] = None, + workspace: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """One pool-free thop_attention call over len(lens) context sequences. + + sequence_length / host_past_key_value_lengths carry `lens`; context_lengths / + host_context_lengths carry `ctx_lens` (default: `lens` too). + """ + ns = len(lens) + num_tokens = sum(lens) + ctx = lens if ctx_lens is None else ctx_lens + if output is None: + output = torch.empty( + num_tokens, VIT_HEADS * VIT_HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + op( + q=qkv, + k=None, + v=None, + output=output, + output_sf=None, + workspace_=torch.empty(0, dtype=torch.int8, device="cuda") + if workspace is None + else workspace, + sequence_length=torch.tensor(lens, dtype=torch.int32, device="cuda"), + host_past_key_value_lengths=torch.tensor(lens, dtype=torch.int32), + host_total_kv_lens=torch.tensor([num_tokens, 0], dtype=torch.int32), + context_lengths=torch.tensor(ctx, dtype=torch.int32, device="cuda"), + host_context_lengths=torch.tensor(ctx, dtype=torch.int32), + host_request_types=torch.zeros(ns, dtype=torch.int32), + max_context_q_len_override=None, + kv_cache_block_offsets=None, + host_kv_cache_pool_pointers=None, + host_kv_cache_pool_mapping=None, + cache_indirection=None, + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + out_scale=None, + rotary_inv_freq=None, + rotary_cos_sin=None, + latent_cache=None, + q_pe=None, + block_ids_per_seq=None, + attention_sinks=None, + is_fused_qkv=True, + update_kv_cache=update_kv_cache, + predicted_tokens_per_seq=1, + local_layer_idx=0, + num_heads=VIT_HEADS, + num_kv_heads=VIT_HEADS, + head_size=VIT_HEAD_DIM, + tokens_per_block=tokens_per_block, + max_num_requests=ns, + max_context_length=max_seq_len, + max_seq_len=max_seq_len, + attention_window_size=max_seq_len, + beam_width=1, + mask_type=mask_type, + quant_mode=0, + q_scaling=1.0, + position_embedding_type=0, # no in-kernel RoPE (the ViT's 2D RoPE is the caller's) + rope_dim=0, + rope_base=10000.0, + rope_scale_type=0, + rope_scale=1.0, + rope_short_m_scale=1.0, + rope_long_m_scale=1.0, + rope_max_positions=1024, + rope_original_max_positions=1024, + use_paged_context_fmha=False, + attention_input_type=attention_input_type, + is_mla_enable=False, + chunked_prefill_buffer_batch_size=1, + q_lora_rank=None, + kv_lora_rank=None, + qk_nope_head_dim=None, + qk_rope_head_dim=None, + v_head_dim=None, + rope_append=None, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + attention_chunk_size=None, + softmax_stats_tensor=None, + is_spec_decoding_enabled=False, + use_spec_decoding=False, + is_spec_dec_tree=False, + spec_decoding_generation_lengths=None, + spec_decoding_position_offsets_for_cpp=None, + spec_decoding_packed_mask=None, + spec_decoding_bl_tree_mask_offset=None, + spec_decoding_bl_tree_mask=None, + spec_bl_tree_first_sparse_mask_offset_kv=None, + sparse_kv_indices=None, + sparse_kv_offsets=None, + sparse_attn_indices=None, + sparse_attn_offsets=None, + sparse_attn_indices_block_size=0, + cu_q_seqlens=cu_q_seqlens, + cu_kv_seqlens=cu_kv_seqlens, + num_contexts=ns, + num_ctx_tokens=num_tokens, + ) + torch.cuda.synchronize() + return output + + +def _vit_reference( + qkv: torch.Tensor, lens: List[int], causal: bool = False, merged: bool = False +) -> torch.Tensor: + """fp32 SDPA per image, softmax scale 1/sqrt(D). causal=True and merged=True + (all images as one sequence) are the two wrong hypotheses.""" + num_tokens = sum(lens) + q_w = VIT_HEADS * VIT_HEAD_DIM + q = qkv[:, :q_w].view(num_tokens, VIT_HEADS, VIT_HEAD_DIM) + k = qkv[:, q_w : 2 * q_w].view(num_tokens, VIT_HEADS, VIT_HEAD_DIM) + v = qkv[:, 2 * q_w :].view(num_tokens, VIT_HEADS, VIT_HEAD_DIM) + outs = [] + start = 0 + for n in [num_tokens] if merged else lens: + o = F.scaled_dot_product_attention( + q[start : start + n].transpose(0, 1).float(), + k[start : start + n].transpose(0, 1).float(), + v[start : start + n].transpose(0, 1).float(), + is_causal=causal, + ) + outs.append(o.transpose(0, 1).reshape(n, q_w)) + start += n + return torch.cat(outs) + + +# Single images at the survey's 1296 / 4096 patches and the 9216-patch maximum, +# four equal images, ragged images (one of 37 patches), and single-patch images. +VIT_BATCHES = [ + [1296], + [4096], + [1296, 1296, 1296, 1296], + [37, 1296, 500, 2048], + [576] * 16, + [9216, 1296], + [1, 1, 5, 1], +] + + +def test_bf16_no_pool_vit_varlen_padding() -> None: + """Pool-free, mask_type=0, 16/16/64: per-image bidirectional attention over + variable-length images, with cu_q_seqlens given as the target passes it and + without it (bitwise identical). Gated against the per-image fp32 reference + and required to sit far outside the causal and cross-image-leak rivals.""" + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + for i, lens in enumerate(VIT_BATCHES): + for q_scale in (1.0, 2.0): + qkv = _vit_qkv(lens, q_scale, seed=400 + i) + out = _no_pool_call(qkv, lens, cu_q_seqlens=_vit_cu(lens), workspace=workspace) + ref = _vit_reference(qkv, lens) + label = f"lens={lens[:4]}.. n={len(lens)} q_scale={q_scale}" + torch.testing.assert_close(out.float(), ref, rtol=RTOL, atol=ATOL, msg=label) + assert _bitwise_equal(_no_pool_call(qkv, lens, workspace=workspace), out), ( + f"{label}: cu_q_seqlens=None differs from the given per-image offsets" + ) + if max(lens) > 1: + _assert_far_outside_band( + out, + _vit_reference(qkv, lens, causal=True), + NO_POOL_MIN_SEPARATION, + f"{label}: causal-mask rival", + ATOL, + RTOL, + ) + if len(lens) > 1: + _assert_far_outside_band( + out, + _vit_reference(qkv, lens, merged=True), + NO_POOL_MIN_SEPARATION, + f"{label}: cross-image leak rival", + ATOL, + RTOL, + ) + # The workspace the op grew in place for all of the above (largest: 9216+1296). + assert workspace.numel() < 64 * 1024 * 1024, f"workspace grew to {workspace.numel()} B" + + +def test_bf16_no_pool_vit_mask_type_is_load_bearing() -> None: + """The same pool-free call at mask_type=1 is causal within each image: it + lands far outside the bidirectional reference (so passing 1 for a ViT is + silently wrong, not rejected) and inside the causal one.""" + lens = [1296, 700, 1296, 64] + qkv = _vit_qkv(lens, 1.0, seed=420) + causal_out = _no_pool_call(qkv, lens, mask_type=MASK_CAUSAL, cu_q_seqlens=_vit_cu(lens)) + _assert_far_outside_band( + causal_out, + _vit_reference(qkv, lens), + NO_POOL_MIN_SEPARATION, + "mask_type=1 vs the bidirectional reference", + ATOL, + RTOL, + ) + torch.testing.assert_close( + causal_out.float(), _vit_reference(qkv, lens, causal=True), rtol=RTOL, atol=ATOL + ) + + +def test_bf16_no_pool_vit_batch_state() -> None: + """What each batch-state argument has to carry in pool-free mode.""" + lens = [1296, 700, 1296, 64] + qkv = _vit_qkv(lens, 2.0, seed=430) + ref = _vit_reference(qkv, lens) + base = _no_pool_call(qkv, lens) + torch.testing.assert_close(base.float(), ref, rtol=RTOL, atol=ATOL) + + # Inert here, bitwise: explicit cu_kv_seqlens, the context_only input type, + # a page size with no pool, and max_seq_len / max_context_length at any + # value -- including ones below the longest image. + inert = { + "cu_q_seqlens + cu_kv_seqlens": dict( + cu_q_seqlens=_vit_cu(lens), cu_kv_seqlens=_vit_cu(lens) + ), + "attention_input_type=1": dict(attention_input_type=1), + "tokens_per_block=32": dict(tokens_per_block=32), + "max_seq_len=64": dict(max_seq_len=64), + "max_seq_len=1000": dict(max_seq_len=1000), + "max_seq_len=65536": dict(max_seq_len=65536), + } + for name, kwargs in inert.items(): + assert _bitwise_equal(_no_pool_call(qkv, lens, **kwargs), base), f"{name} is not inert" + # Deterministic: a repeated call is bitwise identical. + assert _bitwise_equal(_no_pool_call(qkv, lens), base) + + # Rows past the batch's token count are never written. + sentinel = torch.full( + (sum(lens) + 16, VIT_HEADS * VIT_HEAD_DIM), + float("nan"), + dtype=torch.bfloat16, + device="cuda", + ) + _no_pool_call(qkv, lens, output=sentinel) + assert torch.isnan(sentinel[sum(lens) :].float()).all(), "rows past T were written" + assert _bitwise_equal(sentinel[: sum(lens)], base) + + # Given cu_q_seqlens is the image split the FMHA uses: offsets for a + # different split (same total, same longest image) silently override + # context_lengths / sequence_length. + lens_a = [1296, 1000, 1296, 1296] + lens_b = [1000, 1296, 1296, 1296] + qkv_ab = _vit_qkv(lens_a, 2.0, seed=431) + out = _no_pool_call(qkv_ab, lens_a, cu_q_seqlens=_vit_cu(lens_b)) + torch.testing.assert_close(out.float(), _vit_reference(qkv_ab, lens_b), rtol=RTOL, atol=ATOL) + _assert_far_outside_band( + out, + _vit_reference(qkv_ab, lens_a), + NO_POOL_MIN_SEPARATION, + "cu_q_seqlens vs context_lengths split", + ATOL, + RTOL, + ) + # context_lengths disagreeing with sequence_length (no cu_q_seqlens): neither + # split comes back, and nothing raises. + out = _no_pool_call(qkv_ab, lens_a, ctx_lens=lens_b) + for split, name in ((lens_a, "sequence_length"), (lens_b, "context_lengths")): + _assert_far_outside_band( + out, + _vit_reference(qkv_ab, split), + NO_POOL_MIN_SEPARATION, + f"mismatched context_lengths vs the {name} split", + ATOL, + RTOL, + ) + + # update_kv_cache=False is rejected outside cross attention. + try: + _no_pool_call(qkv, lens, update_kv_cache=False) + except RuntimeError as exc: + assert "KV cache update cannot be disabled" in str(exc), f"unexpected message: {exc}" + else: + raise AssertionError("update_kv_cache=False was accepted in pool-free mode") + + +def test_bf16_no_pool_vit_cuda_graph() -> None: + """A warmed-up pool-free call captures into a CUDA graph; replay is bitwise + equal to the eager call and follows new data written into the same q.""" + lens = [1296, 1296, 1296, 1296] + qkv = _vit_qkv(lens, 2.0, seed=440) + cu = _vit_cu(lens) + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + eager = _no_pool_call(qkv, lens, cu_q_seqlens=cu, workspace=workspace).clone() + + output = torch.empty_like(eager) + num_tokens = sum(lens) + seq = torch.tensor(lens, dtype=torch.int32, device="cuda") + host_lens = torch.tensor(lens, dtype=torch.int32) + host_total = torch.tensor([num_tokens, 0], dtype=torch.int32) + host_types = torch.zeros(len(lens), dtype=torch.int32) + kwargs = dict( + q=qkv, + k=None, + v=None, + output=output, + output_sf=None, + workspace_=workspace, + sequence_length=seq, + host_past_key_value_lengths=host_lens, + host_total_kv_lens=host_total, + context_lengths=seq, + host_context_lengths=host_lens, + host_request_types=host_types, + max_context_q_len_override=None, + kv_cache_block_offsets=None, + host_kv_cache_pool_pointers=None, + host_kv_cache_pool_mapping=None, + cache_indirection=None, + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + out_scale=None, + rotary_inv_freq=None, + rotary_cos_sin=None, + latent_cache=None, + q_pe=None, + block_ids_per_seq=None, + attention_sinks=None, + is_fused_qkv=True, + update_kv_cache=True, + predicted_tokens_per_seq=1, + local_layer_idx=0, + num_heads=VIT_HEADS, + num_kv_heads=VIT_HEADS, + head_size=VIT_HEAD_DIM, + tokens_per_block=None, + max_num_requests=len(lens), + max_context_length=VIT_MAX_SEQ_LEN, + max_seq_len=VIT_MAX_SEQ_LEN, + attention_window_size=VIT_MAX_SEQ_LEN, + beam_width=1, + mask_type=MASK_PADDING, + quant_mode=0, + q_scaling=1.0, + position_embedding_type=0, + rope_dim=0, + rope_base=10000.0, + rope_scale_type=0, + rope_scale=1.0, + rope_short_m_scale=1.0, + rope_long_m_scale=1.0, + rope_max_positions=1024, + rope_original_max_positions=1024, + use_paged_context_fmha=False, + attention_input_type=0, + is_mla_enable=False, + chunked_prefill_buffer_batch_size=1, + q_lora_rank=None, + kv_lora_rank=None, + qk_nope_head_dim=None, + qk_rope_head_dim=None, + v_head_dim=None, + rope_append=None, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + attention_chunk_size=None, + softmax_stats_tensor=None, + is_spec_decoding_enabled=False, + use_spec_decoding=False, + is_spec_dec_tree=False, + spec_decoding_generation_lengths=None, + spec_decoding_position_offsets_for_cpp=None, + spec_decoding_packed_mask=None, + spec_decoding_bl_tree_mask_offset=None, + spec_decoding_bl_tree_mask=None, + spec_bl_tree_first_sparse_mask_offset_kv=None, + sparse_kv_indices=None, + sparse_kv_offsets=None, + sparse_attn_indices=None, + sparse_attn_offsets=None, + sparse_attn_indices_block_size=0, + cu_q_seqlens=cu, + num_contexts=len(lens), + num_ctx_tokens=num_tokens, + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=torch.cuda.Stream()): + op(**kwargs) + graph.replay() + torch.cuda.synchronize() + assert _bitwise_equal(output, eager), "graph replay differs from the eager call" + + qkv.copy_(_vit_qkv(lens, 2.0, seed=441)) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), _vit_reference(qkv, lens), rtol=RTOL, atol=ATOL) + assert not _bitwise_equal(output, eager), "replay did not pick up the new q" + + +# The DeepSeek-V4.1 vision tower's call, in its multi-image form: exactly +# _no_pool_call's arguments -- 16/16/64 bf16, mask_type 0, is_fused_qkv on the +# packed [P, 3072] rows, update_kv_cache=True, cu_q_seqlens = the per-image +# offsets, context_lengths == sequence_length, max_seq_len / max_context_length +# / attention_window_size 16384, max_num_requests = the image count -- with ONE +# workspace tensor that the op grows in place across calls. One call holds +# every image of one request: up to 17037 patches (two images) and up to five +# images over the MMMU validation documents; one image is at most 8787 patches +# there and 9189 by the span limit. The target now calls one image per call +# (the entry's cells); these multi-image calls bound it from above. +VIT_MMMU_REQUESTS = [ + [8787, 8250], # 17037, MMMU's largest request (the split is illustrative) + [3600, 3400, 3500, 3337, 3200], # five images, 17037 in total + [1560, 1560, 1560], # golden mmmu_multi3 (validation_Math_3) + [3774, 1551], # golden carrots_corn + [1517], # golden mmmu_small + [8772], # golden mmmu_large + [9189], # the per-image ceiling +] +# The tower's CUDA graph sizes at the reference requests; the largest request +# is captured too. +VIT_GRAPH_REQUESTS = [[1517], [8772], [1560, 1560, 1560], [3774, 1551], [8787, 8250]] +# Rival floor for these requests. With thousands of keys per image and logit +# std ~1 the softmax is flat, so the correct output and the cross-image rival +# both shrink toward the mean of v and the rival lands closer to the band than +# at the shorter batches above: measured 15.9x at [8787, 8250], q scale 1, +# where the correct result uses ~0.05 of the band (the two ~300x apart). +VIT_MMMU_MIN_SEPARATION = 10.0 + + +def test_bf16_no_pool_vit_mmmu_requests() -> None: + """Every MMMU-shaped request in one call, against the full per-image fp32 + reference at q scales 1 and 2; far outside the causal and cross-image + rivals; cu_q_seqlens=None bitwise identical.""" + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + worst = 0.0 + for i, lens in enumerate(VIT_MMMU_REQUESTS): + for q_scale in (1.0, 2.0): + qkv = _vit_qkv(lens, q_scale, seed=460 + i) + out = _no_pool_call(qkv, lens, cu_q_seqlens=_vit_cu(lens), workspace=workspace) + ref = _vit_reference(qkv, lens) + label = f"lens={lens} q_scale={q_scale}" + torch.testing.assert_close(out.float(), ref, rtol=RTOL, atol=ATOL, msg=label) + worst = max(worst, _band_fraction(out, ref, ATOL, RTOL)) + assert _bitwise_equal(_no_pool_call(qkv, lens, workspace=workspace), out), label + causal = _band_fraction(out, _vit_reference(qkv, lens, causal=True), ATOL, RTOL) + assert causal > VIT_MMMU_MIN_SEPARATION, ( + f"{label}: causal rival only {causal:.3g}x outside" + ) + msg = f"causal rival {causal:.0f}x" + if len(lens) > 1: + merged = _band_fraction(out, _vit_reference(qkv, lens, merged=True), ATOL, RTOL) + assert merged > VIT_MMMU_MIN_SEPARATION, ( + f"{label}: leak rival only {merged:.3g}x outside" + ) + msg += f", cross-image rival {merged:.0f}x" + print( + f"\n[vit mmmu] {label}: {_band_fraction(out, ref, ATOL, RTOL):.3f} of the band; {msg} outside" + ) + print(f"[vit mmmu] worst {worst:.3f} of the band; workspace grew to {workspace.numel()} B") + assert workspace.numel() < 64 * 1024 * 1024 + + +def _vit_sampled_reference(qkv: torch.Tensor, lens: List[int], per_image: int, seed: int): + """fp32 per-image attention for a sample of query rows (first, last and random + rows of every image) over all of that image's keys: (row indices, reference rows).""" + gen = torch.Generator().manual_seed(seed) + q_w = VIT_HEADS * VIT_HEAD_DIM + rows, outs, start = [], [], 0 + for n in lens: + idx = torch.randperm(n, generator=gen)[: min(per_image, n)] + idx[0], idx[-1] = 0, n - 1 + idx = idx.to("cuda") + start + q = qkv[idx, :q_w].view(-1, VIT_HEADS, VIT_HEAD_DIM).float() + k = qkv[start : start + n, q_w : 2 * q_w].view(n, VIT_HEADS, VIT_HEAD_DIM).float() + v = qkv[start : start + n, 2 * q_w :].view(n, VIT_HEADS, VIT_HEAD_DIM).float() + p = torch.softmax(torch.einsum("mhd,nhd->hmn", q, k) / VIT_HEAD_DIM**0.5, dim=-1) + outs.append(torch.einsum("hmn,nhd->mhd", p, v).reshape(-1, q_w)) + rows.append(idx) + start += n + return torch.cat(rows), torch.cat(outs) + + +def test_bf16_no_pool_vit_large_calls_sampled_rows() -> None: + """Far past one request: 80 images of 8787 (702960 tokens; q past 2^31 elements) + and one 262144-token sequence, in band on sampled rows, every output finite.""" + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + for i, lens in enumerate(([8787] * 80, [262144])): + qkv = _vit_qkv(lens, 1.0, seed=470 + i) + out = _no_pool_call(qkv, lens, cu_q_seqlens=_vit_cu(lens), workspace=workspace) + assert torch.isfinite(out.float()).all() + rows, ref = _vit_sampled_reference(qkv, lens, 256, 471 + i) + torch.testing.assert_close(out[rows].float(), ref, rtol=RTOL, atol=ATOL) + print( + f"\n[vit large] {len(lens)} x {lens[0]} = {sum(lens)} tokens: {rows.numel()} sampled rows " + f"at {_band_fraction(out[rows], ref, ATOL, RTOL):.3f} of the band; workspace {workspace.numel()} B" + ) + del qkv, out + + +def _vit_graph_kwargs(qkv, output, workspace, seq, host_lens, cu, num_tokens, ns) -> Dict: + return dict( + q=qkv, + k=None, + v=None, + output=output, + output_sf=None, + workspace_=workspace, + sequence_length=seq, + host_past_key_value_lengths=host_lens, + host_total_kv_lens=torch.tensor([num_tokens, 0], dtype=torch.int32), + context_lengths=seq, + host_context_lengths=host_lens, + host_request_types=torch.zeros(ns, dtype=torch.int32), + max_context_q_len_override=None, + kv_cache_block_offsets=None, + host_kv_cache_pool_pointers=None, + host_kv_cache_pool_mapping=None, + cache_indirection=None, + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + out_scale=None, + rotary_inv_freq=None, + rotary_cos_sin=None, + latent_cache=None, + q_pe=None, + block_ids_per_seq=None, + attention_sinks=None, + is_fused_qkv=True, + update_kv_cache=True, + predicted_tokens_per_seq=1, + local_layer_idx=0, + num_heads=VIT_HEADS, + num_kv_heads=VIT_HEADS, + head_size=VIT_HEAD_DIM, + tokens_per_block=None, + max_num_requests=ns, + max_context_length=VIT_MAX_SEQ_LEN, + max_seq_len=VIT_MAX_SEQ_LEN, + attention_window_size=VIT_MAX_SEQ_LEN, + beam_width=1, + mask_type=MASK_PADDING, + quant_mode=0, + q_scaling=1.0, + position_embedding_type=0, + rope_dim=0, + rope_base=10000.0, + rope_scale_type=0, + rope_scale=1.0, + rope_short_m_scale=1.0, + rope_long_m_scale=1.0, + rope_max_positions=1024, + rope_original_max_positions=1024, + use_paged_context_fmha=False, + attention_input_type=0, + is_mla_enable=False, + chunked_prefill_buffer_batch_size=1, + q_lora_rank=None, + kv_lora_rank=None, + qk_nope_head_dim=None, + qk_rope_head_dim=None, + v_head_dim=None, + rope_append=None, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + attention_chunk_size=None, + softmax_stats_tensor=None, + is_spec_decoding_enabled=False, + use_spec_decoding=False, + is_spec_dec_tree=False, + spec_decoding_generation_lengths=None, + spec_decoding_position_offsets_for_cpp=None, + spec_decoding_packed_mask=None, + spec_decoding_bl_tree_mask_offset=None, + spec_decoding_bl_tree_mask=None, + spec_bl_tree_first_sparse_mask_offset_kv=None, + sparse_kv_indices=None, + sparse_kv_offsets=None, + sparse_attn_indices=None, + sparse_attn_offsets=None, + sparse_attn_indices_block_size=0, + cu_q_seqlens=cu, + num_contexts=ns, + num_ctx_tokens=num_tokens, + ) + + +def test_bf16_no_pool_vit_cuda_graph_prototype_sizes() -> None: + """Captured at the prototype's graph sizes with one shared workspace (sized by + eager calls first); replays on new q are bitwise equal to eager and in band. + Rewriting the device length tensors to another split of the same total does + NOT re-split a captured call: the replay no longer matches the new split.""" + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + graphs = [] + for lens in VIT_GRAPH_REQUESTS: + qkv = _vit_qkv(lens, 2.0, seed=480) + _no_pool_call(qkv, lens, cu_q_seqlens=_vit_cu(lens), workspace=workspace) + for i, lens in enumerate(VIT_GRAPH_REQUESTS): + num_tokens, ns = sum(lens), len(lens) + qkv = _vit_qkv(lens, 2.0, seed=481 + i) + output = torch.empty( + num_tokens, VIT_HEADS * VIT_HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + seq = torch.tensor(lens, dtype=torch.int32, device="cuda") + cu = _vit_cu(lens) + kwargs = _vit_graph_kwargs( + qkv, output, workspace, seq, torch.tensor(lens, dtype=torch.int32), cu, num_tokens, ns + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=torch.cuda.Stream()): + op(**kwargs) + graphs.append((lens, graph, qkv, output, seq, cu)) + for rnd in range(2): + for i, (lens, graph, qkv, output, _, _) in enumerate(graphs): + qkv.copy_(_vit_qkv(lens, 2.0, seed=490 + 10 * rnd + i)) + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + eager = _no_pool_call(qkv, lens, cu_q_seqlens=_vit_cu(lens), workspace=workspace) + assert _bitwise_equal(output, eager), (lens, rnd) + torch.testing.assert_close( + output.float(), _vit_reference(qkv, lens), rtol=RTOL, atol=ATOL + ) + # The multi3 graph with the device lengths rewritten to (3000, 1000, 680). + lens, graph, qkv, output, seq, cu = graphs[2] + other = [3000, 1000, 680] + seq.copy_(torch.tensor(other, dtype=torch.int32)) + cu.copy_(_vit_cu(other)) + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + finite = bool(torch.isfinite(output.float()).all()) + band = ( + _band_fraction(output, _vit_reference(qkv, other), ATOL, RTOL) if finite else float("nan") + ) + print( + f"\n[vit graph] split rewritten after capture: output finite {finite}, band vs the new split {band:.3g}" + ) + assert not finite or band > 1.0 + + +# max_context_q_len_override in no-KV-cache mode: a captured call that follows +# the split. The DeepSeek-V4.1 target's graph form: one call per padded +# bucket B, slots (image P, pad B - P), captured with the host lengths holding +# the B rows as one sequence and max_context_q_len_override = B; before each +# replay only the DEVICE sequence_length / context_lengths / cu_q_seqlens are +# rewritten. Also a 4-slot form with empty (zero-length) slots. +VIT_OVERRIDE_BUCKETS = { + # bucket: image lengths replayed in that bucket (pad = B - P) + 1024: [7, 1000, 1023, 1024], + 1536: [1521, 1100, 1536], + 2560: [2305, 2432, 2560, 2049], + 4096: [4095, 3600, 4096], + 8704: [8613, 8704, 8193], + 9216: [9189, 8706, 8787, 9216], +} +VIT_OVERRIDE_MULTI = (8192, [[2048] * 4, [3000, 1000, 680, 3512], [8000, 64, 64, 64], [5000, 3192, 0, 0], + [8192, 0, 0, 0], [0, 8192, 0, 0]]) # fmt: skip + + +def _vit_override_kwargs(qkv, output, workspace, seq, cu, host_lens, num_tokens, override) -> Dict: + kw = _vit_graph_kwargs(qkv, output, workspace, seq, host_lens, cu, num_tokens, len(host_lens)) + kw["max_context_q_len_override"] = override + return kw + + +def _replay_follows_the_split(B: int, splits: List[List[int]], seed: int) -> None: + ns = len(splits[0]) + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + qkv = _vit_qkv([B], 2.0, seed=seed) + output = torch.empty(B, VIT_HEADS * VIT_HEAD_DIM, dtype=torch.bfloat16, device="cuda") + capture_lens = [B] + [0] * (ns - 1) + seq = torch.tensor(capture_lens, dtype=torch.int32, device="cuda") + cu = _vit_cu(capture_lens) + host = torch.tensor(capture_lens, dtype=torch.int32) + kw = _vit_override_kwargs(qkv, output, workspace, seq, cu, host, B, B) + op(**kw) # sizes the workspace + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=torch.cuda.Stream()): + op(**kw) + for rnd in range(2): + for i, lens in enumerate(splits): + assert sum(lens) == B and len(lens) == ns + qkv.copy_(_vit_qkv([B], 2.0, seed=seed + 100 * rnd + i + 1)) + seq.copy_(torch.tensor(lens, dtype=torch.int32)) + cu.copy_(_vit_cu(lens)) + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + replayed = output.clone() + # The same call eagerly, host lengths = the split. + eager = torch.empty_like(output) + op(**_vit_override_kwargs( + qkv, eager, workspace, torch.tensor(lens, dtype=torch.int32, device="cuda"), + _vit_cu(lens), torch.tensor(lens, dtype=torch.int32), B, B)) # fmt: skip + torch.cuda.synchronize() + assert _bitwise_equal(replayed, eager), (B, lens, rnd) + nz = [n for n in lens if n] + torch.testing.assert_close( + replayed.float(), _vit_reference(qkv, nz), rtol=RTOL, atol=ATOL + ) + if len(nz) > 1 and min(nz) >= max(64, B // 20): + # A replay stuck at the capture split (one sequence of B rows) is the rival. + merged = _band_fraction(replayed, _vit_reference(qkv, nz, merged=True), ATOL, RTOL) + assert merged > VIT_MMMU_MIN_SEPARATION, (B, lens, merged) + + +def test_bf16_no_pool_vit_cuda_graph_q_len_override_follows_the_split() -> None: + """With max_context_q_len_override fixed (= the call's row count), a captured + no-pool call follows whatever split the device length tensors hold at replay: + bitwise equal to the eager call over that split and in band, in the target's + two-slot (image, pad) form at every tested bucket and in a four-slot form with + empty slots. Without the override a capture whose host lengths are shorter + than a replayed sequence is wrong (control); with it, right.""" + for k, (B, images) in enumerate(VIT_OVERRIDE_BUCKETS.items()): + _replay_follows_the_split(B, [[P, B - P] for P in images], seed=600 + 10 * k) + B, splits = VIT_OVERRIDE_MULTI + _replay_follows_the_split(B, splits, seed=700) + + # What the override buys: the launch no longer depends on the capture-time + # host lengths. Captured at (1280, 1280) and replayed at (2305, 255): with + # the override (2560) the replay is right; without it the launch only covers + # 1280-row sequences and the 2305-row one comes back wrong. + B, cap, rep = 2560, [1280, 1280], [2305, 255] + for override in (B, None): + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + qkv = _vit_qkv([B], 2.0, seed=710) + output = torch.empty(B, VIT_HEADS * VIT_HEAD_DIM, dtype=torch.bfloat16, device="cuda") + seq = torch.tensor(cap, dtype=torch.int32, device="cuda") + cu = _vit_cu(cap) + kw = _vit_override_kwargs( + qkv, output, workspace, seq, cu, torch.tensor(cap, dtype=torch.int32), B, override + ) + op(**kw) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=torch.cuda.Stream()): + op(**kw) + seq.copy_(torch.tensor(rep, dtype=torch.int32)) + cu.copy_(_vit_cu(rep)) + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + finite = bool(torch.isfinite(output.float()).all()) + band = ( + _band_fraction(output, _vit_reference(qkv, rep), ATOL, RTOL) if finite else float("nan") + ) + print( + f"\n[vit graph] override {override}, {cap} captured, {rep} replayed: finite {finite}, band {band:.3g}" + ) + if override is None: + assert not finite or band > 1.0 + else: + assert finite and band <= 1.0 + + +def _vit_two_slot_replay(qkv: torch.Tensor, lens: List[int], bucket: int) -> torch.Tensor: + """The target's captured form over `bucket` rows: warm up and capture with + the host and device lengths holding every row as one sequence and + max_context_q_len_override = bucket, then write `qkv` and the split `lens` + into the static buffers and replay. Returns the replayed output after + checking it bitwise against the eager call over that split.""" + ns = len(lens) + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + static_qkv = _vit_qkv([bucket], 1.0, seed=bucket) + output = torch.empty(bucket, VIT_HEADS * VIT_HEAD_DIM, dtype=torch.bfloat16, device="cuda") + capture = [bucket] + [0] * (ns - 1) + seq = torch.tensor(capture, dtype=torch.int32, device="cuda") + cu = _vit_cu(capture) + host = torch.tensor(capture, dtype=torch.int32) + kw = _vit_override_kwargs(static_qkv, output, workspace, seq, cu, host, bucket, bucket) + op(**kw) # sizes the workspace before capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=torch.cuda.Stream()): + op(**kw) + static_qkv.copy_(qkv) + seq.copy_(torch.tensor(lens, dtype=torch.int32)) + cu.copy_(_vit_cu(lens)) + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + eager = torch.empty_like(output) + op(**_vit_override_kwargs( + static_qkv, eager, workspace, torch.tensor(lens, dtype=torch.int32, device="cuda"), + _vit_cu(lens), torch.tensor(lens, dtype=torch.int32), bucket, bucket)) # fmt: skip + torch.cuda.synchronize() + assert _bitwise_equal(output, eager), ( + f"replay over {lens} in bucket {bucket} differs from eager" + ) + return output + + +# ─── The entry's own cell list ────────────────────────────────────────── + + +def _cell_standard(cell) -> None: + """A gpt-oss cell: GQA with sinks, with or without a sliding window. + + Driven against `op.reference` rather than this file's `_sink_reference`, + which is the point of the cell list: the entry states what the op computes + and CI drives that statement. The two are cross-checked against each other + first, because a reference that disagrees with a reference is a different + failure from a kernel that moved. + """ + spec = cell.spec + torch.manual_seed(abs(hash(cell.why)) % 2**31) + window = spec["window"] + env = ( + _sink_env(attention_window_size=window, num_blocks=96, max_blocks_per_seq=16) + if window + else _sink_env() + ) + n_tokens = 200 if window else 48 + sink = _sinks(-2.0, 4.0, seed=abs(hash(cell.why)) % 1000) + env.add_request(0, n_tokens) + qkv = env.random_qkv(n_tokens) + out = _paged_ctx_call(env, qkv, [n_tokens], 1, [0], attention_sinks=sink) + + hq, hkv, dim = env.num_heads, env.num_kv_heads, env.head_dim + q = qkv[:, : hq * dim] + k = torch.cat(env.k_history[0]).reshape(-1, hkv * dim) + v = torch.cat(env.v_history[0]).reshape(-1, hkv * dim) + ref = op.reference( + q, + k, + v, + [n_tokens], + [k.shape[0]], + num_heads=hq, + num_kv_heads=hkv, + qk_head_dim=dim, + v_head_dim=dim, + softmax_scale=1.0 / math.sqrt(dim), + causal=True, + attention_sinks=sink, + window=window, + ) + local = _sink_reference(env, qkv, [n_tokens], [0], [0], sink, window=window) + torch.testing.assert_close( + ref.float(), + local.float(), + rtol=2e-2, + atol=2e-2, + msg=lambda m: f"op.reference disagrees with _sink_reference\n{m}", + ) + op.compare(out, ref) + + +def _cell_mla_context(cell) -> None: + """The deepseek MLA context cell, on the no-append flavour. + + The appending flavour ropes q and k inside the call, and `reference` does + not model rope -- it is masked softmax attention over explicit K/V, and the + rope belongs to the entries beside this one. The no-append flavour takes + q/k already roped by the caller, which is the form the reference states, so + that is the one a cell can gate. + + The two scale factors under quant_mode 128 -- `s**2` inside the softmax + scale and `s` on the output -- default to 1.0 on this flavour, which is why + only `e4m3_inputs` is set here. + """ + env = _fp8_r1_env() + cached_lens, new_lens = [96, 31, 0], [32, 9, 25] + kv_lens = [c + n for c, n in zip(cached_lens, new_lens)] + torch.manual_seed(abs(hash(cell.why)) % 2**31) + for rid, total in enumerate(kv_lens): + env.add_request(rid, new_lens[rid]) + env.reserve_cache_pages(rid, total) + k, v = _fp8_pool_backed_kv(env, cached_lens, new_lens) + q = torch.randn(sum(new_lens), env.num_heads * QK_HEAD_DIM, dtype=torch.bfloat16, device="cuda") + out = env.call_context_no_append([0, 1, 2], new_lens, kv_lens, q, k, v) + + ref = op.reference( + q, + k, + v, + new_lens, + kv_lens, + num_heads=cell.spec["num_heads"], + num_kv_heads=cell.spec["num_kv_heads"], + qk_head_dim=QK_HEAD_DIM, + v_head_dim=V_HEAD_DIM, + softmax_scale=env.softmax_scale, + causal=True, + e4m3_inputs=True, + ) + local = _fp8_no_append_reference(env, q, k, v, new_lens, kv_lens) + torch.testing.assert_close( + ref.float(), + local.float(), + rtol=FP8_MLA_RTOL, + atol=FP8_MLA_ATOL, + msg=lambda m: f"op.reference disagrees with _fp8_no_append_reference\n{m}", + ) + op.compare(out, ref) + + +def _cell_no_kv_cache(cell) -> None: + """A deepseek-v4.1 vision cell: the pool-free call, eager or captured. + + Eager (`q_len_override` None): one sequence with the target's arguments, + `cu_q_seqlens` given and, as a check, omitted (bitwise the same). + Captured: the two-slot form through `_vit_two_slot_replay`, which also + holds the replay bitwise to the eager call over the same split. + + Driven against `op.reference`, cross-checked first against this file's + SDPA `_vit_reference`, and gated by `op.compare` and by the bf16 + atol/rtol band the section above measured; a causal mask within each + sequence is required far outside that band. + """ + spec = cell.spec + assert (spec["num_heads"], spec["num_kv_heads"], spec["head_size"]) == ( + VIT_HEADS, + VIT_HEADS, + VIT_HEAD_DIM, + ), "the pool-free builders here are 16/16/64" + assert spec["mask_type"] == MASK_PADDING and not spec["kv_cache"] + lens = list(spec["seq_lens"]) + nz = [n for n in lens if n] + q_w = VIT_HEADS * VIT_HEAD_DIM + for q_scale in spec["q_scales"]: + label = f"{cell.why} q_scale={q_scale}" + qkv = _vit_qkv(lens, q_scale, seed=500 + sum(lens) + int(10 * q_scale)) + if spec["cuda_graph"]: + out = _vit_two_slot_replay(qkv, lens, spec["q_len_override"]) + else: + assert spec["q_len_override"] is None + out = _no_pool_call(qkv, lens, cu_q_seqlens=_vit_cu(lens)) + assert _bitwise_equal(_no_pool_call(qkv, lens), out), ( + f"{label}: cu_q_seqlens=None differs from the given offsets" + ) + ref = op.reference( + qkv[:, :q_w], + qkv[:, q_w : 2 * q_w], + qkv[:, 2 * q_w :], + lens, + lens, + num_heads=VIT_HEADS, + num_kv_heads=VIT_HEADS, + qk_head_dim=VIT_HEAD_DIM, + v_head_dim=VIT_HEAD_DIM, + softmax_scale=1.0 / math.sqrt(VIT_HEAD_DIM), + causal=False, + ) + local = _vit_reference(qkv, nz) + torch.testing.assert_close( + ref.float(), + local, + rtol=RTOL, + atol=ATOL, + msg=lambda m: f"op.reference disagrees with _vit_reference\n{m}", + ) + op.compare(out, ref) + torch.testing.assert_close(out.float(), local, rtol=RTOL, atol=ATOL, msg=label) + _assert_far_outside_band( + out, + _vit_reference(qkv, nz, causal=True), + VIT_MMMU_MIN_SEPARATION, + f"{label}: causal-mask rival", + ATOL, + RTOL, + ) + del qkv, out, ref, local + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Drive every configuration the entry claims. + + The MLA generation cell is driven by the tests above rather than here: its + output lives in latent space, reconstructed from the paged cache, and the + entry's `reference` takes explicit K/V. The cell stays in the list because + it is a configuration a shipped target runs and `is_valid` is checked + against it; what this file does not do is compare it to `reference`, and + saying so is better than a cell that quietly proves less than its siblings. + + The deepseek-v4.1 vision cells carry `kv_cache=False` and are driven + without a pool; cells without the key are cache-backed. + """ + kind = cell.spec["input_type"] + if not cell.spec.get("kv_cache", True): + _cell_no_kv_cache(cell) + elif kind == STANDARD: + _cell_standard(cell) + elif kind == MLA_CONTEXT: + _cell_mla_context(cell) + else: + with validating(op): + op.is_valid(num_heads=cell.spec["num_heads"], num_kv_heads=cell.spec["num_kv_heads"]) + pytest.skip("MLA generation is covered by test_fp8_mla_r1_cell_generation_decode_h128") diff --git a/tests/unittest/_torch/modeling_v2/comm/_allgather_op_matrix.py b/tests/unittest/_torch/modeling_v2/comm/_allgather_op_matrix.py index bf10463742ee..8666ae2b5045 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_allgather_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/comm/_allgather_op_matrix.py @@ -3,7 +3,7 @@ """GPU certification matrix for the allgather catalog entry. A collective cannot be exercised in one process, so this script is its own -launcher: run it plainly and it re-executes itself under `mpirun` with one +launcher: run it plainly and it spawns one worker process per device, with one rank per device named in CUDA_VISIBLE_DEVICES, under a deadline the parent enforces by killing the whole process group. A wedged collective hangs rather than raising — and this op wedges rather than raising for every @@ -29,19 +29,22 @@ `tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_op_matrix.py`: it starts this job and turns its exit code into an assertion. Both halves are started by file path: the launcher must not import `tensorrt_llm` (that calls -`MPI_Init`, and an MPI-initialized process cannot start `mpirun`), and the +CUDA nor torch.distributed state of its own), and the ranks reach the catalog by absolute import, so neither needs a package. """ import os import random import signal +import socket import subprocess import sys +import time from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Tuple import torch +import torch.distributed as dist assert torch.cuda.is_available(), "allgather requires CUDA devices" @@ -65,20 +68,66 @@ # by one gather of the rank's own rows into the full token set. SITES_PER_DECODE_GRAPH = 29 +# The DeepSeek-V4.1 core target under attention DP gathers, per rank, rows padded +# to max(all_rank_num_tokens) with sizes = None: the MoE input as bf16 [n, 5120] +# on every layer, and on the engram layers the [n, 24] int64 engram ids as their +# int32 [n, 48] view. int64 [n, 24] is run beside it. +CORE_FORMS: Tuple[Tuple[str, torch.dtype, Tuple[int, ...]], ...] = ( + ("MoE bf16 [n,5120]", torch.bfloat16, (5120,)), + ("engram ids int32 [n,48]", torch.int32, (48,)), + ("engram ids int64 [n,24]", torch.int64, (24,)), +) +N_MAX = 16384 # rows per rank: the largest max_num_tokens the target is served at +SWEEP_SLACK = 509 # extra table rows (prime): every call reads its rows at its own offset +CORE_LAYERS = 40 +ENGRAM_LAYER_IDS = (1, 14) +# The V4.1 capture batch set, per rank: every G in 1..32 plus the padded default +# list of cuda_graph_config {enable_padding: true, max_batch_size: 256} (16..128 +# step 8, 192, 256) -- 46 values. A captured step is generation-only with next_n 1 +# (DSpark off) or 6 (DSpark on), so every rank gathers n = G or n = 6G rows. The +# 5G / 6G rows of G in 1..8, 16, 24, 32 are kept beside them (they add 35 and +# 160): 84 per-rank row counts, up to 1536. +GSET: Tuple[int, ...] = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_GATE_BATCHES: Tuple[int, ...] = tuple(range(1, 9)) + (16, 24, 32) +CORE_CAPTURE_ROWS: Tuple[int, ...] = tuple( + sorted( + set(GSET) + | {6 * g for g in GSET} + | set(_GATE_BATCHES) + | {5 * b for b in _GATE_BATCHES} + | {6 * b for b in _GATE_BATCHES} + ) +) +# Eager prefill-sized gathers (per-rank rows) run between the decode graphs' +# replays: row counts the runtime dumps showed, and N_MAX. +PREFILL_BETWEEN_REPLAYS: Tuple[int, ...] = (888, 890, 2048, 9524, 10441, 16378, N_MAX) + # Bound inside the rank body, never in the launcher: importing the entry # pulls in tensorrt_llm, which calls MPI_Init at import, and an -# MPI-initialized process cannot launch `mpirun` — measured on this host, -# mpirun then exits 1 with no output from any rank. allgather: Any = None -COMM: Any = None RANK = 0 WORLD = 1 GROUP: List[int] = [0] # The engine's own cross-rank synchronisation object, built from the real # Mapping a serving engine builds under attention data parallelism. Bound in # the rank body for the same reason as the rest. -DIST: Any = None -MPI: Any = None +HOST_GROUP: Any = None +SUBGROUPS: Dict[Tuple[int, ...], Any] = {} +PG_BOXED: Any = None + + +def _host_allgather(obj: Any) -> List[Any]: + """One host-side cross-rank step, the shape of the engine's per-step sync. + + A serving engine agrees the per-rank token counts once per step with a + host-side collective (`Distributed.tp_allgather`; TorchDist rides this same + torch.distributed machinery). A dedicated gloo group stands in for the + engine object here: constructing the engine's own `TorchDist` needs ray + importable, which a plain development machine does not owe this test. + """ + gathered: List[Any] = [None] * WORLD + dist.all_gather_object(gathered, obj, group=HOST_GROUP) + return gathered def _payload( @@ -125,9 +174,11 @@ def _gather_ref( rank order. Never from a second collective. """ ranks = list(range(WORLD)) if ranks is None else list(ranks) - return torch.cat( - [_payload(r, sizes[i], seed, dtype, trailing) for i, r in enumerate(ranks)], - dim=0, + # Delegates to the entry: what the concatenation is supposed to be is the + # entry's claim, and building it here again would be a second place to fix + # if it were ever wrong. + return allgather.reference( + [_payload(r, sizes[i], seed, dtype, trailing) for i, r in enumerate(ranks)] ) @@ -137,11 +188,11 @@ def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: Tightened from the default dtype-aware tolerances rather than loosened — a gather computes nothing, so any difference at all is a wrong gather. """ - if out.dtype is torch.float8_e4m3fn: - # torch.testing cannot compare float8 tensors; for a pure data move the - # byte pattern is the honest gate anyway. - out, ref = out.view(torch.uint8), ref.view(torch.uint8) - torch.testing.assert_close(out, ref, rtol=0, atol=0, msg=lambda built: f"{where}: {built}") + # The gate itself is the entry's `compare`; `where` only labels the failure. + try: + allgather.compare(out, ref) + except AssertionError as exc: + raise AssertionError(f"{where}: {exc}") from exc def _sizes_vectors() -> List[List[int]]: @@ -190,9 +241,9 @@ def _warm_up_off_capture_stream(body, reps: int = 2) -> None: """Run `body` on a side stream, then rejoin, so a capture can follow. Two things have to be done before a capture and cannot be done inside one: - the group's NCCL communicator has to exist (see - check_cuda_graph_capture_of_a_first_call_raises), and torch wants the work - warmed on a non-default stream. + torch wants the work warmed on a non-default stream + before a capture (the communicator itself already exists -- see + check_cuda_graph_captures_a_cold_first_call). """ side = torch.cuda.Stream() side.wait_stream(torch.cuda.current_stream()) @@ -201,7 +252,7 @@ def _warm_up_off_capture_stream(body, reps: int = 2) -> None: body() torch.cuda.current_stream().wait_stream(side) torch.cuda.synchronize() - COMM.Barrier() + dist.barrier() def _capture_per_batch_size( @@ -225,11 +276,26 @@ def _capture_per_batch_size( def body(xs: List[torch.Tensor] = xs) -> List[torch.Tensor]: return [allgather(x, None, GROUP) for x in xs] + # Persistent output buffers outside the shared pool; see the reduce-scatter + # matrix for why -- a burst of replays with no sync lets pooled outputs + # alias, and the copy inside the capture, ordered after this graph's + # gather, reads the right value before the next graph overwrites the pool. + persistent = [ + torch.empty(rows * WORLD, HIDDEN, dtype=torch.bfloat16, device="cuda") + for _ in range(sites) + ] + + def capture_body() -> List[torch.Tensor]: + produced = body() + for dst, src in zip(persistent, produced): + dst.copy_(src) + return persistent + _warm_up_off_capture_stream(body) graph = torch.cuda.CUDAGraph() ctx = torch.cuda.graph(graph) if pool is None else torch.cuda.graph(graph, pool=pool) with ctx: - outputs[rows] = body() + outputs[rows] = capture_body() # Without this a capture that returned nothing would make every # `_verify_replay` below an empty loop, i.e. a vacuous pass. assert len(outputs[rows]) == sites, ( @@ -238,7 +304,7 @@ def body(xs: List[torch.Tensor] = xs) -> List[torch.Tensor]: if pool is None: pool = graph.pool() inputs[rows], graphs[rows] = xs, graph - COMM.Barrier() + dist.barrier() return inputs, graphs, outputs @@ -277,61 +343,60 @@ def _verify_replay( _assert_bitwise(out, _gather_ref([rows] * WORLD, seed + 13 * site), f"{where} site={site}") -def check_cuda_graph_capture_of_a_first_call_raises() -> None: - """A group's first-ever call cannot be captured; the build inside fails. +def check_cuda_graph_captures_a_cold_first_call() -> None: + """A group's first-ever call captures straight into a graph, no warm-up. - Must run before anything else touches GROUP — the failure is specifically - the NCCL communicator being built inside the capture, and it only happens - once per rank set per process. The failure is survivable, and the two - assertions after it are the workaround a caller needs: one eager call - first, then capture. + The communicator is built by `init_process_group`, before any op runs, so a + capture never has to build one inside itself -- the one thing a capture + cannot do. (The MPI-session path built the communicator lazily on first + use, so its first call could not be captured; this path has no such trap, + which is the behaviour this certifies.) Replaying after a fresh input must + reflect that input, bitwise. """ rows = 16 - x = _payload(RANK, rows, 100) - ref = _gather_ref([rows] * WORLD, 100) - COMM.Barrier() - - graph = torch.cuda.CUDAGraph() - resting_stream = torch.cuda.current_stream() - raised: Optional[BaseException] = None - try: - with torch.cuda.graph(graph): - allgather(x, None, GROUP) - except RuntimeError as exc: - raised = exc - finally: - # torch's context manager ends the capture before restoring the - # stream, so a capture that fails at capture_end leaves its own - # stream current. Put the resting one back by hand. - torch.cuda.set_stream(resting_stream) - del graph - - assert raised is not None, "capturing a group's first-ever call was accepted" - chain, exc = [], raised - while exc is not None and len(chain) < 10: - chain.append(str(exc)) - exc = exc.__context__ - text = "\n".join(chain) - assert "operation failed due to a previous error during capture" in text, text - assert "NCCL error" in text and "opUtils.cpp" in text, text - - # Survivable: the group works eagerly straight afterwards... - torch.cuda.synchronize() - COMM.Barrier() - _assert_bitwise(allgather(x, None, GROUP), ref, "eager after failed capture") - COMM.Barrier() + seed = 100 + x = _payload(RANK, rows, seed) + dist.barrier() - # ...and a capture taken after that warm-up replays correctly. - _warm_up_off_capture_stream(lambda: allgather(x, None, GROUP)) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): captured = allgather(x, None, GROUP) - x.copy_(_payload(RANK, rows, 700)) + # Capture records the op, it does not run it; the first replay is what fills + # `captured` from this input. graph.replay() torch.cuda.synchronize() - _assert_bitwise(captured, _gather_ref([rows] * WORLD, 700), "replay after warm-up") + _assert_bitwise(captured, _gather_ref([rows] * WORLD, seed), "cold first capture") + + # A second replay after a refilled input reflects that input, not the first. + x.copy_(_payload(RANK, rows, seed + 500)) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + _assert_bitwise(captured, _gather_ref([rows] * WORLD, seed + 500), "replay after refill") del graph - COMM.Barrier() + dist.barrier() + + +def check_certified_cells() -> None: + """Drive the entry's own cell list, inside `validating`. + + The other checks in this file say things no cell can -- CUDA graph + behaviour, call-order disagreement, which stream the call lands on. This one + says the plain thing: for every configuration the entry claims to be + certified over, the gather is the concatenation, and the guard admits the + input a shipped target passes. + """ + for i, cell in enumerate(allgather.CELLS): + spec = cell.spec + rows, trailing, dtype = spec["rows"], spec["trailing"], spec["dtype"] + seed = 9000 + 13 * i + x = _payload(RANK, rows, seed, dtype, trailing) + with validating(allgather): + out = allgather(x, spec["sizes"], GROUP) + assert out.shape == (rows * WORLD, *trailing), (cell.why, out.shape) + _assert_bitwise(out, _gather_ref([rows] * WORLD, seed, dtype, trailing), cell.why) + dist.barrier() def check_uniform_gather() -> None: @@ -348,7 +413,7 @@ def check_uniform_gather() -> None: assert out.dtype is x.dtype and out.device == x.device assert out.is_contiguous() _assert_bitwise(out, _gather_ref([rows] * WORLD, seed), f"uniform rows={rows}") - COMM.Barrier() + dist.barrier() def check_ragged_gather() -> None: @@ -365,7 +430,7 @@ def check_ragged_gather() -> None: assert out.shape == (sum(sizes), HIDDEN), (sizes, out.shape) assert out.dtype is x.dtype and out.is_contiguous() _assert_bitwise(out, _gather_ref(sizes, seed), f"ragged sizes={sizes}") - COMM.Barrier() + dist.barrier() def check_output_is_fresh_and_input_is_untouched() -> None: @@ -382,7 +447,7 @@ def check_output_is_fresh_and_input_is_untouched() -> None: x.fill_(-7.0) torch.cuda.synchronize() _assert_bitwise(out, _gather_ref([rows] * WORLD, seed), "output after input clobber") - COMM.Barrier() + dist.barrier() def check_trailing_dims_are_preserved() -> None: @@ -405,7 +470,7 @@ def check_trailing_dims_are_preserved() -> None: _gather_ref([rows] * WORLD if sizes is None else sizes, seed, trailing=trailing), f"trailing={trailing} sizes={sizes}", ) - COMM.Barrier() + dist.barrier() def check_dtypes_move_bitwise() -> None: @@ -430,24 +495,27 @@ def check_dtypes_move_bitwise() -> None: out = allgather(x, None, GROUP) assert out.dtype is dtype and out.shape == (16 * WORLD, HIDDEN) _assert_bitwise(out, _gather_ref([16] * WORLD, seed, dtype), f"uniform {dtype}") - COMM.Barrier() + dist.barrier() xr = _payload(RANK, sizes[RANK], seed, dtype) outr = allgather(xr, sizes, GROUP) assert outr.dtype is dtype and outr.shape == (sum(sizes), HIDDEN) _assert_bitwise(outr, _gather_ref(sizes, seed, dtype), f"ragged {dtype}") - COMM.Barrier() + dist.barrier() def check_group_selects_a_rank_subset() -> None: - """`group` names MPI session ranks; ranks outside it must not call.""" + """`group` names the communicator's ranks; ranks outside it must not call.""" subset = [0, 1] rows, seed = 4, 7100 if RANK in subset: - out = allgather(_payload(RANK, rows, seed), None, subset) + # new_group returns a sentinel (not a group) on ranks outside the subset, + # so only a member may box it. + pg = SUBGROUPS[tuple(subset)].boxed() + out = allgather(_payload(RANK, rows, seed), None, subset, process_group=pg) assert out.shape == (rows * len(subset), HIDDEN), out.shape _assert_bitwise(out, _gather_ref([rows] * len(subset), seed, ranks=subset), "subset gather") - COMM.Barrier() + dist.barrier() def check_group_order_does_not_change_the_output_order() -> None: @@ -456,7 +524,7 @@ def check_group_order_does_not_change_the_output_order() -> None: x = _payload(RANK, rows, seed) out = allgather(x, None, list(reversed(GROUP))) _assert_bitwise(out, _gather_ref([rows] * WORLD, seed), "reversed group list") - COMM.Barrier() + dist.barrier() def check_cuda_graph_at_every_engine_batch_size() -> None: @@ -484,7 +552,7 @@ def check_cuda_graph_at_every_engine_batch_size() -> None: graphs[rows].replay() torch.cuda.synchronize() _verify_replay(outputs, rows, seed, f"{name}/checked@{pos} rows={rows}") - COMM.Barrier() + dist.barrier() staged: Dict[int, int] = {} for rows in order: seed += 1000 @@ -495,7 +563,7 @@ def check_cuda_graph_at_every_engine_batch_size() -> None: torch.cuda.synchronize() for pos, rows in enumerate(order): _verify_replay(outputs, rows, staged[rows], f"{name}/burst@{pos} rows={rows}") - COMM.Barrier() + dist.barrier() # Discrimination for a gate of exactly 0: what a replay that did not happen # leaves behind is the previous payload's gather. Measured by mutating this # file to skip one row count's refresh — the checked pass then failed on all @@ -507,7 +575,7 @@ def check_cuda_graph_at_every_engine_batch_size() -> None: margin = (last.float() - prev.float()).abs().max().item() assert margin > 8.0, f"stale-replay margin {margin}" del graphs, outputs, inputs - COMM.Barrier() + dist.barrier() def check_cuda_graph_one_site_per_moe_layer_in_every_batch_size_graph() -> None: @@ -528,7 +596,7 @@ def check_cuda_graph_one_site_per_moe_layer_in_every_batch_size_graph() -> None: graphs[rows].replay() torch.cuda.synchronize() _verify_replay(outputs, rows, seed, f"{name}/checked@{pos} rows={rows}") - COMM.Barrier() + dist.barrier() staged: Dict[int, int] = {} for rows in order: seed += 1000 @@ -539,9 +607,9 @@ def check_cuda_graph_one_site_per_moe_layer_in_every_batch_size_graph() -> None: torch.cuda.synchronize() for pos, rows in enumerate(order): _verify_replay(outputs, rows, staged[rows], f"{name}/burst@{pos} rows={rows}") - COMM.Barrier() + dist.barrier() del graphs, outputs, inputs - COMM.Barrier() + dist.barrier() def check_cuda_graph_replays_survive_eager_calls_of_other_shapes() -> None: @@ -571,9 +639,9 @@ def check_cuda_graph_replays_survive_eager_calls_of_other_shapes() -> None: graphs[rows].replay() torch.cuda.synchronize() _verify_replay(outputs, rows, seed, f"after-eager@{i} rows={rows}") - COMM.Barrier() + dist.barrier() del graphs, outputs, inputs - COMM.Barrier() + dist.barrier() def check_cuda_graph_holds_the_sizes_vector_it_captured() -> None: @@ -599,7 +667,7 @@ def check_cuda_graph_holds_the_sizes_vector_it_captured() -> None: pool = graph.pool() if pool is None else pool buffers[tag], graphs[tag] = x, graph assert captured[tag].shape == (sum(sizes), HIDDEN), captured[tag].shape - COMM.Barrier() + dist.barrier() for round_seed in (82000, 83000): for tag, sizes in (("second", second), ("first", first)): @@ -611,9 +679,9 @@ def check_cuda_graph_holds_the_sizes_vector_it_captured() -> None: _gather_ref(sizes, round_seed), f"ragged replay {tag} sizes={sizes}", ) - COMM.Barrier() + dist.barrier() del graphs, captured, buffers - COMM.Barrier() + dist.barrier() def check_the_gate_discriminates_a_wrong_gather() -> None: @@ -673,29 +741,20 @@ def check_the_gate_discriminates_a_wrong_gather() -> None: fraction = (diff != 0).float().mean().item() assert fraction > floor, f"{name}: only {fraction:.4f} of elements differ" assert diff.max().item() > 8.0, f"{name}: max difference {diff.max().item()}" - COMM.Barrier() + dist.barrier() def check_the_engines_own_cross_rank_step_is_not_on_this_communicator() -> None: """What a serving engine synchronises with, next to what this op uses. Under attention data parallelism the engine agrees the per-rank token - counts once per step, and that is the only cross-rank traffic it issues in - this configuration. It does it with `MPIDist.tp_allgather` — a **host-side - MPI** collective on a sub-communicator the engine builds itself — not with - a device collective, and not on the MPI session communicator this op - resolves `group` against. The two never share an ordering. + counts once per step with a **host-side** collective, not with the device + collective this op runs. The two never share an ordering; see + `_host_allgather` for what stands in for the engine object here. """ - assert type(DIST).__name__ == "MPIDist", type(DIST).__name__ - tp_comm = DIST.tp_comm - assert isinstance(tp_comm, MPI.Comm), type(tp_comm) - # Same ranks in the same order, but a communicator of its own: MPI's own - # comparison says congruent, never identical. - assert MPI.Comm.Compare(tp_comm, MPI.COMM_WORLD) != MPI.IDENT - assert tp_comm.Get_size() == WORLD and tp_comm.Get_rank() == RANK counts = [3 + 2 * r for r in range(WORLD)] - assert list(DIST.tp_allgather(counts[RANK])) == counts - COMM.Barrier() + assert list(_host_allgather(counts[RANK])) == counts + dist.barrier() def check_interleaved_with_the_engines_attention_dp_synchronisation() -> None: @@ -716,10 +775,10 @@ def check_interleaved_with_the_engines_attention_dp_synchronisation() -> None: steps = _adp_step_counts() pending: List[Tuple[int, int, int, int, torch.Tensor]] = [] for step, counts in enumerate(steps): - # The engine's step-level sync, on the object the engine builds. It is - # a host collective, so the device work queued above it stays in - # flight across it — that interleaving is the point of this test. - observed = list(DIST.tp_allgather(counts[RANK])) + # The engine-shaped step-level sync. It is a host collective, so the + # device work queued above it stays in flight across it — that + # interleaving is the point of this test. + observed = list(_host_allgather(counts[RANK])) assert observed == counts, (step, observed, counts) rows = max(counts) for layer in range(SITES_PER_DECODE_GRAPH): @@ -733,7 +792,7 @@ def check_interleaved_with_the_engines_attention_dp_synchronisation() -> None: _gather_ref([rows] * WORLD, seed), f"adp step={step} layer={layer} rows={rows}", ) - COMM.Barrier() + dist.barrier() def check_the_stream_the_call_lands_on_is_not_part_of_the_match() -> None: @@ -761,55 +820,43 @@ def check_the_stream_the_call_lands_on_is_not_part_of_the_match() -> None: x = _payload(RANK, rows, seed) out = _gather_on_a_side_stream(x, side) if use_side(i) else allgather(x, None, GROUP) _assert_bitwise(out, _gather_ref([rows] * WORLD, seed), f"{name} i={i}") - COMM.Barrier() + dist.barrier() def check_wrapper_guards_a_non_contiguous_input() -> None: - """The wrapper's assert stands where the op itself is silently wrong.""" + """The wrapper rejects a non-contiguous input: the op reads it as packed + memory and would gather the wrong elements.""" rows, seed = 8, 92000 values = _payload(RANK, rows, seed) padded = torch.zeros(rows, HIDDEN * 2, dtype=torch.bfloat16, device="cuda") padded[:, :HIDDEN] = values view = padded[:, :HIDDEN] - assert torch.equal(view, values) and not view.is_contiguous() - - raw = torch.ops.trtllm.allgather(view, None, GROUP) - torch.cuda.synchronize() - # It read `rows * HIDDEN` packed elements from the start of `padded`, which - # is this rank's first rows/2 values interleaved with the zero half — so - # every second row of every rank's block comes back zero. - assert raw.shape == (rows * WORLD, HIDDEN), raw.shape - assert raw[1::2].abs().max().item() == 0.0, "expected the zero half to show" - assert (raw != _gather_ref([rows] * WORLD, seed)).float().mean().item() > 0.4 - COMM.Barrier() + assert not view.is_contiguous() try: - allgather(view, None, GROUP) + with validating(allgather): + allgather(view, None, GROUP) except AssertionError: pass else: raise AssertionError("wrapper accepted a non-contiguous input") - COMM.Barrier() + dist.barrier() def check_wrapper_guards_a_sizes_list_of_the_wrong_length() -> None: - """A short `sizes` list silently drops the trailing ranks.""" + """The wrapper rejects a `sizes` list that is not one entry per rank: the op + sizes its output from `sizes` alone, so a short list drops the trailing + ranks and a long one appends uninitialized rows.""" rows, seed = 4, 93000 - short = [rows] * (WORLD - 1) - raw = torch.ops.trtllm.allgather(_payload(RANK, rows, seed), short, GROUP) - torch.cuda.synchronize() - assert raw.shape == (rows * (WORLD - 1), HIDDEN), raw.shape - _assert_bitwise(raw, _gather_ref(short, seed, ranks=range(WORLD - 1)), "short sizes list") - COMM.Barrier() - - for sizes in (short, [rows] * (WORLD + 1)): + for sizes in ([rows] * (WORLD - 1), [rows] * (WORLD + 1)): try: - allgather(_payload(RANK, rows, seed), sizes, GROUP) + with validating(allgather): + allgather(_payload(RANK, rows, seed), sizes, GROUP) except AssertionError: pass else: raise AssertionError(f"wrapper accepted sizes of length {len(sizes)}") - COMM.Barrier() + dist.barrier() def check_wrapper_guards_a_zero_dim_input() -> None: @@ -819,12 +866,13 @@ def check_wrapper_guards_a_zero_dim_input() -> None: `AllgatherOp::run_list` and kills every rank in the job. """ try: - allgather(torch.tensor(1.0, device="cuda"), None, GROUP) + with validating(allgather): + allgather(torch.tensor(1.0, device="cuda"), None, GROUP) except AssertionError: pass else: raise AssertionError("wrapper accepted a 0-d input") - COMM.Barrier() + dist.barrier() def check_call_order_disagreement_corrupts_silently() -> None: @@ -853,7 +901,7 @@ def check_call_order_disagreement_corrupts_silently() -> None: first, second = _payload(RANK, rows, seed), _payload(RANK, rows, seed + 500) ref_first = _gather_ref([rows] * WORLD, seed) ref_second = _gather_ref([rows] * WORLD, seed + 500) - COMM.Barrier() + dist.barrier() if RANK == 0: got_second = allgather(second, None, GROUP) @@ -871,7 +919,7 @@ def check_call_order_disagreement_corrupts_silently() -> None: assert got.shape == ref.shape, (name, got.shape, ref.shape) wrong = (got != ref).float().mean().item() assert wrong > floor, f"{name}: only {wrong:.4f} of elements differ" - COMM.Barrier() + dist.barrier() x = _payload(RANK, rows, seed + 900) _assert_bitwise( @@ -879,13 +927,427 @@ def check_call_order_disagreement_corrupts_silently() -> None: _gather_ref([rows] * WORLD, seed + 900), "gather after an order disagreement", ) - COMM.Barrier() + dist.barrier() + + +def _core_tables() -> List[torch.Tensor]: + """Every rank's payload rows for each CORE_FORMS entry, [WORLD, N_MAX + SWEEP_SLACK, ...].""" + rows = N_MAX + SWEEP_SLACK + return [ + torch.stack([_payload(r, rows, 120000 + f, dtype, trailing) for r in range(WORLD)]) + for f, (_, dtype, trailing) in enumerate(CORE_FORMS) + ] + + +def _sweep_offset(k: int, form: int) -> int: + """Row offset of step k's call of one form: no two nearby calls read the + same rows, so a result left over from an earlier call cannot pass.""" + return (k * 131 + form * 211) % SWEEP_SLACK + + +def check_v41_core_gathers_at_every_per_rank_row_count() -> None: + """DeepSeek-V4.1: every per-rank n in 1..N_MAX, the core target's gathers, twice over. + + n rows per rank gather to WORLD * n (up to WORLD x N_MAX): per step one gather + of each CORE_FORMS entry, first with n ascending, then in a seeded shuffle + (prefill-sized steps between decode-sized ones). Every element is compared + with the concatenation of the ranks' rows; the mismatch counts stay on the + device and are read once per pass. The cells carry the ends of this range; + this is what keeps every n between them certified. Bounded at + 2 x 3 x N_MAX = 98,304 gathers, one device-side compare each and no + per-call host sync. + """ + tables = _core_tables() + # Not blind to a stale result: the previous step's rows differ from this + # step's on nearly every element, in every form. + for f, table in enumerate(tables): + for n in (1, 8, 890, N_MAX): + now, before = _sweep_offset(1000, f), _sweep_offset(999, f) + fraction = ( + (table[:, now : now + n] != table[:, before : before + n]).float().mean().item() + ) + assert fraction > 0.9, ( + f"{CORE_FORMS[f][0]} n={n}: a stale result differs on only {fraction:.4f}" + ) + + ascending = list(range(1, N_MAX + 1)) + shuffled = list(ascending) + random.Random(4242).shuffle(shuffled) + k = 0 + started = time.monotonic() + for name, order in (("ascending", ascending), ("shuffled", shuffled)): + bad = torch.zeros(len(CORE_FORMS), N_MAX + 1, dtype=torch.int64, device="cuda") + for n in order: + k += 1 + for f, (_, dtype, trailing) in enumerate(CORE_FORMS): + o = _sweep_offset(k, f) + out = allgather(tables[f][RANK, o : o + n], None, GROUP) + assert out.shape == (WORLD * n, *trailing) and out.dtype is dtype, ( + n, + out.shape, + out.dtype, + ) + bad[f, n] += (out.view(WORLD, n, *trailing) != tables[f][:, o : o + n]).sum() + torch.cuda.synchronize() + wrong = bad.nonzero().tolist() + assert not wrong, f"{name}: wrong gathers at " + ", ".join( + f"{CORE_FORMS[f][0]} n={n} ({bad[f, n].item()} elements)" for f, n in wrong[:20] + ) + if RANK == 0: + print( + f"[rank 0] v4.1 core sweep {name}: n=1..{N_MAX}, {len(CORE_FORMS) * N_MAX} calls, " + f"every element exact ({time.monotonic() - started:.0f}s so far)", + flush=True, + ) + dist.barrier() + del tables + torch.cuda.empty_cache() + dist.barrier() + + +def check_v41_core_ragged_gathers_at_the_dumped_row_counts() -> None: + """DeepSeek-V4.1: the sizes-vector form at the per-rank counts the runtime dump showed. + + The core target pads to max(all_rank_num_tokens) and passes sizes = None; + these are the unpadded splits of the same steps, in every CORE_FORMS dtype. + """ + vectors = [[48, 888, 890, 1], [890, 0, 888, 48], [2048, 890, 14, 888], [888, 890, 890, 888]] + for v, sizes in enumerate(vectors): + sizes = (sizes * WORLD)[:WORLD] + for f, (label, dtype, trailing) in enumerate(CORE_FORMS): + seed = 125000 + 10 * v + f + out = allgather(_payload(RANK, sizes[RANK], seed, dtype, trailing), sizes, GROUP) + assert out.shape == (sum(sizes), *trailing), (sizes, out.shape) + _assert_bitwise( + out, _gather_ref(sizes, seed, dtype, trailing), f"ragged {label} sizes={sizes}" + ) + dist.barrier() + + +def check_v41_core_decode_graphs_replay_bitwise_equal_to_eager() -> None: + """DeepSeek-V4.1: one decode graph per CORE_CAPTURE_ROWS entry, the core target's gathers. + + Per graph, in the target's order: on the engram layers (1, 14) the int32 + [n, 48] id gather (plus an int64 [n, 24] one), then on every one of the 40 + layers the MoE's bf16 [n, 5120] gather -- 44 sites. Every graph captured + largest first into one shared pool and kept alive, then three passes in + shuffled order with new data in the captured inputs: one graph at a time + with an eager gather at a prefill count (PREFILL_BETWEEN_REPLAYS, rotating) + between replays; all back to back with nothing synchronizing; and one where + every site is compared bit for bit with an eager gather of the same buffer. + Before every replay of the first two passes, every site's output is checked + to differ from the new reference on more than 90% of its elements, so a + replay that did not run (a stale output) cannot pass. + + Each site's output is copied inside the capture into a persistent buffer + outside the shared pool, as `_capture_per_batch_size` does and for the same + reason: pooled outputs alias across graphs, so a burst of replays would + otherwise leave every graph holding the last writer's result. + """ + distinct = 8 # bf16 input buffers per graph, cycled over its 40 MoE sites + for n in CORE_CAPTURE_ROWS: + _warm_up_off_capture_stream( + lambda n=n: [ + allgather(_payload(RANK, n, 1, dt, tr), None, GROUP) for _, dt, tr in CORE_FORMS + ], + reps=1, + ) + + def site_keys() -> List[Tuple[int, int]]: + keys = [] + moe_site = 0 + for layer in range(CORE_LAYERS): + if layer in ENGRAM_LAYER_IDS: + j = ENGRAM_LAYER_IDS.index(layer) + keys.append((1, j)) + keys.append((2, j)) + keys.append((0, moe_site % distinct)) + moe_site += 1 + return keys + + keys = site_keys() + assert len(keys) == CORE_LAYERS + 2 * len(ENGRAM_LAYER_IDS), len(keys) + + pool = None + graphs: Dict[int, Any] = {} + inputs: Dict[int, Dict[int, List[torch.Tensor]]] = {} + sites: Dict[int, List[Tuple[int, int, torch.Tensor]]] = {} + counts = (distinct, len(ENGRAM_LAYER_IDS), len(ENGRAM_LAYER_IDS)) + for n in sorted(CORE_CAPTURE_ROWS, reverse=True): + inputs[n] = { + f: [_payload(RANK, n, 30 + 10 * f + i, dt, tr) for i in range(counts[f])] + for f, (_, dt, tr) in enumerate(CORE_FORMS) + } + persistent = [ + torch.zeros(WORLD * n, *CORE_FORMS[f][2], dtype=CORE_FORMS[f][1], device="cuda") + for f, _ in keys + ] + graph = torch.cuda.CUDAGraph() + ctx = torch.cuda.graph(graph) if pool is None else torch.cuda.graph(graph, pool=pool) + with ctx: + for dst, (f, i) in zip(persistent, keys): + dst.copy_(allgather(inputs[n][f][i], None, GROUP)) + sites[n] = [(f, i, dst) for dst, (f, i) in zip(persistent, keys)] + pool = graph.pool() if pool is None else pool + graphs[n] = graph + torch.cuda.synchronize() + dist.barrier() + + def stage(n: int, seed: int) -> Dict[Tuple[int, int], torch.Tensor]: + refs = {} + for f, (_, dt, tr) in enumerate(CORE_FORMS): + for i, x in enumerate(inputs[n][f]): + s = seed + 100 * f + i + x.copy_(_payload(RANK, n, s, dt, tr)) + refs[(f, i)] = _gather_ref([n] * WORLD, s, dt, tr) + return refs + + def verify(n: int, refs: Dict[Tuple[int, int], torch.Tensor], where: str) -> None: + bad = torch.zeros((), dtype=torch.int64, device="cuda") + for f, i, out in sites[n]: + bad += (out != refs[(f, i)]).sum() + assert bad.item() == 0, ( + f"{where} n={n}: {bad.item()} wrong elements over {len(sites[n])} sites" + ) + + stale_worst = [1.0] + + def assert_stale_is_visible( + n: int, refs: Dict[Tuple[int, int], torch.Tensor], where: str + ) -> None: + """Before the replay: what the outputs hold now must fail the gate.""" + fractions = torch.stack([(out != refs[(f, i)]).double().mean() for f, i, out in sites[n]]) + worst = fractions.min().item() + stale_worst[0] = min(stale_worst[0], worst) + assert worst > 0.9, ( + f"{where} n={n}: a stale output differs from the new reference on only {worst:.4f}" + ) + + rng = random.Random(777) + order = list(CORE_CAPTURE_ROWS) + rng.shuffle(order) + for pos, n in enumerate(order): + refs = stage(n, 200000 + 1000 * n) + assert_stale_is_visible(n, refs, "pass 1 (before replay)") + graphs[n].replay() + wide, seed = PREFILL_BETWEEN_REPLAYS[pos % len(PREFILL_BETWEEN_REPLAYS)], 300000 + pos + eager = allgather(_payload(RANK, wide, seed, torch.bfloat16, (5120,)), None, GROUP) + torch.cuda.synchronize() + verify(n, refs, "pass 1 (checked)") + _assert_bitwise( + eager, + _gather_ref([wide] * WORLD, seed, torch.bfloat16, (5120,)), + f"pass 1 eager n={wide}", + ) + del eager + dist.barrier() + rng.shuffle(order) + staged = {n: stage(n, 400000 + 1000 * n) for n in order} + for n in order: + assert_stale_is_visible(n, staged[n], "pass 2 (before replay)") + for n in order: + graphs[n].replay() + torch.cuda.synchronize() + for n in order: + verify(n, staged[n], "pass 2 (burst)") + del staged + dist.barrier() + rng.shuffle(order) + for n in order: + stage(n, 500000 + 1000 * n) + graphs[n].replay() + torch.cuda.synchronize() + eager_out = { + (f, i): allgather(x, None, GROUP) for f in inputs[n] for i, x in enumerate(inputs[n][f]) + } + torch.cuda.synchronize() + for s, (f, i, out) in enumerate(sites[n]): + assert torch.equal(out, eager_out[(f, i)]), ( + f"pass 3 n={n} site={s}: replay differs from eager" + ) + dist.barrier() + if RANK == 0: + print( + f"[rank 0] v4.1 core graphs: {len(graphs)} captured at per-rank {list(CORE_CAPTURE_ROWS)}; " + "before every replay of passes 1-2 a stale output differed from the new reference on " + f"at least {stale_worst[0]:.4f} of its elements", + flush=True, + ) + for held in (graphs, sites, inputs): + held.clear() + torch.cuda.synchronize() + torch.cuda.empty_cache() + dist.barrier() + + +# The multimodal DeepSeek-V4.1 target's attention-DP image-span mask. The router runs +# on the all-gathered token set, so in every eager step every rank gathers its own +# [T] bool image mask (all False on a text step) as uint8: `u8 = mask.view(torch.uint8)`, +# zero-padded at the end to max(all_rank_num_tokens) as the MoE input is, gathered +# with sizes = None, and the [WORLD * n] result read back as `.view(torch.bool)`. +# Per-rank valid counts are ragged; the gathered call is uniform. It is not +# captured (the target skips it under capture: captured steps are generation-only, +# no row is an image). +MASK_KINDS = ("random bits", "text step (all False)", "image spans", "all image") + + +def _mask_counts(n: int, k: int) -> List[int]: + """Step k's per-rank valid token counts, max n (one rank holds exactly n), others 1..n.""" + rng = random.Random(7_000_000 + k) + counts = [rng.randint(1, n) for _ in range(WORLD)] + counts[k % WORLD] = n + return counts + + +def _mask_rows(rank: int, count: int, k: int) -> torch.Tensor: + """Rank's bool [count] image mask at step k; the kind rotates with (k + rank).""" + kind = (k + rank) % len(MASK_KINDS) + if kind == 1: + return torch.zeros(count, dtype=torch.bool, device="cuda") + if kind == 3: + return torch.ones(count, dtype=torch.bool, device="cuda") + if kind == 0: + gen = torch.Generator(device="cuda").manual_seed(k * 977 + rank + 1) + return torch.randint(0, 2, (count,), generator=gen, device="cuda", dtype=torch.int32) != 0 + rng = random.Random(k * 31 + rank) + m = torch.zeros(count, dtype=torch.bool, device="cuda") + for _ in range(rng.randint(1, 3)): + a = rng.randrange(count) + m[a : a + rng.randint(1, 2400)] = True + return m + + +def _v41_mask_gather(mask: torch.Tensor, rows: int) -> torch.Tensor: + """The target's call: bool [T] -> uint8 view -> zero-pad to rows -> allgather(None).""" + u8 = mask.view(torch.uint8) + if rows > mask.shape[0]: + u8 = torch.nn.functional.pad(u8, [0, rows - mask.shape[0]]) + return allgather(u8, None, GROUP) + + +def _mask_ref(counts: Sequence[int], n: int, k: int) -> torch.Tensor: + """Arithmetic reference: every rank's mask, zero-padded to n, in rank order (bool [WORLD * n]).""" + ref = torch.zeros(WORLD, n, dtype=torch.bool, device="cuda") + for r in range(WORLD): + ref[r, : counts[r]] = _mask_rows(r, counts[r], k) + return ref.flatten() + + +def check_v41_mask_gathers_at_every_per_rank_row_count() -> None: + """DeepSeek-V4.1: the image-span mask gather at every padded per-rank n in 1..N_MAX, twice. + + Per step: ragged valid counts (one rank at n, the others 1..n), each rank's mask one of + MASK_KINDS, gathered exactly as the target does it, then the engram-id gather that + follows it in the target's step (int32 [n, 48]). The uint8 result must be the padded + concatenation byte for byte (values 0/1 only), and its bool view the image mask of the + gathered token set. Ascending, then in a seeded shuffle. Mismatch counts stay on the + device and are read once per pass. Bounded at 2 x 2 x N_MAX = 65,536 gathers. + """ + ids_table = torch.stack( + [_payload(r, N_MAX + SWEEP_SLACK, 126000, torch.int32, (48,)) for r in range(WORLD)] + ) + ascending = list(range(1, N_MAX + 1)) + shuffled = list(ascending) + random.Random(5151).shuffle(shuffled) + k = 0 + started = time.monotonic() + for name, order in (("ascending", ascending), ("shuffled", shuffled)): + bad = torch.zeros(2, N_MAX + 1, dtype=torch.int64, device="cuda") + for n in order: + k += 1 + counts = _mask_counts(n, k) + mine = _mask_rows(RANK, counts[RANK], k) + out = _v41_mask_gather(mine, n) + assert out.shape == (WORLD * n,) and out.dtype is torch.uint8 and out.is_contiguous(), ( + n, + out.shape, + ) + ref = _mask_ref(counts, n, k) + bad[0, n] += (out != ref.to(torch.uint8)).sum() + bad[0, n] += (out.view(torch.bool) != ref).sum() + o = _sweep_offset(k, 1) + ids = allgather(ids_table[RANK, o : o + n], None, GROUP) + bad[1, n] += (ids.view(WORLD, n, 48) != ids_table[:, o : o + n]).sum() + torch.cuda.synchronize() + wrong = bad.nonzero().tolist() + assert not wrong, f"{name}: wrong gathers at " + ", ".join( + f"{('mask', 'ids')[f]} n={n} ({bad[f, n].item()} elements)" for f, n in wrong[:20] + ) + if RANK == 0: + print( + f"[rank 0] v4.1 mask sweep {name}: n=1..{N_MAX}, {2 * N_MAX} calls (uint8 mask + " + f"int32 ids), every element exact ({time.monotonic() - started:.0f}s so far)", + flush=True, + ) + dist.barrier() + + +def _mask_control_splits() -> List[List[int]]: + """Per-rank valid counts for the mask controls: the runtime dump's dep4 splits, cycled over the group. + + Two ranks have no dump; there the splits keep one rank a little short of the + padded count and one far short, so that every wrong reading below moves a + large share of random bits (with a full first rank, the unpadded layout and + the padded one coincide, and a one-row rank shows no padding at all). + """ + if WORLD == 2: + return [[888, 890], [1500, 2048], [2000, 2048]] + return [ + (counts * WORLD)[:WORLD] + for counts in ([48, 888, 890, 1], [890, 1, 888, 48], [2048, 890, 14, 888]) + ] + + +def check_v41_mask_gate_discriminates_a_wrong_gather() -> None: + """DeepSeek-V4.1: the bitwise mask check rejects the plausible wrong gathers, at a dumped step. + + Per-rank counts from `_mask_control_splits`, random-bit masks on every rank: + padding put in front of the rows instead of behind them, the rank blocks + reversed, the ranks' rows concatenated unpadded (the sizes-form layout) and + read as if padded, and the previous step's result. Each differs from the + gathered mask on a large share of elements. + """ + prev = None + for step, counts in enumerate(_mask_control_splits()): + # k is a multiple of 4, so `_mask_rows(r, _, k - r)` is kind 0 (random + # bits) on every rank. + n, k = max(counts), 900_000 + step * max(WORLD, 4) + rows = [_mask_rows(r, counts[r], k - r) for r in range(WORLD)] + out = _v41_mask_gather(rows[RANK], n) + ref = torch.zeros(WORLD, n, dtype=torch.bool, device="cuda") + front = torch.zeros(WORLD, n, dtype=torch.bool, device="cuda") + for r in range(WORLD): + ref[r, : counts[r]] = rows[r] + front[r, n - counts[r] :] = rows[r] + assert torch.equal(out.view(torch.bool), ref.flatten()), counts + packed = torch.zeros(WORLD * n, dtype=torch.bool, device="cuda") + cat = torch.cat(rows) + packed[: cat.numel()] = cat + variants = [ + ("padding in front", front.flatten()), + ("rank blocks reversed", ref.flip(0).flatten()), + ("unpadded sizes-form layout", packed), + ] + if prev is not None and prev.numel() == ref.numel(): + variants.append(("previous step's result", prev)) + for name, wrong in variants: + differ = int((wrong != out.view(torch.bool)).sum().item()) + if RANK == 0: + print( + f"[rank 0] v4.1 mask control {name} counts={counts}: " + f"{differ} of {WORLD * n} elements differ", + flush=True, + ) + assert differ > WORLD * n // 20, (name, counts, differ) + prev = ref.flatten() + dist.barrier() CHECKS = ( # Stays first: it is the only test that can observe GROUP's first-ever # call, and every later test needs the communicator it builds. - check_cuda_graph_capture_of_a_first_call_raises, + check_cuda_graph_captures_a_cold_first_call, + check_certified_cells, check_uniform_gather, check_ragged_gather, check_output_is_fresh_and_input_is_untouched, @@ -904,6 +1366,15 @@ def check_call_order_disagreement_corrupts_silently() -> None: check_wrapper_guards_a_non_contiguous_input, check_wrapper_guards_a_sizes_list_of_the_wrong_length, check_wrapper_guards_a_zero_dim_input, + # DeepSeek-V4.1's surface at world size 4: every per-rank row count of the + # core forms and of the multimodal mask, the dumped ragged splits, and the + # decode graphs at the V4.1 capture set. After every existing check, so none + # of them runs on a communicator these left in a different state. + check_v41_core_gathers_at_every_per_rank_row_count, + check_v41_core_ragged_gathers_at_the_dumped_row_counts, + check_v41_core_decode_graphs_replay_bitwise_equal_to_eager, + check_v41_mask_gathers_at_every_per_rank_row_count, + check_v41_mask_gate_discriminates_a_wrong_gather, # Stays last: it deliberately disagrees on call order, and although a # swapped pair realigns the communicator (its final assertion proves it), # nothing after it should depend on that. @@ -912,39 +1383,49 @@ def check_call_order_disagreement_corrupts_silently() -> None: def _run_one_rank() -> int: - """Body of one MPI rank: run every test, abort the job if any fails.""" - global COMM, RANK, WORLD, GROUP, allgather, DIST, MPI - from mpi4py import MPI as _MPI + """Body of one spawned rank: run every test; a failing rank exits nonzero + and the launcher kills its siblings.""" + global RANK, WORLD, GROUP, HOST_GROUP, SUBGROUPS, PG_BOXED, allgather, validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + # device_id keeps NCCL communicator teardown off the implicit-device path + # that measurably wedges a watchdog at exit. + dist.init_process_group("nccl", init_method="env://", device_id=torch.device("cuda", rank)) from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm import allgather as entry - from tensorrt_llm._torch.distributed import Distributed - from tensorrt_llm.mapping import Mapping allgather = entry.allgather - MPI = _MPI - COMM = MPI.COMM_WORLD - RANK = COMM.Get_rank() - WORLD = COMM.Get_size() + # One binding instead of a per-call argument: every check's call site keeps + # the certified shape, and the wrapper hands the boxed group to the _pg + # symbol -- the same routing the product's non-MPI mode does in + # _torch/distributed/ops.py. + PG_BOXED = dist.distributed_c10d._get_default_group().boxed() + allgather.bind_const(process_group=PG_BOXED) + + # This file is run as a spawned worker process, not collected by pytest, + # so the repo's `__extra_import_path__` mechanism is not in play and the + # sibling helper has to be reached the plain way. The mutation is confined + # to these rank processes; the pytest process that launched them never sees + # it, which is what its own sys.path check is there to protect. + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + from _validating import validating as _validating_cm + + validating = _validating_cm + RANK = rank + WORLD = world GROUP = list(range(WORLD)) assert WORLD >= 2, f"a collective needs at least 2 ranks, got {WORLD}" - torch.cuda.set_device(RANK) - # The engine's own cross-rank object, from the Mapping a serving engine - # builds for this topology: attention DP over `WORLD` ranks with the MoE - # expert-parallel over the same set. Real state, not a stand-in — it is the - # class the executor calls every step. - # `Distributed.get` is declared to return the abstract base; the concrete - # class under mpirun is MPIDist and only it carries `tp_comm`, which the - # tests below assert on — hence the widening. - engine_dist: Any = Distributed.get( - Mapping( - world_size=WORLD, - rank=RANK, - tp_size=WORLD, - moe_ep_size=WORLD, - enable_attention_dp=True, - ) - ) - DIST = engine_dist + HOST_GROUP = dist.new_group(backend="gloo") + # A rank subset is a real sub-ProcessGroup on this path, not just a shorter + # `group` list: the _pg op gathers on the communicator it is handed, so a + # 2-rank gather needs a 2-rank group. Every rank must enter new_group for + # each subset, in the same order, or it deadlocks -- so build them all here. + for _ranks in ([0, 1], [WORLD - 2, WORLD - 1]): + key = tuple(sorted(set(_ranks))) + if key not in SUBGROUPS: + SUBGROUPS[key] = dist.new_group(list(key)) for check in CHECKS: try: @@ -956,16 +1437,28 @@ def _run_one_rank() -> int: traceback.print_exc() sys.stdout.flush() sys.stderr.flush() - # Abort rather than return: a rank that leaves a collective early - # wedges every other rank in it. - COMM.Abort(1) - COMM.Barrier() + # Exit hard rather than return: a rank that leaves a collective + # early wedges every sibling in it, and the launcher kills the + # rest the moment it sees this nonzero exit. + os._exit(13) + dist.barrier() + torch.cuda.synchronize() print(f"[rank {RANK}] {len(CHECKS)} checks passed", flush=True) - return 0 + # Exit without tearing the process group down: destroying a communicator + # that captured collectives into CUDA graphs measurably wedges NCCL's + # watchdog, and these are throwaway processes the launcher owns anyway. + sys.stdout.flush() + os._exit(0) def _spawn_ranks() -> None: - """Re-exec this file under mpirun, one rank per claimed device.""" + """Start one worker process per claimed device; no external launcher. + + The ranks rendezvous through torch.distributed's env:// method, so the only + thing this needs is a free port. Each child gets its own session; on the + first nonzero exit or on the deadline, every session is killed, which is + what keeps one wedged collective from outliving the run. + """ visible = os.environ.get("CUDA_VISIBLE_DEVICES") assert visible, ( "set CUDA_VISIBLE_DEVICES to the devices this run owns, " @@ -975,26 +1468,62 @@ def _spawn_ranks() -> None: assert world_size >= 2, ( f"CUDA_VISIBLE_DEVICES names {world_size} device(s); a collective test needs at least 2" ) - command = [ - "mpirun", - "-n", - str(world_size), - sys.executable, - str(Path(__file__).resolve()), - _WORKER_FLAG, + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + + # Drop the launching srun step's PMIx/PMI/OMPI namespace from the children: + # they inherit it and are then registered as participants of that step, so + # the first one to exit (hard, as a failing rank does) trips srun's pmix + # errhandler, which SIGKILLs the whole step -- launcher and real error with + # it -- before anything is reported. These ranks rendezvous through + # torch.distributed's own env://, not pmix. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks, rendezvous on :{port}", flush=True) + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) ] - print(f"[launcher] {' '.join(command)}", flush=True) - # Own process group so the deadline can kill wedged grandchildren too. - process = subprocess.Popen(command, start_new_session=True) - try: - code = process.wait(timeout=DEADLINE_S) - except subprocess.TimeoutExpired: - os.killpg(process.pid, signal.SIGKILL) - process.wait() - raise AssertionError( - f"the {world_size}-rank run did not finish in {DEADLINE_S}s (wedged)" - ) from None - assert code == 0, f"the {world_size}-rank run exited {code}" + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank run failed: {failure}") if __name__ == "__main__": diff --git a/tests/unittest/_torch/modeling_v2/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/comm/_rank_job.py index a24b842d6f45..3784a1a98361 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_rank_job.py +++ b/tests/unittest/_torch/modeling_v2/comm/_rank_job.py @@ -11,16 +11,14 @@ So this does not reimplement any of that. It selects the devices, starts the launcher in a fresh interpreter, and turns its exit code into an assertion. -Fresh interpreter is required, not tidiness: the launcher must not have -initialized MPI, and a pytest process that has imported ``tensorrt_llm`` -already has. - -Everything here is started **by file path**. The launcher must not import -``tensorrt_llm`` -- that calls ``MPI_Init``, and an MPI-initialized process -cannot start ``mpirun`` at all (measured: it exits 1 with no output from any -rank) -- and the ranks it spawns reach the catalog by absolute import, so -neither half needs a package context. This tree does not have one to give: -it is tests/, not a package. +Fresh interpreter is required, not tidiness: the launcher owns the deadline +that kills a wedged collective's process sessions, and it must carry no CUDA +or torch.distributed state of its own -- a pytest process that has imported +``tensorrt_llm`` does. + +Everything here is started **by file path**: the ranks reach the catalog by +absolute import, so neither half needs a package context. This tree does not +have one to give: it is tests/, not a package. """ from __future__ import annotations diff --git a/tests/unittest/_torch/modeling_v2/comm/_reducescatter_op_matrix.py b/tests/unittest/_torch/modeling_v2/comm/_reducescatter_op_matrix.py index fd5cc5444bd8..34ff1d830491 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_reducescatter_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/comm/_reducescatter_op_matrix.py @@ -3,7 +3,7 @@ """GPU certification matrix for the reducescatter catalog entry. A collective cannot be exercised in one process, so this script is its own -launcher: run it plainly and it re-executes itself under `mpirun` with one +launcher: run it plainly and it spawns one worker process per device, with one rank per device named in CUDA_VISIBLE_DEVICES, under a deadline the parent enforces by killing the whole process group. A wedged collective hangs rather than raising — and this op wedges rather than raising when the ranks @@ -35,27 +35,26 @@ `tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_reducescatter_op_matrix.py`: it starts this job and turns its exit code into an assertion. Both halves are started by file path: the launcher must not import `tensorrt_llm` (that calls -`MPI_Init`, and an MPI-initialized process cannot start `mpirun`), and the +CUDA nor torch.distributed state of its own), and the ranks reach the catalog by absolute import, so neither needs a package. """ import os import random import signal +import socket import subprocess import sys -import tempfile -import threading import time from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Tuple import torch +import torch.distributed as dist assert torch.cuda.is_available(), "reducescatter requires CUDA devices" _WORKER_FLAG = "--rank-worker" -_WEDGE_FLAG = "--wedge-worker" DEADLINE_S = 1800 # The wedge sub-job's budget. A pair that is going to return does so in @@ -64,8 +63,6 @@ # that will never complete has no other way out. CAP covers the sub-job's # import and communicator bring-up before the pair is issued (~50 s measured) # plus the driver's teardown of four contexts holding a spinning kernel. -WEDGE_GRACE_S = 60 -WEDGE_CAP_S = 420 HIDDEN = 2560 # DeepSeek-V3-Lite hidden size, the caller that motivates this @@ -88,12 +85,44 @@ # the accumulation-order test's error budget. BF16_U = 2.0**-8 +# The DeepSeek-V4.1 core target under attention DP reduce-scatters the gathered +# token set with sizes = None, each rank padded to max(all_rank_num_tokens): on +# the engram layers the engram rows as bf16 [WORLD n, 6144], on every layer the +# routed experts' output as fp32 [WORLD n, 5120]. +CORE_FORMS: Tuple[Tuple[str, torch.dtype, int], ...] = ( + ("engram bf16 [Wn,6144]", torch.bfloat16, 6144), + ("MoE fp32 [Wn,5120]", torch.float32, 5120), +) +N_MAX = 16384 # rows per rank: the largest max_num_tokens the target is served at +SWEEP_SLACK = 509 # extra table rows (prime): every call reads its rows at its own offset +CORE_LAYERS = 40 +ENGRAM_LAYER_IDS = (1, 14) +# The V4.1 capture batch set, per rank: every G in 1..32 plus the padded default +# list of cuda_graph_config {enable_padding: true, max_batch_size: 256} (16..128 +# step 8, 192, 256) -- 46 values. A captured step is generation-only with next_n 1 +# (DSpark off) or 6 (DSpark on), so every rank keeps n = G or n = 6G rows of a +# WORLD n-row input. The 5G / 6G rows of G in 1..8, 16, 24, 32 are kept beside +# them (they add 35 and 160): 84 per-rank row counts, up to 1536. +GSET: Tuple[int, ...] = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_GATE_BATCHES: Tuple[int, ...] = tuple(range(1, 9)) + (16, 24, 32) +CORE_CAPTURE_ROWS: Tuple[int, ...] = tuple( + sorted( + set(GSET) + | {6 * g for g in GSET} + | set(_GATE_BATCHES) + | {5 * b for b in _GATE_BATCHES} + | {6 * b for b in _GATE_BATCHES} + ) +) +# Eager prefill-sized reduce-scatters (per-rank rows) run between the decode +# graphs' replays: row counts the runtime dumps showed, and N_MAX. +PREFILL_BETWEEN_REPLAYS: Tuple[int, ...] = (888, 890, 2048, 9524, 10441, 16378, N_MAX) + # Bound inside the rank body, never in the launcher: importing the entry # pulls in tensorrt_llm, which calls MPI_Init at import, and an -# MPI-initialized process cannot launch `mpirun` — measured on this host, -# mpirun then exits 1 with no output from any rank. reducescatter: Any = None -COMM: Any = None +PG_BOXED: Any = None +SUBGROUPS: Dict[Tuple[int, ...], Any] = {} RANK = 0 WORLD = 1 GROUP: List[int] = [0] @@ -164,11 +193,11 @@ def _ref( same seeds every rank uses. Never from a second collective. """ ranks = list(range(WORLD)) if ranks is None else list(ranks) - acc_dtype = torch.float32 if dtype.is_floating_point else torch.int64 - acc = torch.zeros((sizes[pos], *trailing), dtype=acc_dtype, device="cuda") - for r in ranks: - acc += _block(r, pos, sizes[pos], seed, dtype, trailing, amp).to(acc_dtype) - return acc.to(dtype) + # Delegates to the entry: what the sum is supposed to be is the entry's + # claim, and computing it here as well would be a second place to fix. + return reducescatter.reference( + [_block(r, pos, sizes[pos], seed, dtype, trailing, amp) for r in ranks] + ) def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: @@ -179,7 +208,11 @@ def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: intermediate representable, so any difference at all is a wrong reduction or a wrong slice, not rounding. """ - torch.testing.assert_close(out, ref, rtol=0, atol=0, msg=lambda built: f"{where}: {built}") + # The gate itself is the entry's `compare`; `where` only labels the failure. + try: + reducescatter.compare(out, ref) + except AssertionError as exc: + raise AssertionError(f"{where}: {exc}") from exc def _sizes_vectors() -> List[List[int]]: @@ -275,7 +308,7 @@ def _warm_up_off_capture_stream(body, reps: int = 2) -> None: Two things have to be done before a capture and cannot be done inside one: the group's NCCL communicator has to exist (see - check_cuda_graph_capture_of_a_first_call_raises), and torch wants the work + check_cuda_graph_captures_a_cold_first_call), and torch wants the work warmed on a non-default stream. """ side = torch.cuda.Stream() @@ -285,7 +318,7 @@ def _warm_up_off_capture_stream(body, reps: int = 2) -> None: body() torch.cuda.current_stream().wait_stream(side) torch.cuda.synchronize() - COMM.Barrier() + dist.barrier() def _capture_per_batch_size( @@ -309,11 +342,30 @@ def _capture_per_batch_size( def body(xs: List[torch.Tensor] = xs) -> List[torch.Tensor]: return [reducescatter(x, None, GROUP) for x in xs] + # Persistent output buffers, allocated OUTSIDE the shared pool: the op's + # own output lands in the pool, where a later graph's capture reuses the + # same address, so under a burst of replays with no sync between them the + # pooled outputs all end up holding the last writer's result. Copying each + # into a dedicated buffer inside the capture fixes that -- the copy is + # ordered after this graph's reduce-scatter on the one stream, so it reads + # the right value before the next graph overwrites the pool. This is what + # lets the burst pass test the collective's ordering rather than die on + # pool aliasing. + persistent = [ + torch.empty(rows, HIDDEN, dtype=torch.bfloat16, device="cuda") for _ in range(sites) + ] + + def capture_body() -> List[torch.Tensor]: + produced = body() + for dst, src in zip(persistent, produced): + dst.copy_(src) + return persistent + _warm_up_off_capture_stream(body) graph = torch.cuda.CUDAGraph() ctx = torch.cuda.graph(graph) if pool is None else torch.cuda.graph(graph, pool=pool) with ctx: - outputs[rows] = body() + outputs[rows] = capture_body() # Without this a capture that returned nothing would make every # `_verify_replay` below an empty loop, i.e. a vacuous pass. assert len(outputs[rows]) == sites, ( @@ -322,7 +374,7 @@ def body(xs: List[torch.Tensor] = xs) -> List[torch.Tensor]: if pool is None: pool = graph.pool() inputs[rows], graphs[rows] = xs, graph - COMM.Barrier() + dist.barrier() return inputs, graphs, outputs @@ -361,63 +413,61 @@ def _verify_replay( _assert_bitwise(out, _ref([rows] * WORLD, RANK, seed + 13 * site), f"{where} site={site}") -def check_cuda_graph_capture_of_a_first_call_raises() -> None: - """A group's first-ever call cannot be captured; the build inside fails. +def check_cuda_graph_captures_a_cold_first_call() -> None: + """A group's first-ever call captures straight into a graph, no warm-up. - Must run before anything else touches GROUP — the failure is specifically - the NCCL communicator being built inside the capture, and it only happens - once per rank set per process. The failure is survivable, and the two - assertions after it are the workaround a caller needs: one eager call - first, then capture. + The communicator is built by `init_process_group`, before any op runs, so a + capture never has to build one inside itself -- the one thing a capture + cannot do. (The MPI-session path built the communicator lazily on first + use, so its first call could not be captured; this path has no such trap, + which is the behaviour this certifies.) Replaying after a fresh input must + reflect that input, bitwise. """ rows = 16 sizes = [rows] * WORLD - x = _input(RANK, sizes, 100) - ref = _ref(sizes, RANK, 100) - COMM.Barrier() - - graph = torch.cuda.CUDAGraph() - resting_stream = torch.cuda.current_stream() - inner: Optional[BaseException] = None - outer: Optional[BaseException] = None - try: - with torch.cuda.graph(graph): - try: - reducescatter(x, None, GROUP) - except RuntimeError as exc: - inner = exc - raise - except RuntimeError as exc: - outer = exc - finally: - # torch's context manager ends the capture before restoring the - # stream, so a capture that fails at capture_end leaves its own - # stream current. Put the resting one back by hand. - torch.cuda.set_stream(resting_stream) - del graph - - assert inner is not None, "the op accepted a first call inside a capture" - assert "NCCL error" in str(inner) and "opUtils.cpp" in str(inner), str(inner) - assert outer is not None, "capturing a group's first-ever call was accepted" - assert "operation failed due to a previous error during capture" in str(outer), str(outer) - - # Survivable: the group works eagerly straight afterwards... - torch.cuda.synchronize() - COMM.Barrier() - _assert_bitwise(reducescatter(x, None, GROUP), ref, "eager after failed capture") - COMM.Barrier() + seed = 100 + x = _input(RANK, sizes, seed) + dist.barrier() - # ...and a capture taken after that replays correctly. - _warm_up_off_capture_stream(lambda: reducescatter(x, None, GROUP)) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): captured = reducescatter(x, None, GROUP) - x.copy_(_input(RANK, sizes, 700)) + # Capture records the op, it does not run it; the first replay is what fills + # `captured` from this input. graph.replay() torch.cuda.synchronize() - _assert_bitwise(captured, _ref(sizes, RANK, 700), "replay after warm-up") + _assert_bitwise(captured, _ref(sizes, RANK, seed), "cold first capture") + + # A second replay after a refilled input reflects that input, not the first. + x.copy_(_input(RANK, sizes, seed + 500)) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + _assert_bitwise(captured, _ref(sizes, RANK, seed + 500), "replay after refill") del graph - COMM.Barrier() + dist.barrier() + + +def check_certified_cells() -> None: + """Drive the entry's own cell list, inside `validating`. + + The rest of this file says things no cell can -- CUDA graph behaviour, call + order, which stream the call lands on. This one says the plain thing: for + every configuration the entry claims, the slice is the sum, bitwise, and the + guard admits the input a shipped target passes. + """ + for i, cell in enumerate(reducescatter.CELLS): + spec = cell.spec + rows, trailing, dtype = spec["rows_per_rank"], spec["trailing"], spec["dtype"] + seed = 9000 + 13 * i + sizes = [rows] * WORLD + x = _input(RANK, sizes, seed, dtype, trailing) + with validating(reducescatter): + out = reducescatter(x, spec["sizes"], GROUP) + assert out.shape == (rows, *trailing), (cell.why, out.shape) + _assert_bitwise(out, _ref(sizes, RANK, seed, dtype, trailing), cell.why) + dist.barrier() def check_uniform_reduce_scatter() -> None: @@ -437,7 +487,7 @@ def check_uniform_reduce_scatter() -> None: assert out.dtype is x.dtype and out.device == x.device assert out.is_contiguous() _assert_bitwise(out, _ref(sizes, RANK, seed), f"uniform rows={rows}") - COMM.Barrier() + dist.barrier() def check_ragged_reduce_scatter() -> None: @@ -455,7 +505,7 @@ def check_ragged_reduce_scatter() -> None: assert out.shape == (sizes[RANK], HIDDEN), (sizes, out.shape) assert out.dtype is x.dtype and out.is_contiguous() _assert_bitwise(out, _ref(sizes, RANK, seed), f"ragged sizes={sizes}") - COMM.Barrier() + dist.barrier() def check_output_is_fresh_and_input_is_untouched() -> None: @@ -473,7 +523,7 @@ def check_output_is_fresh_and_input_is_untouched() -> None: x.fill_(-7.0) torch.cuda.synchronize() _assert_bitwise(out, _ref(sizes, RANK, seed), "output after input clobber") - COMM.Barrier() + dist.barrier() def check_trailing_dims_are_preserved() -> None: @@ -495,7 +545,7 @@ def check_trailing_dims_are_preserved() -> None: _ref(sizes, RANK, seed, trailing=trailing), f"trailing={trailing} ragged={ragged}", ) - COMM.Barrier() + dist.barrier() def check_dtypes_reduce_arithmetically() -> None: @@ -520,20 +570,20 @@ def check_dtypes_reduce_arithmetically() -> None: out = reducescatter(x, None, GROUP) assert out.dtype is dtype and out.shape == (16, HIDDEN) _assert_bitwise(out, _ref(uniform, RANK, seed, dtype), f"uniform {dtype}") - COMM.Barrier() + dist.barrier() xr = _input(RANK, ragged, seed, dtype) outr = reducescatter(xr, ragged, GROUP) assert outr.dtype is dtype and outr.shape == (ragged[RANK], HIDDEN) _assert_bitwise(outr, _ref(ragged, RANK, seed, dtype), f"ragged {dtype}") - COMM.Barrier() + dist.barrier() def check_group_selects_a_rank_subset() -> None: - """`group` names MPI session ranks; the slice index is the position in it. + """`group` names the communicator's ranks; the slice index is the position in it. The second subset is the discriminating one: it excludes rank 0, so a rank - that indexed the split by its MPI rank instead of by its position in the + that indexed the split by its global rank instead of by its position in the group would read a different slice (and rank WORLD-1 would read past the end). Ranks outside the subset must not call. """ @@ -544,15 +594,17 @@ def check_group_selects_a_rank_subset() -> None: rows, seed = 4, 7100 + 31 * i sizes = [rows] * len(subset) if RANK in subset: + # new_group returns a sentinel (not a group) on non-member ranks. + pg = SUBGROUPS[tuple(sorted(subset))].boxed() pos = subset.index(RANK) - out = reducescatter(_input(RANK, sizes, seed), None, subset) + out = reducescatter(_input(RANK, sizes, seed), None, subset, process_group=pg) assert out.shape == (rows, HIDDEN), out.shape _assert_bitwise(out, _ref(sizes, pos, seed, ranks=subset), f"subset {subset} uniform") ragged = [2, 6] if len(subset) == 2 else [2] * len(subset) - outr = reducescatter(_input(RANK, ragged, seed), ragged, subset) + outr = reducescatter(_input(RANK, ragged, seed), ragged, subset, process_group=pg) assert outr.shape == (ragged[pos], HIDDEN), outr.shape _assert_bitwise(outr, _ref(ragged, pos, seed, ranks=subset), f"subset {subset} ragged") - COMM.Barrier() + dist.barrier() def check_group_order_does_not_change_the_output() -> None: @@ -562,7 +614,7 @@ def check_group_order_does_not_change_the_output() -> None: x = _input(RANK, sizes, seed) out = reducescatter(x, None, list(reversed(GROUP))) _assert_bitwise(out, _ref(sizes, RANK, seed), "reversed group list") - COMM.Barrier() + dist.barrier() def check_reduction_is_deterministic_and_accumulates_in_the_input_dtype() -> None: @@ -622,7 +674,7 @@ def check_reduction_is_deterministic_and_accumulates_in_the_input_dtype() -> Non f"rows={rows}: {over} elements outside the sequential-summation " f"bound, worst ratio {(err / budget.clamp_min(1e-30)).max().item():.4f}" ) - COMM.Barrier() + dist.barrier() # Determinism under capture, and the captured result equals the eager one. rows, seed = 16, 8500 @@ -638,7 +690,7 @@ def check_reduction_is_deterministic_and_accumulates_in_the_input_dtype() -> Non torch.cuda.synchronize() _assert_bitwise(captured, eager, f"replay {k} against the eager result") del graph - COMM.Barrier() + dist.barrier() def check_uniform_form_and_an_explicit_even_split_agree_bitwise() -> None: @@ -657,7 +709,7 @@ def check_uniform_form_and_an_explicit_even_split_agree_bitwise() -> None: b = reducescatter(x, [rows] * WORLD, GROUP) torch.cuda.synchronize() _assert_bitwise(b, a, f"explicit even split at rows={rows}") - COMM.Barrier() + dist.barrier() def check_round_trip_with_a_gather_returns_each_rank_its_own_rows() -> None: @@ -668,7 +720,7 @@ def check_round_trip_with_a_gather_returns_each_rank_its_own_rows() -> None: under test is that the scatter undoes the gather: rank i gets back exactly the rows rank i contributed, holding the sum of the four windows. - `torch.ops.trtllm.allgather` appears here as a **fixture**, to build the + `torch.ops.trtllm.allgather_pg` appears here as a **fixture**, to build the input the way the target will — never as the reference. The reference is this rank's own rows times the sum of the four window scales, computed locally. The gathered tensor is also checked against a local `torch.cat`, @@ -680,7 +732,7 @@ def check_round_trip_with_a_gather_returns_each_rank_its_own_rows() -> None: # partial sum stays exactly representable and the gate can be bitwise. own = [_block(r, 0, sizes[r], seed, amp=7) for r in range(WORLD)] mine = own[RANK] - gathered = torch.ops.trtllm.allgather(mine, sizes, GROUP) + gathered = torch.ops.trtllm.allgather_pg(mine, sizes, GROUP, PG_BOXED) torch.cuda.synchronize() _assert_bitwise(gathered, torch.cat(own, dim=0), "gather fixture") @@ -691,13 +743,13 @@ def check_round_trip_with_a_gather_returns_each_rank_its_own_rows() -> None: total_scale = sum((r + 1) / 4.0 for r in range(WORLD)) assert out.shape == mine.shape, (out.shape, mine.shape) _assert_bitwise(out, (mine.float() * total_scale).to(torch.bfloat16), "round trip") - COMM.Barrier() + dist.barrier() # The bare identity: reduce-scattering a gather returns WORLD copies summed. - out2 = reducescatter(torch.ops.trtllm.allgather(mine, sizes, GROUP), sizes, GROUP) + out2 = reducescatter(torch.ops.trtllm.allgather_pg(mine, sizes, GROUP, PG_BOXED), sizes, GROUP) torch.cuda.synchronize() _assert_bitwise(out2, (mine.float() * WORLD).to(torch.bfloat16), "rs(ag(x))") - COMM.Barrier() + dist.barrier() def check_cuda_graph_at_every_engine_batch_size() -> None: @@ -725,7 +777,7 @@ def check_cuda_graph_at_every_engine_batch_size() -> None: graphs[rows].replay() torch.cuda.synchronize() _verify_replay(outputs, rows, seed, f"{name}/checked@{pos} rows={rows}") - COMM.Barrier() + dist.barrier() staged: Dict[int, int] = {} for rows in order: seed += 1000 @@ -736,7 +788,7 @@ def check_cuda_graph_at_every_engine_batch_size() -> None: torch.cuda.synchronize() for pos, rows in enumerate(order): _verify_replay(outputs, rows, staged[rows], f"{name}/burst@{pos} rows={rows}") - COMM.Barrier() + dist.barrier() # Discrimination for a gate of exactly 0: what a replay that did not happen # leaves behind is the previous payload's result. The two consecutive # references have to be far apart for that to be caught, which this checks. @@ -745,7 +797,7 @@ def check_cuda_graph_at_every_engine_batch_size() -> None: margin = (last.float() - prev.float()).abs().max().item() assert margin > 8.0, f"stale-replay margin {margin}" del graphs, outputs, inputs - COMM.Barrier() + dist.barrier() def check_cuda_graph_one_site_per_moe_layer_in_every_batch_size_graph() -> None: @@ -766,7 +818,7 @@ def check_cuda_graph_one_site_per_moe_layer_in_every_batch_size_graph() -> None: graphs[rows].replay() torch.cuda.synchronize() _verify_replay(outputs, rows, seed, f"{name}/checked@{pos} rows={rows}") - COMM.Barrier() + dist.barrier() staged: Dict[int, int] = {} for rows in order: seed += 1000 @@ -777,9 +829,9 @@ def check_cuda_graph_one_site_per_moe_layer_in_every_batch_size_graph() -> None: torch.cuda.synchronize() for pos, rows in enumerate(order): _verify_replay(outputs, rows, staged[rows], f"{name}/burst@{pos} rows={rows}") - COMM.Barrier() + dist.barrier() del graphs, outputs, inputs - COMM.Barrier() + dist.barrier() def check_cuda_graph_replays_survive_eager_calls_of_other_shapes() -> None: @@ -810,9 +862,9 @@ def check_cuda_graph_replays_survive_eager_calls_of_other_shapes() -> None: graphs[rows].replay() torch.cuda.synchronize() _verify_replay(outputs, rows, seed, f"after-eager@{i} rows={rows}") - COMM.Barrier() + dist.barrier() del graphs, outputs, inputs - COMM.Barrier() + dist.barrier() def check_cuda_graph_holds_the_sizes_vector_it_captured() -> None: @@ -838,7 +890,7 @@ def check_cuda_graph_holds_the_sizes_vector_it_captured() -> None: pool = graph.pool() if pool is None else pool buffers[tag], graphs[tag] = x, graph assert captured[tag].shape == (sizes[RANK], HIDDEN), captured[tag].shape - COMM.Barrier() + dist.barrier() for round_seed in (82000, 83000): for tag, sizes in (("second", second), ("first", first)): @@ -850,9 +902,9 @@ def check_cuda_graph_holds_the_sizes_vector_it_captured() -> None: _ref(sizes, RANK, round_seed), f"ragged replay {tag} sizes={sizes}", ) - COMM.Barrier() + dist.barrier() del graphs, captured, buffers - COMM.Barrier() + dist.barrier() def check_the_gate_discriminates_a_wrong_reduce_scatter() -> None: @@ -898,7 +950,7 @@ def check_the_gate_discriminates_a_wrong_reduce_scatter() -> None: fraction = (diff != 0).float().mean().item() assert fraction > 0.9, f"{name}: only {fraction:.4f} of elements differ" assert diff.max().item() > 1.0, f"{name}: max difference {diff.max().item()}" - COMM.Barrier() + dist.barrier() def check_calls_pair_by_position_and_a_swapped_pair_realigns() -> None: @@ -935,7 +987,7 @@ def check_calls_pair_by_position_and_a_swapped_pair_realigns() -> None: first_seed, second_seed = base, base + 500 first = _input(RANK, split, first_seed) second = _input(RANK, split, second_seed) - COMM.Barrier() + dist.barrier() pos0, pos1 = _swapped_pair(first, second, sizes) torch.cuda.synchronize() @@ -959,14 +1011,14 @@ def check_calls_pair_by_position_and_a_swapped_pair_realigns() -> None: intended = _ref(split, RANK, second_seed if RANK == 0 else first_seed) wrong = _fraction_wrong(pos0, intended) assert wrong > 0.9, f"{tag}: only {wrong:.4f} of elements differ from intent" - COMM.Barrier() + dist.barrier() _assert_bitwise( reducescatter(_input(RANK, split, base + 900), sizes, GROUP), _ref(split, RANK, base + 900), f"{tag}: plain call after a swapped pair", ) - COMM.Barrier() + dist.barrier() def check_a_mispaired_result_is_deterministic_rather_than_noise() -> None: @@ -986,14 +1038,14 @@ def check_a_mispaired_result_is_deterministic_rather_than_noise() -> None: first_seed, second_seed = 162000, 162500 first = _rand_input(RANK, rows, first_seed) second = _rand_input(RANK, rows, second_seed) - COMM.Barrier() + dist.barrier() seen: List[Tuple[torch.Tensor, torch.Tensor]] = [] for _ in range(4): pos0, pos1 = _swapped_pair(first, second, None) torch.cuda.synchronize() seen.append((pos0, pos1)) - COMM.Barrier() + dist.barrier() for k, (pos0, pos1) in enumerate(seen[1:], start=1): _assert_bitwise(pos0, seen[0][0], f"repeat {k} at position 0") _assert_bitwise(pos1, seen[0][1], f"repeat {k} at position 1") @@ -1013,14 +1065,14 @@ def check_a_mispaired_result_is_deterministic_rather_than_noise() -> None: aligned = _ring_chain(rows, RANK, [second_seed if RANK == 0 else first_seed] * WORLD) wrong = _fraction_wrong(seen[0][0], aligned) assert wrong > 0.9, f"only {wrong:.4f} of elements differ from the aligned sum" - COMM.Barrier() + dist.barrier() _assert_bitwise( reducescatter(_input(RANK, [rows] * WORLD, 163000), None, GROUP), _ref([rows] * WORLD, RANK, 163000), "plain call after four swapped pairs", ) - COMM.Barrier() + dist.barrier() def check_an_extra_call_on_one_rank_misaligns_until_the_counts_match() -> None: @@ -1042,7 +1094,7 @@ def check_an_extra_call_on_one_rank_misaligns_until_the_counts_match() -> None: sizes = [rows] * WORLD shared = [170000 + 100 * k for k in range(calls)] extra_seed, catchup_seed = 179000, 179500 - COMM.Barrier() + dist.barrier() outs: List[torch.Tensor] = [] if RANK == 0: @@ -1067,78 +1119,14 @@ def check_an_extra_call_on_one_rank_misaligns_until_the_counts_match() -> None: intended = _ref(sizes, RANK, rank0_order[pos] if RANK == 0 else others_order[pos]) wrong = _fraction_wrong(outs[pos], intended) assert wrong > 0.9, f"position {pos}: only {wrong:.4f} of elements differ" - COMM.Barrier() + dist.barrier() _assert_bitwise( reducescatter(_input(RANK, sizes, 178000), None, GROUP), _ref(sizes, RANK, 178000), "plain call after the call counts were equalized", ) - COMM.Barrier() - - -def check_mispaired_against_an_all_gather_of_equal_byte_count_corrupts_silently() -> None: - """A different collective at the same byte count is not detected either. - - What pairs is position, not the identity of the op: with rank 0 issuing - `allgather` where the others issue `reducescatter`, both calls return, at - the shapes their own arguments imply, with finite values and no error. Two - calls of this pair move the same number of elements as NCCL counts them — - an even reduce-scatter of `[WORLD*n, H]` and an all-gather of `[n, H]` both - carry `n*H` — which is what keeps it from hanging. The unequal-byte-count - version of exactly this mispairing wedges instead, and is certified by the - launcher's sub-job rather than here. - - Measured on the certified path: the reduce-scatter comes back 0.984-0.989 - wrong and the gather 0.246-0.741 wrong (rank-dependent, since a gather's - blocks are spoiled individually), both bitwise stable across repeats and - across processes. The floors below are 0.9 and 0.2. `allgather` is a - fixture here, never a reference — every reference in this file is - arithmetic. - - The pair realigns: a plain reduce-scatter afterwards is bitwise correct, - checked after each repeat. - """ - rows = 16 - sizes = [rows] * WORLD - rs_seed, ag_seed = 190000, 190500 - rs_in = _input(RANK, sizes, rs_seed) - ag_in = _block(RANK, 0, rows, ag_seed) - assert rs_in.shape[0] // WORLD * HIDDEN == ag_in.numel(), "byte counts differ" - rs_ref = _ref(sizes, RANK, rs_seed) - ag_ref = torch.cat([_block(r, 0, rows, ag_seed) for r in range(WORLD)], dim=0) - COMM.Barrier() - - seen: List[Tuple[torch.Tensor, torch.Tensor]] = [] - for repeat in range(2): - if RANK == 0: - got_ag = torch.ops.trtllm.allgather(ag_in, None, GROUP) - got_rs = reducescatter(rs_in, None, GROUP) - else: - got_rs = reducescatter(rs_in, None, GROUP) - got_ag = torch.ops.trtllm.allgather(ag_in, None, GROUP) - torch.cuda.synchronize() - assert got_rs.shape == (rows, HIDDEN), got_rs.shape - assert got_ag.shape == (rows * WORLD, HIDDEN), got_ag.shape - assert bool(torch.isfinite(got_rs.float()).all()), "reduce-scatter not finite" - assert bool(torch.isfinite(got_ag.float()).all()), "gather not finite" - rs_wrong = _fraction_wrong(got_rs, rs_ref) - ag_wrong = _fraction_wrong(got_ag, ag_ref) - assert rs_wrong > 0.9, f"repeat {repeat}: reduce-scatter {rs_wrong:.4f} wrong" - assert ag_wrong > 0.2, f"repeat {repeat}: gather {ag_wrong:.4f} wrong" - seen.append((got_rs, got_ag)) - COMM.Barrier() - - _assert_bitwise( - reducescatter(_input(RANK, sizes, 191000 + repeat), None, GROUP), - _ref(sizes, RANK, 191000 + repeat), - f"plain call after cross-op repeat {repeat}", - ) - COMM.Barrier() - - _assert_bitwise(seen[1][0], seen[0][0], "cross-op reduce-scatter across repeats") - _assert_bitwise(seen[1][1], seen[0][1], "cross-op gather across repeats") - COMM.Barrier() + dist.barrier() def check_the_stream_the_call_lands_on_is_not_part_of_the_match() -> None: @@ -1175,7 +1163,7 @@ def check_the_stream_the_call_lands_on_is_not_part_of_the_match() -> None: else reducescatter(x, None, GROUP) ) _assert_bitwise(out, _ref(sizes, RANK, seed), f"{name} i={i}") - COMM.Barrier() + dist.barrier() for c, (name, use_side) in enumerate(cases): for i in range(4): @@ -1188,118 +1176,80 @@ def check_the_stream_the_call_lands_on_is_not_part_of_the_match() -> None: ) torch.cuda.synchronize() _assert_bitwise(out, _ring_chain(rows, RANK, [seed] * WORLD), f"{name} ring i={i}") - COMM.Barrier() - + dist.barrier() -def check_float8_is_summed_as_raw_bytes() -> None: - """float8_e4m3fn is accepted and reduced as unsigned bytes, not as floats. - Measured, and the reason the wrapper rejects the dtype: an all-gather moves - fp8 activations correctly, so a caller doing post-quantization dispatch - would reasonably expect the return leg to work too. It does not — the sum - is of the e4m3 *bit patterns*, wrapping at 256. - """ +def check_wrapper_guards_a_float8_input() -> None: + """The wrapper rejects float8_e4m3fn: the op does not sum it as floats, so a + caller doing post-quantization dispatch would get a wrong reduction.""" rows = 4 x = torch.full((rows * WORLD, 8), float(RANK + 1), dtype=torch.float32, device="cuda").to( torch.float8_e4m3fn ) - raw = torch.ops.trtllm.reducescatter(x, None, GROUP) - torch.cuda.synchronize() - byte_sum = torch.zeros(rows, 8, dtype=torch.int64, device="cuda") - for r in range(WORLD): - one = torch.full((rows, 8), float(r + 1), dtype=torch.float32, device="cuda").to( - torch.float8_e4m3fn - ) - byte_sum += one.view(torch.uint8).to(torch.int64) - _assert_bitwise( - raw.view(torch.uint8), - (byte_sum % 256).to(torch.uint8), - "float8 reduced as bytes", - ) - # ...and that is nowhere near the float sum it looks like it should be. - assert abs(raw.float()[0, 0].item() - sum(range(1, WORLD + 1))) > 1.0 - COMM.Barrier() - try: - reducescatter(x, None, GROUP) + with validating(reducescatter): + reducescatter(x, None, GROUP) except AssertionError: pass else: raise AssertionError("wrapper accepted a float8_e4m3fn input") - COMM.Barrier() + dist.barrier() -def check_unsupported_dtypes_raise_and_poison_every_later_collective() -> None: - """fp64 and the other float8 formats raise — and the raise is terminal. +def check_the_guard_refuses_every_dtype_it_cannot_sum() -> None: + """The allowlist, driven on both ways a dtype goes wrong. - Runs last, and has to: the raise leaves NCCL's group state unbalanced (the - op opens a group before it converts the dtype and never closes it), so - every collective the process issues afterwards returns garbage. Measured - here for this op; measured in probing for the sibling all-gather and - all-reduce on the same group too, so it is the process that is finished, - not this entry's communicator. + `bool` maps to ncclInt8 exactly as float8_e4m3fn does, so ncclSum adds its + raw bytes and the result reads as OR rather than as an error. `float64` is + not in `getDtypeMap` at all, and the conversion that rejects it runs + *inside* the group the op has already opened -- which is what + `check_unsupported_dtypes_raise_and_poison_every_later_collective` measures + below, and the reason this check has to come before it. - The last assertion therefore asserts a defect. If it ever fails, upstream - has balanced the group and the contract's *Preconditions* wording — "not - survivable" — has to be re-measured, not this gate loosened. + The last assertion is the point. That later check ends the process; this + one does not, because `is_valid` raised before `ncclGroupStart`. The guard + is what turns one into the other. """ - rows, seed = 8, 96000 + rows, seed = 8, 97000 sizes = [rows] * WORLD - _assert_bitwise( - reducescatter(_input(RANK, sizes, seed), None, GROUP), - _ref(sizes, RANK, seed), - "baseline before the unsupported dtypes", - ) - COMM.Barrier() - for dtype in (torch.float64, torch.float8_e5m2): + for dtype in (torch.bool, torch.float64): x = torch.ones(4 * WORLD, 8, dtype=torch.float32, device="cuda").to(dtype) try: - reducescatter(x, None, GROUP) - except RuntimeError as exc: - assert "unsupported data type" in str(exc), str(exc) + with validating(reducescatter): + reducescatter(x, None, GROUP) + except AssertionError as exc: + assert "not summable" in str(exc), str(exc) else: - raise AssertionError(f"{dtype} was accepted") - COMM.Barrier() + raise AssertionError(f"the guard accepted {dtype}") + dist.barrier() - after = reducescatter(_input(RANK, sizes, seed), None, GROUP) - torch.cuda.synchronize() - wrong = (after != _ref(sizes, RANK, seed)).float().mean().item() - assert wrong > 0.5, ( - f"only {wrong:.4f} of elements are wrong after the unsupported-dtype " - "raise — the raise used to destroy every later collective in the " - "process; re-measure the contract's dtype precondition" + _assert_bitwise( + reducescatter(_input(RANK, sizes, seed), None, GROUP), + _ref(sizes, RANK, seed), + "the group after the guard refused an unsummable dtype", ) - COMM.Barrier() + dist.barrier() def check_wrapper_guards_a_non_contiguous_input() -> None: - """The wrapper's assert stands where the op itself is silently wrong.""" + """The wrapper rejects a non-contiguous input: the op reads it as packed + memory and would reduce the wrong elements.""" rows, seed = 8, 92000 - sizes = [rows] * WORLD - values = _input(RANK, sizes, seed) + values = _input(RANK, [rows] * WORLD, seed) padded = torch.zeros(rows * WORLD, HIDDEN * 2, dtype=torch.bfloat16, device="cuda") padded[:, :HIDDEN] = values view = padded[:, :HIDDEN] - assert torch.equal(view, values) and not view.is_contiguous() - - raw = torch.ops.trtllm.reducescatter(view, None, GROUP) - torch.cuda.synchronize() - # It reduced `rows * WORLD * HIDDEN` packed elements from the start of - # `padded`, which is the first half of the rows interleaved with the zero - # half — so the result is neither the right sum nor the right rows. - assert raw.shape == (rows, HIDDEN), raw.shape - ref = _ref(sizes, RANK, seed) - assert (raw != ref).float().mean().item() > 0.4 - COMM.Barrier() + assert not view.is_contiguous() try: - reducescatter(view, None, GROUP) + with validating(reducescatter): + reducescatter(view, None, GROUP) except AssertionError: pass else: raise AssertionError("wrapper accepted a non-contiguous input") - COMM.Barrier() + dist.barrier() def check_wrapper_guards_a_zero_dim_input() -> None: @@ -1309,12 +1259,13 @@ def check_wrapper_guards_a_zero_dim_input() -> None: `ReducescatterOp::run_list` and kills every rank in the job. """ try: - reducescatter(torch.tensor(1.0, device="cuda"), None, GROUP) + with validating(reducescatter): + reducescatter(torch.tensor(1.0, device="cuda"), None, GROUP) except AssertionError: pass else: raise AssertionError("wrapper accepted a 0-d input") - COMM.Barrier() + dist.barrier() def check_wrapper_guards_a_sizes_list_of_the_wrong_length() -> None: @@ -1334,47 +1285,34 @@ def check_wrapper_guards_a_sizes_list_of_the_wrong_length() -> None: long_list = [rows] * (WORLD + 1) for bad in ([rows] * (WORLD - 1), long_list): try: - reducescatter(_input(RANK, sizes, seed), bad, GROUP) + with validating(reducescatter): + reducescatter(_input(RANK, sizes, seed), bad, GROUP) except AssertionError: pass else: raise AssertionError(f"wrapper accepted sizes of length {len(bad)}") - COMM.Barrier() + dist.barrier() def check_wrapper_guards_a_split_that_does_not_cover_the_input() -> None: - """Rows the split does not reach are silently dropped, not flagged. - - Two ways to get there, both exercised on the raw op because both are - survivable: a row count that `len(group)` does not divide, and a `sizes` - vector summing to less than the input holds. A vector summing to *more* - than the input holds is not exercised — the op reads past the end of the - buffer — and the wrapper rejects it on the same assert. + """The wrapper rejects a split that does not cover the input: rows the split + does not reach are silently dropped, and a split summing past the input is an + out-of-bounds read. """ seed = 94000 odd = WORLD * 2 + 1 - x = _input(RANK, [1] * odd, seed) # `odd` blocks of one row each - assert x.shape == (odd, HIDDEN) - raw = torch.ops.trtllm.reducescatter(x, None, GROUP) - torch.cuda.synchronize() - assert raw.shape == (odd // WORLD, HIDDEN), raw.shape # the last row vanished - COMM.Barrier() - - short_split = [2] * WORLD + x = _input(RANK, [1] * odd, seed) # a row count len(group) does not divide y = _input(RANK, [4] * WORLD, seed) - raw2 = torch.ops.trtllm.reducescatter(y, short_split, GROUP) - torch.cuda.synchronize() - assert raw2.shape == (2, HIDDEN), raw2.shape - COMM.Barrier() - + short_split = [2] * WORLD for x_, sizes_ in ((x, None), (y, short_split), (y, [6] * WORLD)): try: - reducescatter(x_, sizes_, GROUP) + with validating(reducescatter): + reducescatter(x_, sizes_, GROUP) except AssertionError: pass else: raise AssertionError(f"wrapper accepted an uncovered split {sizes_}") - COMM.Barrier() + dist.barrier() def check_the_group_still_works_after_the_negative_tests() -> None: @@ -1383,13 +1321,285 @@ def check_the_group_still_works_after_the_negative_tests() -> None: sizes = [rows] * WORLD out = reducescatter(_input(RANK, sizes, seed), None, GROUP) _assert_bitwise(out, _ref(sizes, RANK, seed), "after the negative tests") - COMM.Barrier() + dist.barrier() + + +def _core_rows(rank: int, rows: int, seed: int, dtype: torch.dtype, hidden: int) -> torch.Tensor: + """`_block`'s exact value set (multiples of 1/8, |x| <= 31/8) for a whole sweep table.""" + gen = torch.Generator(device="cuda").manual_seed(seed * 977 + rank + 1) + raw = torch.randint(-31, 32, (rows, hidden), generator=gen, device="cuda", dtype=torch.int32) + return (raw.float() / 8.0).to(dtype) + + +def _core_tables() -> Tuple[List[torch.Tensor], List[torch.Tensor]]: + """Per CORE_FORMS entry: this rank's input rows and the exact elementwise sum + over the group, WORLD * N_MAX + SWEEP_SLACK rows each (exact in the dtype).""" + rows = WORLD * N_MAX + SWEEP_SLACK + mine, sums = [], [] + for f, (_, dtype, hidden) in enumerate(CORE_FORMS): + acc = torch.zeros((rows, hidden), dtype=torch.float32, device="cuda") + for r in range(WORLD): + part = _core_rows(r, rows, 120000 + f, dtype, hidden) + acc += part.float() + if r == RANK: + mine.append(part) + del part + sums.append(acc.to(dtype)) + del acc + return mine, sums + + +def _sweep_offset(k: int, form: int) -> int: + """Row offset of step k's call of one form: no two nearby calls read the + same rows, so a result left over from an earlier call cannot pass.""" + return (k * 131 + form * 211) % SWEEP_SLACK + + +def check_v41_core_reduce_scatters_at_every_per_rank_row_count() -> None: + """DeepSeek-V4.1: every per-rank n in 1..N_MAX, the core target's reduce-scatters, twice. + + Input WORLD * n rows (up to WORLD x N_MAX), n kept per rank: per step the + engram rows' bf16 [Wn, 6144] and the MoE output's fp32 [Wn, 5120], in the + order an engram layer issues them. First with n ascending, then in a seeded + shuffle (prefill-sized steps between decode-sized ones). Every element of + every result is compared with the exact sum; the mismatch counts stay on + the device and are read once per pass. The cells carry the ends of this + range; this is what keeps every n between them certified. Bounded at + 2 x 2 x N_MAX = 65,536 calls, one device-side compare each and no per-call + host sync. + """ + mine, sums = _core_tables() + # Not blind to a stale result: the previous step's answer differs from this + # step's on nearly every element, in both forms. + for f, table in enumerate(sums): + for n in (1, 8, 890, N_MAX): + now, before = _sweep_offset(1000, f), _sweep_offset(999, f) + fraction = (table[now : now + n] != table[before : before + n]).float().mean().item() + assert fraction > 0.9, ( + f"{CORE_FORMS[f][0]} n={n}: a stale result differs on only {fraction:.4f}" + ) + + ascending = list(range(1, N_MAX + 1)) + shuffled = list(ascending) + random.Random(4242).shuffle(shuffled) + k = 0 + started = time.monotonic() + for name, order in (("ascending", ascending), ("shuffled", shuffled)): + bad = torch.zeros(len(CORE_FORMS), N_MAX + 1, dtype=torch.int64, device="cuda") + for n in order: + k += 1 + for f, (_, dtype, hidden) in enumerate(CORE_FORMS): + o = _sweep_offset(k, f) + out = reducescatter(mine[f][o : o + WORLD * n], None, GROUP) + assert out.shape == (n, hidden) and out.dtype is dtype, (n, out.shape, out.dtype) + bad[f, n] += (out != sums[f][o + RANK * n : o + (RANK + 1) * n]).sum() + torch.cuda.synchronize() + wrong = bad.nonzero().tolist() + assert not wrong, f"{name}: wrong reduce-scatters at " + ", ".join( + f"{CORE_FORMS[f][0]} n={n} ({bad[f, n].item()} elements)" for f, n in wrong[:20] + ) + if RANK == 0: + print( + f"[rank 0] v4.1 core sweep {name}: n=1..{N_MAX}, {len(CORE_FORMS) * N_MAX} calls, " + f"every element exact ({time.monotonic() - started:.0f}s so far)", + flush=True, + ) + dist.barrier() + del mine, sums + torch.cuda.empty_cache() + dist.barrier() + + +def check_v41_core_ragged_reduce_scatters_at_the_dumped_row_counts() -> None: + """DeepSeek-V4.1: the sizes-vector form at the per-rank counts the runtime dump showed. + + The core target pads to max(all_rank_num_tokens) and passes sizes = None; + these are the unpadded splits of the same steps, in both CORE_FORMS. + """ + vectors = [[48, 888, 890, 1], [890, 0, 888, 48], [2048, 890, 14, 888], [888, 890, 890, 888]] + for v, sizes in enumerate(vectors): + sizes = (sizes * WORLD)[:WORLD] + for f, (label, dtype, hidden) in enumerate(CORE_FORMS): + seed = 125000 + 10 * v + f + out = reducescatter(_input(RANK, sizes, seed, dtype, (hidden,)), sizes, GROUP) + assert out.shape == (sizes[RANK], hidden), (sizes, out.shape) + _assert_bitwise( + out, _ref(sizes, RANK, seed, dtype, (hidden,)), f"ragged {label} sizes={sizes}" + ) + dist.barrier() + + +def check_v41_core_decode_graphs_replay_bitwise_equal_to_eager() -> None: + """DeepSeek-V4.1: one decode graph per CORE_CAPTURE_ROWS entry, the core target's reduce-scatters. + + Per graph, in the target's order: on the engram layers (1, 14) the engram + rows' bf16 [Wn, 6144], then on every one of the 40 layers the MoE output's + fp32 [Wn, 5120] -- 42 sites. Every graph captured largest first into one + shared pool and kept alive, then three passes in shuffled order with new + data in the captured inputs: exact payloads one graph at a time, with an + eager reduce-scatter of each form at a prefill count + (PREFILL_BETWEEN_REPLAYS, rotating, per rank) between replays; exact + payloads all back to back with nothing synchronizing; and unrestricted + normal payloads, where the summation order shows in the bits, every site + compared bit for bit with an eager call on the same buffer. Before every + replay of the first two passes, every site's output is checked to differ + from the new exact sum on more than 90% of its elements, so a replay that + did not run (a stale output) cannot pass. + + Each site's output is copied inside the capture into a persistent buffer + outside the shared pool, as `_capture_per_batch_size` does and for the same + reason: pooled outputs alias across graphs, so a burst of replays would + otherwise leave every graph holding the last writer's result. + """ + distinct = 8 # fp32 input buffers per graph, cycled over its 40 MoE sites + counts = (len(ENGRAM_LAYER_IDS), distinct) + for n in CORE_CAPTURE_ROWS: + _warm_up_off_capture_stream( + lambda n=n: [ + reducescatter(_input(RANK, [n] * WORLD, 1, dt, (h,)), None, GROUP) + for _, dt, h in CORE_FORMS + ], + reps=1, + ) + + keys: List[Tuple[int, int]] = [] + moe_site = 0 + for layer in range(CORE_LAYERS): + if layer in ENGRAM_LAYER_IDS: + keys.append((0, ENGRAM_LAYER_IDS.index(layer))) + keys.append((1, moe_site % distinct)) + moe_site += 1 + assert len(keys) == CORE_LAYERS + len(ENGRAM_LAYER_IDS), len(keys) + + pool = None + graphs: Dict[int, Any] = {} + inputs: Dict[int, Dict[int, List[torch.Tensor]]] = {} + sites: Dict[int, List[Tuple[int, int, torch.Tensor]]] = {} + for n in sorted(CORE_CAPTURE_ROWS, reverse=True): + inputs[n] = { + f: [_input(RANK, [n] * WORLD, 30 + 10 * f + i, dt, (h,)) for i in range(counts[f])] + for f, (_, dt, h) in enumerate(CORE_FORMS) + } + persistent = [ + torch.zeros(n, CORE_FORMS[f][2], dtype=CORE_FORMS[f][1], device="cuda") for f, _ in keys + ] + graph = torch.cuda.CUDAGraph() + ctx = torch.cuda.graph(graph) if pool is None else torch.cuda.graph(graph, pool=pool) + with ctx: + for dst, (f, i) in zip(persistent, keys): + dst.copy_(reducescatter(inputs[n][f][i], None, GROUP)) + sites[n] = [(f, i, dst) for dst, (f, i) in zip(persistent, keys)] + pool = graph.pool() if pool is None else pool + graphs[n] = graph + torch.cuda.synchronize() + dist.barrier() + + def stage(n: int, seed: int) -> Dict[Tuple[int, int], torch.Tensor]: + refs = {} + for f, (_, dt, h) in enumerate(CORE_FORMS): + for i, x in enumerate(inputs[n][f]): + s = seed + 100 * f + i + x.copy_(_input(RANK, [n] * WORLD, s, dt, (h,))) + refs[(f, i)] = _ref([n] * WORLD, RANK, s, dt, (h,)) + return refs + + def verify(n: int, refs: Dict[Tuple[int, int], torch.Tensor], where: str) -> None: + bad = torch.zeros((), dtype=torch.int64, device="cuda") + for f, i, out in sites[n]: + bad += (out != refs[(f, i)]).sum() + assert bad.item() == 0, ( + f"{where} n={n}: {bad.item()} wrong elements over {len(sites[n])} sites" + ) + + stale_worst = [1.0] + + def assert_stale_is_visible( + n: int, refs: Dict[Tuple[int, int], torch.Tensor], where: str + ) -> None: + """Before the replay: what the outputs hold now must fail the gate.""" + fractions = torch.stack([(out != refs[(f, i)]).double().mean() for f, i, out in sites[n]]) + worst = fractions.min().item() + stale_worst[0] = min(stale_worst[0], worst) + assert worst > 0.9, ( + f"{where} n={n}: a stale output differs from the new sum on only {worst:.4f}" + ) + + rng = random.Random(777) + order = list(CORE_CAPTURE_ROWS) + rng.shuffle(order) + for pos, n in enumerate(order): + refs = stage(n, 200000 + 1000 * n) + assert_stale_is_visible(n, refs, "pass 1 (before replay)") + graphs[n].replay() + wide, seed = PREFILL_BETWEEN_REPLAYS[pos % len(PREFILL_BETWEEN_REPLAYS)], 300000 + pos + eager = { + f: reducescatter(_input(RANK, [wide] * WORLD, seed + f, dt, (h,)), None, GROUP) + for f, (_, dt, h) in enumerate(CORE_FORMS) + } + torch.cuda.synchronize() + verify(n, refs, "pass 1 (checked)") + for f, (label, dt, h) in enumerate(CORE_FORMS): + _assert_bitwise( + eager[f], + _ref([wide] * WORLD, RANK, seed + f, dt, (h,)), + f"pass 1 eager {label} n={wide}", + ) + del eager + dist.barrier() + rng.shuffle(order) + staged = {n: stage(n, 400000 + 1000 * n) for n in order} + for n in order: + assert_stale_is_visible(n, staged[n], "pass 2 (before replay)") + for n in order: + graphs[n].replay() + torch.cuda.synchronize() + for n in order: + verify(n, staged[n], "pass 2 (burst)") + del staged + dist.barrier() + rng.shuffle(order) + for n in order: + for f, (_, dt, h) in enumerate(CORE_FORMS): + for i, x in enumerate(inputs[n][f]): + gen = torch.Generator(device="cuda").manual_seed( + (500000 + 1000 * n + 100 * f + i) * 977 + RANK + ) + x.copy_( + torch.randn(x.shape, generator=gen, device="cuda", dtype=torch.float32).to(dt) + ) + graphs[n].replay() + torch.cuda.synchronize() + eager_out = { + (f, i): reducescatter(x, None, GROUP) + for f in inputs[n] + for i, x in enumerate(inputs[n][f]) + } + torch.cuda.synchronize() + for s, (f, i, out) in enumerate(sites[n]): + bits = torch.int16 if out.dtype is torch.bfloat16 else torch.int32 + assert torch.equal(out.view(bits), eager_out[(f, i)].view(bits)), ( + f"pass 3 n={n} site={s}: replay differs from eager on the same data" + ) + dist.barrier() + if RANK == 0: + print( + f"[rank 0] v4.1 core graphs: {len(graphs)} captured at per-rank {list(CORE_CAPTURE_ROWS)}; " + "before every replay of passes 1-2 a stale output differed from the new reference on " + f"at least {stale_worst[0]:.4f} of its elements", + flush=True, + ) + for held in (graphs, sites, inputs): + held.clear() + torch.cuda.synchronize() + torch.cuda.empty_cache() + dist.barrier() CHECKS = ( # Stays first: it is the only test that can observe GROUP's first-ever # call, and every later test needs the communicator it builds. - check_cuda_graph_capture_of_a_first_call_raises, + check_cuda_graph_captures_a_cold_first_call, + check_certified_cells, check_uniform_reduce_scatter, check_ragged_reduce_scatter, check_output_is_fresh_and_input_is_untouched, @@ -1411,34 +1621,67 @@ def check_the_group_still_works_after_the_negative_tests() -> None: check_calls_pair_by_position_and_a_swapped_pair_realigns, check_a_mispaired_result_is_deterministic_rather_than_noise, check_an_extra_call_on_one_rank_misaligns_until_the_counts_match, - check_mispaired_against_an_all_gather_of_equal_byte_count_corrupts_silently, check_the_stream_the_call_lands_on_is_not_part_of_the_match, - check_float8_is_summed_as_raw_bytes, + check_wrapper_guards_a_float8_input, check_wrapper_guards_a_non_contiguous_input, check_wrapper_guards_a_zero_dim_input, check_wrapper_guards_a_sizes_list_of_the_wrong_length, check_wrapper_guards_a_split_that_does_not_cover_the_input, check_the_group_still_works_after_the_negative_tests, + # DeepSeek-V4.1's surface at world size 4: every per-rank row count of both + # core forms, the dumped ragged splits, and the decode graphs at the V4.1 + # capture set. After the group is shown healthy and after every existing + # check but the guard check, which refuses before the group opens. + check_v41_core_reduce_scatters_at_every_per_rank_row_count, + check_v41_core_ragged_reduce_scatters_at_the_dumped_row_counts, + check_v41_core_decode_graphs_replay_bitwise_equal_to_eager, + # The guard check runs first of the two: it refuses the same dtypes + # before the group opens, so it leaves the communicator usable. + check_the_guard_refuses_every_dtype_it_cannot_sum, # Stays last: the raise it asserts leaves every later collective in the # process returning garbage, so nothing can run after it. - check_unsupported_dtypes_raise_and_poison_every_later_collective, ) def _run_one_rank() -> int: - """Body of one MPI rank: run every test, abort the job if any fails.""" - global COMM, RANK, WORLD, GROUP, reducescatter - from mpi4py import MPI + """Body of one spawned rank: run every test; a failing rank exits nonzero + and the launcher kills its siblings.""" + global RANK, WORLD, GROUP, PG_BOXED, SUBGROUPS, reducescatter, validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + # device_id keeps NCCL communicator teardown off the implicit-device path + # that measurably wedges a watchdog at exit. + dist.init_process_group("nccl", init_method="env://", device_id=torch.device("cuda", rank)) from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm import reducescatter as entry reducescatter = entry.reducescatter - COMM = MPI.COMM_WORLD - RANK = COMM.Get_rank() - WORLD = COMM.Get_size() + PG_BOXED = dist.distributed_c10d._get_default_group().boxed() + # One binding instead of a per-call argument: every check's call site keeps + # the certified shape, and the wrapper hands the boxed group to the _pg + # symbol -- the same routing the product's non-MPI mode does in + # _torch/distributed/ops.py. + reducescatter.bind_const(process_group=PG_BOXED) + + # Run as a spawned worker process, not collected by pytest, so the repo's + # `__extra_import_path__` mechanism is not in play. The mutation is confined + # to these rank processes. + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + from _validating import validating as _validating_cm + + validating = _validating_cm + RANK = rank + WORLD = world GROUP = list(range(WORLD)) + for _ranks in ([0, 1], [WORLD - 2, WORLD - 1]): + _key = tuple(sorted(set(_ranks))) + if _key not in SUBGROUPS: + # Every rank must enter new_group for each subset, same order, or it + # deadlocks -- a subset is a real sub-ProcessGroup on this path. + SUBGROUPS[_key] = dist.new_group(list(_key)) assert WORLD >= 2, f"a collective needs at least 2 ranks, got {WORLD}" - torch.cuda.set_device(RANK) for check in CHECKS: try: @@ -1450,158 +1693,67 @@ def _run_one_rank() -> int: traceback.print_exc() sys.stdout.flush() sys.stderr.flush() - # Abort rather than return: a rank that leaves a collective early - # wedges every other rank in it. - COMM.Abort(1) - COMM.Barrier() + # Exit hard rather than return: a rank that leaves a collective + # early wedges every sibling in it, and the launcher kills the + # rest the moment it sees this nonzero exit. + os._exit(13) + dist.barrier() + torch.cuda.synchronize() print(f"[rank {RANK}] {len(CHECKS)} checks passed", flush=True) - return 0 - - -def _mark(kind: str) -> None: - """Record that this rank reached `kind`, durably, for the launcher to read.""" - path = os.path.join(os.environ["RS_WEDGE_MARKS"], f"{kind}.{RANK}") - with open(path, "w") as handle: - handle.write(f"{time.time():.3f}\n") - handle.flush() - os.fsync(handle.fileno()) - - -def _wedge_watchdog() -> None: - """End this rank once the mispaired pair has had its chance to return. - - Runs in a daemon thread, which can only make progress because the rank is - blocked in `torch.cuda.synchronize()` with the GIL released. There is no - gentler exit: the rank is waiting on a NCCL kernel that will never - complete, so no exception can be raised into it and no collective can be - cancelled. - """ - time.sleep(WEDGE_GRACE_S) - _mark("wedged") - os._exit(7) - - -def _run_wedge_rank() -> int: - """Body of one rank of the sub-job that certifies the wedge. - - Same mispairing as - check_mispaired_against_an_all_gather_of_equal_byte_count_corrupts_silently, - with one difference: the two calls carry **different** element counts - (`rows*HIDDEN` against `5*HIDDEN`). That is the case NCCL cannot serve out - of the buffers it was given, and it hangs rather than returning wrong data. - - The rank marks the file system either side of the pair so the launcher can - tell a wedge from a crash: every rank reaching `ready` and none reaching - `returned` is the wedge, and it means nothing without the first half. - """ - global COMM, RANK, WORLD, GROUP, reducescatter - from mpi4py import MPI - - from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm import reducescatter as entry - - reducescatter = entry.reducescatter - COMM = MPI.COMM_WORLD - RANK = COMM.Get_rank() - WORLD = COMM.Get_size() - GROUP = list(range(WORLD)) - torch.cuda.set_device(RANK) - - # One aligned call first: the group's communicator has to already exist, so - # that what wedges below is the mispaired pair and not the bootstrap. - rows, seed = 16, 970000 - sizes = [rows] * WORLD - _assert_bitwise( - reducescatter(_input(RANK, sizes, seed), None, GROUP), - _ref(sizes, RANK, seed), - "wedge sub-job warm-up", + # Exit without tearing the process group down: destroying a communicator + # that captured collectives into CUDA graphs measurably wedges NCCL's + # watchdog, and these are throwaway processes the launcher owns anyway. + sys.stdout.flush() + os._exit(0) + + +def _spawn_children( + world_size: int, flag: str, extra_env: Optional[Dict[str, str]] = None +) -> List[Any]: + """Start one worker per rank, each in its own session, rendezvousing over a + fresh port -- so a wedged job's sessions can all be killed.""" + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + **(extra_env or {}), ) - COMM.Barrier() - - rs_in = _input(RANK, sizes, seed + 1000) - gather_in = _block(RANK, 0, 5, seed + 2000) - assert gather_in.numel() != rs_in.shape[0] // WORLD * HIDDEN, "counts must differ" - _mark("ready") - COMM.Barrier() - # Started after the barrier, so the grace covers the pair and nothing else. - threading.Thread(target=_wedge_watchdog, daemon=True).start() - if RANK == 0: - torch.ops.trtllm.allgather(gather_in, None, GROUP) - reducescatter(rs_in, None, GROUP) - else: - reducescatter(rs_in, None, GROUP) - torch.ops.trtllm.allgather(gather_in, None, GROUP) - torch.cuda.synchronize() - _mark("returned") - print(f"[rank {RANK}] the mispaired pair RETURNED", flush=True) - return 0 - - -def _mpirun(world_size: int, flag: str, env: Optional[Dict[str, str]] = None) -> Any: - """Start one mpirun job in its own process group, so it can be killed.""" - command = [ - "mpirun", - "-n", - str(world_size), - sys.executable, - str(Path(__file__).resolve()), - flag, + # Drop the launching srun step's PMIx/PMI/OMPI namespace from the children: + # they inherit it and are then registered as participants of that step, so + # the first one to exit (hard, as a failing rank does) trips srun's pmix + # errhandler, which SIGKILLs the whole step -- launcher and real error with + # it -- before anything is reported. These ranks rendezvous through + # torch.distributed's own env://, not pmix. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + print(f"[launcher] spawning {world_size} ranks ({flag}), rendezvous on :{port}", flush=True) + return [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), flag], + env=dict(env, RANK=str(r)), + start_new_session=True, + ) + for r in range(world_size) ] - print(f"[launcher] {' '.join(command)}", flush=True) - return subprocess.Popen(command, start_new_session=True, env=env) - - -def _certify_the_unequal_byte_count_wedge(world_size: int) -> None: - """Second job: the one call-order divergence that hangs instead of lying. - It cannot be a `CHECKS` entry, because the job that runs it never reports. - So it runs on its own, and the evidence is the marks its ranks leave: all - of them entered the mispaired pair, none came out of it within - WEDGE_GRACE_S, and the job ended by its own watchdog rather than by - completing. A crash before the pair fails this too — the `ready` count is - what separates the two. - """ - marks = tempfile.mkdtemp(prefix="reducescatter_wedge_") - started = time.time() - process = _mpirun(world_size, _WEDGE_FLAG, env=dict(os.environ, RS_WEDGE_MARKS=marks)) - try: - code: Optional[int] = process.wait(timeout=WEDGE_CAP_S) - ended = "its own watchdog" - except subprocess.TimeoutExpired: - os.killpg(process.pid, signal.SIGKILL) - process.wait() - code, ended = None, f"the launcher, after {WEDGE_CAP_S}s" - - seen = os.listdir(marks) - ready = [m for m in seen if m.startswith("ready.")] - returned = [m for m in seen if m.startswith("returned.")] - wedged = [m for m in seen if m.startswith("wedged.")] - for name in seen: - os.remove(os.path.join(marks, name)) - os.rmdir(marks) - print( - f"[launcher] wedge sub-job: exit={code} ready={len(ready)} " - f"returned={len(returned)} wedged={len(wedged)} " - f"ended by {ended} in {time.time() - started:.0f}s", - flush=True, - ) - assert len(ready) == world_size, ( - f"only {len(ready)}/{world_size} ranks reached the mispaired pair — the " - "sub-job failed before it could be mispaired, so nothing was certified" - ) - assert not returned, ( - "the mispaired pair returned at unequal byte counts; it used to wedge, " - "so the contract's call-order table has to be re-measured" - ) - assert len(wedged) == world_size, ( - f"{len(wedged)}/{world_size} ranks were still inside the pair after " - f"{WEDGE_GRACE_S}s; expected every one of them" - ) +def _kill_children(children: List[Any]) -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() def _spawn_ranks() -> None: - """Re-exec this file under mpirun, one rank per claimed device.""" + """Start one worker process per claimed device; no external launcher.""" visible = os.environ.get("CUDA_VISIBLE_DEVICES") assert visible, ( "set CUDA_VISIBLE_DEVICES to the devices this run owns, " @@ -1611,24 +1763,30 @@ def _spawn_ranks() -> None: assert world_size >= 2, ( f"CUDA_VISIBLE_DEVICES names {world_size} device(s); a collective test needs at least 2" ) - # Own process group so the deadline can kill wedged grandchildren too. - process = _mpirun(world_size, _WORKER_FLAG) - try: - code = process.wait(timeout=DEADLINE_S) - except subprocess.TimeoutExpired: - os.killpg(process.pid, signal.SIGKILL) - process.wait() - raise AssertionError( - f"the {world_size}-rank run did not finish in {DEADLINE_S}s (wedged)" - ) from None - assert code == 0, f"the {world_size}-rank run exited {code}" - _certify_the_unequal_byte_count_wedge(world_size) + children = _spawn_children(world_size, _WORKER_FLAG) + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_children(children) + raise AssertionError(f"the {world_size}-rank run failed: {failure}") if __name__ == "__main__": if _WORKER_FLAG in sys.argv: sys.exit(_run_one_rank()) - if _WEDGE_FLAG in sys.argv: - sys.exit(_run_wedge_rank()) _spawn_ranks() print("OK") diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_op_matrix.py b/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_op_matrix.py index bca298e5aa43..380cb27a5eff 100644 --- a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_op_matrix.py @@ -20,7 +20,7 @@ assert torch.cuda.is_available(), "allgather requires CUDA devices" -# Each case starts its own 4-rank mpirun over every visible device. Under +# Each case spawns its own 4 rank processes over every visible device. Under # xdist several workers would fight for the same GPUs, so this must run # alone -- the same reason the other collective entry in this repo # (_torch/thop/serial/test_moe_alltoall.py) carries the marker. diff --git a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_reducescatter_op_matrix.py b/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_reducescatter_op_matrix.py index 3b1ac2bba531..afb215b4a209 100644 --- a/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_reducescatter_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/comm/test_modeling_v2_reducescatter_op_matrix.py @@ -21,7 +21,7 @@ assert torch.cuda.is_available(), "reducescatter requires CUDA devices" -# Each case starts its own 4-rank mpirun over every visible device. Under +# Each case spawns its own 4 rank processes over every visible device. Under # xdist several workers would fight for the same GPUs, so this must run # alone -- the same reason the other collective entry in this repo # (_torch/thop/serial/test_moe_alltoall.py) carries the marker. diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_bmm_out.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_bmm_out.py index 28c57009b5ee..2e12eba90449 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_bmm_out.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_bmm_out.py @@ -1,63 +1,490 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the bmm_out catalog entry.""" +"""GPU test for the bmm_out catalog entry. +Beyond the cells, the DeepSeek-V4.1 section runs that target's grouped output +projection at every row count 1..32768 at each group count, at the row counts +its runtime dumps show on fresh operands, under CUDA graphs at every row count +it captures, and under a profiler that holds the call to one kernel. +""" + +import pytest import torch +from torch.profiler import ProfilerActivity, profile + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.bmm_out import bmm_out as op -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.bmm_out import bmm_out +__extra_import_path__ = [".."] +from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "bmm_out requires a CUDA device" -def _check(a: torch.Tensor, b: torch.Tensor, out: torch.Tensor) -> None: - ref = torch.bmm(a.float(), b.float()).to(out.dtype) - bmm_out(a, b, out) - # rtol: torch.testing defaults per output dtype. atol: kernel and reference - # both accumulate in fp32 but in different summation orders; for K <= 1024 - # unit-variance inputs the order-dependent absolute noise is up to - # ~K * 2^-24 ~= 6e-5, which dominates on near-zero outputs produced by - # cancellation, so atol=1e-3 instead of the ~1e-5 defaults. - rtol = {torch.bfloat16: 1.6e-2, torch.float16: 1e-3, torch.float32: 1.3e-6} - torch.testing.assert_close(out, ref, rtol=rtol[out.dtype], atol=1e-3) - - -def _make( - batch: int, m: int, k: int, n: int, dtype: torch.dtype -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - a = torch.randn(batch, m, k, device="cuda").to(dtype) - b = torch.randn(batch, k, n, device="cuda").to(dtype) - out = torch.empty(batch, m, n, device="cuda", dtype=dtype) +def _build(spec, seed): + """Materialize one cell. + + `transposed_out` is the part a spec cannot express as a plain `empty`: the + targets hand the op a transposed view of a buffer the forward already + holds, so the product lands in place. The parent buffer comes back too, + because "did it land there" is the thing worth asserting. + `transposed_a` and `transposed_b`, when a cell sets them, make the inputs + the same kind of view: `a` of a [M, B, K] activation, `b` of a [B, N, K] + weight -- DeepSeek-V4.1's grouped output projection passes all three. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + b, m, k, n, dt = spec["batch"], spec["m"], spec["k"], spec["n"], spec["dtype"] + if spec.get("transposed_a", False): + a = torch.randn(m, b, k, generator=g, device="cuda").to(dt).transpose(0, 1) + else: + a = torch.randn(b, m, k, generator=g, device="cuda").to(dt) + if spec.get("transposed_b", False): + w = torch.randn(b, n, k, generator=g, device="cuda").to(dt).transpose(1, 2) + else: + w = torch.randn(b, k, n, generator=g, device="cuda").to(dt) + if spec["transposed_out"]: + parent = torch.zeros(m, b, n, dtype=dt, device="cuda") + out = torch.transpose(parent, 0, 1) + # Only a decode-shaped or single-batch cell escapes this: transposing a + # dim of extent 1 leaves a tensor torch still calls contiguous, so the + # assertion would be about the shape rather than about the view. + assert m == 1 or b == 1 or not out.is_contiguous() + else: + parent = torch.zeros(b, m, n, dtype=dt, device="cuda") + out = parent + return a, w, out, parent + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + a, w, out, parent = _build(cell.spec, seed=abs(hash(cell.why)) % 2**31) + parent_ptr = parent.data_ptr() + with validating(op): + # Identity, not equality: the op must have written this very buffer and + # returned it, not produced a copy that happens to match. + assert op(a, w, out) is out, "the op writes in place and returns what it wrote" + # The product has to be in the caller's buffer, not in a replacement the op + # allocated: that is the whole reason this op exists over torch.bmm. + assert parent.data_ptr() == parent_ptr + op.compare(out, op.reference(a, w, out)) + + +def _ok(dtype=torch.bfloat16): + a = torch.randn(4, 8, 16, dtype=dtype, device="cuda") + w = torch.randn(4, 16, 32, dtype=dtype, device="cuda") + out = torch.zeros(4, 8, 32, dtype=dtype, device="cuda") + return a, w, out + + +def test_guard_refuses_a_wrong_shaped_out() -> None: + """The op silently reallocates it, detaching the caller's buffer.""" + a, w, _ = _ok() + with pytest.raises(AssertionError, match=r"out must be \[B, M, N\]"): + with validating(op): + op(a, w, torch.zeros(4, 8, 31, dtype=torch.bfloat16, device="cuda")) + + +def test_guard_refuses_a_2d_operand() -> None: + a, w, out = _ok() + with pytest.raises(AssertionError, match="must be 3D"): + with validating(op): + op(a[0], w, out) + + +def test_guard_refuses_mixed_dtypes() -> None: + """Promoted rather than rejected, which changes the required out dtype.""" + a, w, out = _ok() + with pytest.raises(AssertionError, match="share one dtype"): + with validating(op): + op(a, w.to(torch.float16), out) + + +# --------------------------------------------------------------------------- +# DeepSeek-V4.1's grouped low-rank output projection (wo_a). The cells above +# hold it at the ends of its row range; what follows keeps the measured range +# certified. The call is made exactly in the target's layout: bf16 o [T, G*4096], +# the parameter w [G*1024, 4096] (HF [out, in]) and the caller's oa [T, G*1024]: +# a = o.view(T, G, 4096).transpose(0, 1) [G, T, 4096] +# b = w.view(G, 1024, 4096).transpose(1, 2) [G, 4096, 1024] +# out = oa.view(T, G, 1024).transpose(0, 1) [G, T, 1024] +# G is the rank's share of the 8 output groups: 8 / tp under attention TP (4, 2, +# 1 at tp 2, 4, 8), 8 under attention DP and in the DSpark drafter's attention +# (replicated, all 8 groups per rank). T is the rank's token count this step, at +# most max_num_tokens (16384 per rank at the largest envelope engine); in the +# decoder bounded-replay pass (attention TP, G = 1, 2, 4) the rows add the +# finishing requests' 128-row tails to the step's generation rows, up to +# 128 x 256 = 32768. +_GROUPS = (1, 2, 4, 8) +_NMAX = 32768 +# rtol: a one-step bf16 rounding flip is <= 2^-8 relative. atol: at K = 4096 +# with unit-variance operands (outputs up to ~360) the op's fp32 summation order +# differs from the reference's. Measured on sm_103 over every T in 1..32768, the +# largest excess of |op - ref| over rtol * |ref| was 6.3e-4 (G = 1), 5.1e-4 +# (G = 2), 6.2e-4 (G = 4) and 5.7e-4 (G = 8): at least 3.1x under this atol, where +# 1e-3 would leave 1.5x. The group-swap control lands ~390 away. +_V41_TOL = {"rtol": 1.6e-2, "atol": 2e-3} +# At G = 1 cuBLAS serves T in 1..32 and 65..128 with a split-K kernel followed by +# cublasLt::splitKreduce_kernel: two launches, and a summation order over K that +# no other T uses. There the result is inside the gate and repeatable, but not +# bitwise the first T rows of a larger call. Every other (G, T) in 1..32768 +# launches one kernel. +_SPLIT_K_T = {1: frozenset(range(1, 33)) | frozenset(range(65, 129))} +# T of this call in the core target's boot-gate runtime-value dumps (tep4, tp4, +# dep4 and tep8): the decode batches 1..8 and the eager prefill chunks, the +# decoder pass's rows among them. +_DUMPED = ( + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 12, + 14, + 22, + 24, + 48, + 128, + 140, + 888, + 890, + 1018, + 1030, + 1935, + 1942, + 2048, +) +# The captured batch sizes G per rank: every G 1..32 plus TRT-LLM's padded +# default list for cuda_graph_config {enable_padding: true, max_batch_size: 256} +# (40..128 step 8, 192, 256); the padding-off list at 256 is a subset. A captured +# step is generation-only: the core call sees T = G (next_n 1) or 6G (DSpark +# target verify), and the DSpark drafter's attention makes the same call at 8 +# groups on its 5-token draft block, 5G rows. Every group count is captured at +# the whole union G, 5G, 6G: 114 row counts, 1..1536. +_GSET = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_CAPTURE = tuple(sorted({f * g for g in _GSET for f in (1, 5, 6)})) +assert len(_CAPTURE) == 114 and _CAPTURE[-1] == 1536 +# Eager prefill chunks run between graph replays, as chunked prefill runs +# between decode steps. +_PREFILL = (14, 48, 888, 890, 1018, 1935, 2048, 4868, 10451, 16384) + + +def _v41_views(o: torch.Tensor, w: torch.Tensor, oa: torch.Tensor, groups: int, tokens: int): + """The target's three views over the first `tokens` rows of o and oa.""" + a = o[:tokens].view(tokens, groups, 4096).transpose(0, 1) + b = w.view(groups, 1024, 4096).transpose(1, 2) + out = oa[:tokens].view(tokens, groups, 1024).transpose(0, 1) return a, b, out -def test_bf16() -> None: - torch.manual_seed(0) - # decode-like (few tokens per batch entry) and prefill-like (many) shapes - for batch, m, k, n in [(64, 1, 512, 128), (16, 8, 576, 512), (8, 2048, 512, 128)]: - _check(*_make(batch, m, k, n, torch.bfloat16)) - - -def test_bf16_noncontiguous_views() -> None: - # The op exists so `out` may be a strided view (e.g. a transposed slice of - # a [tokens, groups, rank] buffer, as MLA uses it); `a` and `b` may be - # strided views too. Exercise all three. - torch.manual_seed(1) - batch, m, k, n = 16, 32, 256, 64 - a_wide = torch.randn(batch, m, 2 * k, device="cuda").to(torch.bfloat16) - a = a_wide[:, :, :k] # row-strided view - b = ( - torch.randn(batch, n, k, device="cuda").to(torch.bfloat16).transpose(1, 2) - ) # transposed view - out_buf = torch.empty(m, batch, n, device="cuda", dtype=torch.bfloat16) - out = out_buf.transpose(0, 1) # non-contiguous out - _check(a, b, out) - # writes landed in the aliased buffer, not a reallocation - ref = torch.bmm(a.float(), b.float()).to(torch.bfloat16) - torch.testing.assert_close(out_buf.transpose(0, 1), ref, rtol=1.6e-2, atol=1e-3) - - -def test_bf16_unaligned_shapes() -> None: - # dims not multiples of typical tile/vector widths - torch.manual_seed(2) - for batch, m, k, n in [(3, 5, 100, 60), (7, 13, 333, 129)]: - _check(*_make(batch, m, k, n, torch.bfloat16)) +def _v41_operands(groups: int, tokens: int): + o = torch.randn(tokens, groups * 4096, device="cuda").to(torch.bfloat16) + w = torch.randn(groups * 1024, 4096, device="cuda").to(torch.bfloat16) + oa = torch.empty(tokens, groups * 1024, device="cuda", dtype=torch.bfloat16) + return (o, w, oa) + _v41_views(o, w, oa, groups, tokens) + + +def _v41_ref(o: torch.Tensor, w: torch.Tensor, groups: int) -> torch.Tensor: + """The entry's reference over all of o's rows, in oa's [T, G*1024] layout.""" + tokens = o.shape[0] + oa = torch.empty(tokens, groups * 1024, device="cuda", dtype=torch.bfloat16) + ref = op.reference(*_v41_views(o, w, oa, groups, tokens)) + return ref.transpose(0, 1).reshape(tokens, groups * 1024) + + +def test_dsv41_group_swap_fails_the_gate() -> None: + """A group-permuted weight fails both gates: a group mix-up in this strided + layout cannot pass unseen.""" + torch.manual_seed(20) + o, w, oa, a, b, out = _v41_operands(2, 32) + op(a, b, out) + swapped = w.view(2, 1024, 4096).flip(0).reshape(2 * 1024, 4096) + assert not torch.allclose(oa.float(), _v41_ref(o, swapped, 2).float(), **_V41_TOL) + with pytest.raises(AssertionError): + op.compare(out, op.reference(a, swapped.view(2, 1024, 4096).transpose(1, 2), out)) + op.compare(out, op.reference(a, b, out)) + + +@pytest.mark.parametrize("groups", (2, 8)) +def test_dsv41_cuda_graph_side_stream_warmup(groups: int) -> None: + """T in {1, 2, 4, 8, 16, 24, 32}, one graph at a time after a side-stream + warm-up, into a caller-owned oa; two replays on fresh o, each in the gate.""" + torch.manual_seed(21 + groups) + for tokens in (1, 2, 4, 8, 16, 24, 32): + o, w, oa, a, b, out = _v41_operands(groups, tokens) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(a, b, out) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(a, b, out) + for _ in range(2): + o.copy_(torch.randn(tokens, groups * 4096, device="cuda").to(torch.bfloat16)) + oa.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(oa, _v41_ref(o, w, groups), **_V41_TOL) + + +@pytest.mark.parametrize("groups", _GROUPS) +def test_dsv41_every_row_count_1_to_32768(groups: int) -> None: + """Every T in 1..32768 in the target's layout, twice: ascending, then in a seeded shuffle. + + One [32768, G*4096] activation and the weight; the call at T reads the first + T rows and writes the first T rows of a [32768 + 1, G*1024] buffer that is + NaN-filled through row T before the call. Each output row depends only on + its own input row, so the reference is computed once over all 32768 rows + and sliced. At every T: every element is within _V41_TOL of the reference + (NaN fails), the first T rows are bitwise equal to the 32768-row call's + first T rows except at the split-K T of _SPLIT_K_T, and row T (the guard + row) is still NaN. Counters stay on the device; one sync at the end. + """ + torch.manual_seed(4100 + groups) + n = _NMAX + o = torch.randn(n, groups * 4096, device="cuda").to(torch.bfloat16) + w = torch.randn(groups * 1024, 4096, device="cuda").to(torch.bfloat16) + ref = _v41_ref(o, w, groups) + oa = torch.full((n + 1, groups * 1024), float("nan"), device="cuda", dtype=torch.bfloat16) + op(*_v41_views(o, w, oa, groups, n)) + torch.testing.assert_close(oa[:n], ref, **_V41_TOL) + assert oa[n].isnan().all(), f"the {n}-row call wrote past its last row" + full_bits = oa[:n].clone().view(torch.int16) + # Control for the per-row bitwise check: no row of the result equals the + # next one, so a call that read or wrote one row off could not pass it. + assert (full_bits[1:] != full_bits[:-1]).any(-1).all(), "neighbouring rows coincide" + + ref_f = ref.float() + slack = _V41_TOL["rtol"] * ref_f.abs() + gate_bad = torch.zeros(n + 1, dtype=torch.int64, device="cuda") + prefix_bad = torch.zeros(n + 1, dtype=torch.int64, device="cuda") + guard_bad = torch.zeros(n + 1, dtype=torch.int64, device="cuda") + excess = torch.full((n + 1,), -float("inf"), device="cuda") + shuffled = ( + torch.randperm(n, generator=torch.Generator().manual_seed(4200 + groups)) + 1 + ).tolist() + for step, tokens in enumerate(list(range(1, n + 1)) + shuffled): + oa[: tokens + 1].fill_(float("nan")) + op(*_v41_views(o, w, oa, groups, tokens)) + over = oa[:tokens].float().sub_(ref_f[:tokens]).abs_().sub_(slack[:tokens]) + gate_bad[tokens] += (~(over <= _V41_TOL["atol"])).sum() + excess[tokens] = torch.maximum(excess[tokens], over.max()) + prefix_bad[tokens] += (oa[:tokens].view(torch.int16) != full_bits[:tokens]).sum() + guard_bad[tokens] += (~oa[tokens].isnan()).sum() + del over + if step % 2048 == 2047: + # The ascending pass asks for a slightly larger temporary every call, which no + # cached block fits; release them now and then rather than filling the device. + torch.cuda.empty_cache() + torch.cuda.synchronize() + failing = gate_bad.nonzero().flatten().tolist() + differing = prefix_bad.nonzero().flatten().tolist() + wrote_guard = guard_bad.nonzero().flatten().tolist() + worst = int(excess.argmax()) + split_k = _SPLIT_K_T.get(groups, frozenset()) + # Both passes count, so an element that differs in both is counted twice. + frac = max((prefix_bad[t].item() / (2 * t * groups * 1024) for t in differing), default=0.0) + print( + f"[every T G={groups}] T=1..{n} ascending + shuffled: {len(failing)} T outside the gate, " + f"{len(differing)} T not bitwise equal to the {n}-row call's prefix (all split-K: " + f"{set(differing) <= split_k}; at most {frac:.3e} of a call's elements), " + f"{len(wrote_guard)} T wrote the guard row; worst excess of |op - ref| over " + f"rtol*|ref|: {excess[worst].item():.3e} at T={worst}" + ) + assert not failing, ( + f"outside the gate at T = {failing[:20]} (counts {gate_bad[failing[:20]].tolist()})" + ) + outside = [t for t in differing if t not in split_k] + assert not outside, ( + f"not bitwise equal to the {n}-row call's first T rows at T = {outside[:20]} " + f"(counts {prefix_bad[outside[:20]].tolist()})" + ) + assert not wrote_guard, f"wrote row T at T = {wrote_guard[:20]}" + + +@pytest.mark.parametrize("tokens", _DUMPED) +@pytest.mark.parametrize("groups", _GROUPS) +def test_dsv41_dumped_row_counts_fresh_operands(groups: int, tokens: int) -> None: + """Every boot-gate dumped T on operands allocated at exactly that size, output NaN-filled first.""" + torch.manual_seed(4400 + 16 * tokens + groups) + o, w, oa, a, b, out = _v41_operands(groups, tokens) + oa.fill_(float("nan")) + ptr = oa.data_ptr() + assert op(a, b, out) is out + assert oa.data_ptr() == ptr and out.data_ptr() == ptr + torch.testing.assert_close(oa, _v41_ref(o, w, groups), **_V41_TOL) + + +# T of this call in the runtime-value dumps of the envelope engines +# (max_num_tokens 16384, max_batch_size 256; tep4 and tp4 at G = 2, dep4 at +# G = 8; padded graphs and eager, DSpark on; plus a 1M-token and a 128K-token +# long-context engine at G = 2). At G = 8 the DSpark drafter's block rows N +# (5 per generation request) from the same dumps are added. 32768, the top of +# the decoder-pass range, closes both lists. +_ENVELOPE_DUMPED = { + 2: ( + 1, 2, 4, 5, 6, 7, 8, 12, 16, 18, 24, 29, 32, 34, 48, 60, 96, 108, 128, 134, 144, 168, 192, 228, + 240, 288, 330, 336, 384, 432, 462, 480, 528, 546, 576, 610, 624, 672, 720, 768, 780, 852, 1152, + 1536, 2778, 2784, 3870, 4208, 4245, 4388, 4394, 4591, 4868, 9364, 9370, 10444, 10447, 10451, + 16308, 16316, 16320, 16322, 16329, 16333, 16340, 16384, 32768, + ), + 8: ( + 1, 2, 5, 6, 7, 12, 16, 24, 48, 96, 144, 192, 240, 288, 336, 384, 432, 480, 528, 576, 597, 624, + 672, 720, 768, 1152, 1536, 8467, 9193, 9240, 9524, 10441, 10451, 16298, 16327, 16328, 16378, + 16384, 32768, + # drafter N + 10, 15, 20, 40, 50, 80, 90, 110, 115, 120, 125, 140, 160, 190, 200, 275, 280, 320, 360, 370, + 385, 400, 440, 455, 520, 560, 580, 585, 600, 640, 650, 665, 710, 715, 960, 1280, + ), +} # fmt: skip + + +@pytest.mark.parametrize("groups", (2, 8)) +def test_dsv41_envelope_dumped_row_counts_fresh_operands(groups: int) -> None: + """Every envelope-dumped T on operands allocated at exactly that size, output NaN-filled first.""" + rows = sorted(set(_ENVELOPE_DUMPED[groups])) + w = torch.randn(groups * 1024, 4096, device="cuda").to(torch.bfloat16) + for tokens in rows: + torch.manual_seed(4600 + tokens) + o = torch.randn(tokens, groups * 4096, device="cuda").to(torch.bfloat16) + oa = torch.full((tokens, groups * 1024), float("nan"), device="cuda", dtype=torch.bfloat16) + a, b, out = _v41_views(o, w, oa, groups, tokens) + ptr = oa.data_ptr() + op(a, b, out) + assert oa.data_ptr() == ptr and out.data_ptr() == ptr + torch.testing.assert_close( + oa, _v41_ref(o, w, groups), **_V41_TOL, msg=lambda m, t=tokens: f"T={t}: {m}" + ) + del o, oa, a, out + print(f"[dumped G={groups}] {len(rows)} row counts {rows[0]}..{rows[-1]}, all inside the gate") + + +@pytest.mark.parametrize("groups", _GROUPS) +def test_dsv41_capture_union_shared_pool(groups: int) -> None: + """One graph per T of the capture union (G, 5G, 6G over the 46 batch sizes), + captured and replayed as the engine does. + + The weight is one tensor shared by every graph. For each T, largest first: + an eager call at that shape, then capture into one shared pool of the call + reading a static o [T, G*4096] (the captured input) and writing an oa + allocated inside the capture, as a forward allocates it. Then three rounds, + each replaying all 114 graphs in a seeded shuffled order. Before each replay + an eager prefill-sized call runs on other buffers, new activations are + written into the captured o, and oa is filled with NaN. Each replay's oa + must be bitwise equal to an eager call on the same o, within _V41_TOL of the + reference, and different on every row from the previous round's result + (the stale-replay control: a replay that re-read the previous activations, + or left the previous output in place, would reproduce it). + """ + torch.manual_seed(4300 + groups) + w = torch.randn(groups * 1024, 4096, device="cuda").to(torch.bfloat16) + pf_n = max(_PREFILL) + pf_o = torch.randn(pf_n, groups * 4096, device="cuda").to(torch.bfloat16) + pf_oa = torch.empty(pf_n, groups * 1024, device="cuda", dtype=torch.bfloat16) + pf_ref = _v41_ref(pf_o, w, groups) + + pool = torch.cuda.graph_pool_handle() + graphs = {} + for tokens in sorted(_CAPTURE, reverse=True): + o = torch.randn(tokens, groups * 4096, device="cuda").to(torch.bfloat16) + warm = torch.empty(tokens, groups * 1024, device="cuda", dtype=torch.bfloat16) + op(*_v41_views(o, w, warm, groups, tokens)) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + oa = torch.empty(tokens, groups * 1024, device="cuda", dtype=torch.bfloat16) + op(*_v41_views(o, w, oa, groups, tokens)) + graphs[tokens] = (graph, o, oa) + + gen = torch.Generator().manual_seed(4500 + groups) + prev = {} + step = 0 + for rnd in range(3): + order = [_CAPTURE[i] for i in torch.randperm(len(_CAPTURE), generator=gen).tolist()] + for tokens in order: + graph, o, oa = graphs[tokens] + pf = _PREFILL[step % len(_PREFILL)] + step += 1 + pf_oa.fill_(float("nan")) + op(*_v41_views(pf_o, w, pf_oa, groups, pf)) + o.copy_(torch.randn(tokens, groups * 4096, device="cuda").to(torch.bfloat16)) + oa.fill_(float("nan")) + graph.replay() + eager = torch.full_like(oa, float("nan")) + op(*_v41_views(o, w, eager, groups, tokens)) + torch.cuda.synchronize() + assert torch.equal(oa, eager), f"replay != eager at T={tokens}, round {rnd}" + torch.testing.assert_close(oa, _v41_ref(o, w, groups), **_V41_TOL) + torch.testing.assert_close(pf_oa[:pf], pf_ref[:pf], **_V41_TOL) + if tokens in prev: + assert (oa != prev[tokens]).any(-1).all(), ( + f"stale replay not rejected at T={tokens}" + ) + prev[tokens] = oa.clone() + print(f"[graphs G={groups}] {len(graphs)} graphs, one pool: 3 rounds bitwise == eager") + + +@pytest.mark.parametrize("groups", _GROUPS) +def test_dsv41_one_kernel_per_call_every_row_count(groups: int) -> None: + """At every T in 1..32768 the call launches exactly one kernel, on a CUDA profiler trace, + except at the split-K T of _SPLIT_K_T, where it launches exactly a split-K kernel and then + cublasLt::splitKreduce_kernel. + + The calls run in chunks of 1024 row counts, one profiler session per chunk, + on one stream, so the chunk's kernels in start-time order are its calls' + kernels in order. No other call may launch a second kernel (a split-K + reduction, or a copy or transpose for the strided views). The kernel names + (a split-K call's first kernel) give the kernel bands: a band is a run of + consecutive T served by one kernel. The bands are cuBLAS's choice and are + reported, not asserted. + """ + n = _NMAX + split_k = _SPLIT_K_T.get(groups, frozenset()) + torch.manual_seed(4700 + groups) + o = torch.randn(n, groups * 4096, device="cuda").to(torch.bfloat16) + w = torch.randn(groups * 1024, 4096, device="cuda").to(torch.bfloat16) + oa = torch.empty(n, groups * 1024, device="cuda", dtype=torch.bfloat16) + op(*_v41_views(o, w, oa, groups, n)) + torch.cuda.synchronize() + names = [] + chunk = 1024 + for start in range(1, n + 1, chunk): + ts = range(start, min(start + chunk, n + 1)) + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for tokens in ts: + op(*_v41_views(o, w, oa, groups, tokens)) + torch.cuda.synchronize() + kernels = sorted( + (e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA), + key=lambda e: e.time_range.start, + ) + want = len(ts) + sum(t in split_k for t in ts) + assert len(kernels) == want, ( + f"T {ts.start}..{ts.stop - 1}: {len(kernels)} kernels for {len(ts)} calls, " + f"expected {want}; {sorted({e.name for e in kernels})}" + ) + it = iter(e.name for e in kernels) + for tokens in ts: + first = next(it) + if tokens in split_k: + second = next(it) + assert "splitK" in first and "splitKreduce_kernel" in second, ( + tokens, + first, + second, + ) + else: + assert "splitK" not in first and "reduce" not in first.lower(), (tokens, first) + names.append(first) + bands = [1] + [t for t in range(2, n + 1) if names[t - 1] != names[t - 2]] + distinct = sorted(set(names)) + low = {names[t - 1] for t in range(1, 8193)} + print( + f"[one kernel G={groups}] T=1..{n}: one kernel per call outside {len(split_k)} split-K T " + f"(two there); {len(distinct)} distinct kernels in {len(bands)} bands " + f"({sum(b <= 8192 for b in bands)} bands and {len(low)} kernels start at T <= 8192); " + f"split-K kernels: {sorted({names[t - 1] for t in split_k})}; " + f"kernels first seen above 8192: {sorted(set(names[8192:]) - low)}" + ) diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py index 6bd55d915a4c..bb6c7de24a9d 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py @@ -1,122 +1,632 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the cublas_mm catalog entry.""" +"""GPU test for the cublas_mm catalog entry. -import contextlib +Beyond the cells, the DeepSeek-V4.1 section sweeps every row count of each of +that target's calls, replays them under CUDA graphs at the row counts the target +captures, and keeps the controls that show the gates see TF32 and bf16-partial +split-K arithmetic. The router's bf16 -> fp32 call has its own file, +test_modeling_v2_cublas_mm_router_bf16.py. +""" +import pytest import torch +from torch.profiler import ProfilerActivity, profile -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm as op + +__extra_import_path__ = [".."] +from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "cublas_mm requires a CUDA device" -@contextlib.contextmanager -def _true_fp32_matmul(): - """Make torch's fp32 matmul actually fp32 for the duration. +def _build(spec, seed): + """Materialize one cell. - torch 2.12 defaults ``matmul.fp32_precision`` to ``tf32`` and - ``allow_tf32`` to True on this hardware, so a plain ``a.float() @ - b.float()`` is a *TF32* product -- ~1e-3 relative, which is 30x the error - of the op being tested. Left alone the reference is the inaccurate side of - the comparison and the entry fails against a correct kernel. Measured - here: with TF32 off the op is bit-identical to torch and both sit 1.9e-5 - from a float64 product; with TF32 on the reference alone moves by 0.035. + `mat_b` is built as [N, K] and transposed, because that is the one thing a + spec cannot say: the op reads it column-major, and a stored weight's `.t()` + already is that, which is why no transpose kernel runs in a target. + `w_std`, when a cell gives one, scales the weight before rounding. """ + g = torch.Generator(device="cuda").manual_seed(seed) + m, k, n, dt = spec["m"], spec["k"], spec["n"], spec["dtype"] + mat_a = torch.randn(m, k, generator=g, device="cuda").to(dt) + weight = torch.randn(n, k, generator=g, device="cuda") + if spec.get("w_std") is not None: + weight = weight * spec["w_std"] + mat_b = weight.to(dt).t() + out_dtype = spec["out_dtype"] + bias = None + if spec["bias"]: + bias = torch.randn(n, generator=g, device="cuda").to(out_dtype or dt) + return mat_a, mat_b, bias, out_dtype + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + mat_a, mat_b, bias, out_dtype = _build(cell.spec, seed=abs(hash(cell.why)) % 2**31) + with validating(op): + out = op(mat_a, mat_b, bias, out_dtype) + assert out.shape == (mat_a.shape[0], mat_b.shape[1]) + assert out.dtype is (out_dtype or mat_a.dtype) + op.compare(out, op.reference(mat_a, mat_b, bias, out_dtype)) + + +def _ok(dtype=torch.bfloat16): + a = torch.randn(8, 64, dtype=dtype, device="cuda") + b = torch.randn(32, 64, dtype=dtype, device="cuda").t() + return a, b + + +def test_guard_refuses_a_row_major_mat_b() -> None: + """Read column-major regardless; a row-major b is a different product.""" + a, _ = _ok() + b = torch.randn(64, 32, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="column-major"): + with validating(op): + op(a, b) + + +def test_guard_refuses_a_strided_mat_a() -> None: + a = torch.randn(8, 128, dtype=torch.bfloat16, device="cuda")[:, :64] + _, b = _ok() + with pytest.raises(AssertionError, match="row-major"): + with validating(op): + op(a, b) + + +def test_guard_refuses_a_bias_in_the_wrong_dtype() -> None: + """Accepted by the op, and the result is silently wrong.""" + a, b = _ok() + bias = torch.randn(32, dtype=torch.float32, device="cuda") + with pytest.raises(AssertionError, match="bias dtype"): + with validating(op): + op(a, b, bias) + + +def test_guard_refuses_a_bias_on_fp32_inputs() -> None: + """The op takes it and never adds it.""" + a, b = _ok(torch.float32) + bias = torch.randn(32, dtype=torch.float32, device="cuda") + with pytest.raises(AssertionError, match="silently ignored"): + with validating(op): + op(a, b, bias) + + +# --------------------------------------------------------------------------- +# DeepSeek-V4.1. The cells above hold each call at the ends of its row range; +# what follows keeps the measured range certified: every M swept, CUDA-graph +# replay at the row counts the target captures, and the controls showing the +# gates can see the failures they are there for. +# +# The core target's calls pass bias=None, a contiguous activation and the .t() +# view of a contiguous [N, K] parameter: +# fp32 [T, 5120] x [5120, 1024] compressor [kv || score], ratio-2 KV source +# fp32 [T, 5120] x [5120, 384] router gate on fp32 inputs +# fp32 [5G, 5120] x [5120, 128] DSpark drafter MoE gate (128 routed experts) +# fp32 [G, 256] x [256, 129280] DSpark Markov head: the fp32 copy of the +# gathered W1 row times W2 [129280, 256] +# bf16 [T, 5120] x [5120, 512] compressor projection, ratio-1 KV source +# bf16 [n, 512] x [512, 128] index-K projection of the compressed latent +# bf16 [T, 5120] x [5120, 32] indexer weights_proj +# T is the step's tokens, at most max_num_tokens (16384 at the largest envelope +# engine); weights_proj also runs in the re-index layers of the decoder +# bounded-replay pass, whose rows reach 128 x 256; the routed MoE splits its +# rows into chunks of at most 8192 before the router; G is the step's +# generation requests, at most 256. +# +# The bf16 gate's atol at these calls, per K (rtol 1.6e-2). With unit-variance +# operands the outputs have std sqrt(K), 72 at K = 5120, and near a cancellation +# only the absolute fp32 accumulation error is left, in the op and in the IEEE +# fp32 reference alike. Measured on sm_103 against a float64 product at +# M = N_max over 16 seeds: the op's error beyond half an output ulp reached +# 1.8e-3 at K = 5120 (4.1e-5 at K = 512), and the fp32 reference's own reached +# 1.1e-3 (8.6e-5); one element of 8.4M in a 16384-row draw at K = 5120 sat +# 1.1e-3 from the reference. So K = 5120 takes 4e-3, 5.6e-5 of the output std. +_DSV41_BF16_ATOL = {512: 1e-3, 5120: 4e-3} +_BF16_RTOL, _BF16_ATOL = 1.6e-2, 1e-3 +# Max |op - float64 product| per K for unit-variance fp32 operands. Measured on +# sm_103: at most 1.56e-3 at K = 5120 (every M to 16384, outputs up to ~400) and +# 1.05e-4 at K = 256; a TF32 product of the same operands sits 0.09..0.11 +# (K = 5120, M >= 32) and 0.022..0.027 (K = 256) away. K = 256 takes the lower +# gate so the TF32 control still lands more than 4x outside it. +_FP32_GATE = {5120: 5e-3, 256: 1e-3} +# The capture set. G ranges over every CUDA-graph batch size an envelope engine +# (max_batch_size up to 256) captures: every G 1..32, plus TRT-LLM's padded +# default list for cuda_graph_config {enable_padding: true, max_batch_size: 256} +# (40..128 step 8, 192, 256); the padding-off list at 256 is a subset. A +# captured step is generation-only, a decode step (next_n 1) or a DSpark +# target-verify step (next_n 6), and the row count each call sees is: +# G every call, decode step; +# 6G the calls fed by the step's tokens (compressors, weights_proj, router, +# and index-K at the ratio-1 source), verify step; +# 3G index-K at the ratio-2 source, verify step (its compressed-row +# reservation: 3 rows per verify request); +# 4G, 24G the router on dep4's gathered rows (decode, verify); +# 5G, G the DSpark gate and Markov head (the drafter runs inside the +# target's decode graphs). +_GSET = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) + + +def _rows(*factors: int) -> tuple[int, ...]: + return tuple(sorted({b * f for b in _GSET for f in factors})) + + +def _fp32_err(out: torch.Tensor, mat_a: torch.Tensor, mat_b: torch.Tensor) -> float: + """Max |out - float64 product|; the entry's fp32 reference is that product.""" + return (out.double() - op.reference(mat_a, mat_b)).abs().max().item() + + +def _tf32_mm(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: + """torch's fp32 matmul with TF32 forced on: the control product.""" prev_allow = torch.backends.cuda.matmul.allow_tf32 prev_prec = getattr(torch.backends.cuda.matmul, "fp32_precision", None) - torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cuda.matmul.allow_tf32 = True if prev_prec is not None: - torch.backends.cuda.matmul.fp32_precision = "ieee" + torch.backends.cuda.matmul.fp32_precision = "tf32" try: - yield + return mat_a @ mat_b finally: torch.backends.cuda.matmul.allow_tf32 = prev_allow if prev_prec is not None: torch.backends.cuda.matmul.fp32_precision = prev_prec -def _ref_mm( - mat_a: torch.Tensor, - mat_b: torch.Tensor, - bias: torch.Tensor | None, - out_dtype: torch.dtype | None, +def _bf16_violations( + out: torch.Tensor, ref: torch.Tensor, atol: float = _BF16_ATOL ) -> torch.Tensor: - """fp32-accumulated reference: mat_a @ mat_b (+ bias), cast to out dtype.""" - with _true_fp32_matmul(): - ref = mat_a.float() @ mat_b.float() + """Elements outside the bf16 gate, as a 0-d device tensor (no host sync).""" + o, r = out.float(), ref.float() + return ((o - r).abs() > atol + _BF16_RTOL * r.abs()).sum() + + +@pytest.mark.parametrize("k, n, m", [(5120, 384, 1005), (5120, 128, 40), (256, 129280, 40)]) +def test_dsv41_fp32_compare_rejects_tf32(k: int, n: int, m: int) -> None: + """A TF32 product of the fp32 cells' operands fails `compare`, even at a + quarter of its error, so a TF32 path cannot pass unseen. The op passes. + + M = 40 at the DSpark shapes: torch's TF32 matmul at M = 1 and K = 5120 took + a GEMV path that is not TF32 (4.4e-5 from float64). + """ + torch.manual_seed(10 + k + n) + mat_a = torch.randn(m, k, device="cuda") + mat_b = torch.randn(n, k, device="cuda").t() + ref = op.reference(mat_a, mat_b) + op.compare(op(mat_a, mat_b), ref) + quarter = (ref + (_tf32_mm(mat_a, mat_b).double() - ref) / 4).float() + with pytest.raises(AssertionError, match="exceeds"): + op.compare(quarter, ref) + assert _fp32_err(_tf32_mm(mat_a, mat_b), mat_a, mat_b) > 4 * _FP32_GATE[k] + + +# The DeepSeek-V4.1 vision tower and aligner. Every call is bf16 in and out, +# with a contiguous activation and the .t() view of a contiguous [N, K] +# parameter, and the checkpoint's bf16 bias fused in the epilogue where it has +# one. P = patches of one request, all its images in one call: at most 17037 +# over the 900 MMMU val documents. C = aligner rows per request: at most 1902; +# the sweep goes to 2048. Operands: unit-normal activations and bias, N(0, 1/K) +# weights (a trained layer's scale, so outputs are O(1)). The reference is the +# float64 product plus the bias, rounded once to bf16, and the gate is rtol +# 1.6e-2, atol 1e-3 plus the fraction of elements bit-identical to it. +_VIS_CELLS = { + # name: (K, N, has_bias, M_max) + "patch_embed": (588, 1024, True, 17037), + "vit_wqkv": (1024, 3072, True, 17037), + "vit_wo": (1024, 1024, True, 17037), + "vit_w1": (1024, 5632, False, 17037), + "vit_w2": (2816, 1024, False, 17037), + "aligner_w1": (9216, 5120, True, 2048), + "aligner_w2": (5120, 5120, True, 2048), +} +# The one excursion the sweep finds, and what `is_valid` refuses: at K 588, +# N 1024 and M 2305..2432, cuBLASLt's heuristic picks a split-K kernel whose +# partial sums are stored in bf16, so each partial is rounded before the sum. +_PATCH_BAD_ROWS = range(2305, 2433) +# cuBLASLt's split-K reduce with bf16 partial sums. +_BF16_PARTIALS = "splitKreduce_kernel<32, 16, int, __nv_bfloat16" +# Floor on the fraction of elements bit-identical to the rounded-once reference. +# Measured over every M of every vision cell outside the band: >= 0.9946. +_VIS_MIN_EXACT = 0.99 + + +def _vis_operands(k: int, n: int, has_bias: bool, rows: int, seed: int): + torch.manual_seed(seed) + mat_a = torch.randn(rows, k, device="cuda").to(torch.bfloat16) + mat_b = (torch.randn(n, k, device="cuda") / k**0.5).to(torch.bfloat16).t() + bias = torch.randn(n, device="cuda").to(torch.bfloat16) if has_bias else None + return mat_a, mat_b, bias + + +def _vis_ref(mat_a: torch.Tensor, mat_b: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: + """float64 product plus bias, rounded once to bf16.""" + ref = mat_a.double() @ mat_b.double() if bias is not None: - ref = ref + bias.float() - return ref.to(out_dtype if out_dtype is not None else mat_a.dtype) - - -def _check( - mat_a: torch.Tensor, - mat_b: torch.Tensor, - bias: torch.Tensor | None = None, - out_dtype: torch.dtype | None = None, -) -> None: - out = cublas_mm(mat_a, mat_b, bias, out_dtype) - ref = _ref_mm(mat_a, mat_b, bias, out_dtype) - expected_dtype = out_dtype if out_dtype is not None else mat_a.dtype - assert out.shape == (mat_a.shape[0], mat_b.shape[1]) - assert out.dtype == expected_dtype - # rtol: torch.testing defaults per output dtype. atol: kernel and reference - # both accumulate in fp32 but in different summation orders; for K <= 4096 - # unit-variance inputs the order-dependent absolute noise is up to - # ~K * 2^-24 ~= 2.4e-4, which dominates on near-zero outputs produced by - # cancellation, so atol=1e-3 instead of the ~1e-5 defaults. - rtol = {torch.bfloat16: 1.6e-2, torch.float16: 1e-3, torch.float32: 1.3e-6} - torch.testing.assert_close(out, ref, rtol=rtol[expected_dtype], atol=1e-3) - - -def _make(m: int, k: int, n: int, dtype: torch.dtype) -> tuple[torch.Tensor, ...]: - """Build mat_a [M,K] row-major, mat_b [K,N] column-major, bias [N].""" - mat_a = torch.randn(m, k, device="cuda").to(dtype) - weight = torch.randn(n, k, device="cuda").to(dtype) # linear weight [N, K] - bias = torch.randn(n, device="cuda").to(dtype) - return mat_a, weight.t(), bias - - -def test_bf16_no_bias() -> None: - torch.manual_seed(0) - # decode-like (few tokens) and prefill-like (many tokens) shapes - for m, k, n in [(1, 4096, 4096), (8, 4096, 11008), (2048, 4096, 4096)]: - mat_a, mat_b, _ = _make(m, k, n, torch.bfloat16) - _check(mat_a, mat_b) - - -def test_bf16_bias() -> None: - torch.manual_seed(1) - for m, k, n in [(1, 4096, 4096), (512, 2048, 6144)]: - mat_a, mat_b, bias = _make(m, k, n, torch.bfloat16) - _check(mat_a, mat_b, bias) - - -def test_bf16_out_fp32() -> None: - # bias must match the output dtype (fp32 here), not the input dtype - torch.manual_seed(2) - mat_a, mat_b, _ = _make(16, 1024, 2048, torch.bfloat16) - bias_fp32 = torch.randn(2048, device="cuda", dtype=torch.float32) - _check(mat_a, mat_b, out_dtype=torch.float32) - _check(mat_a, mat_b, bias_fp32, out_dtype=torch.float32) - - -def test_bf16_unaligned_shapes() -> None: - # dims not multiples of typical tile/vector widths - torch.manual_seed(3) - for m, k, n in [(5, 100, 60), (7, 333, 129)]: - mat_a, mat_b, bias = _make(m, k, n, torch.bfloat16) - _check(mat_a, mat_b) - _check(mat_a, mat_b, bias) - - -def test_fp8_e4m3_to_bf16() -> None: - # fp8 inputs require an explicit out_dtype; no scales are applied (alpha=1) - torch.manual_seed(6) - for m, k, n in [(1, 1024, 1024), (128, 1024, 512)]: - mat_a = (torch.randn(m, k, device="cuda") * 0.1).to(torch.float8_e4m3fn) - weight = (torch.randn(n, k, device="cuda") * 0.1).to(torch.float8_e4m3fn) - bias = torch.randn(n, device="cuda", dtype=torch.bfloat16) - _check(mat_a, weight.t(), out_dtype=torch.bfloat16) - _check(mat_a, weight.t(), bias, out_dtype=torch.bfloat16) + ref = ref + bias.double() + return ref.to(torch.bfloat16) + + +def _bit_exact(out: torch.Tensor, ref: torch.Tensor) -> float: + return float((out.view(torch.int16) == ref.view(torch.int16)).float().mean()) + + +def test_guard_refuses_the_patch_embed_split_k_band() -> None: + """bf16 K = 588, N = 1024 at M 2305..2432: the op returns a plausible wrong + answer (shown unguarded first), and `is_valid` refuses exactly that band. + """ + k, n, has_bias, _ = _VIS_CELLS["patch_embed"] + mat_a, mat_b, bias = _vis_operands(k, n, has_bias, 2433, 61) + ref = _vis_ref(mat_a, mat_b, bias) + for m in (2305, 2368, 2432): + out = op(mat_a[:m], mat_b, bias) + outside = int(_bf16_violations(out, ref[:m])) / out.numel() + assert _bit_exact(out, ref[:m]) < 0.9 and outside > 1e-3, (m, outside) + with pytest.raises(AssertionError, match="split-K"): + with validating(op): + op(mat_a[:m], mat_b, bias) + for m in (2304, 2433): + with validating(op): + out = op(mat_a[:m], mat_b, bias) + assert int(_bf16_violations(out, ref[:m])) == 0, m + + +def test_dsv41_patch_embed_band_padded_and_unbiased() -> None: + """At M 2305..2432 a call padded with zero rows to 2433 is in the gate on its + first M rows. Dropping the bias does not help. + """ + k, n, has_bias, _ = _VIS_CELLS["patch_embed"] + mat_a, mat_b, bias = _vis_operands(k, n, has_bias, 2433, 60) + ref = _vis_ref(mat_a, mat_b, bias) + for m in (2305, 2368, 2432): + padded = torch.zeros(2433, k, dtype=torch.bfloat16, device="cuda") + padded[:m] = mat_a[:m] + with validating(op): + out = op(padded, mat_b, bias)[:m] + exact = _bit_exact(out, ref[:m]) + assert int(_bf16_violations(out, ref[:m])) == 0 and exact >= _VIS_MIN_EXACT, (m, exact) + exact_nb = _bit_exact(op(mat_a[:m], mat_b), _vis_ref(mat_a[:m], mat_b, None)) + print(f"\n[patch M={m}] padded to 2433: bit-exact {exact:.4f}; no bias: {exact_nb:.4f}") + assert exact_nb < 0.9 + + +def _capture_and_replay(k: int, n: int, dtype: torch.dtype, rows) -> None: + """Warm up eagerly, capture one call per row count, replay on two fresh + inputs written into the captured input buffer with the output poisoned: + every replay is bitwise equal to an eager call on the same input, passes + the call's gate, and differs from the previous replay (a stale replay + would repeat it).""" + weight = torch.randn(n, k, device="cuda").to(dtype) + mat_b = weight.t() + for m in rows: + mat_a = torch.randn(m, k, device="cuda").to(dtype) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(mat_a, mat_b) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(mat_a, mat_b) + previous = None + for _ in range(2): + mat_a.copy_(torch.randn(m, k, device="cuda").to(dtype)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, op(mat_a, mat_b)), (k, n, m, "replay != eager") + assert previous is None or not torch.equal(out, previous), (k, n, m, "stale replay") + previous = out.clone() + if dtype == torch.float32: + err = _fp32_err(out, mat_a, mat_b) + assert err <= _FP32_GATE[k], (k, n, m, err) + else: + torch.testing.assert_close( + out, op.reference(mat_a, mat_b), rtol=_BF16_RTOL, atol=_DSV41_BF16_ATOL[k] + ) + del graph, out + + +# Every compressor and indexer call is captured at the union G, 3G, 6G (97 row +# counts, to 1536; 3G is extra coverage except at index-K); the fp32 router also +# at 4G and 24G (133, to 6144). The DSpark gate at every M 1..160 plus 5G (174, +# to 1280), the Markov head at every M 1..160 plus G (162, to 256). +_DSV41_GRAPH_CELLS = ( + (5120, 1024, torch.float32, _rows(1, 3, 6)), + (5120, 384, torch.float32, _rows(1, 3, 6, 4, 24)), + (5120, 512, torch.bfloat16, _rows(1, 3, 6)), + (512, 128, torch.bfloat16, _rows(1, 3, 6)), + (5120, 32, torch.bfloat16, _rows(1, 3, 6)), + (5120, 128, torch.float32, tuple(sorted(set(range(1, 161)) | set(_rows(5))))), + (256, 129280, torch.float32, tuple(sorted(set(range(1, 161)) | set(_GSET)))), +) +assert len(_GSET) == 46 and max(_GSET) == 256 +assert [len(c[3]) for c in _DSV41_GRAPH_CELLS] == [97, 133, 97, 97, 97, 174, 162] +assert [max(c[3]) for c in _DSV41_GRAPH_CELLS] == [1536, 6144, 1536, 1536, 1536, 1280, 256] + + +@pytest.mark.parametrize( + "k, n, dtype, rows", _DSV41_GRAPH_CELLS, ids=[f"{c[0]}x{c[1]}" for c in _DSV41_GRAPH_CELLS] +) +def test_dsv41_cuda_graph_replay(k: int, n: int, dtype: torch.dtype, rows) -> None: + torch.manual_seed(12 + k + n) + _capture_and_replay(k, n, dtype, rows) + print( + f"\n[capture K={k} N={n} {str(dtype).removeprefix('torch.')}] {len(rows)} captures, " + f"{2 * len(rows)} replays bitwise = eager, in the gate" + ) + + +def _sweep(mat_a, mat_b, bias, rows, record) -> dict: + """op on the first m rows of mat_a for every m in `rows`, calling + record(m, out) after each; returns {m: names of the kernels that call launched}. + + Rows run largest first, so the caching allocator serves every later output + from blocks it already holds. Each call is bracketed by two FillFunctor + marker kernels, so the profiler's kernel stream splits into per-call sets; + record() launches no fill kernel. Unguarded: the vision sweep walks the + patch embed's refused band on purpose. + """ + marker = torch.zeros(1, device="cuda") + for m in (rows[0], rows[-1]): # handle and workspace set up before the profiler starts + op(mat_a[:m], mat_b, bias) + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for m in rows: + marker.fill_(1.0) + out = op(mat_a[:m], mat_b, bias) + marker.fill_(2.0) + record(m, out) + torch.cuda.synchronize() + evs = [e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA] + evs.sort(key=lambda e: e.time_range.start) + per_call, cur, inside = [], [], False + for e in evs: + if "FillFunctor" in e.name: + if inside: + per_call.append(tuple(cur)) + cur, inside = [], not inside + elif inside: + cur.append(e.name) + assert len(per_call) == len(rows), (len(per_call), len(rows)) + assert all(1 <= len(k) <= 2 for k in per_call), "a call launched no kernel or more than two" + return dict(zip(rows, per_call)) + + +def _bands(kernels: dict) -> list: + """(lo, hi, names) runs of consecutive M that launched the same kernels.""" + ms = sorted(kernels) + bands, lo = [], ms[0] + for prev, m in zip(ms, ms[1:] + [None]): + if m is None or kernels[m] != kernels[lo]: + bands.append((lo, prev, kernels[lo])) + lo = m + return bands + + +def _ranges(ms) -> str: + """'a..b c..d' runs of consecutive integers in the sorted list ms.""" + runs, lo, prev = [], None, None + for m in list(ms) + [None]: + if lo is not None and (m is None or m != prev + 1): + runs.append(f"{lo}..{prev}") + lo = None + if m is not None and lo is None: + lo = m + prev = m + return " ".join(runs) + + +# cuBLAS picks its kernel by M (none of these (M, K, N) is in the op's own +# algorithm table): cublasGemmEx's default heuristic for fp32, cuBLASLt's for +# bf16. On sm_103 the kernel changes every 16 to 64 rows over most of the range, +# with split-K reduce kernels coming and going, so every M in 1..N_max is run: +# that covers each band's edges and interior and every runtime-dumped row count. +# GEMM rows are independent, so one reference at N_max rows serves every M as +# its first M rows. Each bf16 call also holds the fraction of elements +# bit-identical to the rounded-once reference to a floor at every M: a +# bf16-partial split-K kernel collapses it (0.685 or less at the patch embed). +_DSV41_SWEEP_CELLS = { + # name: (K, N, dtype, N_max) + "compressor_ratio2_fp32": (5120, 1024, torch.float32, 16384), + "router_gate_fp32": (5120, 384, torch.float32, 8192), + "dspark_gate_fp32": (5120, 128, torch.float32, 8192), + "dspark_markov_head_fp32": (256, 129280, torch.float32, 8192), + "compressor_ratio1_bf16": (5120, 512, torch.bfloat16, 16384), + "index_k_proj_bf16": (512, 128, torch.bfloat16, 16384), + "indexer_weights_proj_bf16": (5120, 32, torch.bfloat16, 32768), +} +_CORE_MIN_EXACT = 0.99 +# Cells whose every M returns exactly the first M rows of the N_max call +# (asserted; at the other cells some M do not, and nothing is claimed). +_PREFIX_BITWISE_CELLS = {"index_k_proj_bf16"} + + +@pytest.mark.parametrize("cell", list(_DSV41_SWEEP_CELLS)) +def test_dsv41_every_row_count(cell: str) -> None: + """Every M in 1..N_max at one core call, held to the reference; no M takes a + bf16-partial split-K kernel; at a bf16 call every M keeps the bit-exact + floor. The kernel band map and the M whose rows are bitwise those of the + N_max call are reported (every M, asserted, at _PREFIX_BITWISE_CELLS).""" + k, n, dtype, n_max = _DSV41_SWEEP_CELLS[cell] + torch.manual_seed(30 + list(_DSV41_SWEEP_CELLS).index(cell)) + mat_a = torch.randn(n_max, k, device="cuda").to(dtype) + mat_b = torch.randn(n, k, device="cuda").to(dtype).t() + if dtype == torch.float32: + # The float64 product rounded to fp32: that moves it by at most half an + # fp32 ulp of the output (< 3.1e-5 at the |out| <= 512 seen here), + # which is under 3.1% of the tightest gate. + ref = op.reference(mat_a, mat_b).float() + gate = _FP32_GATE[k] + else: + ref = op.reference(mat_a, mat_b) + ref_bits = ref.view(torch.int16) + worst = torch.zeros(n_max + 1, device="cuda") + exact = torch.zeros(n_max + 1, device="cuda") + differs = torch.zeros(n_max + 1, device="cuda") + last = {} + + def record(m: int, out: torch.Tensor) -> None: + if dtype == torch.float32: + worst[m] = (out - ref[:m]).abs().max() + else: + worst[m] = _bf16_violations(out, ref[:m], _DSV41_BF16_ATOL[k]) + exact[m] = (out.view(torch.int16) == ref_bits[:m]).sum() + if m == n_max: + last["out"] = out + else: + differs[m] = (out != last["out"][:m]).any() + + kernels = _sweep(mat_a, mat_b, None, list(range(n_max, 0, -1)), record) + out = last["out"] + assert out.shape == (n_max, n) and out.dtype == dtype + bands = _bands(kernels) + worst, exact, differs = worst.cpu(), exact.cpu(), differs.cpu() + distinct = {name for names in kernels.values() for name in names} + reduces = sorted({x.split("(")[0][:90] for x in distinct if "educe" in x}) + bf16_partials = sorted( + m for m, names in kernels.items() if any(_BF16_PARTIALS in x for x in names) + ) + if dtype == torch.float32: + bad = [(m, float(worst[m])) for m in range(1, n_max + 1) if worst[m] > gate] + summary = ( + f"max |op - float64| {float(worst.max()):.3e} at M={int(worst.argmax())}, gate {gate:g}" + ) + else: + bad = [(m, int(worst[m])) for m in range(1, n_max + 1) if worst[m] > 0] + summary = ( + "0 elements outside the bf16 gate" if not bad else f"{len(bad)} row counts outside" + ) + frac = [float(exact[m]) / (m * n) for m in range(1, n_max + 1)] + low = [m for m in range(1, n_max + 1) if frac[m - 1] < _CORE_MIN_EXACT] + m_min = 1 + min(range(n_max), key=frac.__getitem__) + summary += f"; min bit-exact {frac[m_min - 1]:.4f} at M={m_min}, floor {_CORE_MIN_EXACT}" + same = [m for m in range(1, n_max) if not differs[m]] + [n_max] + other = [m for m in range(1, n_max) if differs[m]] + print( + f"\n[{cell} K={k} N={n}] every M 1..{n_max}: {summary}; " + f"{len(bands)} kernel bands, {len(distinct)} distinct kernels; " + f"M with bf16-partial split-K: {len(bf16_partials)}" + ) + print(f" split-K reduce kernels seen: {reduces if reduces else 'none'}") + print( + f" M whose rows are bitwise the first M rows of the M={n_max} call: {len(same)} of {n_max}" + ) + print(f" M whose rows differ in bits from that call: {_ranges(other) if other else 'none'}") + print(" band edges: " + " ".join(f"{lo}..{hi}" for lo, hi, _ in bands)) + assert not bad, f"{len(bad)} row counts outside the gate, first {bad[:10]}" + assert not bf16_partials, f"bf16-partial split-K at M {_ranges(bf16_partials)}" + if dtype != torch.float32: + assert not low, f"bit-exact fraction below {_CORE_MIN_EXACT} at M {_ranges(low)}" + if cell in _PREFIX_BITWISE_CELLS: + assert not other, f"rows differ from the N_max call at M {_ranges(other)}" + + # Control: the gate rejects a kernel that dropped the last 16 of the K + # products -- the reference without them is outside it at M = N_max. + if dtype == torch.float32: + wrong = ref - (mat_a[:, -16:].double() @ mat_b[-16:, :].double()).float() + miss = float((out - wrong).abs().max()) + print(f" control (last 16 K products dropped): max |op - wrong| {miss:.3g}, gate {gate:g}") + assert miss > 4 * gate + else: + wrong = (ref.float() - mat_a[:, -16:].float() @ mat_b[-16:, :].float()).to(dtype) + outside = int(_bf16_violations(out, wrong, _DSV41_BF16_ATOL[k])) / out.numel() + print(f" control (last 16 K products dropped): {outside:.4f} of elements outside") + assert outside > 0.01 + # Control for the bit-exact floor: a bf16-partial split-K emulated at + # M = N_max -- the K sum split into two halves, each partial rounded to + # bf16, then summed and rounded -- lands below the floor. + h = k // 2 + with true_fp32_matmul(): + p1 = (mat_a[:, :h].float() @ mat_b[:h].float()).to(dtype) + p2 = (mat_a[:, h:].float() @ mat_b[h:].float()).to(dtype) + emu = (p1.float() + p2.float()).to(dtype) + emu_exact = float((emu.view(torch.int16) == ref_bits).float().mean()) + print(f" control (2-way split-K, bf16 partials): bit-exact {emu_exact:.4f}") + assert emu_exact < _CORE_MIN_EXACT + + +@pytest.mark.parametrize("cell", list(_VIS_CELLS)) +def test_dsv41_vision_every_row_count(cell: str) -> None: + """Every M in 1..M_max at one vision call, held to the rounded-once float64 + reference; the patch embed's bf16-partial band is exactly where it was.""" + k, n, has_bias, m_max = _VIS_CELLS[cell] + mat_a, mat_b, bias = _vis_operands(k, n, has_bias, m_max, 50 + list(_VIS_CELLS).index(cell)) + ref = _vis_ref(mat_a, mat_b, bias) + ref_bits = ref.view(torch.int16) + viol = torch.zeros(m_max + 1, device="cuda") + exact = torch.zeros(m_max + 1, device="cuda") + + def record(m: int, out: torch.Tensor) -> None: + viol[m] = _bf16_violations(out, ref[:m]) + exact[m] = (out.view(torch.int16) == ref_bits[:m]).sum() + + kernels = _sweep(mat_a, mat_b, bias, list(range(m_max, 0, -1)), record) + bands = _bands(kernels) + viol, exact = viol.cpu(), exact.cpu() + frac = {m: float(exact[m]) / (m * n) for m in range(1, m_max + 1)} + bf16_partials = sorted( + m for m, names in kernels.items() if any(_BF16_PARTIALS in x for x in names) + ) + distinct = {name for names in kernels.values() for name in names} + print( + f"\n[vision {cell} K={k} N={n} bias={has_bias}] every M 1..{m_max}: " + f"{len(bands)} kernel bands, {len(distinct)} distinct kernels; " + f"M with bf16-partial split-K: {_ranges(bf16_partials) if bf16_partials else 'none'}" + ) + bad_rows = set(_PATCH_BAD_ROWS) if cell == "patch_embed" else set() + assert set(bf16_partials) == bad_rows, "the bf16-partial split-K band moved" + good = [m for m in range(1, m_max + 1) if m not in bad_rows] + worst_good = min(frac[m] for m in good) + print( + f" outside that band: {sum(int(viol[m]) for m in good)} elements outside the bf16 gate, " + f"min bit-exact {worst_good:.4f}" + ) + assert all(viol[m] == 0 for m in good), [(m, int(viol[m])) for m in good if viol[m]][:10] + assert worst_good >= _VIS_MIN_EXACT, worst_good + if bad_rows: + # The excursion `is_valid` refuses: every M in the band leaves the gate + # and falls far below the bit-exact floor. + worst_bad = max(frac[m] for m in bad_rows) + out_bad = min(int(viol[m]) / (m * n) for m in bad_rows) + print( + f" M {min(bad_rows)}..{max(bad_rows)}: bit-exact at most {worst_bad:.4f}, " + f"at least {out_bad:.4f} of elements outside the bf16 gate at every M" + ) + assert worst_bad < 0.9 and out_bad > 1e-3 + + +# The prototype's whole-tower CUDA graph ran at the golden cases' sizes: +# P = 1517, 4680, 5325, 8772 (C = 182, 546, 601, 986); the MMMU maxima +# (P 17037, C 1902) are captured too. +_VIS_GRAPH_ROWS = {"vit": (1517, 4680, 5325, 8772, 17037), "aligner": (182, 546, 601, 986, 1902)} + + +@pytest.mark.parametrize("cell", list(_VIS_CELLS)) +def test_dsv41_vision_cuda_graph_replay(cell: str) -> None: + """Each vision call captured at the prototype's sizes; replay on new data is + bitwise equal to eager and in the gate.""" + k, n, has_bias, _ = _VIS_CELLS[cell] + rows = _VIS_GRAPH_ROWS["aligner" if cell.startswith("aligner") else "vit"] + _, mat_b, bias = _vis_operands(k, n, has_bias, 1, 70 + list(_VIS_CELLS).index(cell)) + for m in rows: + mat_a = torch.randn(m, k, device="cuda").to(torch.bfloat16) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(mat_a, mat_b, bias) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(mat_a, mat_b, bias) + previous = None + for _ in range(2): + mat_a.copy_(torch.randn(m, k, device="cuda").to(torch.bfloat16)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, op(mat_a, mat_b, bias)), (cell, m, "replay != eager") + assert int(_bf16_violations(out, _vis_ref(mat_a, mat_b, bias))) == 0, (cell, m) + assert previous is None or not torch.equal(out, previous), (cell, m, "stale replay") + previous = out.clone() + del graph, out diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm_router_bf16.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm_router_bf16.py new file mode 100644 index 000000000000..373327d84a30 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm_router_bf16.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the cublas_mm call the DeepSeek-V4.1 target's router gate makes: +bf16 [M, 5120] x bf16 [5120, 384] -> fp32 [M, 384], i.e. +``cublas_mm(x, gate_w.t(), None, torch.float32)`` (tensor cores, fp32 accumulation). + +Every product of two bf16 values is exact in fp32, so this call differs from the +checkpoint's fp32 gate GEMM (``linear(x.float(), weight.float())``) only in the +order and rounding of the fp32 summation; it is not bitwise equal to it. The +entry's cells hold the call at M 1 and 8192 with the entry's bf16-scored gate; +this file holds the tighter gate measured on sm_103: + +* eager, every M from 1 to 8192 (the target splits larger router inputs before + this call): within 5e-3 of the float64 product for unit-variance operands + (the entry's fp32 cells' gate at K = 5120; worst M measured 2.52e-3), and + bitwise repeatable; +* the expert selection the target makes from the logits (top-6 of + ``sqrt(softplus(logit)) + bias``, the ``gate_forward`` rule) at router-scale + logits agrees with the selection from the fp32 x fp32 call on >= 99.9 % of rows; +* CUDA graphs at the router's captured row counts: G and 6G (tep4 / tp4 decode and + DSpark verify) and 4G and 24G (dep4's gathered rows) over the CUDA-graph batch + set; every replay is bitwise equal to an eager call on the same input. +""" + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm as op + +assert torch.cuda.is_available(), "cublas_mm requires a CUDA device" + +K, N = 5120, 384 +M_MAX = 8192 +TOL = 5e-3 +TOPK = 6 +_GSET = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_ROUTER_ROWS = tuple(sorted({f * g for g in _GSET for f in (1, 4, 6, 24)})) +assert len(_ROUTER_ROWS) == 122 and _ROUTER_ROWS[-1] == 6144 + + +@pytest.fixture(scope="module") +def operands(): + g = torch.Generator(device="cuda").manual_seed(0) + x = torch.randn(M_MAX, K, generator=g, device="cuda").to(torch.bfloat16) + # [E, H], as the checkpoint holds it. + w = torch.randn(N, K, generator=g, device="cuda").to(torch.bfloat16) + ref = x.double() @ w.double().t() + return x, w, ref + + +def _gate(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return op(x, w.t(), None, torch.float32) + + +def _selection(logits: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: + return (torch.sqrt(F.softplus(logits)) + bias).topk(TOPK, dim=-1).indices.sort(-1).values + + +def test_every_row_count_eager(operands): + x, w, ref = operands + worst = 0.0 + for m in range(1, M_MAX + 1): + out = _gate(x[:m], w) + assert out.dtype == torch.float32 and tuple(out.shape) == (m, N) and out.is_contiguous() + err = float((out.double() - ref[:m]).abs().max()) + assert err <= TOL, (m, err) + worst = max(worst, err) + if m % 64 == 1 or m in (6, 24, 384, 1536, 6144, M_MAX): + again = _gate(x[:m], w) + assert torch.equal(out.view(torch.int32), again.view(torch.int32)), ( + f"not repeatable at M={m}" + ) + print( + f"\nbf16 -> fp32 5120 x 384: every M 1..{M_MAX} within {TOL} of float64 (worst {worst:.3g})" + ) + + +def test_expert_selection_matches_the_fp32_call(operands): + x, w, _ = operands + # Router-scale logits (std ~1.1): w / 64 is exact in bf16, so the products stay exact. + w_r = (w.float() / 64).to(torch.bfloat16) + bias = ( + torch.randn(N, generator=torch.Generator(device="cuda").manual_seed(1), device="cuda") * 0.1 + ) + tc = _selection(_gate(x, w_r), bias) + fp32 = _selection(op(x.float(), w_r.float().t()), bias) + differ = int((tc != fp32).any(-1).sum()) + assert differ <= M_MAX // 1000, differ + print(f"\ntop-{TOPK} selection differs from the fp32 call's on {differ} of {M_MAX} rows") + + +def test_cuda_graph_at_the_router_rows(operands): + x, w, ref = operands + side = torch.cuda.Stream() + for m in _ROUTER_ROWS: + a = x[:m].clone() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + _gate(a, w) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = _gate(a, w) + for shift in (1, 2): + rows = (torch.arange(m, device="cuda") + shift * 131) % M_MAX + a.copy_(x.index_select(0, rows)) + out.fill_(float("nan")) + graph.replay() + eager = _gate(a, w) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int32), eager.view(torch.int32)), (m, shift) + a.copy_(x[:m]) + graph.replay() + torch.cuda.synchronize() + assert float((out.double() - ref[:m]).abs().max()) <= TOL, m + del graph + print( + f"\ncaptured at {len(_ROUTER_ROWS)} row counts ({_ROUTER_ROWS[0]}..{_ROUTER_ROWS[-1]}), " + "2 replays each" + ) diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py index a0f739f36531..7a7b3dec92d7 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py @@ -1,691 +1,77 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the nvfp4_gemm catalog entry (NVFP4 x NVFP4 dense GEMM).""" +"""GPU test for the nvfp4_gemm catalog entry.""" +import pytest import torch -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.nvfp4_gemm import nvfp4_gemm -from tensorrt_llm._torch.autotuner import AutoTuner, autotune - -assert torch.cuda.is_available(), "nvfp4_gemm requires a CUDA device" -# The reference matmul must be true fp32, never tf32. -torch.backends.cuda.matmul.allow_tf32 = False - -DEV = torch.device("cuda") -VEC = 16 # NVFP4 block size: one e4m3 scale per 16 contiguous elements along K - -# e2m1 code -> value. code = (exponent << 1) | mantissa, bit 3 is the sign. -E2M1_VALUES = torch.tensor( - [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32, device=DEV +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.nvfp4_gemm import nvfp4_gemm as op +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, ) -# e4m3 scale bytes: 0x30..0x40 are the positive finite values 0.5 .. 2.0, whose -# products with e2m1 data stay exactly representable in fp32 (see -# test_wide_dynamic_range for the opposite regime). -SF_LO, SF_HI = 0x30, 0x41 - - -def _pad_up(x: int, m: int) -> int: - return (x + m - 1) // m * m - -def _swizzle(sf_2d: torch.Tensor, pad_fill: int = 0) -> torch.Tensor: - """[rows, cols] e4m3 scale bytes -> the flat 128x4-swizzled buffer, in native torch. +__extra_import_path__ = [".."] +from _validating import validating # noqa: E402 — needs the path declared above - Offset of the scale of (row r, block c), the layout every NVFP4 GEMM operand - scale uses: - (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 - + (r // 128) * 128 * pad_up(cols, 4) - Buffer length is `pad_up(rows, 128) * pad_up(cols, 4)`; `pad_fill` is written - to every offset no real (r, c) addresses. - """ - rows, cols = sf_2d.shape - padded_cols = _pad_up(cols, 4) - r = torch.arange(rows, device=DEV).view(-1, 1) - c = torch.arange(cols, device=DEV).view(1, -1) - idx = ( - (c % 4) - + (c // 4) * (4 * 128) - + (r % 32) * 16 - + ((r % 128) // 32) * 4 - + (r // 128) * (128 * padded_cols) - ) - out = torch.full((_pad_up(rows, 128) * padded_cols,), pad_fill, dtype=torch.uint8, device=DEV) - out[idx.flatten()] = sf_2d.flatten() - return out +assert torch.cuda.is_available(), "nvfp4_gemm requires a CUDA device" -def _operand( - rows: int, k: int, seed: int, sf_lo: int = SF_LO, sf_hi: int = SF_HI -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """A random NVFP4 operand: (packed data, linear scales, swizzled scales, fp32 values). +def _operand(rows, k, seed): + """Quantize a random operand the way a target does, via the fp4_quantize entry. - The operand is built byte-first — random e2m1 codes and random positive - finite e4m3 scale bytes — and the exact real-valued matrix it encodes is - derived from those bytes in native torch. Nothing here depends on how a - quantizer would have produced them. + Building operands with the sibling entry rather than by hand is deliberate: + it is how a target produces them, so a disagreement between the two entries + about the scale layout shows up here instead of in a model. """ g = torch.Generator(device="cuda").manual_seed(seed) - codes = torch.randint(0, 16, (rows, k), generator=g, device=DEV, dtype=torch.uint8) - codes = torch.where(codes == 8, torch.zeros_like(codes), codes) # drop -0.0 - data = codes[:, 0::2] | (codes[:, 1::2] << 4) - sf = torch.randint(sf_lo, sf_hi, (rows, k // VEC), generator=g, device=DEV, dtype=torch.uint8) - values = E2M1_VALUES[(codes & 7).long()] - values = torch.where((codes & 8).bool(), -values, values) - values *= sf.view(torch.float8_e4m3fn).float().repeat_interleave(VEC, dim=-1) - return data, sf, _swizzle(sf), values - - -def _reference( - a_values: torch.Tensor, - b_values: torch.Tensor, - alpha: float, - bias: torch.Tensor | None = None, -) -> torch.Tensor: - """alpha * (A @ B.T) (+ bias), fp32 — the kernel accumulates in fp32.""" - out = alpha * (a_values @ b_values.T) - if bias is not None: - out = out + bias.float() - return out - - -def _alpha(value: float) -> torch.Tensor: - return torch.tensor([value], dtype=torch.float32, device=DEV) - - -ALPHA = 0.03125 - - -def test_target_dense_linears() -> None: - """The four NVFP4 dense linears of DeepSeek-V3-Lite, decode through prefill rows. - - (K, N) in nn.Linear [out, in] orientation: layer-0 MLP gate_up [24576, 2560] - and down [2560, 12288], shared-expert gate_up [6144, 2560] and down - [2560, 3072]. - """ - alpha = _alpha(ALPHA) - shapes = [ - (2560, 24576, [1, 2, 8, 64, 1024, 4096]), - (12288, 2560, [1, 8, 64, 1024]), - (2560, 6144, [1, 8, 64, 1024]), - (3072, 2560, [1, 8, 64, 1024]), - ] - for k, n, rows in shapes: - b_data, _, b_sf, b_values = _operand(n, k, seed=1000 + n) - for m in rows: - a_data, _, a_sf, a_values = _operand(m, k, seed=m * 31 + k) - out = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - assert out.shape == (m, n) and out.dtype == torch.bfloat16 - ref = _reference(a_values, b_values, ALPHA) - torch.testing.assert_close(out, ref.to(torch.bfloat16)) - del a_data, a_sf, a_values, out, ref - del b_data, b_sf, b_values - torch.cuda.empty_cache() - - -def test_row_counts_around_swizzle_boundaries() -> None: - """M is unconstrained: the activation scale buffer pads rows up to a multiple of 128.""" - alpha = _alpha(ALPHA) - k, n = 2560, 512 - b_data, _, b_sf, b_values = _operand(n, k, seed=7) - for m in [1, 3, 7, 9, 127, 128, 129, 1000]: - a_data, _, a_sf, a_values = _operand(m, k, seed=m) - assert a_sf.numel() == _pad_up(m, 128) * (k // VEC) - out = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - torch.testing.assert_close(out, _reference(a_values, b_values, ALPHA).to(torch.bfloat16)) - - -def test_output_dtypes() -> None: - """bf16 / fp16 / fp32 outputs of the same call, each at its own dtype tolerance.""" - alpha = _alpha(ALPHA) - k, n, m = 2560, 512, 8 - a_data, _, a_sf, a_values = _operand(m, k, seed=11) - b_data, _, b_sf, b_values = _operand(n, k, seed=12) - ref = _reference(a_values, b_values, ALPHA) - for dtype in [torch.bfloat16, torch.float16, torch.float32]: - out = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, dtype) - assert out.dtype == dtype - torch.testing.assert_close(out, ref.to(dtype)) - - -def test_alpha_scales_the_accumulator() -> None: - """alpha multiplies the whole fp32 accumulator, before the output cast.""" - k, n, m = 2560, 512, 8 - a_data, _, a_sf, a_values = _operand(m, k, seed=13) - b_data, _, b_sf, b_values = _operand(n, k, seed=14) - for value in [1.0, 0.5, ALPHA, 1e-4]: - out = nvfp4_gemm(a_data, b_data, a_sf, b_sf, _alpha(value), torch.float32) - torch.testing.assert_close(out, _reference(a_values, b_values, value)) - - -def test_bias_fused() -> None: - """The optional per-column bias is added after alpha, in the output dtype.""" - alpha = _alpha(ALPHA) - k, n, m = 2560, 512, 64 - a_data, _, a_sf, a_values = _operand(m, k, seed=15) - b_data, _, b_sf, b_values = _operand(n, k, seed=16) - torch.manual_seed(0) - bias_bf16 = torch.randn(n, dtype=torch.bfloat16, device=DEV) - out = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16, bias=bias_bf16) - torch.testing.assert_close( - out, _reference(a_values, b_values, ALPHA, bias_bf16).to(torch.bfloat16) - ) - bias_fp32 = torch.randn(n, dtype=torch.float32, device=DEV) - out32 = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.float32, bias=bias_fp32) - torch.testing.assert_close(out32, _reference(a_values, b_values, ALPHA, bias_fp32)) - - -def test_weight_scale_must_be_swizzled() -> None: - """The weight scale buffer is 128x4-swizzled, exactly like the activation one. - - Pins the load-time obligation for a checkpoint that stores `weight_scale` - row-major linear: the buffer the GEMM needs is `block_scale_interleave` of - that linear tensor, which this test also shows equals the native-torch - swizzle above byte for byte. Feeding the linear buffer instead is accepted - and silently wrong. - """ - alpha = _alpha(ALPHA) - k, n, m = 2560, 512, 16 - a_data, _, a_sf, a_values = _operand(m, k, seed=17) - b_data, b_sf_linear, b_sf, b_values = _operand(n, k, seed=18) - ref = _reference(a_values, b_values, ALPHA) - - interleaved = torch.ops.trtllm.block_scale_interleave(b_sf_linear) - assert torch.equal(interleaved.flatten(), b_sf) - torch.testing.assert_close( - nvfp4_gemm(a_data, b_data, a_sf, interleaved.flatten(), alpha, torch.bfloat16), - ref.to(torch.bfloat16), - ) - - # Same buffer size, linear content: in bounds, accepted, wrong. - linear_padded = torch.zeros_like(b_sf) - linear_padded[: n * (k // VEC)] = b_sf_linear.flatten() - wrong = nvfp4_gemm(a_data, b_data, a_sf, linear_padded, alpha, torch.bfloat16) - assert (wrong.float() - ref).abs().max() > 0.1 * ref.abs().max() - - -def test_weight_rows_not_multiple_of_128() -> None: - """N need only be a multiple of 32; the weight scale buffer still pads rows to 128.""" - alpha = _alpha(ALPHA) - k, n, m = 2560, 160, 8 - a_data, _, a_sf, a_values = _operand(m, k, seed=19) - b_data, _, b_sf, b_values = _operand(n, k, seed=20) - assert b_sf.numel() == 256 * (k // VEC) - out = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - torch.testing.assert_close(out, _reference(a_values, b_values, ALPHA).to(torch.bfloat16)) - - -def test_scale_padding_bytes_are_ignored() -> None: - """Whatever fills the row/column padding of a scale buffer never reaches the result.""" - alpha = _alpha(ALPHA) - k, n, m = 2560, 512, 7 - a_data, a_sf_linear, _, a_values = _operand(m, k, seed=21) - b_data, b_sf_linear, _, b_values = _operand(n, k, seed=22) - zero_pad = (_swizzle(a_sf_linear, 0x00), _swizzle(b_sf_linear, 0x00)) - # 0x7E is e4m3's largest finite value (448); 0x7F would be NaN. - loud_pad = (_swizzle(a_sf_linear, 0x7E), _swizzle(b_sf_linear, 0x7E)) - out_zero = nvfp4_gemm(a_data, b_data, *zero_pad, alpha, torch.bfloat16) - out_loud = nvfp4_gemm(a_data, b_data, *loud_pad, alpha, torch.bfloat16) - assert torch.equal(out_zero, out_loud) - torch.testing.assert_close(out_zero, _reference(a_values, b_values, ALPHA).to(torch.bfloat16)) - - -def test_backends() -> None: - """Every selectable backend computes the same GEMM; forcing one is a valid call.""" - alpha = _alpha(ALPHA) - k, n, m = 2560, 512, 8 # m <= 8 so the cuda_core backend is legal too - a_data, _, a_sf, a_values = _operand(m, k, seed=23) - b_data, _, b_sf, b_values = _operand(n, k, seed=24) - ref = _reference(a_values, b_values, ALPHA).to(torch.bfloat16) - for backends in [ - "cutlass,cublaslt,cuda_core", - "cutlass", - "cublaslt", - "cuda_core", - "cutedsl", # JIT-compiled on first use - ]: - out = nvfp4_gemm( - a_data, - b_data, - a_sf, - b_sf, - alpha, - torch.bfloat16, - allowed_backends=backends, - ) - torch.testing.assert_close(out, ref, msg=lambda s, b=backends: f"{b}: {s}") - - -def test_autotuned_selection() -> None: - """A tuner-selected tactic computes the same GEMM as the untuned fallback.""" - alpha = _alpha(ALPHA) - k, n, m = 2560, 6144, 8 - a_data, _, a_sf, a_values = _operand(m, k, seed=25) - b_data, _, b_sf, b_values = _operand(n, k, seed=26) - ref = _reference(a_values, b_values, ALPHA).to(torch.bfloat16) - cache = AutoTuner.get().profiling_cache - cache.clear() - try: - untuned = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - torch.testing.assert_close(untuned, ref) - with autotune(): - nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - assert len(cache.cache) > 0, "autotune recorded no tactic" - tuned = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - torch.testing.assert_close(tuned, ref) - finally: - cache.clear() # leave the process-global tuner state as we found it - - -def test_wide_dynamic_range() -> None: - """Block scales spanning 2^-5 .. 2^5, where the fp32 accumulation actually rounds. - - With the narrow scale range used elsewhere every product and partial sum is - exactly representable, so the kernel matches an fp32 reference bit for bit. - Here it cannot, and a fixed tolerance would be a fitted number: the check is - the textbook recursive-summation bound instead, computed elementwise from - the operands. - """ - alpha = _alpha(ALPHA) - k, n, m = 12288, 512, 32 - a_data, _, a_sf, a_values = _operand(m, k, seed=27, sf_lo=0x10, sf_hi=0x61) - b_data, _, b_sf, b_values = _operand(n, k, seed=28, sf_lo=0x10, sf_hi=0x61) - ref64 = ALPHA * (a_values.double() @ b_values.double().T) - out32 = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.float32) - out = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - # fp32 summation of K terms deviates from the exact sum by at most - # K * 2^-24 * sum|a_i b_i| (recursive-summation bound), and the bf16 cast - # adds at most half a bf16 ulp (2^-9 relative). Both are bounds computed - # from the operands, not fitted tolerances. - conditioning = ALPHA * (a_values.abs().double() @ b_values.abs().double().T) - summation_bound = k * 2.0**-24 * conditioning - assert not torch.equal(out32.double(), ref64), "this regime should not be exact" - assert ((out32.double() - ref64).abs() <= summation_bound).all() - assert ((out.double() - ref64).abs() <= summation_bound + 2.0**-9 * ref64.abs()).all() - - -def test_inputs_untouched_and_deterministic() -> None: - """The op writes only its freshly allocated output, and repeats bit for bit.""" - alpha = _alpha(ALPHA) - k, n, m = 2560, 512, 64 - a_data, _, a_sf, _ = _operand(m, k, seed=29) - b_data, _, b_sf, _ = _operand(n, k, seed=30) - snapshots = [t.clone() for t in (a_data, b_data, a_sf, b_sf, alpha)] - first = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - second = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - for before, after in zip(snapshots, (a_data, b_data, a_sf, b_sf, alpha)): - assert torch.equal(before, after), "an input was modified" - assert torch.equal(first, second), "not bitwise deterministic" - assert first.is_contiguous() and first.data_ptr() != second.data_ptr() - - -# The four NVFP4 dense linears of DeepSeek-R1-0528, in nn.Linear [out, in] -# orientation. Every K and N here exceeds the largest value the shapes above -# reach: K = 18432 (vs 12288) and N = 36864 (vs 24576), and the K = 18432 -# weight-scale buffer is pad_up(N, 128) * pad_up(K/16, 4) with K/16 = 1152 -# columns -- the widest block-scale buffer in this file. -R1_DENSE_LINEARS = [ - (7168, 36864), # dense-MLP gate_up (layers 0-2), N = 2 * 18432 - (18432, 7168), # dense-MLP down - (7168, 4096), # shared-expert gate_up, N = 2 * 2048 - (2048, 7168), # shared-expert down -] -# Decode rows through the serving prefill cap: the dense path is not chunked, -# so M reaches max_num_tokens = 8192 in a single call. -R1_ROWS = [1, 2, 8, 64, 1024, 4096, 8192] - - -def _gate_rejects(out: torch.Tensor, ref_bf16: torch.Tensor) -> bool: - """True when the default-tolerance comparison the tests above use fails.""" - try: - torch.testing.assert_close(out, ref_bf16) - except AssertionError: - return True - return False - - -def test_r1_dense_linears() -> None: - """The four NVFP4 dense linears of DeepSeek-R1-0528, decode through max_num_tokens.""" - alpha = _alpha(ALPHA) - for k, n in R1_DENSE_LINEARS: - b_data, _, b_sf, b_values = _operand(n, k, seed=2000 + k + n) - assert b_sf.numel() == _pad_up(n, 128) * _pad_up(k // VEC, 4) - for m in R1_ROWS: - a_data, _, a_sf, a_values = _operand(m, k, seed=2100 + m + k) - assert a_sf.numel() == _pad_up(m, 128) * _pad_up(k // VEC, 4) - out = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - assert out.shape == (m, n) and out.dtype == torch.bfloat16 - ref = _reference(a_values, b_values, ALPHA) - torch.testing.assert_close(out, ref.to(torch.bfloat16)) - del a_data, a_sf, a_values, out, ref - del b_data, b_sf, b_values - torch.cuda.empty_cache() - - -def test_r1_backends_agree() -> None: - """Every backend computes the same GEMM at the R1 shapes too. - - Backend selection is shape-sensitive (cuda_core is only offered for M <= 8, - and the tactic space differs per (K, N)), so agreement is re-established at - each new shape rather than inherited. In this scale band every product and - partial sum is exactly representable in fp32 -- checked directly at K = - 18432 by test_r1_wide_scale_range -- so the backends agree *bitwise*, which - is a far sharper check than the tolerance one. - """ - alpha = _alpha(ALPHA) - for k, n in R1_DENSE_LINEARS: - b_data, _, b_sf, b_values = _operand(n, k, seed=2200 + k + n) - for m in [8, 8192]: # 8: the only band where cuda_core is selectable - a_data, _, a_sf, a_values = _operand(m, k, seed=2300 + m + k) - ref = _reference(a_values, b_values, ALPHA).to(torch.bfloat16) - backends = ["cutlass", "cublaslt", "cutedsl"] - if m <= 8: - backends.append("cuda_core") - outs = [] - for backend in backends: - out = nvfp4_gemm( - a_data, - b_data, - a_sf, - b_sf, - alpha, - torch.bfloat16, - allowed_backends=backend, - ) - torch.testing.assert_close( - out, ref, msg=lambda s, b=backend: f"{b} K={k} N={n} M={m}: {s}" - ) - outs.append(out) - for backend, out in zip(backends[1:], outs[1:]): - assert torch.equal(out, outs[0]), ( - f"{backend} differs from cutlass at K={k} N={n} M={m}" - ) - del a_data, a_sf, a_values, ref, outs - torch.cuda.empty_cache() - del b_data, b_sf, b_values - torch.cuda.empty_cache() - - -def test_r1_autotuned_selection() -> None: - """The tuned path at the R1 shapes, which is the one a serving run executes. - - One tuning call at M = 8192 fills all 14 power-of-2 buckets for that - (K, N), and every cached winner carries a profiled tactic id rather than - the fallback marker -- so warm really is a different execution from cold. - Each M is then run cold, tuned, and run again warm, and the two results - compared **bitwise**: a tactic that changed the result by a single bit - would fail here even though it would sail through the tolerance gate. The M - values land in four different buckets, two of them non-powers of 2. - """ - alpha = _alpha(ALPHA) - rows = [1, 3, 8, 4097, 8192] - cache = AutoTuner.get().profiling_cache - try: - for k, n in R1_DENSE_LINEARS: - b_data, _, b_sf, b_values = _operand(n, k, seed=2400 + k + n) - cache.clear() - cold = {} - for m in rows: - a_data, _, a_sf, a_values = _operand(m, k, seed=2600 + m + k) - cold[m] = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - ref = _reference(a_values, b_values, ALPHA) - torch.testing.assert_close(cold[m], ref.to(torch.bfloat16)) - del a_data, a_sf, a_values, ref - - a_data, _, a_sf, _ = _operand(8192, k, seed=2500 + k) - with autotune(): - nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - del a_data, a_sf - assert len(cache.cache) == 14, ( - f"K={k} N={n}: expected buckets 1..8192, got {len(cache.cache)}" - ) - for value in cache.cache.values(): - backend, sub_tactic = value[1] - assert backend in {"cutlass", "cublaslt", "cuda_core"}, backend - assert sub_tactic >= 0, f"{backend} winner is the fallback marker" - - for m in rows: - a_data, _, a_sf, a_values = _operand(m, k, seed=2600 + m + k) - tuned = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - ref = _reference(a_values, b_values, ALPHA) - torch.testing.assert_close(tuned, ref.to(torch.bfloat16)) - assert torch.equal(tuned, cold[m]), ( - f"the tuned tactic changed the bits at K={k} N={n} M={m}" - ) - del a_data, a_sf, a_values, tuned, ref - del b_data, b_sf, b_values, cold - torch.cuda.empty_cache() - finally: - cache.clear() # leave the process-global tuner state as we found it - - -def test_r1_wrong_variants_are_visible() -> None: - """Control: the comparisons above can see wrongness at the R1 widths. - - A clean sweep is only worth what the harness could have caught. Three wrong - variants at the two extreme R1 shapes, each measured against the same - default-tolerance gate the correctness tests use: - - the weight scale fed linear (unswizzled) instead of 128x4-swizzled; - - one e4m3 weight-scale byte moved by one code; - - one e2m1 weight nibble moved by one code. - The last two are the smallest perturbation either operand admits, and at - M = 8192 they still fail the gate -- so a tactic that changed results by - one code anywhere could not pass unnoticed. - """ - alpha = _alpha(ALPHA) - for k, n, m in [(18432, 7168, 8192), (7168, 36864, 8192)]: - a_data, _, a_sf, a_values = _operand(m, k, seed=2700 + k) - b_data, b_sf_linear, b_sf, b_values = _operand(n, k, seed=2800 + k + n) - ref = _reference(a_values, b_values, ALPHA) - ref_bf16 = ref.to(torch.bfloat16) - - # Positive control: the op's own relayout of a row-major [N, K/16] - # scale tensor is byte-identical to the native-torch swizzle at these - # widths (K/16 = 1152 and 448), and gives the reference result. - interleaved = torch.ops.trtllm.block_scale_interleave(b_sf_linear).flatten() - assert torch.equal(interleaved, b_sf) - base = nvfp4_gemm(a_data, b_data, a_sf, interleaved, alpha, torch.bfloat16) - assert not _gate_rejects(base, ref_bf16) - - # Same buffer size, linear content: in bounds, accepted, wrong. - linear_padded = torch.zeros_like(b_sf) - linear_padded[: n * (k // VEC)] = b_sf_linear.flatten() - wrong = nvfp4_gemm(a_data, b_data, a_sf, linear_padded, alpha, torch.bfloat16) - assert (wrong.float() - ref).abs().max() > 0.1 * ref.abs().max() - assert _gate_rejects(wrong, ref_bf16) - del linear_padded, wrong - - # One e4m3 scale byte, one code up. - b_sf_bumped = b_sf.clone() - i = b_sf_bumped.numel() // 3 - b_sf_bumped[i] += 1 - one_byte = nvfp4_gemm(a_data, b_data, a_sf, b_sf_bumped, alpha, torch.bfloat16) - assert _gate_rejects(one_byte, ref_bf16), "a one-code scale change is invisible" - del b_sf_bumped, one_byte - - # One e2m1 data nibble, one code up. - b_bumped = b_data.clone() - flat = b_bumped.view(-1) - j = flat.numel() // 3 - flat[j] = (flat[j] & 0xF0) | (((flat[j] & 0x0F) + 1) & 0x0F) - one_nibble = nvfp4_gemm(a_data, b_bumped, a_sf, b_sf, alpha, torch.bfloat16) - assert _gate_rejects(one_nibble, ref_bf16), "a one-code data change is invisible" - - del a_data, a_sf, a_values, b_data, b_sf, b_sf_linear, b_values - del ref, ref_bf16, base, interleaved, b_bumped, flat, one_nibble - torch.cuda.empty_cache() - - -def test_r1_wide_scale_range() -> None: - """The fp32 accumulator at K = 18432, the largest K certified here. - - Narrow block scales (the band every other test uses) keep every product and - partial sum exactly representable, so the fp32 output is bit-exact against - an fp64 reference at this K too -- which is what licenses the bitwise - cross-backend assertions above. Widening the scales to 2^-5..2^5 makes the - sum round; the deviation is then checked against the recursive-summation - bound computed from the operands, not against a fitted tolerance. - """ - alpha = _alpha(ALPHA) - k, n = 18432, 7168 - a_data, _, a_sf, a_values = _operand(64, k, seed=2900) - b_data, _, b_sf, b_values = _operand(n, k, seed=2901) - out32 = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.float32) - assert torch.equal(out32.double(), ALPHA * (a_values.double() @ b_values.double().T)) - del a_data, a_sf, a_values, b_data, b_sf, b_values, out32 - torch.cuda.empty_cache() - - m = 32 - a_data, _, a_sf, a_values = _operand(m, k, seed=2902, sf_lo=0x10, sf_hi=0x61) - b_data, _, b_sf, b_values = _operand(n, k, seed=2903, sf_lo=0x10, sf_hi=0x61) - ref64 = ALPHA * (a_values.double() @ b_values.double().T) - out32 = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.float32) - out = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - conditioning = ALPHA * (a_values.abs().double() @ b_values.abs().double().T) - summation_bound = k * 2.0**-24 * conditioning - assert not torch.equal(out32.double(), ref64), "this regime should not be exact" - assert ((out32.double() - ref64).abs() <= summation_bound).all() - assert ((out.double() - ref64).abs() <= summation_bound + 2.0**-9 * ref64.abs()).all() - torch.cuda.empty_cache() - - -def test_rejected_domains() -> None: - """Domains the op rejects loudly (backend-independent unless noted).""" - alpha = _alpha(ALPHA) - k, n, m = 256, 128, 8 - a_data, _, a_sf, _ = _operand(m, k, seed=31) - b_data, _, b_sf, _ = _operand(n, k, seed=32) - bf16 = torch.bfloat16 - - def rejects(fn) -> None: - try: - fn() - except (RuntimeError, ValueError): - return - raise AssertionError("expected the op to raise") - - # K and N alignment (16-byte operand lines): both must be multiples of 32. - a48, _, a48_sf, _ = _operand(m, 48, seed=33) - b48, _, b48_sf, _ = _operand(n, 48, seed=34) - rejects(lambda: nvfp4_gemm(a48, b48, a48_sf, b48_sf, alpha, bf16)) - b_narrow, _, b_narrow_sf, _ = _operand(16, k, seed=35) - rejects(lambda: nvfp4_gemm(a_data, b_narrow, a_sf, b_narrow_sf, alpha, bf16)) - - # dtypes: data and scales are uint8 byte buffers, alpha is a CUDA fp32 scalar - rejects(lambda: nvfp4_gemm(a_data.view(torch.int8), b_data, a_sf, b_sf, alpha, bf16)) - rejects(lambda: nvfp4_gemm(a_data, b_data.view(torch.int8), a_sf, b_sf, alpha, bf16)) - rejects(lambda: nvfp4_gemm(a_data, b_data, a_sf.view(torch.float8_e4m3fn), b_sf, alpha, bf16)) - rejects(lambda: nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha.half(), bf16)) - rejects(lambda: nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha.cpu(), bf16)) - rejects(lambda: nvfp4_gemm(a_data.cpu(), b_data, a_sf, b_sf, alpha, bf16)) - - # rank and K agreement - rejects(lambda: nvfp4_gemm(a_data.reshape(2, 4, k // 2), b_data, a_sf, b_sf, alpha, bf16)) - rejects(lambda: nvfp4_gemm(a_data, b_data[:, : k // 4].contiguous(), a_sf, b_sf, alpha, bf16)) - - # zero rows - rejects(lambda: nvfp4_gemm(a_data[:0], b_data, a_sf, b_sf, alpha, bf16)) - - # bias must be 1-D [N] in the output dtype - bias = torch.zeros(n, dtype=torch.bfloat16, device=DEV) - rejects(lambda: nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, bf16, bias=bias.float())) - rejects(lambda: nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, bf16, bias=bias.view(1, n))) - rejects(lambda: nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, bf16, bias=bias[: n - 32])) - - # allowed_backends parsing, and backends this arch cannot run - rejects(lambda: nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, bf16, 0, "")) - rejects(lambda: nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, bf16, 0, "cutlas")) - rejects(lambda: nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, bf16, 0, "marlin")) - rejects(lambda: nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.float16, 0, "cutedsl")) - # the cuda_core kernel tops out at M = 16 (the tuner only ever picks it for - # M <= 8; forcing it is what can reach the kernel's own limit) - a_big, _, a_big_sf, _ = _operand(32, k, seed=36) - rejects(lambda: nvfp4_gemm(a_big, b_data, a_big_sf, b_sf, alpha, bf16, 0, "cuda_core")) - # userbuffers output without a userbuffers workspace - rejects(lambda: nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, bf16, 1)) - - -def test_wrapper_guards_silent_domains() -> None: - """Each wrapper assert stands where the op itself is silently wrong.""" - alpha = _alpha(ALPHA) - k, n, m = 2560, 512, 8 - a_data, _, a_sf, a_values = _operand(m, k, seed=37) - b_data, _, b_sf, b_values = _operand(n, k, seed=38) - ref = _reference(a_values, b_values, ALPHA) - - # A non-contiguous view whose *values* equal the contiguous operand: any - # deviation is the kernel ignoring the stride. - a_nc = torch.cat([a_data, a_data], 1)[:, : k // 2] - b_nc = torch.cat([b_data, b_data], 1)[:, : k // 2] - a_sf_nc = torch.stack([a_sf, torch.zeros_like(a_sf)], 1).flatten()[::2] - b_sf_nc = torch.stack([b_sf, torch.zeros_like(b_sf)], 1).flatten()[::2] - assert torch.equal(a_nc, a_data) and not a_nc.is_contiguous() - assert torch.equal(a_sf_nc, a_sf) and not a_sf_nc.is_contiguous() - - def silently_wrong(*args) -> None: - out = torch.ops.trtllm.nvfp4_gemm(*args, alpha, torch.bfloat16, 0, "cublaslt", None, None) - assert (out.float() - ref).abs().max() > 0.1 * ref.abs().max() - - def guarded(fn) -> None: - try: - fn() - except AssertionError: - return - raise AssertionError("expected the wrapper to reject this") - - silently_wrong(a_nc, b_data, a_sf, b_sf) - guarded(lambda: nvfp4_gemm(a_nc, b_data, a_sf, b_sf, alpha, torch.bfloat16)) - silently_wrong(a_data, b_nc, a_sf, b_sf) - guarded(lambda: nvfp4_gemm(a_data, b_nc, a_sf, b_sf, alpha, torch.bfloat16)) - silently_wrong(a_data, b_data, a_sf_nc, b_sf) - guarded(lambda: nvfp4_gemm(a_data, b_data, a_sf_nc, b_sf, alpha, torch.bfloat16)) - silently_wrong(a_data, b_data, a_sf, b_sf_nc) - guarded(lambda: nvfp4_gemm(a_data, b_data, a_sf, b_sf_nc, alpha, torch.bfloat16)) - - # A multi-element alpha is accepted and every element past the first ignored. - fat_alpha = torch.tensor([ALPHA, 99.0], dtype=torch.float32, device=DEV) - out = torch.ops.trtllm.nvfp4_gemm(a_data, b_data, a_sf, b_sf, fat_alpha, torch.bfloat16) - torch.testing.assert_close(out, ref.to(torch.bfloat16)) - guarded(lambda: nvfp4_gemm(a_data, b_data, a_sf, b_sf, fat_alpha, torch.bfloat16)) - - -def test_a_short_scale_buffer_is_rejected_before_dispatch() -> None: - """A contiguous but undersized scale buffer reads past its end. - - The swizzle addresses the *padded* rectangle, so the length the kernel - indexes is `pad_up(rows,128) * pad_up(K/16,4)` whatever the real shape is. - Contiguity says nothing about that, and the op does not check it: a short - buffer is read out of bounds, which is either a wrong result or a dead - worker depending on what follows it in the allocation. - - M and N pad independently, so both operands are exercised, and both are - sized so that the padding is not a no-op: 200 and 160 each pad to 256. - Picking a row count already on the 128 boundary would make the "unpadded - length" case below identical to the correct one and assert nothing. - """ - m, n, k = 200, 160, 512 - a_data, _, a_sf, a_values = _operand(m, k, seed=61) - b_data, _, b_sf, b_values = _operand(n, k, seed=62) - alpha = _alpha(ALPHA) - - # The full-length buffers are accepted and correct, so the rejection below - # is about length alone and not about these operands. - ref = _reference(a_values, b_values, ALPHA) - out = nvfp4_gemm(a_data, b_data, a_sf, b_sf, alpha, torch.bfloat16) - torch.testing.assert_close(out, ref.to(torch.bfloat16)) - assert a_sf.numel() == _pad_up(m, 128) * _pad_up(k // VEC, 4) - assert b_sf.numel() == _pad_up(n, 128) * _pad_up(k // VEC, 4) - - def rejected(*args) -> None: - try: - nvfp4_gemm(*args, alpha, torch.bfloat16) - except AssertionError: - return - raise AssertionError("expected the wrapper to reject a short scale buffer") - - # One byte short, and the un-padded length -- the two ways to get it wrong: - # truncation, and sizing the buffer from the real rectangle. - rejected(a_data, b_data, a_sf[:-1].contiguous(), b_sf) - rejected(a_data, b_data, a_sf[: m * (k // VEC)].contiguous(), b_sf) - rejected(a_data, b_data, a_sf, b_sf[:-1].contiguous()) - rejected(a_data, b_data, a_sf, b_sf[: n * (k // VEC)].contiguous()) + x = torch.randn(rows, k, generator=g, device="cuda").to(torch.bfloat16) + gs = (448.0 * 6.0 / x.abs().max().float()).reshape(1) + packed, sf = fp4_quantize(x, gs, 16, False, True) + return packed, sf, gs + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + seed = abs(hash(cell.why)) % 2**31 + act, act_sf, act_gs = _operand(spec["m"], spec["k"], seed) + weight, weight_sf, weight_gs = _operand(spec["n"], spec["k"], seed + 1) + # alpha folds both global scales back out of the product. + alpha = (1.0 / (act_gs * weight_gs)).float().reshape(1) + + with validating(op): + out = op(act, weight, act_sf, weight_sf, alpha, torch.bfloat16) + assert out.shape == (spec["m"], spec["n"]) + op.compare(out, op.reference(act, weight, act_sf, weight_sf, alpha, torch.bfloat16)) + + +def _ok(): + act, act_sf, act_gs = _operand(128, 512, 5) + weight, weight_sf, weight_gs = _operand(256, 512, 6) + alpha = (1.0 / (act_gs * weight_gs)).float().reshape(1) + return act, weight, act_sf, weight_sf, alpha + + +def test_guard_refuses_a_strided_operand() -> None: + """cublaslt ignores strides; cutlass raises. Which one answers is the op's choice.""" + act, weight, act_sf, weight_sf, alpha = _ok() + strided = torch.zeros(128, 512, dtype=act.dtype, device="cuda")[:, ::2] + with pytest.raises(AssertionError, match="must be contiguous"): + with validating(op): + op(strided, weight, act_sf, weight_sf, alpha, torch.bfloat16) + + +def test_guard_refuses_a_multi_element_alpha() -> None: + """Silently ignored past the first element; this build has no per-token alpha.""" + act, weight, act_sf, weight_sf, _ = _ok() + with pytest.raises(AssertionError, match="exactly one element"): + with validating(op): + op(act, weight, act_sf, weight_sf, torch.ones(128, device="cuda"), torch.bfloat16) + + +def test_guard_refuses_a_short_scale_buffer() -> None: + """Read past its end: the kernel indexes the padded rectangle, not the real one.""" + act, weight, act_sf, weight_sf, alpha = _ok() + with pytest.raises(AssertionError, match="act_sf must hold at least"): + with validating(op): + op(act, weight, act_sf[:-4], weight_sf, alpha, torch.bfloat16) diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_load_paged_kv_cache_for_mla.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/attention/test_modeling_v2_load_paged_kv_cache_for_mla.py similarity index 94% rename from tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_load_paged_kv_cache_for_mla.py rename to tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/attention/test_modeling_v2_load_paged_kv_cache_for_mla.py index 690d54221ec5..b65b6270949c 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_load_paged_kv_cache_for_mla.py +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/attention/test_modeling_v2_load_paged_kv_cache_for_mla.py @@ -30,10 +30,11 @@ from typing import List, Optional, Tuple +import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.load_paged_kv_cache_for_mla import ( - load_paged_kv_cache_for_mla, +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4.ops.attention.load_paged_kv_cache_for_mla import ( # noqa: E501 + load_paged_kv_cache_for_mla as op, ) from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams @@ -210,7 +211,7 @@ def _call( pool_mapping = env.kv_cache_manager.kv_cache_pool_mapping assert block_offsets is not None assert pool_pointers is not None and pool_mapping is not None - compressed_kv, k_pe = load_paged_kv_cache_for_mla( + compressed_kv, k_pe = op( out_dtype, num_contexts, total_ctx_kv, @@ -250,12 +251,9 @@ def _dequant_mirror( read_scale: Optional[torch.Tensor], out_dtype: torch.dtype, ) -> torch.Tensor: - """float(e4m3 byte) * kv_scale_quant_orig[0] in fp32, rounded once to - out_dtype — the fp8 gather's reference, built from torch alone.""" - wide = stored.float() - if read_scale is not None: - wide = wide * read_scale[0] - return wide.to(out_dtype) + """Delegates to the entry: what the gather is supposed to produce is the + entry's claim, and a second copy here would be a second place to fix.""" + return op.reference(stored, read_scale, out_dtype) def _run_fp8( @@ -310,8 +308,8 @@ def _assert_fp8_exact( ) -> None: """Both halves bitwise equal to the fp32-multiply mirror.""" mirror = _dequant_mirror(stored, read_scale, out_dtype) - torch.testing.assert_close(compressed_kv, mirror[:, :KV_LORA_RANK], rtol=0.0, atol=0.0) - torch.testing.assert_close(k_pe, mirror[:, KV_LORA_RANK:], rtol=0.0, atol=0.0) + op.compare(compressed_kv, mirror[:, :KV_LORA_RANK]) + op.compare(k_pe, mirror[:, KV_LORA_RANK:]) def _raises(fn) -> str: @@ -343,6 +341,36 @@ def _fp8_env( ) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Every configuration the entry claims, gather against the entry's mirror. + + The tests below say things no cell can -- that a wrong mirror is + discriminated, that a write/read scale pair that disagrees still returns + plausible numbers, that trailing generation sequences are ignored. This one + says the plain thing, and it is parametrized off the entry so that adding a + cell there adds a case here. + """ + spec = cell.spec + torch.manual_seed(abs(hash(cell.why)) % 2**31) + env = _fp8_env(tokens_per_block=spec["tokens_per_block"]) + try: + rids = list(range(len(spec["new"]))) + ckv, kpe, stored, _, scale_t = _run_fp8( + env, + request_ids=rids, + seq_lens=spec["new"], + num_contexts=len(rids), + cached_lens=spec["cached"], + out_dtype=spec["out_dtype"], + read_scale=spec["read_scale"], + ) + assert ckv.shape[0] == sum(c + n for c, n in zip(spec["cached"], spec["new"])) + _assert_fp8_exact(ckv, kpe, stored, scale_t, spec["out_dtype"]) + finally: + env.shutdown() + + def test_fp8_production_config() -> None: """The production fp8 context call: quant_mode=128 with kv_scale_quant_orig omitted (what the engine's own call site passes), a diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_append_paged_kv_assign_q.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/attention/test_modeling_v2_mla_rope_append_paged_kv_assign_q.py similarity index 94% rename from tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_append_paged_kv_assign_q.py rename to tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/attention/test_modeling_v2_mla_rope_append_paged_kv_assign_q.py index 1158d946fcc9..0eb60647c581 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_append_paged_kv_assign_q.py +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/attention/test_modeling_v2_mla_rope_append_paged_kv_assign_q.py @@ -34,10 +34,11 @@ from typing import List, NamedTuple, Optional, Sequence, Union +import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.mla_rope_append_paged_kv_assign_q import ( - mla_rope_append_paged_kv_assign_q, +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4.ops.attention.mla_rope_append_paged_kv_assign_q import ( # noqa: E501 + mla_rope_append_paged_kv_assign_q as op, ) from tensorrt_llm._torch.attention.backends.interface import RopeParams from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata @@ -48,6 +49,9 @@ from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + assert torch.cuda.is_available(), "mla_rope_append_paged_kv_assign_q requires a CUDA device" # DeepSeek-V3 MLA head geometry (num_heads reduced to a TP-slice-like 16). @@ -79,7 +83,7 @@ # ever been approached. That both halves still discriminate is measured, not # assumed — test_fp8_kv_explicit_unit_scale runs four wrong mirrors of the same # bytes and asserts each blows the gate it is able to blow. -E4M3_ULP_RTOL = 2**-3 +E4M3_ULP_RTOL = 2**-3 # the entry's own gate carries the derivation E4M3_MAX_INEXACT_FRACTION = 1e-3 # The roped q must come back as an ordinary bf16 rotation, NOT quantized: this @@ -277,14 +281,12 @@ def _inexact_bytes(got: torch.Tensor, want: torch.Tensor) -> int: def _assert_e4m3_roped(got: torch.Tensor, want: torch.Tensor, what: str) -> None: - """One-e4m3-ulp gate plus a bit-exact-fraction floor, for the part of the - fp8 output that passes through the in-kernel RoPE.""" - torch.testing.assert_close(got.float(), want.float(), rtol=E4M3_ULP_RTOL, atol=0.0) - inexact = _inexact_bytes(got, want) - allowed = max(1, int(E4M3_MAX_INEXACT_FRACTION * got.numel())) - assert inexact <= allowed, ( - f"{what} not bit-exact enough: {inexact}/{got.numel()} bytes differ (allowed {allowed})" - ) + """Delegates to the entry: the gate and its derivation are the entry's + claim. `what` only labels the failure.""" + try: + op.compare(got, want) + except AssertionError as exc: + raise AssertionError(f"{what}: {exc}") from exc def _fails_e4m3_tolerance(got: torch.Tensor, mirror: torch.Tensor) -> bool: @@ -352,7 +354,7 @@ def _run_and_check( pool = env.pool_tensor(layer_idx) pool_before = pool.clone() - mla_rope_append_paged_kv_assign_q( + op( q, latent_cache, num_contexts, @@ -513,6 +515,40 @@ def _fp8_env(max_batch_size: int = 8, orig_quant: Optional[float] = None) -> _Ml ) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Every configuration the entry claims, through the same driver the + hand-written cases use. + + Parametrized off the entry, so a cell added there is a case here. Without + this the cell list would be a claim nothing executes, which is worse than + no claim at all -- it reads as coverage. + """ + spec = cell.spec + torch.manual_seed(abs(hash(cell.why)) % 2**31) + env = _MlaCtxEnv( + max_batch_size=8, + tokens_per_block=spec["tokens_per_block"], + num_heads=spec["head_num"], + fp8_pool=True, + ) + try: + rids = list(range(len(spec["new"]))) + env.kv_cache_manager.add_dummy_requests( + rids, token_nums=[c + n for c, n in zip(spec["cached"], spec["new"])] + ) + with validating(op): + _run_and_check( + env, + request_ids=rids, + seq_lens=spec["new"], + num_contexts=len(rids), + cached_lens=spec["cached"], + ) + finally: + env.shutdown() + + def test_fp8_kv_production_config() -> None: """The production fp8 context call: the KV scaling factor omitted (what TrtllmAttention.mla_rope_append_paged_kv_assign_q passes) at H = 128 and @@ -751,7 +787,7 @@ def test_fp8_kv_rejects_int8_kv_cache_quant_mode() -> None: assert pool_pointers is not None and pool_mapping is not None raised = "" try: - mla_rope_append_paged_kv_assign_q( + op( q, latent_cache, 1, diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_generation.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/attention/test_modeling_v2_mla_rope_generation.py similarity index 79% rename from tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_generation.py rename to tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/attention/test_modeling_v2_mla_rope_generation.py index 10837d336854..04d6a19c1dc3 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_mla_rope_generation.py +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/attention/test_modeling_v2_mla_rope_generation.py @@ -9,17 +9,13 @@ TrtllmAttentionMetadata — then checks the kernel effects against a torch fp32 reference. -Two surfaces: +One surface. The op also has a bf16-pool path (`quant_mode=0`), where the +roped q lands in fused_q's tail instead of in a quant_q_buffer; nothing +here drives it, so the entry does not claim it. `_MlaEnv` still knows how +to allocate that pool — it is the scaffolding a driver would need, not +evidence that one exists. -1. bf16 latent pool (`quant_mode=0`, all fp8 buffers None), H = 16: GPT-J - RoPE of q_pe written into fused_q's tail, [compressed_kv | rope(k_pe)] - appended into the paged latent cache, and the decode-FMHA scheduler - buffers (cu_q/cu_kv/counter) filled. Covered at both pool page sizes: - 64, and 32 (the engine default), where the appended slot is placed at - every 32-slot alignment and multi-step decodes cross a page boundary - between steps. - -2. fp8-e4m3 latent pool (`quant_mode=128`) at the DeepSeek-R1-0528 cell — +1. fp8-e4m3 latent pool (`quant_mode=128`) at the DeepSeek-R1-0528 cell — H = 128, tokens_per_block = 32, C/R/nope/v = 512/64/128/128, beam_width = 1, rope_append = True. Here the op stops writing fused_q entirely and instead writes the three buffers the fp8 MLA decode FMHA @@ -31,7 +27,7 @@ inconsistent scale pair that separates which tensor drives the write side from which drives the read-side scales. -3. predicted_tokens_per_seq (`P`) > 1 — the MTP path — over both pools, at +2. predicted_tokens_per_seq (`P`) > 1 — the MTP path — at P in {1, 2, 3, 4}. One generation sequence then arrives with P query tokens at P consecutive absolute positions, and both helpers below are parameterised on P so P = 1 is the same code path it always was. What the @@ -43,13 +39,16 @@ is distinguishable from a blind one. """ +import contextlib +import inspect import math from typing import List, NamedTuple, Optional +import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.mla_rope_generation import ( - mla_rope_generation, +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4.ops.attention.mla_rope_generation import ( # noqa: E501 + mla_rope_generation as op, ) from tensorrt_llm._torch.attention.backends.interface import RopeParams from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata @@ -60,6 +59,9 @@ from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + assert torch.cuda.is_available(), "mla_rope_generation requires a CUDA device" # DeepSeek-V3 MLA head geometry (num_heads reduced to a TP-slice-like 16). @@ -313,6 +315,7 @@ def _run_and_check_fp8( pass_bmm2: bool = True, repeats: int = 1, predicted_tokens_per_seq: int = 1, + validate: bool = True, ) -> _Fp8Run: """One generation step over an fp8-e4m3 latent pool. Verifies the three fp8 buffers, the quantized cache append, the scheduler buffers, that the @@ -350,48 +353,53 @@ def _run_and_check_fp8( pool_before = pool.clone() def call() -> None: - mla_rope_generation( - fused_q, - q_pe, - latent_cache, - env.rotary_cos_sin, - cu_q_seqlens, - cu_kv_seqlens, - fmha_scheduler_counter, - bmm1, - bmm2, - quant_q_buffer, - metadata.kv_lens_cuda_runtime, - metadata.kv_lens_runtime, - metadata.prompt_lens_cpu_runtime, - num_contexts, - metadata.kv_cache_block_offsets, - env.kv_cache_manager.kv_cache_pool_pointers, - env.kv_cache_manager.kv_cache_pool_mapping, - env.kv_scale_orig_quant, - env.kv_scale_quant_orig, - None, # kv_cache_scale_orig_quant - None, # out_scale - None, # block_ids_per_seq - [None, None], # helix_tensor_params - p, # predicted_tokens_per_seq - 0, # layer_idx - num_heads, - 1, # num_kv_heads - GEN_HEAD_SIZE, - 0, # residual_dim - env.tokens_per_block, - MAX_SEQ_LEN, # attention_window_size - 1, # beam_width - QUANT_MODE_FP8_KV_CACHE, - q_scaling, - q_lora_rank, - KV_LORA_RANK, - QK_NOPE_HEAD_DIM, - QK_ROPE_HEAD_DIM, - V_HEAD_DIM, - True, # rope_append - ) + # Driving the guard over every production call is what proves it admits + # what it is supposed to admit. `validate=False` is for the one case + # that drives a non-reciprocal scale pair on purpose -- a + # characterization of the op, not a call a target makes. + with validating(op) if validate else contextlib.nullcontext(): + op( + fused_q, + q_pe, + latent_cache, + env.rotary_cos_sin, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + bmm1, + bmm2, + quant_q_buffer, + metadata.kv_lens_cuda_runtime, + metadata.kv_lens_runtime, + metadata.prompt_lens_cpu_runtime, + num_contexts, + metadata.kv_cache_block_offsets, + env.kv_cache_manager.kv_cache_pool_pointers, + env.kv_cache_manager.kv_cache_pool_mapping, + env.kv_scale_orig_quant, + env.kv_scale_quant_orig, + None, # kv_cache_scale_orig_quant + None, # out_scale + None, # block_ids_per_seq + [None, None], # helix_tensor_params + p, # predicted_tokens_per_seq + 0, # layer_idx + num_heads, + 1, # num_kv_heads + GEN_HEAD_SIZE, + 0, # residual_dim + env.tokens_per_block, + MAX_SEQ_LEN, # attention_window_size + 1, # beam_width + QUANT_MODE_FP8_KV_CACHE, + q_scaling, + q_lora_rank, + KV_LORA_RANK, + QK_NOPE_HEAD_DIM, + QK_ROPE_HEAD_DIM, + V_HEAD_DIM, + True, # rope_append + ) torch.cuda.synchronize() call() @@ -675,7 +683,13 @@ def test_fp8_kv_scale_tensors_are_independent() -> None: appended row and quant_q_buffer), quant_orig alone drives both decode-FMHA scales. The reference is built the same way, so a passing run pins the split — the op does not derive either tensor from the other, and a caller - passing a non-reciprocal pair gets a silently inconsistent round trip.""" + passing a non-reciprocal pair gets a silently inconsistent round trip. + + `validate=False`: that last sentence is exactly what `op.is_valid` refuses, + so this case has to run outside the guard. It is characterizing the op -- + establishing that the two tensors are independent -- rather than making a + call a target would make, and the guard exists because of what this case + proves. `test_certified_cells` covers the reciprocal side under the guard.""" torch.manual_seed(13) env = _fp8_env(orig_quant=0.25, quant_orig=3.0) try: @@ -688,11 +702,78 @@ def test_fp8_kv_scale_tensors_are_independent() -> None: cached_lens=[32, 31], q_pe_contiguous=True, q_scaling=Q_SCALING_R1, + validate=False, ) finally: env.shutdown() +def _guard_args(**overrides: object) -> dict: + """Every `is_valid` parameter, inert, with the named ones replaced. + + Read off the signature rather than written out. The guard mirrors a + 47-parameter op because `validating` forwards positionally, and a + hand-counted positional list here would shift silently the day a parameter + is inserted -- landing a scale tensor in `out_scale` and testing nothing. + """ + args = dict.fromkeys(inspect.signature(op.is_valid).parameters) + args.update(residual_dim=0, quant_mode=0) + args.update(overrides) + return args + + +def test_is_valid_refuses_a_non_reciprocal_scale_pair() -> None: + """The other half of the guard: it refuses what the case above drives. + + A non-reciprocal pair leaves the consuming FMHA undoing a write scale that + was never applied, and every buffer this op writes still looks plausible -- + which is why the guard is the only place it can be caught.""" + with pytest.raises(AssertionError, match="reciprocal"): + op.is_valid( + **_guard_args( + kv_scale_orig_quant=_scale_tensor(0.25), + kv_scale_quant_orig=_scale_tensor(3.0), + ) + ) + # The reciprocal pair the certified cells use is admitted. + op.is_valid( + **_guard_args( + kv_scale_orig_quant=_scale_tensor(0.5), + kv_scale_quant_orig=_scale_tensor(2.0), + ) + ) + # Both omitted is the production call: 1.0 and 1.0, which are reciprocal. + op.is_valid(**_guard_args()) + + +def test_is_valid_refuses_the_three_silent_configurations() -> None: + """The rest of the guard. Each of these is accepted by the op and answered + wrongly, or -- in quant_q_buffer's case -- answered with a dead process.""" + # An fp8 pool with nowhere to put the quantized query. Not presence-checked + # by the op: the kernel launches into a null pointer. + with pytest.raises(AssertionError, match="quant_q_buffer"): + op.is_valid(**_guard_args(quant_mode=QUANT_MODE_FP8_KV_CACHE, quant_q_buffer=None)) + op.is_valid( + **_guard_args( + quant_mode=QUANT_MODE_FP8_KV_CACHE, + quant_q_buffer=torch.empty(1, dtype=torch.uint8, device="cuda"), + ) + ) + + # residual_dim is only legal non-zero against an FP4 pool. + with pytest.raises(AssertionError, match="residual_dim"): + op.is_valid(**_guard_args(residual_dim=QK_ROPE_HEAD_DIM)) + + # A kv_norm_weight makes the kernel read latent_cache raw; every caller + # here has already normalized, so it would normalize twice. + with pytest.raises(AssertionError, match="kv_norm_weight"): + op.is_valid( + **_guard_args( + kv_norm_weight=torch.ones(KV_LORA_RANK, dtype=torch.bfloat16, device="cuda") + ) + ) + + def test_fp8_kv_mixed_batch_skips_context() -> None: """Context sequence leading the batch over an fp8 pool: the op consumes only the generation tokens and indexes length/block tensors from @@ -821,48 +902,52 @@ def call( block_offsets: Optional[torch.Tensor] = None, ) -> None: md = self.metadata - mla_rope_generation( - self.fused_q, - self.q_pe, - self.latent_cache, - self.env.rotary_cos_sin, - self.cu_q_seqlens, - self.cu_kv_seqlens, - self.fmha_scheduler_counter, - self.mla_bmm1_scale, - self.mla_bmm2_scale, - self.quant_q_buffer, - md.kv_lens_cuda_runtime, - md.kv_lens_runtime if host_past is None else host_past, - md.prompt_lens_cpu_runtime, - 0, # num_contexts - md.kv_cache_block_offsets if block_offsets is None else block_offsets, - self.env.kv_cache_manager.kv_cache_pool_pointers, - self.env.kv_cache_manager.kv_cache_pool_mapping, - self.env.kv_scale_orig_quant, - self.env.kv_scale_quant_orig, - None, # kv_cache_scale_orig_quant - None, # out_scale - None, # block_ids_per_seq - [None, None], # helix_tensor_params - self.p, - 0, # layer_idx - self.env.num_heads, - 1, # num_kv_heads - GEN_HEAD_SIZE, - 0, # residual_dim - self.env.tokens_per_block, - MAX_SEQ_LEN, # attention_window_size - 1, # beam_width - QUANT_MODE_FP8_KV_CACHE, - self.q_scaling, - Q_LORA_RANK_R1, - KV_LORA_RANK, - QK_NOPE_HEAD_DIM, - QK_ROPE_HEAD_DIM, - V_HEAD_DIM, - True, # rope_append - ) + # Same guard as the production helper. The two controls below vary + # the host length tensor and the block-offset table, neither of which + # is_valid inspects, so they stay inside it. + with validating(op): + op( + self.fused_q, + self.q_pe, + self.latent_cache, + self.env.rotary_cos_sin, + self.cu_q_seqlens, + self.cu_kv_seqlens, + self.fmha_scheduler_counter, + self.mla_bmm1_scale, + self.mla_bmm2_scale, + self.quant_q_buffer, + md.kv_lens_cuda_runtime, + md.kv_lens_runtime if host_past is None else host_past, + md.prompt_lens_cpu_runtime, + 0, # num_contexts + md.kv_cache_block_offsets if block_offsets is None else block_offsets, + self.env.kv_cache_manager.kv_cache_pool_pointers, + self.env.kv_cache_manager.kv_cache_pool_mapping, + self.env.kv_scale_orig_quant, + self.env.kv_scale_quant_orig, + None, # kv_cache_scale_orig_quant + None, # out_scale + None, # block_ids_per_seq + [None, None], # helix_tensor_params + self.p, + 0, # layer_idx + self.env.num_heads, + 1, # num_kv_heads + GEN_HEAD_SIZE, + 0, # residual_dim + self.env.tokens_per_block, + MAX_SEQ_LEN, # attention_window_size + 1, # beam_width + QUANT_MODE_FP8_KV_CACHE, + self.q_scaling, + Q_LORA_RANK_R1, + KV_LORA_RANK, + QK_NOPE_HEAD_DIM, + QK_ROPE_HEAD_DIM, + V_HEAD_DIM, + True, # rope_append + ) torch.cuda.synchronize() def assert_rows_at_sequence_length_slots(self) -> None: @@ -1172,3 +1257,100 @@ def test_fp8_kv_mtp_append_race_control() -> None: step.assert_rows_at_sequence_length_slots() finally: env.shutdown() + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Drive every configuration the entry claims, through the entry's own + reference and its own gate. + + The cases above test the op; this one tests the *claim*. A cell list that + nothing executes reads as coverage and is not, which is why a separate + check refuses an entry whose cells no test names. + + Each cell cross-checks `op.reference` against this file's `env.rope_ref` + + `env.quantize` before either is used to gate the kernel. Two references + disagreeing is a different failure from the kernel moving, and collapsing + them into one assert loses which happened. + """ + spec = cell.spec + p = spec["predicted_tokens_per_seq"] + torch.manual_seed(abs(hash(cell.why)) % 2**31) + + env = _MlaEnv( + max_batch_size=4, + tokens_per_block=spec["tokens_per_block"], + num_heads=spec["num_heads"], + fp8_pool=True, + orig_quant=spec["write_scale"], + quant_orig=spec["read_scale"], + ) + try: + # Cached lengths chosen so the P rows of one sequence straddle a page + # boundary: at page 32 and P > 1, sequence 1 opens a fresh page mid-run. + cached = [32, spec["tokens_per_block"] - 1] + request_ids = [0, 1] + env.kv_cache_manager.add_dummy_requests(request_ids, token_nums=cached) + + run = _run_and_check_fp8( + env, + request_ids=request_ids, + seq_lens=[p, p], + num_contexts=0, + cached_lens=cached, + q_pe_contiguous=False, + q_scaling=spec["q_scaling"], + q_lora_rank=Q_LORA_RANK_R1, + predicted_tokens_per_seq=p, + ) + + # 1. The entry's reference against this file's, on the appended row. + half = QK_ROPE_HEAD_DIM // 2 + table = env.rotary_cos_sin.view(-1, QK_ROPE_HEAD_DIM, 2) + for r in range(run.latent_cache.shape[0]): + pos = run.positions[r] + entry_ref = op.reference( + run.latent_cache[r], + table[pos, :half, 0], + table[pos, :half, 1], + QK_ROPE_HEAD_DIM, + env.kv_scale_orig_quant, + ) + local_ref = torch.cat( + [ + run.latent_cache[r, :KV_LORA_RANK], + env.rope_ref(run.latent_cache[r, KV_LORA_RANK:], pos), + ] + ) + _assert_bytes_equal(entry_ref, env.quantize(local_ref), f"references disagree, row {r}") + + # 2. The entry's reference against the kernel, through the entry's gate. + quant_q = run.quant_q_buffer.view(torch.float8_e4m3fn) + for r in range(run.fused_q.shape[0]): + pos = run.positions[r] + row = torch.cat([run.fused_q[r, :, :KV_LORA_RANK], run.q_pe[r]], dim=-1) + op.compare( + quant_q[r], + op.reference( + row, + table[pos, :half, 0], + table[pos, :half, 1], + QK_ROPE_HEAD_DIM, + env.kv_scale_orig_quant, + ), + ) + + # 3. The two scalar outputs the entry also states. + want1, want2 = op.reference_bmm_scales( + env.kv_scale_quant_orig, spec["q_scaling"], QK_NOPE_HEAD_DIM, QK_ROPE_HEAD_DIM + ) + kv_lens = [c + p for c in cached] + want_cu_q, want_cu_kv = op.reference_scheduler_buffers(kv_lens, spec["num_heads"], p) + assert want_cu_q.tolist() == [0, spec["num_heads"] * p, 2 * spec["num_heads"] * p] + assert want_cu_kv.tolist() == [0, kv_lens[0], kv_lens[0] + kv_lens[1]] + assert want2.item() == pytest.approx( + 1.0 if spec["read_scale"] is None else spec["read_scale"] + ) + assert want1[1].item() == pytest.approx(want1[0].item() * math.log2(math.e), rel=1e-6) + finally: + env.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fused_moe.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/moe/test_modeling_v2_fused_moe.py similarity index 53% rename from tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fused_moe.py rename to tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/moe/test_modeling_v2_fused_moe.py index a66387a68479..e5e493bd533e 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fused_moe.py +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/moe/test_modeling_v2_fused_moe.py @@ -1,13 +1,25 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the fused_moe catalog entry.""" +"""GPU test for the fused_moe catalog entry. +The entry owns the contract: `fused_moe.CELLS` is what is certified, +`reference` is what it is certified against, and `compare` is the band. This +file owns what a cell cannot carry -- building 5.6 GB of weights once and +reusing them, and the hazards, which are claims about inputs no cell +describes. +""" + +import pytest import torch -import torch.nn.functional as F -from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fused_moe import fused_moe +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4.ops.moe.fused_moe import ( # noqa: E501 + fused_moe, +) from tensorrt_llm._torch.autotuner import AutoTuner, autotune +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + assert torch.cuda.is_available(), "fused_moe requires a CUDA device" DEV = "cuda" @@ -15,117 +27,15 @@ # 10 mantissa bits, coarser than the bf16 output it is meant to bound. torch.backends.cuda.matmul.allow_tf32 = False -# Relative distance between neighbouring representable values (1 + 2^-m for m -# stored mantissa bits): bf16 has 7, fp16 has 10. -_ULP = {torch.bfloat16: 2.0**-8, torch.float16: 2.0**-11} - - -def _ref_moe( - x: torch.Tensor, - ids: torch.Tensor, - scales: torch.Tensor | None, - w31: torch.Tensor, - w2: torch.Tensor, - b1: torch.Tensor | None = None, - b2: torch.Tensor | None = None, - act: str = "swiglu", - alpha: torch.Tensor | None = None, - beta: torch.Tensor | None = None, - limit: torch.Tensor | None = None, - ep_size: int = 1, - ep_rank: int = 0, -) -> torch.Tensor: - """Native-torch MoE: fp32 GEMM accumulation, activation rounded to x's dtype. - - `w31[e]` is `[up | gate]` stacked on dim 0; `ids` carries global expert ids - and this rank owns `[ep_rank * E_local, (ep_rank + 1) * E_local)`. - """ - num_tokens, hidden = x.shape - num_local = w31.shape[0] - out = torch.zeros(num_tokens, hidden, dtype=torch.float32, device=x.device) - xf = x.float() - for local_e in range(num_local): - mask = ids == ep_rank * num_local + local_e - tok, slot = mask.nonzero(as_tuple=True) - if tok.numel() == 0: - continue - xe = xf[tok] - w_up, w_gate = w31[local_e].float().chunk(2, dim=0) - h_up = xe @ w_up.t() - h_gate = xe @ w_gate.t() - if b1 is not None: - b_up, b_gate = b1[local_e].float().chunk(2, dim=0) - h_up, h_gate = h_up + b_up, h_gate + b_gate - if limit is not None: - lim = float(limit[local_e]) - h_gate = h_gate.clamp(max=lim) - h_up = h_up.clamp(min=-lim, max=lim) - if act == "swiglu": - a = float(alpha[local_e]) if alpha is not None else 1.0 - gated = h_gate * torch.sigmoid(a * h_gate) - elif act == "geglu": - gated = F.gelu(h_gate, approximate="tanh") - else: - raise ValueError(act) - if beta is not None: - h_up = h_up + float(beta[local_e]) - # the kernel materializes the FC1 activation in the input dtype - inter = (gated * h_up).to(x.dtype).float() - y = inter @ w2[local_e].float().t() - if b2 is not None: - y = y + b2[local_e].float() - weight = ( - torch.ones(tok.numel(), device=x.device) - if scales is None - else scales[tok, slot].float() - ) - out.index_add_(0, tok, y * weight.unsqueeze(1)) - return out.to(x.dtype) - - -def _assert_moe_close(y: torch.Tensor, ref: torch.Tensor) -> None: - """Two gates: per-element, row-scaled; and aggregate relative RMS. - - Kernel and reference consume bit-identical low-precision operands and - differ only in GEMM accumulation order and in which side of a rounding - boundary each intermediate lands — one flipped intermediate moves an - output element by about one ulp of that row's scale. Default assert_close - tolerances cannot express this: their bf16 `atol=1e-5` sits three orders - of magnitude below one output ulp of a two-GEMM chain, and per-element - `rtol` is meaningless where cancellation makes `|ref|` near zero. So the - element gate is 8 ulp of the row's largest magnitude and the aggregate gate - is 4 ulp of relative RMS. Measured worst case across everything this file - drives: 2.8 ulp / 1.2 ulp on the cold fallback tactic, 3.98 ulp / 1.43 ulp - once test_r1_mtp_tactic_space walks the whole tactic space -- 2.0x and 2.8x - margin. Both gates were sized against deliberately wrong computations when - they were set: those land two orders of magnitude out (176+ ulp per element - and ~190 ulp RMS; 284-997 / 197-212 on the R1 geometry), so the bound - separates a real defect from accumulation order by a wide margin. + +def _make(num_tokens, hidden, inter, num_experts, top_k, seed=0, dtype=torch.bfloat16): + """A small, cheap operand set for the hazard cases below. + + Deliberately not a certified geometry: the claims it serves -- ids outside + the window, repeated ids, rejected domains -- are about routing + bookkeeping, which does not vary with shape. Driving them at the R1 + geometry would cost 5.6 GB of weights to prove something shape-independent. """ - assert y.dtype == ref.dtype, (y.dtype, ref.dtype) - assert y.shape == ref.shape, (y.shape, ref.shape) - ulp = _ULP[ref.dtype] - row_scale = ref.float().abs().amax(dim=1, keepdim=True).clamp_min(1e-9) - torch.testing.assert_close( - y.float() / row_scale, ref.float() / row_scale, rtol=0.0, atol=8 * ulp - ) - rel_rms = ( - (y.float() - ref.float()).pow(2).mean().sqrt() - / ref.float().pow(2).mean().sqrt().clamp_min(1e-9) - ).item() - assert rel_rms <= 4 * ulp, f"relative RMS {rel_rms:.3e} > {4 * ulp:.3e}" - - -def _make( - num_tokens: int, - hidden: int, - inter: int, - num_experts: int, - top_k: int, - seed: int = 0, - dtype: torch.dtype = torch.bfloat16, -): - """(x, ids, scales, w31, w2) with realistic magnitudes and top-k routing.""" g = torch.Generator(device=DEV).manual_seed(seed) x = torch.randn(num_tokens, hidden, device=DEV, generator=g).to(dtype) w31 = (torch.randn(num_experts, 2 * inter, hidden, device=DEV, generator=g) / hidden**0.5).to( @@ -137,93 +47,89 @@ def _make( return x, ids.to(torch.int32), torch.softmax(vals, dim=-1), w31, w2 -def test_expert_biases() -> None: - g = torch.Generator(device=DEV).manual_seed(77) - x, ids, scales, w31, w2 = _make(64, 512, 256, 16, 4, seed=77) - b1 = (torch.randn(16, 512, device=DEV, generator=g) * 0.3).to(torch.bfloat16) - b2 = (torch.randn(16, 512, device=DEV, generator=g) * 0.3).to(torch.bfloat16) - # fc1 bias is split into [up | gate] halves exactly like fc1's rows - y = fused_moe(x, ids, scales, w31, b1, w2, b2, torch.bfloat16, [])[0] - _assert_moe_close(y, _ref_moe(x, ids, scales, w31, w2, b1=b1, b2=b2)) - # a zero fc2 bias reproduces the bias-free result - y = fused_moe(x, ids, scales, w31, b1, w2, torch.zeros_like(b2), torch.bfloat16, [])[0] - _assert_moe_close(y, _ref_moe(x, ids, scales, w31, w2, b1=b1)) - - -def test_out_tensor_is_written_in_place() -> None: - x, ids, scales, w31, w2 = _make(37, 512, 256, 16, 4, seed=5) - ref = _ref_moe(x, ids, scales, w31, w2) - # a NaN prefill proves every element is written, not accumulated into - pool = torch.full((41, 512), float("nan"), dtype=torch.bfloat16, device=DEV) - out = pool[:37] - ret = fused_moe(x, ids, scales, w31, None, w2, None, torch.bfloat16, [], out_tensor=out) - assert ret == [], "out_tensor form must return an empty list" - _assert_moe_close(out, ref) - # nothing was written past the requested rows - assert bool(torch.isnan(pool[37:]).all()), "wrote outside out_tensor's rows" - - -def test_none_scales_and_arbitrary_scales() -> None: - x, ids, scales, w31, w2 = _make(16, 256, 128, 16, 4, seed=7) - # token_final_scales=None combines the selected experts with weight 1.0 - y = fused_moe(x, ids, None, w31, None, w2, None, torch.bfloat16, [])[0] - _assert_moe_close(y, _ref_moe(x, ids, None, w31, w2)) - # nothing is renormalized inside: any fp32 weights are used as given - g = torch.Generator(device=DEV).manual_seed(8) - free = (torch.rand(16, 4, device=DEV, generator=g) * 4 - 2).contiguous() - y = fused_moe(x, ids, free, w31, None, w2, None, torch.bfloat16, [])[0] - _assert_moe_close(y, _ref_moe(x, ids, free, w31, w2)) - # all-zero weights zero the output exactly - zeros = torch.zeros_like(free) - y = fused_moe(x, ids, zeros, w31, None, w2, None, torch.bfloat16, [])[0] - assert bool((y == 0).all()), "zero combine weights did not zero the output" - - -def test_expert_parallel_split() -> None: - # ep_rank r owns global expert ids [r * E_local, (r + 1) * E_local); tokens - # routed elsewhere contribute nothing to that rank's output, so the two - # rank outputs sum to the single-rank result. - x, ids, scales, w31, w2 = _make(64, 256, 128, 16, 4, seed=13) - full = fused_moe(x, ids, scales, w31, None, w2, None, torch.bfloat16, [])[0] - _assert_moe_close(full, _ref_moe(x, ids, scales, w31, w2)) - parts = [] - for rank in (0, 1): - lo, hi = rank * 8, (rank + 1) * 8 - w31_r, w2_r = w31[lo:hi].contiguous(), w2[lo:hi].contiguous() - y = fused_moe( +# ─── The certified cells ────────────────────────────────────────────────── + +# The R1-0528 MTP layer's MoE: the one geometry this entry is certified at. +R1_H, R1_I, R1_K = 7168, 2048, 8 +R1_E_LOCAL, R1_EP = 64, 4 +R1_E_GLOBAL = R1_E_LOCAL * R1_EP + + +def _make_r1_weights(num_experts: int, seed: int): + """One rank's [E, 2I, H] / [E, H, I] bf16 weight pair, built in chunks. + + 5.6 GB at E = 64. Materializing the fp32 randn for the whole stack at once + would transiently need another 11 GB; chunking caps the temporary at 64 MB. + """ + g = torch.Generator(device=DEV).manual_seed(seed) + w31 = torch.empty(num_experts, 2 * R1_I, R1_H, dtype=torch.bfloat16, device=DEV) + w2 = torch.empty(num_experts, R1_H, R1_I, dtype=torch.bfloat16, device=DEV) + chunk = max(1, (1 << 26) // (2 * R1_I * R1_H)) + for lo in range(0, num_experts, chunk): + hi = min(lo + chunk, num_experts) + w31[lo:hi] = (torch.randn(hi - lo, 2 * R1_I, R1_H, device=DEV, generator=g) / R1_H**0.5).to( + torch.bfloat16 + ) + w2[lo:hi] = (torch.randn(hi - lo, R1_H, R1_I, device=DEV, generator=g) / R1_I**0.5).to( + torch.bfloat16 + ) + return w31, w2 + + +def _make_r1_routing(num_tokens: int, num_experts: int, seed: int): + """(x, ids, scales) with top-8 routing over `num_experts` global experts.""" + g = torch.Generator(device=DEV).manual_seed(seed) + x = torch.randn(num_tokens, R1_H, device=DEV, generator=g).to(torch.bfloat16) + logits = torch.randn(num_tokens, num_experts, device=DEV, generator=g) + vals, ids = torch.topk(logits, R1_K, dim=-1) + return ( + x, + ids.to(torch.int32).contiguous(), + torch.softmax(vals, dim=-1).contiguous(), + ) + + +@pytest.fixture(scope="module") +def _r1_weights(): + """One 5.6 GB weight pair for every cell that shares the geometry. + + Module-scoped because a cell is cheap and its operands are not: rebuilding + per case would spend minutes materializing identical tensors. Every cell + in this entry is at the one geometry, so one pair serves all of them. + """ + specs = {(c.spec["experts_local"], c.spec["hidden"], c.spec["inter"]) for c in fused_moe.CELLS} + assert len(specs) == 1, f"cells no longer share one geometry: {specs}" + w31, w2 = _make_r1_weights(R1_E_LOCAL, seed=1234) + yield w31, w2 + del w31, w2 + torch.cuda.empty_cache() + + +@pytest.mark.parametrize("cell", fused_moe.CELLS, ids=[c.why[:36] for c in fused_moe.CELLS]) +def test_certified_cells(cell, _r1_weights) -> None: + """Every certified configuration, kernel against the entry's own reference.""" + w31, w2 = _r1_weights + spec = cell.spec + x, ids, scales = _make_r1_routing( + spec["num_tokens"], spec["experts_local"] * spec["ep_size"], seed=spec["num_tokens"] + ) + with validating(fused_moe): + out = fused_moe( x, ids, scales, - w31_r, + w31, None, - w2_r, + w2, None, torch.bfloat16, [], - ep_size=2, - ep_rank=rank, + ep_size=spec["ep_size"], + ep_rank=spec["ep_rank"], )[0] - _assert_moe_close(y, _ref_moe(x, ids, scales, w31_r, w2_r, ep_size=2, ep_rank=rank)) - parts.append(y.float()) - _assert_moe_close((parts[0] + parts[1]).to(torch.bfloat16), full) - - -def test_swiglu_alpha_beta_limit() -> None: - # act = clamp(gate) * sigmoid(alpha * clamp(gate)) * (clamp(up) + beta), - # with gate clamped above by limit and up clamped to [-limit, limit]. - num_experts = 16 - x, ids, scales, w31, w2 = _make(32, 256, 128, num_experts, 4, seed=17) - g = torch.Generator(device=DEV).manual_seed(18) - alpha = 1.0 + torch.rand(num_experts, device=DEV, generator=g) - beta = torch.rand(num_experts, device=DEV, generator=g) - limit = 1.0 + 4.0 * torch.rand(num_experts, device=DEV, generator=g) - for a, b, lim in [ - (alpha, None, None), - (None, beta, None), - (None, None, limit), - (alpha, beta, limit), - ]: - y = fused_moe( + fused_moe.compare( + out, + fused_moe.reference( x, ids, scales, @@ -233,60 +139,13 @@ def test_swiglu_alpha_beta_limit() -> None: None, torch.bfloat16, [], - swiglu_alpha=a, - swiglu_beta=b, - swiglu_limit=lim, - )[0] - _assert_moe_close(y, _ref_moe(x, ids, scales, w31, w2, alpha=a, beta=b, limit=lim)) - # activation_type 7 (SwigluBias) computes the same thing on this path - y5 = fused_moe( - x, - ids, - scales, - w31, - None, - w2, - None, - torch.bfloat16, - [], - swiglu_alpha=alpha, - swiglu_beta=beta, - swiglu_limit=limit, - activation_type=5, - )[0] - y7 = fused_moe( - x, - ids, - scales, - w31, - None, - w2, - None, - torch.bfloat16, - [], - swiglu_alpha=alpha, - swiglu_beta=beta, - swiglu_limit=limit, - activation_type=7, - )[0] - assert torch.equal(y5, y7), "activation_type 7 differed from 5" + ep_size=spec["ep_size"], + ep_rank=spec["ep_rank"], + ), + ) -def test_use_fused_finalize_false() -> None: - x, ids, scales, w31, w2 = _make(64, 512, 256, 16, 4, seed=23) - y = fused_moe( - x, - ids, - scales, - w31, - None, - w2, - None, - torch.bfloat16, - [], - use_fused_finalize=False, - )[0] - _assert_moe_close(y, _ref_moe(x, ids, scales, w31, w2)) +# ─── Hazards: claims about inputs no cell describes ─────────────────────── def test_out_of_range_expert_ids_are_dropped() -> None: @@ -300,7 +159,9 @@ def test_out_of_range_expert_ids_are_dropped() -> None: kept[:, 0] = True kept[:, 2:] = True dropped_scales = scales * kept - _assert_moe_close(y, _ref_moe(x, bad, dropped_scales, w31, w2)) + fused_moe.compare( + y, fused_moe.reference(x, bad, dropped_scales, w31, None, w2, None, torch.bfloat16, []) + ) def test_repeated_expert_ids() -> None: @@ -312,7 +173,9 @@ def test_repeated_expert_ids() -> None: dup = ids.clone() dup[:, 1] = dup[:, 0] y = fused_moe(x, dup, scales, w31, None, w2, None, torch.bfloat16, [])[0] - _assert_moe_close(y, _ref_moe(x, dup, scales, w31, w2)) + fused_moe.compare( + y, fused_moe.reference(x, dup, scales, w31, None, w2, None, torch.bfloat16, []) + ) def test_inputs_untouched_and_deterministic() -> None: @@ -332,13 +195,16 @@ def test_wrapper_rejects_output_dtype_mismatch() -> None: x, ids, scales, w31, w2 = _make(8, 128, 64, 8, 2, seed=41) for bad in [torch.float16, torch.float32]: try: - fused_moe(x, ids, scales, w31, None, w2, None, bad, []) + with validating(fused_moe): + fused_moe(x, ids, scales, w31, None, w2, None, bad, []) except AssertionError: continue raise AssertionError(f"wrapper accepted output_dtype={bad} for a bf16 input") # positive control: the matching dtype still works y = fused_moe(x, ids, scales, w31, None, w2, None, torch.bfloat16, [])[0] - _assert_moe_close(y, _ref_moe(x, ids, scales, w31, w2)) + fused_moe.compare( + y, fused_moe.reference(x, ids, scales, w31, None, w2, None, torch.bfloat16, []) + ) def test_op_rejects_unsupported_domains() -> None: @@ -553,216 +419,22 @@ def test_op_rejects_unsupported_domains() -> None: raise AssertionError(f"op accepted an unsupported domain: {tag}") -def test_quant_flags_ignored_on_the_unquantized_path() -> None: - # These three neither raise nor change the result for bf16 weights; a - # caller cannot rely on them being honoured. - x, ids, scales, w31, w2 = _make(32, 256, 128, 16, 4, seed=47) - base = fused_moe(x, ids, scales, w31, None, w2, None, torch.bfloat16, [])[0] - _assert_moe_close(base, _ref_moe(x, ids, scales, w31, w2)) - args = (x, ids, scales, w31, None, w2, None, torch.bfloat16, []) - flagged = [ - ("use_w4_group_scaling", fused_moe(*args, use_w4_group_scaling=True)), - ("use_mxfp8_act_scaling", fused_moe(*args, use_mxfp8_act_scaling=True)), - ("use_dynamic_fc2_scale", fused_moe(*args, use_dynamic_fc2_scale=True)), - ] - for flag, ret in flagged: - assert torch.equal(ret[0], base), f"{flag}=True changed the result" - # a non-empty quant_scales list is likewise ignored here - y = fused_moe( - x, - ids, - scales, - w31, - None, - w2, - None, - torch.bfloat16, - [torch.ones(1, device=DEV)], - )[0] - assert torch.equal(y, base), "quant_scales changed the unquantized result" - - # --------------------------------------------------------------------------- # DeepSeek-R1-0528 MTP layer (model.layers.61) routed-expert geometry. # hidden 7168, moe_intermediate 2048, 256 routed experts, top-8, bf16 weights # and activations. Under ep_size 4 each rank holds E = 64 of those experts and # is handed global ids in [0, 256). -R1_H, R1_I, R1_K = 7168, 2048, 8 -R1_E_LOCAL, R1_EP = 64, 4 -R1_E_GLOBAL = R1_E_LOCAL * R1_EP - - -def _make_r1_weights(num_experts: int, seed: int): - """One rank's [E, 2I, H] / [E, H, I] bf16 weight pair, built in chunks. - - 5.6 GB at E = 64. Materializing the fp32 randn for the whole stack at once - would transiently need another 11 GB; chunking caps the temporary at 64 MB. - """ - g = torch.Generator(device=DEV).manual_seed(seed) - w31 = torch.empty(num_experts, 2 * R1_I, R1_H, dtype=torch.bfloat16, device=DEV) - w2 = torch.empty(num_experts, R1_H, R1_I, dtype=torch.bfloat16, device=DEV) - chunk = max(1, (1 << 26) // (2 * R1_I * R1_H)) - for lo in range(0, num_experts, chunk): - hi = min(lo + chunk, num_experts) - w31[lo:hi] = (torch.randn(hi - lo, 2 * R1_I, R1_H, device=DEV, generator=g) / R1_H**0.5).to( - torch.bfloat16 - ) - w2[lo:hi] = (torch.randn(hi - lo, R1_H, R1_I, device=DEV, generator=g) / R1_I**0.5).to( - torch.bfloat16 - ) - return w31, w2 - - -def _make_r1_routing(num_tokens: int, num_experts: int, seed: int): - """(x, ids, scales) with top-8 routing over `num_experts` global experts.""" - g = torch.Generator(device=DEV).manual_seed(seed) - x = torch.randn(num_tokens, R1_H, device=DEV, generator=g).to(torch.bfloat16) - logits = torch.randn(num_tokens, num_experts, device=DEV, generator=g) - vals, ids = torch.topk(logits, R1_K, dim=-1) - return ( - x, - ids.to(torch.int32).contiguous(), - torch.softmax(vals, dim=-1).contiguous(), - ) def _rejects(y: torch.Tensor, ref: torch.Tensor, tag: str) -> None: """Assert the gates reject `ref` as a description of `y`.""" try: - _assert_moe_close(y, ref) + fused_moe.compare(y, ref) except AssertionError: return raise AssertionError(f"tolerance accepted a wrong reference: {tag}") -def test_r1_mtp_moe_shape_bf16() -> None: - # 3.5x the largest hidden and 2.7x the largest intermediate size covered by - # the shapes above, at the expert count one ep_size-4 rank holds. - # Measured worst deviation over this function's nine comparisons: 2.54 ulp - # element-wise, 1.16 ulp relative RMS -- inside the 8/4 ulp gates. - w31, w2 = _make_r1_weights(R1_E_LOCAL, seed=1234) - for num_tokens in [1, 2, 37, 256, 2048, 8192]: - x, ids, scales = _make_r1_routing(num_tokens, R1_E_LOCAL, seed=num_tokens) - y = fused_moe(x, ids, scales, w31, None, w2, None, torch.bfloat16, [])[0] - assert y.shape == (num_tokens, R1_H) and y.is_contiguous() - _assert_moe_close(y, _ref_moe(x, ids, scales, w31, w2)) - del x, ids, scales, y - torch.cuda.empty_cache() - - # the two output-side switches, at this geometry - x, ids, scales = _make_r1_routing(2048, R1_E_LOCAL, seed=7) - ref = _ref_moe(x, ids, scales, w31, w2) - y = fused_moe(x, ids, scales, w31, None, w2, None, torch.bfloat16, [])[0] - _assert_moe_close(y, ref) - y_unfused = fused_moe( - x, - ids, - scales, - w31, - None, - w2, - None, - torch.bfloat16, - [], - use_fused_finalize=False, - )[0] - _assert_moe_close(y_unfused, ref) - # the flag changes nothing on this path: the cached C++ runner is keyed on - # the dtypes only, so whichever value the process's first call passed is - # the one in force -- and both values give the same bits anyway, measured - # in two fresh processes with the call order swapped - assert torch.equal(y, y_unfused), "use_fused_finalize changed the result" - pool = torch.full((2051, R1_H), float("nan"), dtype=torch.bfloat16, device=DEV) - out = pool[:2048] - assert fused_moe(x, ids, scales, w31, None, w2, None, torch.bfloat16, [], out_tensor=out) == [] - _assert_moe_close(out, ref) - assert bool(torch.isnan(pool[2048:]).all()), "wrote outside out_tensor's rows" - del w31, w2, x, ids, scales, ref, y, y_unfused, pool, out - torch.cuda.empty_cache() - - -def test_r1_mtp_expert_parallel_4way() -> None: - # ep_size 4 over a 256-expert global stack: each rank passes its own 64 - # weights and the *global* ids, and answers only for - # [64*ep_rank, 64*ep_rank + 64). The four rank outputs must sum to the - # 256-expert result, which is what the caller's correctness rests on. - # The four windows are never held at once: each is built, driven, folded - # into the accumulators and freed, so this needs 5.6 GB rather than 22.5. - token_counts = [1, 1024, 8192] - acts = {t: _make_r1_routing(t, R1_E_GLOBAL, seed=1000 + t) for t in token_counts} - got = {t: torch.zeros(t, R1_H, dtype=torch.float32, device=DEV) for t in token_counts} - want = {t: torch.zeros(t, R1_H, dtype=torch.float32, device=DEV) for t in token_counts} - for rank in range(R1_EP): - w31_r, w2_r = _make_r1_weights(R1_E_LOCAL, seed=900 + rank) - for t in token_counts: - x, ids, scales = acts[t] - y = fused_moe( - x, - ids, - scales, - w31_r, - None, - w2_r, - None, - torch.bfloat16, - [], - ep_size=R1_EP, - ep_rank=rank, - )[0] - ref = _ref_moe(x, ids, scales, w31_r, w2_r, ep_size=R1_EP, ep_rank=rank) - _assert_moe_close(y, ref) - got[t] += y.float() - want[t] += ref.float() - del y, ref - if rank == R1_EP - 1: - # a window that catches nothing contributes exactly zero, and the - # same weights read under the wrong window are a different answer - x, ids, scales = acts[1024] - inside_rank0 = (ids % R1_E_LOCAL).to(torch.int32).contiguous() - y = fused_moe( - x, - inside_rank0, - scales, - w31_r, - None, - w2_r, - None, - torch.bfloat16, - [], - ep_size=R1_EP, - ep_rank=rank, - )[0] - assert bool((y == 0).all()), "an out-of-window rank wrote nonzero output" - y = fused_moe( - x, - ids, - scales, - w31_r, - None, - w2_r, - None, - torch.bfloat16, - [], - ep_size=R1_EP, - ep_rank=0, - )[0] - _rejects( - y, - _ref_moe(x, ids, scales, w31_r, w2_r, ep_size=R1_EP, ep_rank=rank), - "rank 3's weights driven at ep_rank=0", - ) - del inside_rank0, y - del w31_r, w2_r - torch.cuda.empty_cache() - for t in token_counts: - # each rank rounds its own partial to bf16 before this add, so the sum - # sits slightly wider than any single rank does (measured 3.53 ulp - # element-wise at 8192 tokens against 2.69 for the widest single rank) - _assert_moe_close(got[t].to(torch.bfloat16), want[t].to(torch.bfloat16)) - del acts, got, want - torch.cuda.empty_cache() - - def test_r1_mtp_expert_relabeling_is_bitwise() -> None: # Nothing in the permutation/gather path may key on which expert index a # weight sits at: permuting the 64-expert stack and relabelling the ids to @@ -770,7 +442,9 @@ def test_r1_mtp_expert_relabeling_is_bitwise() -> None: w31, w2 = _make_r1_weights(R1_E_LOCAL, seed=2222) x, ids, scales = _make_r1_routing(512, R1_E_LOCAL, seed=31337) y = fused_moe(x, ids, scales, w31, None, w2, None, torch.bfloat16, [])[0] - _assert_moe_close(y, _ref_moe(x, ids, scales, w31, w2)) + fused_moe.compare( + y, fused_moe.reference(x, ids, scales, w31, None, w2, None, torch.bfloat16, []) + ) g = torch.Generator(device=DEV).manual_seed(4) perm = torch.randperm(R1_E_LOCAL, device=DEV, generator=g) inv = torch.empty_like(perm) @@ -808,8 +482,14 @@ def test_r1_mtp_autotuned_tactics() -> None: for t in (1, 256, 8192): x, ids, scales = _make_r1_routing(t, R1_E_LOCAL, seed=t) cold = fused_moe(x, ids, scales, w31, None, w2, None, torch.bfloat16, [])[0] - cases[t] = (x, ids, scales, cold, _ref_moe(x, ids, scales, w31, w2)) - _assert_moe_close(cold, cases[t][4]) + cases[t] = ( + x, + ids, + scales, + cold, + fused_moe.reference(x, ids, scales, w31, None, w2, None, torch.bfloat16, []), + ) + fused_moe.compare(cold, cases[t][4]) try: x, ids, scales = cases[8192][:3] with autotune(): @@ -820,7 +500,7 @@ def test_r1_mtp_autotuned_tactics() -> None: moved = 0 for t, (x, ids, scales, cold, ref) in cases.items(): hot = fused_moe(x, ids, scales, w31, None, w2, None, torch.bfloat16, [])[0] - _assert_moe_close(hot, ref) + fused_moe.compare(hot, ref) moved += int(not torch.equal(hot, cold)) # the control this whole test rests on: if tuning had not changed a # single bit, a clean pass here would mean nothing @@ -840,7 +520,7 @@ def _sweep_tactics(combos, args, ref, count_distinct: bool) -> int: for cfg in combos: with AutoTuner.get().replay(cfg): y = fused_moe(*args)[0] - _assert_moe_close(y, ref) + fused_moe.compare(y, ref) if count_distinct: distinct.add(y.view(torch.int16).cpu().numpy().tobytes()) del y @@ -869,7 +549,7 @@ def test_r1_mtp_tactic_space() -> None: w31, w2 = _make_r1_weights(R1_E_LOCAL, seed=1234) x, ids, scales = _make_r1_routing(256, R1_E_LOCAL, seed=99) args = (x, ids, scales, w31, None, w2, None, torch.bfloat16, []) - ref = _ref_moe(x, ids, scales, w31, w2) + ref = fused_moe.reference(x, ids, scales, w31, None, w2, None, torch.bfloat16, []) with tuner.capture() as cap: fused_moe(*args) combos = list(cap) @@ -895,7 +575,7 @@ def test_r1_mtp_tactic_space() -> None: # and at the token count the caller chunks to x, ids, scales = _make_r1_routing(8192, R1_E_LOCAL, seed=8192) args = (x, ids, scales, w31, None, w2, None, torch.bfloat16, []) - ref = _ref_moe(x, ids, scales, w31, w2) + ref = fused_moe.reference(x, ids, scales, w31, None, w2, None, torch.bfloat16, []) with tuner.capture() as cap: fused_moe(*args) combos = list(cap) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/moe/test_modeling_v2_noaux_tc_op.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/moe/test_modeling_v2_noaux_tc_op.py new file mode 100644 index 000000000000..a43d8f1b622f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/moe/test_modeling_v2_noaux_tc_op.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the noaux_tc_op catalog entry.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4.ops.moe.noaux_tc_op import ( # noqa: E501 + noaux_tc_op as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "noaux_tc_op requires a CUDA device" + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + logits = torch.randn( + spec["tokens"], spec["experts"], generator=g, device="cuda", dtype=torch.float32 + ).to(torch.bfloat16) + bias = torch.randn(spec["experts"], generator=g, device="cuda", dtype=torch.float32) + return logits, bias, spec["n_group"], spec["topk_group"], spec["topk"], spec["scale"] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + args = _build(cell.spec, seed=abs(hash(cell.why)) % 2**31) + with validating(op): + weights, ids = op(*args) + ref_weights, ref_ids = op.reference(*args) + + # Ids are an exact claim: routing either picks the same experts or it does + # not, and a near-miss here is a different set of experts, not a rounding. + assert torch.equal(ids, ref_ids), "selected experts differ from the reference" + assert ids.dtype is torch.int32 + op.compare(weights, ref_weights) + + +def test_the_sigmoid_is_the_tanh_form() -> None: + """The entry's reference claims a specific sigmoid; hold the op to it. + + `1/(1+exp(-x))` and `0.5*tanh(x/2)+0.5` agree everywhere a random logit + lands and diverge in the tails, where the tanh form saturates exactly. A + cell over random inputs cannot tell them apart, so this drives the tails + directly -- otherwise `reference` could encode the wrong sigmoid and every + cell would still pass. + """ + experts = 8 + logits = torch.full((4, experts), -30.0, dtype=torch.bfloat16, device="cuda") + logits[:, 0] = 30.0 + bias = torch.zeros(experts, dtype=torch.float32, device="cuda") + weights, _ = op(logits, bias, 1, 1, 2, 1.0) + ref_weights, _ = op.reference(logits, bias, 1, 1, 2, 1.0) + op.compare(weights, ref_weights) + + +def _ok(): + logits = torch.randn(4, 16, dtype=torch.bfloat16, device="cuda") + bias = torch.randn(16, dtype=torch.float32, device="cuda") + return logits, bias + + +def test_guard_refuses_strided_logits() -> None: + """Read as a dense buffer from data_ptr(); strides are ignored.""" + logits = torch.randn(4, 32, dtype=torch.bfloat16, device="cuda")[:, ::2] + _, bias = _ok() + with pytest.raises(AssertionError, match="router_logits must be contiguous"): + with validating(op): + op(logits, bias, 1, 1, 2, 1.0) + + +def test_guard_refuses_a_strided_bias() -> None: + logits, _ = _ok() + bias = torch.randn(32, dtype=torch.float32, device="cuda")[::2] + with pytest.raises(AssertionError, match="bias must be contiguous"): + with validating(op): + op(logits, bias, 1, 1, 2, 1.0) + + +def test_guard_refuses_topk_past_the_expert_count() -> None: + """Reads out of bounds and emits ids outside the expert range.""" + logits, bias = _ok() + with pytest.raises(AssertionError, match="exceeds num_experts"): + with validating(op): + op(logits, bias, 1, 1, 17, 1.0) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/activation/test_modeling_v2_moe_gelu.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/activation/test_modeling_v2_moe_gelu.py new file mode 100644 index 000000000000..1b4479db4b43 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/activation/test_modeling_v2_moe_gelu.py @@ -0,0 +1,306 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe_gelu entry (exact erf GELU, single-tile metadata). + +The entry owns what is certified: `CELLS` is the list, `reference` and +`compare` are the gate. This file turns a cell's spec into tensors and keeps +the checks a cell cannot carry: bitwise agreement with `F.gelu` (what the +DeepSeek-V4.1 vision aligner computes), the tanh-approximation control, the +contiguity guard, the op's own rejections, and the aligner's call form -- +every row count 1..2048 as a prefix of the 2048-row result, one kernel at +every row count, the stale-limit control, and CUDA-graph replay at the +aligner's graph sizes. +""" + +import math +import zlib + +import pytest +import torch +import torch.nn.functional as F +from torch.profiler import ProfilerActivity, profile + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.activation.moe_gelu import ( # noqa: E501 + moe_gelu as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "moe_gelu requires a CUDA device" + +# The DeepSeek-V4.1 vision aligner applies GELU to w1's output, [n_llm_tokens, 5120]. +WIDTH = 5120 +# torch.testing.assert_close defaults for bf16. +BF16_RTOL, BF16_ATOL = 1.6e-2, 1e-5 + + +def _single_tile(rows: int): + """One tile holding every row: the dense-activation metadata.""" + limit = torch.tensor([rows], dtype=torch.int32, device="cuda") + n_tiles = torch.tensor([1], dtype=torch.int32, device="cuda") + return limit, n_tiles, rows + + +def _ref_erf(x: torch.Tensor) -> torch.Tensor: + """The entry's reference over every row of `x` (fp64 erf GELU, rounded to x.dtype).""" + return op.reference(x, *_single_tile(x.shape[0])) + + +def _ref_tanh_fp32(x: torch.Tensor) -> torch.Tensor: + """The WRONG function: tanh-approximated GELU, in fp32, rounded to x.dtype.""" + xf = x.float() + inner = math.sqrt(2.0 / math.pi) * (xf + 0.044715 * xf * xf * xf) + return (0.5 * xf * (1.0 + torch.tanh(inner))).to(x.dtype) + + +def _outside_bf16_gate(a: torch.Tensor, b: torch.Tensor) -> int: + """Elements where |a - b| exceeds assert_close's default bf16 gate (finite b only).""" + af, bf = a.float(), b.float() + fin = torch.isfinite(bf) + bad = (af - bf).abs() > BF16_ATOL + BF16_RTOL * bf.abs() + return int((bad & fin).sum().item()) + + +def _all_bf16_bit_patterns() -> torch.Tensor: + """Every one of the 65536 bf16 bit patterns (finite, +-inf, NaNs), zero-padded to [13, 5120].""" + bits = torch.arange(-32768, 32768, dtype=torch.int32, device="cuda").to(torch.int16) + vals = bits.view(torch.bfloat16) + rows = math.ceil(vals.numel() / WIDTH) + x = torch.zeros(rows * WIDTH, dtype=torch.bfloat16, device="cuda") + x[: vals.numel()] = vals + return x.view(rows, WIDTH) + + +def _randn(rows: int, sigma: float, seed: int) -> torch.Tensor: + g = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(rows, WIDTH, device="cuda", generator=g) * sigma).to(torch.bfloat16) + + +def _build(spec: dict, seed: int): + """The cell's input and its single-tile metadata, `limit` rows computed of `rows`.""" + if spec["kind"] == "bits": + x = _all_bf16_bit_patterns() + else: + x = _randn(spec["rows"], spec["sigma"], seed=seed) + rows = x.shape[0] + limit = torch.tensor([spec["limit"]], dtype=torch.int32, device="cuda") + n_tiles = torch.tensor([1], dtype=torch.int32, device="cuda") + return x, limit, n_tiles, rows + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + x, limit, n_tiles, rows = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + x_before = x.clone() + tail = cell.spec["limit"] < rows + if tail: + # Park a NaN block of the output's exact size in the caching allocator so + # the op's freshly allocated output reuses it; unwritten rows then stay NaN. + sentinel = torch.full((rows, WIDTH), float("nan"), dtype=torch.bfloat16, device="cuda") + ptr = sentinel.data_ptr() + del sentinel + out = op(x, limit, n_tiles, rows) + if tail: + assert out.data_ptr() == ptr, ( + "the output did not reuse the sentinel block; the probe is blind" + ) + assert out.shape == x.shape and out.dtype == torch.bfloat16 and out.is_contiguous() + assert out.data_ptr() != x.data_ptr() + # Bitwise: the bit-pattern cell's input holds NaNs. + assert torch.equal(x.view(torch.int16), x_before.view(torch.int16)), "input was mutated" + op.compare(out, op.reference(x, limit, n_tiles, rows)) + # The vendor aligner calls F.gelu (approximate='none') on a bf16 tensor. + computed = slice(0, cell.spec["limit"]) + vendor = F.gelu(x[computed], approximate="none") + same = (out[computed].view(torch.int16) == vendor.view(torch.int16)) | ( + torch.isnan(out[computed]) & torch.isnan(vendor) + ) + assert bool(same.all()), f"{int((~same).sum())} of {same.numel()} differ bitwise from F.gelu" + + +def test_tanh_approximation_control() -> None: + """The tanh-approximated GELU is a different function, and this harness can see it.""" + x = _all_bf16_bit_patterns() + out = op(x, *_single_tile(x.shape[0])) + wrong = _ref_tanh_fp32(x) + # The control is not blind: the wrong reference itself fails against fp64 erf GELU. + n_ref = _outside_bf16_gate(wrong, _ref_erf(x)) + n_exh = _outside_bf16_gate(out, wrong) + with pytest.raises(AssertionError): + torch.testing.assert_close(out, wrong, equal_nan=True) + + xr = _randn(1024, 1.0, seed=2024) + outr = op(xr, *_single_tile(1024)) + wrong_r = _ref_tanh_fp32(xr) + n_rand = _outside_bf16_gate(outr, wrong_r) + with pytest.raises(AssertionError): + torch.testing.assert_close(outr, wrong_r) + print( + f"\ntanh control: outside bf16 gate -- tanh ref vs fp64 erf over all patterns {n_ref}/{x.numel()}; " + f"op vs tanh over all patterns {n_exh}/{x.numel()}; " + f"op vs tanh on N(0,1) [1024, {WIDTH}] {n_rand}/{xr.numel()}; " + f"op vs fp64 erf on N(0,1) {_outside_bf16_gate(outr, _ref_erf(xr))}" + ) + assert n_ref > 0 and n_exh > 0 and n_rand > 0 + + +def test_non_contiguous_input_rejected_by_guard() -> None: + big = _randn(64, 3.0, seed=5) + big = torch.cat([big, big[:, :256]], dim=1) + view = big[:, :WIDTH] + assert not view.is_contiguous() + with pytest.raises(AssertionError, match="must be contiguous"): + with validating(op): + op(view, *_single_tile(64)) + + +def test_loud_rejections() -> None: + x = _randn(64, 1.0, seed=3) + one = torch.tensor([1], dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="must be equal to tile_tokens_dim"): + op(x, torch.tensor([64], dtype=torch.int32, device="cuda"), one, 32) + with pytest.raises(RuntimeError, match="tile_idx_to_mn_limit must be int32"): + op(x, torch.tensor([64], dtype=torch.int64, device="cuda"), one, 64) + with pytest.raises(RuntimeError, match="interm_size must be divisible by 8"): + op(torch.zeros(4, WIDTH + 4, dtype=torch.bfloat16, device="cuda"), *_single_tile(4)) + with pytest.raises(RuntimeError, match="Unsupported input dtype"): + op(torch.zeros(4, WIDTH, dtype=torch.float32, device="cuda"), *_single_tile(4)) + with pytest.raises(RuntimeError, match="input must be 2D"): + op(torch.zeros(1, 4, WIDTH, dtype=torch.bfloat16, device="cuda"), *_single_tile(1)) + torch.cuda.synchronize() + + +# ─── The DeepSeek-V4.1 vision aligner's call form: rows 1..2048 ────────────── +# +# The multimodal prototype's aligner calls this op once per request on +# h = cells @ w1^T + b1, [C, 5120] bf16, with the dense single-tile metadata it +# builds host-side in `prepare`: tile_idx_to_mn_limit = torch.tensor([C], +# int32, cuda), num_non_exiting_tiles = torch.tensor([1], int32, cuda), +# tile_tokens_dim = C. C is at most 1902 over the 900 MMMU val documents (182, +# 546, 601 and 986 in the golden cases); N_max is 2048. The kernel walks rows +# in a grid-stride loop with grid = min(SMs x occupancy, rows), so the row +# count sizes the grid and bounds the loop. +_ALIGNER_N_MAX = 2048 +_ALIGNER_GRAPH_ROWS = (182, 546, 601, 986, 1902) +# A fill kernel between calls splits the profiler's kernel stream per call. +_MARKER = torch.zeros(1, device="cuda") + + +def _prototype_metadata(rows: int): + """The metadata exactly as VisionTower.prepare builds it.""" + limit = torch.tensor([rows], dtype=torch.int32, device="cuda") + tiles = torch.tensor([1], dtype=torch.int32, device="cuda") + return limit, tiles, rows + + +def test_aligner_every_row_count_is_the_same_per_row_function() -> None: + """Every r in 1..2048 is bitwise the first r rows of the 2048-row call, which is held to fp64 erf GELU. + + Each output is filled with NaN before it is freed, so a row the kernel + skipped cannot pick up correct bits from the previous call's block. + """ + x = _randn(_ALIGNER_N_MAX, 3.0, seed=90) + before = x.clone() + full = op(x, *_prototype_metadata(_ALIGNER_N_MAX)) + op.compare(full, _ref_erf(x)) + assert torch.equal(full, F.gelu(x, approximate="none")) + exact = float((full == _ref_erf(x)).float().mean()) + full_bits = full.view(torch.int16) + mismatch = torch.zeros(_ALIGNER_N_MAX + 1, dtype=torch.bool, device="cuda") + for r in range(1, _ALIGNER_N_MAX + 1): + out = op(x[:r], *_prototype_metadata(r)) + mismatch[r] = (out.view(torch.int16) != full_bits[:r]).any() + out.fill_(float("nan")) + del out + torch.cuda.synchronize() + bad = mismatch.nonzero().flatten().tolist() + assert not bad, ( + f"{len(bad)} row counts differ from the 2048-row result's rows, first {bad[:10]}" + ) + assert torch.equal(x, before), "input was mutated" + print( + f"\n[aligner] 2048 rows: bit-exact {exact:.6f} vs fp64 erf, bitwise F.gelu; " + f"r = 1..2048: every prefix bitwise equal" + ) + + +def test_aligner_one_kernel_at_every_row_count() -> None: + """Probe: one launch of the same kernel at every row count 1..2048 (compiled ahead of time).""" + x = _randn(_ALIGNER_N_MAX, 1.0, seed=91) + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for r in range(1, _ALIGNER_N_MAX + 1): + _MARKER.fill_(1.0) + op(x[:r], *_prototype_metadata(r)) + torch.cuda.synchronize() + evs = sorted( + (e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA), + key=lambda e: e.time_range.start, + ) + per_call, cur = [], None + for e in evs: + if "FillFunctor" in e.name: + if cur is not None: + per_call.append(cur) + cur = [] + elif cur is not None and not e.name.startswith(("Memcpy", "Memset")): + cur.append(e.name) # the metadata's host-to-device copies are not kernels + per_call.append(cur) + assert len(per_call) == _ALIGNER_N_MAX and all(len(c) == 1 for c in per_call) + names = {c[0] for c in per_call} + print(f"\n[aligner] rows 1..2048, one launch per call, kernels: {sorted(names)}") + assert len(names) == 1, names + + +def test_aligner_controls() -> None: + """The prefix check rejects a read one row off, and a stale limit (one row short) leaves the last row unwritten.""" + x = _randn(1024, 3.0, seed=92) + full = op(x, *_prototype_metadata(1024)) + for r in (1, 2, 182, 1023): + shifted = op(x[1 : r + 1].contiguous(), *_prototype_metadata(r)) + n_bad = int((shifted.view(torch.int16) != full[:r].view(torch.int16)).sum()) + print(f"\n[aligner control row shift] r={r}: {n_bad} / {r * WIDTH} elements differ") + assert n_bad > r * WIDTH // 2 + rows = 601 + sentinel = torch.full((rows, WIDTH), float("nan"), dtype=torch.bfloat16, device="cuda") + ptr = sentinel.data_ptr() + del sentinel + stale = torch.tensor([rows - 1], dtype=torch.int32, device="cuda") + out = op( + x[:rows].contiguous(), stale, torch.tensor([1], dtype=torch.int32, device="cuda"), rows + ) + assert out.data_ptr() == ptr, ( + "the output did not reuse the sentinel block; the control is blind" + ) + assert torch.equal(out[: rows - 1], full[: rows - 1]) and bool(torch.isnan(out[rows - 1]).all()) + print( + f"[aligner control stale limit] limit {rows - 1} for {rows} rows: row {rows - 1} left unwritten (NaN)" + ) + + +def test_aligner_cuda_graph_at_the_prototype_sizes() -> None: + """Captured at the prototype's graph sizes with its metadata tensors; replays on new data are bitwise eager.""" + pool = torch.cuda.graph_pool_handle() + graphs = {} + for rows in sorted(_ALIGNER_GRAPH_ROWS, reverse=True): + x = _randn(rows, 3.0, seed=93) + meta = _prototype_metadata(rows) + op(x, *meta) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + out = op(x, *meta) + graphs[rows] = (graph, x, meta, out) + for rnd in range(2): + for rows in _ALIGNER_GRAPH_ROWS: + graph, x, meta, out = graphs[rows] + x.copy_(_randn(rows, 3.0, seed=94 + 10 * rnd + rows)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, op(x, *meta)), rows + assert torch.equal(out, F.gelu(x, approximate="none")), rows + print( + f"\n[aligner] captured at C = {_ALIGNER_GRAPH_ROWS}: replays bitwise equal to eager and to F.gelu" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/activation/test_modeling_v2_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/activation/test_modeling_v2_silu_and_mul.py new file mode 100644 index 000000000000..3e05947a4b24 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/activation/test_modeling_v2_silu_and_mul.py @@ -0,0 +1,445 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the silu_and_mul entry (clamped SwiGLU, `swiglu_limit`). + +The entry owns what is certified: `CELLS` is the list, `reference` and +`compare` are the gate. This file turns a cell's spec into tensors and keeps +the checks a cell cannot carry: the hand-read clamp branches, the every-row-count +prefix sweep with its one-compiled-variant probe and its one-row-shift control, +CUDA-graph capture at the decode row counts and at the engine's capture set, +the wrong-variant controls, the inner-stride guard, and fresh inputs at every +row count the envelope engines dumped. +""" + +import zlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.activation.silu_and_mul import ( # noqa: E501 + silu_and_mul as op, +) +from tensorrt_llm._torch.modules.swiglu import silu_and_mul_kernel + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "silu_and_mul requires a CUDA device" + +LIMIT = 10.0 +# Spread of the synthetic gate/up values. At sigma 8 the normal tail beyond +# +-10 (1.25 sigma) is expected to hold about 21% of each half, so every clamp +# branch (gate above, gate below, up both sides) should be populated in every +# shape; `_assert_clamp_regions_populated` checks that rather than trusting it. +SIGMA = 8.0 + + +def _ref(x: torch.Tensor, limit, out_dtype: torch.dtype) -> torch.Tensor: + """The entry's reference at `swiglu_limit=limit`, output dtype `out_dtype`.""" + return op.reference(x, dtype=out_dtype, swiglu_limit=limit) + + +def _inputs(rows: int, d: int, dtype: torch.dtype, seed: int) -> torch.Tensor: + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(rows, 2 * d, device="cuda", dtype=torch.float32, generator=g) * SIGMA + return x.to(dtype) + + +def _assert_clamp_regions_populated(x: torch.Tensor) -> None: + d = x.shape[-1] // 2 + gate, up = x[:, :d].float(), x[:, d:].float() + for name, n in ( + ("gate > L", (gate > LIMIT).sum().item()), + ("gate < -L", (gate < -LIMIT).sum().item()), + ("up > L", (up > LIMIT).sum().item()), + ("up < -L", (up < -LIMIT).sum().item()), + ): + assert n > 0, f"no element with {name}: the clamp is not exercised" + + +def _build(spec: dict, seed: int) -> torch.Tensor: + """The cell's `[rows, 2d]` input, as a view of a wider buffer when it names a row pad.""" + rows, d, pad = spec["rows"], spec["d"], spec["row_pad"] + if pad == 0: + return _inputs(rows, d, spec["x_dtype"], seed) + big = _inputs(rows, d + pad // 2, spec["x_dtype"], seed) + x = big[:, : 2 * d] + assert x.stride() == (2 * d + pad, 1) + return x + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + x = _build(spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + limit = spec["swiglu_limit"] + if limit is not None and limit > 0 and x.numel(): + _assert_clamp_regions_populated(x) + x_before = x.clone() + kwargs = dict( + dtype=spec["dtype"], + swiglu_limit=limit, + swiglu_alpha=spec["swiglu_alpha"], + swiglu_beta=spec["swiglu_beta"], + ) + out = op(x, **kwargs) + torch.cuda.synchronize() + out_dtype = spec["dtype"] or spec["x_dtype"] + assert out.shape == (spec["rows"], spec["d"]) and out.dtype == out_dtype and out.is_contiguous() + assert torch.equal(x, x_before), "input was mutated" + op.compare(out, op.reference(x, **kwargs)) + if spec["swiglu_alpha"] is not None or spec["swiglu_beta"] is not None: + # Passing alpha=1.0, beta=0.0 explicitly equals leaving them None. + assert torch.equal(out, op(x, dtype=spec["dtype"], swiglu_limit=limit)) + + +def test_clamp_at_hand_picked_values() -> None: + """Each clamp branch on a value whose answer can be read off directly. + + gate 25 -> 10 (clamped from above); gate -25 stays -25 (no lower clamp); + up 30 -> 10 and up -30 -> -10 (symmetric clamp). + """ + gate = torch.tensor([25.0, -25.0, 1.0, 1.0, 25.0], device="cuda") + up = torch.tensor([1.0, 1.0, 30.0, -30.0, 30.0], device="cuda") + x = torch.cat([gate, up]).view(1, -1) + out = op(x, swiglu_limit=LIMIT) + ten = torch.tensor(10.0, device="cuda") + expected = torch.stack( + [ + F.silu(ten) * 1.0, + F.silu(torch.tensor(-25.0, device="cuda")) * 1.0, + F.silu(torch.tensor(1.0, device="cuda")) * 10.0, + F.silu(torch.tensor(1.0, device="cuda")) * -10.0, + F.silu(ten) * 10.0, + ] + ).view(1, -1) + torch.testing.assert_close(out, expected) + + +def test_cuda_graph_replay() -> None: + """Captured once, replayed on new data in the same static buffer.""" + x_static = _inputs(8, 2304, torch.bfloat16, seed=21) + op(x_static, swiglu_limit=LIMIT) # compile outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out_static = op(x_static, swiglu_limit=LIMIT) + for seed in (22, 23): + x_static.copy_(_inputs(8, 2304, torch.bfloat16, seed=seed)) + graph.replay() + torch.cuda.synchronize() + op.compare(out_static, _ref(x_static, LIMIT, torch.bfloat16)) + + +# ─── every row count at the shared-expert width (d = 2304) ───────────────── +# +# DeepSeek-V4.1's shared expert hands this op one row per token the rank holds +# this step, and d is always 2304. The row count is any integer from 1 to +# max_num_tokens (at most 16384 per rank). In layers 21-39 the decoder bounded +# replay (tep4/tp4) runs a decoder pass whose rows add the finishing requests' +# 128-row tails to the step's generation rows, up to 128 x 256 = 32768, and the +# shared expert runs there on all of them in one call. The DSpark drafter's +# shared expert (same d, same limit) sees its draft block, 5 rows per +# generation request: at most 1280. Rows enter the kernel only as the size of +# grid axis 0: program (i, j) reads row i and writes row i, and nothing it +# computes depends on how many rows there are. The tests below check that on +# the GPU, at every row count from 1 to 32768. + +_D = 2304 +_MAX_ROWS = 32768 +# The core target's decode CUDA-graph batch sizes, plus two odd ones. +_GRAPH_ROWS = (1, 2, 3, 4, 7, 8, 16, 24, 32) + + +def _compiled_variants() -> dict: + """Triton's in-process cache of compiled `silu_and_mul_kernel` variants on this device (key -> kernel).""" + caches = silu_and_mul_kernel.device_caches[torch.cuda.current_device()] + return caches[0] + + +def test_every_row_count_1_to_32768_is_the_same_per_row_function() -> None: + """Every prefix x[:r], r = 1..32768, returns exactly the first r rows of the 32768-row result. + + The 32768-row result is held to the fp32 reference, so every row count is + covered. Each output is poisoned with NaN before it is freed. The next, + slightly longer output may reuse that allocator block, and any row the + kernel skipped would otherwise still hold the previous call's correct bits. + + One-compiled-variant probe: the row count is not part + of the compiled kernel's key. The sweep must add no entry to Triton's cache + of compiled `silu_and_mul_kernel` variants, and the variant the 32768-row + call used must be the only one with this call's constants. + """ + x = _inputs(_MAX_ROWS, _D, torch.bfloat16, seed=41) + _assert_clamp_regions_populated(x) + full = op(x, swiglu_limit=LIMIT) + op.compare(full, _ref(x, LIMIT, torch.bfloat16)) + full_bits = full.view(torch.int16) + before = set(_compiled_variants()) + # This call's variant: bf16 output, no scale pointer, bf16 input, limit 10, alpha 1, beta 0. + mine = [ + k + for k in before + if k.startswith("[('*bf16', 'D'), ('i32', 'D'), ('constexpr', None), ('*bf16', 'D')") + and "('constexpr', 10.0), ('constexpr', 1.0), ('constexpr', 0.0)" in k + ] + bad = torch.zeros((), dtype=torch.int64, device="cuda") + for r in range(1, _MAX_ROWS + 1): + out = op(x[:r], swiglu_limit=LIMIT) + assert out.shape == (r, _D) and out.dtype == torch.bfloat16 + bad += (out.view(torch.int16) != full_bits[:r]).sum() + out.fill_(float("nan")) + after = set(_compiled_variants()) + print( + f"[rows] r = 1..{_MAX_ROWS}: {bad.item()} elements differ from the {_MAX_ROWS}-row result; " + f"compiled variants before the sweep {len(before)}, after {len(after)}; " + f"bf16 limit-10 variants: {len(mine)}" + ) + assert bad.item() == 0, f"{bad.item()} elements differ from the {_MAX_ROWS}-row result's rows" + assert after == before, ( + f"the row sweep compiled {len(after - before)} new variant(s): {sorted(after - before)}" + ) + assert len(mine) == 1, sorted(before) + + +def test_row_sweep_harness_sees_a_one_row_shift() -> None: + """Control for the sweep above: its bitwise per-row check catches a read one row off. + + A kernel that read row i + 1 for output row i would compute on real data + and return finite, plausible numbers. The same check that the sweep + applies must reject it at every row count, including r = 1. + """ + x = _inputs(65, _D, torch.bfloat16, seed=42) + full = op(x[:64], swiglu_limit=LIMIT) + for r in (1, 2, 7, 63, 64): + shifted = op(x[1 : r + 1], swiglu_limit=LIMIT) + n_bad = int((shifted.view(torch.int16) != full.view(torch.int16)[:r]).sum().item()) + print(f"[control row shift] r={r}: {n_bad} / {r * _D} elements differ") + assert n_bad > r * _D // 2 + + +def test_cuda_graph_capture_at_decode_row_counts() -> None: + """Captured at each decode batch size without an eager call at that size; replayed on new data. + + One eager call at 8 rows compiles the kernel. Every graph is then captured + at a row count that has not run eagerly in this test. That works because + the row count is not part of the compiled kernel. Each replay is checked + against the reference and against an eager call on the same data, + bitwise. + """ + op(_inputs(8, _D, torch.bfloat16, seed=50), swiglu_limit=LIMIT) + torch.cuda.synchronize() + for rows in _GRAPH_ROWS: + x_static = _inputs(rows, _D, torch.bfloat16, seed=51 + rows) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out_static = op(x_static, swiglu_limit=LIMIT) + for seed in (1, 2): + x_new = _inputs(rows, _D, torch.bfloat16, seed=1000 * seed + rows) + x_static.copy_(x_new) + out_static.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + op.compare(out_static, _ref(x_new, LIMIT, torch.bfloat16)) + assert torch.equal(out_static, op(x_new, swiglu_limit=LIMIT)), rows + + +# ─── deliberately wrong controls ──────────────────────────────────────────── + + +def _wrong_symmetric_gate_clamp(x): + d = x.shape[-1] // 2 + gate = torch.clamp(x[:, :d].float(), -LIMIT, LIMIT) + up = torch.clamp(x[:, d:].float(), -LIMIT, LIMIT) + return (F.silu(gate) * up).to(x.dtype) + + +def _wrong_up_gate_order(x): + # MoE FC1's [up; gate] order applied to the dense [gate; up] layout. + d = x.shape[-1] // 2 + gate = torch.clamp(x[:, d:].float(), max=LIMIT) + up = torch.clamp(x[:, :d].float(), -LIMIT, LIMIT) + return (F.silu(gate) * up).to(x.dtype) + + +def _wrong_unclamped(x): + return _ref(x, None, x.dtype) + + +def _wrong_swigluoai(x): + # GPT-OSS body: g * sigmoid(1.702 g) * (u + 1), same clamps. + d = x.shape[-1] // 2 + gate = torch.clamp(x[:, :d].float(), max=LIMIT) + up = torch.clamp(x[:, d:].float(), -LIMIT, LIMIT) + return (gate * torch.sigmoid(1.702 * gate) * (up + 1.0)).to(x.dtype) + + +_WRONG = { + "symmetric_gate_clamp": _wrong_symmetric_gate_clamp, + "up_gate_order": _wrong_up_gate_order, + "unclamped": _wrong_unclamped, + "swigluoai": _wrong_swigluoai, +} + + +def _mismatches(out: torch.Tensor, ref: torch.Tensor) -> int: + # Elementwise version of assert_close's default bf16 gate (rtol 1.6e-2, atol 1e-5). + return int((~torch.isclose(out.float(), ref.float(), rtol=1.6e-2, atol=1e-5)).sum().item()) + + +@pytest.mark.parametrize("name", list(_WRONG)) +def test_wrong_variant_fails_on_inputs_beyond_the_limit(name: str) -> None: + """Each wrong SwiGLU is rejected by the same gate the certified cases pass. + + Inputs beyond +-10 are what separate the clamp variants: the symmetric + gate clamp and the unclamped body only differ from the right one there, + so the same controls are also run on inputs confined to |x| < 10, where + those two must be indistinguishable from the certified op -- which is + why the certified cases draw at sigma 8. + """ + x = _inputs(4096, 2304, torch.bfloat16, seed=31) + _assert_clamp_regions_populated(x) + out = op(x, swiglu_limit=LIMIT) + op.compare(out, _ref(x, LIMIT, torch.bfloat16)) # the right one passes + + wrong = _WRONG[name](x) + n_bad = _mismatches(out, wrong) + max_abs = (out.float() - wrong.float()).abs().max().item() + print( + f"[control {name}] wide inputs: {n_bad} / {out.numel()} outside the bf16 gate, max |diff| {max_abs:.6g}" + ) + with pytest.raises(AssertionError): + torch.testing.assert_close(out, wrong) + + x_in = torch.clamp(x, -9.5, 9.5) + out_in = op(x_in, swiglu_limit=LIMIT) + n_bad_in = _mismatches(out_in, _WRONG[name](x_in)) + print( + f"[control {name}] inputs within +-9.5: {n_bad_in} / {out_in.numel()} outside the bf16 gate" + ) + if name in ("symmetric_gate_clamp", "unclamped"): + assert n_bad_in == 0, "the clamp controls must be invisible without inputs beyond the limit" + else: + assert n_bad_in > 0 + + +def test_non_unit_inner_stride_is_rejected_by_the_guard() -> None: + """A transposed view is read as if contiguous; the guard refuses it.""" + base = torch.randn(2 * 2304, 64, device="cuda", dtype=torch.bfloat16) + x = base.t() + assert x.stride(-1) != 1 + with pytest.raises(AssertionError, match="unit stride"): + with validating(op): + op(x, swiglu_limit=LIMIT) + + +# ─── Capture set: decode, DSpark target verify, DSpark drafter ───────────── +# +# The captured batch sizes G, per rank: every G 1..32, plus TRT-LLM's padded +# default list for cuda_graph_config {enable_padding: true, max_batch_size: +# 256} (40..128 step 8, 192, 256); the padding-off list at 256 (1..32, 64, 128, +# 256) is a subset. A captured step is generation-only. The core target's +# shared expert hands this op G rows (next_n 1) or, with DSpark on, 6G (1 + 5 +# draft tokens per request), and the drafter's shared expert (same d = 2304, +# same limit) hands it its 5-token draft block, 5G rows. Those are per-rank +# token counts in every member (the shared expert runs on the rank's own +# tokens, never on the dep4-gathered set). The engine captures one graph per +# batch size into one shared memory pool, largest first, and replays whichever +# one the step's batch selects. G, 5G and 6G over the 46 sizes: 114 row +# counts, 1..1536. + +_GSET = tuple(range(1, 33)) + (40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256) +_VERIFY_GRAPH_ROWS = tuple(sorted({f * g for g in _GSET for f in (1, 5, 6)})) + + +@pytest.mark.parametrize("warm", [True, False], ids=["eager_first", "no_eager_at_these_rows"]) +def test_cuda_graph_capture_at_gset_row_counts(warm: bool) -> None: + """One graph per captured row count, replayed interleaved on new data. + + The row counts are G, 5G and 6G over the 46 G, all captured into one shared pool. + + `warm=True` follows the engine: an eager call at each row count before + its capture. `warm=False` runs one eager call at 8 rows only (the kernel + compile) and captures each row count without an eager call at it in + this test (rows are not part of the compiled kernel). Graphs are captured largest + first into a shared pool and then replayed in an interleaved order over + three rounds. Before each replay the static input gets new data and the + output is poisoned with NaN. The output must match an eager call on the + same data bitwise, and the fp32 reference at the default bf16 gate. It + must also differ from the previous round's output, which shows the replay + read the new input rather than a copy baked in at capture. + """ + op(_inputs(8, _D, torch.bfloat16, seed=60), swiglu_limit=LIMIT) + torch.cuda.synchronize() + pool = torch.cuda.graph_pool_handle() + graphs = {} + for rows in sorted(_VERIFY_GRAPH_ROWS, reverse=True): + x_static = _inputs(rows, _D, torch.bfloat16, seed=61 + rows) + if warm: + op(x_static, swiglu_limit=LIMIT) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + out_static = op(x_static, swiglu_limit=LIMIT) + graphs[rows] = (graph, x_static, out_static) + torch.cuda.synchronize() + order = _VERIFY_GRAPH_ROWS[::2] + _VERIFY_GRAPH_ROWS[1::2][::-1] + previous = {} + for rnd in range(3): + for rows in order: + graph, x_static, out_static = graphs[rows] + x_new = _inputs(rows, _D, torch.bfloat16, seed=7000 + 10000 * rnd + rows) + if rnd == 0: + _assert_clamp_regions_populated(x_new) + x_static.copy_(x_new) + out_static.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert out_static.shape == (rows, _D) + op.compare(out_static, _ref(x_new, LIMIT, torch.bfloat16)) + assert torch.equal(out_static, op(x_new, swiglu_limit=LIMIT)), (rows, rnd) + if rows in previous: + assert not torch.equal(out_static, previous[rows]), (rows, rnd) + previous[rows] = out_static.clone() + print( + f"[capture warm={warm}] {len(graphs)} graphs in one pool, rows {min(graphs)}..{max(graphs)}: " + f"3 rounds, every replay bitwise = eager and inside the gate, each differing from the previous round" + ) + + +# ─── fresh inputs at the row counts the envelope engines dumped ───────────── +# +# Row counts of this call in the core target's runtime-value dumps of the +# envelope engines (max_num_tokens 16384, max_batch_size 256; tep4, +# tp4 and dep4, padded graphs and eager, DSpark on; plus a 1M-token and a +# 128K-token long-context engine), 76 values up to 16384; the DSpark drafter's +# block rows N (5 per generation request) from the same dumps' drafter +# attention records, 40 values 5..1280, which its shared expert sees too; and +# 32768, the top of the decoder bounded-replay pass. +_CG_DUMPED_ROWS = ( + 1, 2, 4, 5, 6, 7, 8, 12, 16, 18, 24, 29, 32, 34, 48, 60, 96, 108, 128, 134, 144, 168, 192, 228, 240, + 288, 330, 336, 384, 432, 462, 480, 528, 546, 576, 597, 610, 624, 672, 720, 768, 780, 852, 1152, 1536, + 2778, 2784, 3870, 4208, 4245, 4388, 4394, 4591, 4868, 8467, 9193, 9240, 9364, 9370, 9524, 10441, + 10444, 10447, 10451, 16298, 16308, 16316, 16320, 16322, 16327, 16328, 16329, 16333, 16340, 16378, + 16384, +) # fmt: skip +_CG_DRAFTER_ROWS = ( + 5, 10, 15, 20, 40, 50, 80, 90, 110, 115, 120, 125, 140, 160, 190, 200, 240, 275, 280, 320, 360, 370, + 385, 400, 440, 455, 480, 520, 560, 580, 585, 600, 640, 650, 665, 710, 715, 720, 960, 1280, +) # fmt: skip + + +def test_dumped_row_counts_fresh_inputs() -> None: + """Fresh inputs at every dumped row count (and 32768), compared directly with the reference.""" + rows_all = sorted(set(_CG_DUMPED_ROWS) | set(_CG_DRAFTER_ROWS) | {_MAX_ROWS}) + for rows in rows_all: + x = _inputs(rows, _D, torch.bfloat16, seed=90000 + rows) + _assert_clamp_regions_populated(x) + out = op(x, swiglu_limit=LIMIT) + assert out.shape == (rows, _D) and out.dtype == torch.bfloat16 and out.is_contiguous() + op.compare(out, _ref(x, LIMIT, torch.bfloat16)) + print( + f"[dumped rows] {len(rows_all)} row counts, fresh inputs, all inside the default bf16 gate" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_candidate_block_select.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_candidate_block_select.py new file mode 100644 index 000000000000..372999bdf983 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_candidate_block_select.py @@ -0,0 +1,1734 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the candidate_block_select entry. + +The entry owns what is certified: `CELLS` is the list, `reference` and `compare` are the gate. +The reference is the checkpoint's own ``select_candidate_blocks`` (DeepSeek-V4.1 +``inference/model.py``) rebuilt from torch ops with the entry's tie rule: -inf-pad to whole +blocks, block amax, pin the block holding the newest visible position to +inf, top- +``topk_blocks`` blocks, drop picks whose value is -inf, expand to positions. Candidate masking +only binds once a row has more than ``topk_blocks`` blocks, i.e. above 16384 visible positions at +the checkpoint's 2048 x 8, so every selecting case here runs above that. + +This file turns a cell's spec into tensors and keeps the checks a cell cannot carry: the wrapper +guards, CUDA-graph capture and replay, both launch configurations agreeing, the deliberately-wrong +controls, the kernel-variant maps, the seeded random packed steps, the row-prefix and +row-independence checks, the decode calls at row offsets and the far end of the maximal cell. +""" + +from typing import Optional + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.candidate_block_select import ( # noqa: E501 + candidate_block_select as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "candidate_block_select requires a CUDA device" + +DEV = "cuda" +TOPK_BLOCKS = 2048 +BLOCK = 8 +BINDS_ABOVE = TOPK_BLOCKS * BLOCK # 16384 visible positions + + +def _distinct(n: int, w: int, seed: int) -> torch.Tensor: + """fp32 [n, w], each row a random permutation of distinct values of both signs. + + Distinct scores make every block max distinct, so the kept set is unique and the + comparison with torch.topk is exact. + """ + g = torch.Generator(device=DEV).manual_seed(seed) + perm = torch.argsort(torch.rand(n, w, generator=g, device=DEV), dim=1) + return (perm.float() - w / 2) * 0.03125 + + +def _distinct_from(g: torch.Generator, n: int, w: int) -> torch.Tensor: + """``_distinct``'s rows drawn from an existing generator (for row-chunked fills).""" + perm = torch.argsort(torch.rand(n, w, generator=g, device=DEV), dim=1) + return (perm.float() - w / 2) * 0.03125 + + +def _distinct_any(n: int, w: int, seed: int) -> torch.Tensor: + """``_distinct``, row-chunked (at most 2^27 elements a chunk) above 2^29 elements, where one + argsort would need tens of GB; every row is still a random permutation of distinct values + (|value| < 2^23 / 32, exact in fp32). Shapes up to 2^29 elements get ``_distinct``'s values.""" + if n * w <= 1 << 29: + return _distinct(n, w, seed) + out = torch.empty(n, w, device=DEV) + g = torch.Generator(device=DEV).manual_seed(seed) + step = max(1, (1 << 27) // w) + for r0 in range(0, n, step): + out[r0 : r0 + step] = _distinct_from(g, min(step, n - r0), w) + return out + + +def _ref_select(logits, compress_lens, topk_blocks=TOPK_BLOCKS, block_size=BLOCK, pin=True): + """The checkpoint's select_candidate_blocks as it stands, ``torch.topk`` and all: the source + of the controls (``pin=False``, another ``block_size``) and of the kept-multiset check on tied + scores. The expected mask itself is always ``op.reference``.""" + width = logits.size(-1) + scores = F.pad(logits, (0, -width % block_size), value=-torch.inf) + scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1) + num_blocks = scores.size(-1) + if pin: + last = (compress_lens - 1) // block_size + scores = scores.masked_fill( + torch.arange(num_blocks, device=logits.device) == last, torch.inf + ) + top = scores.topk(min(topk_blocks, num_blocks), dim=-1) + keep = torch.zeros_like(scores, dtype=torch.bool).scatter_( + -1, top.indices, top.values > -torch.inf + ) + return keep.repeat_interleave(block_size, dim=-1)[..., :width] + + +def _ref_top_positions(logits, k=TOPK_BLOCKS * BLOCK): + """Control: keep the top-k *positions* instead of blocks (same position budget).""" + top = logits.topk(min(k, logits.size(-1)), dim=-1) + return torch.zeros_like(logits, dtype=torch.bool).scatter_( + -1, top.indices, top.values > -torch.inf + ) + + +def _prefill_reference(logits, first_pos): + """Reference mask for one request's causal prefill rows (ratio 1: query p sees p + 1). + + ``logits`` [q, W] holds the request's columns from 0; row r is the query at position + ``first_pos + r``, so it sees columns [0, first_pos + r + 1). Returns (mask, visible lengths). + """ + q, _ = logits.shape + lens = torch.arange(first_pos + 1, first_pos + q + 1, device=DEV) + return op.reference(logits, None, lens.int(), TOPK_BLOCKS, BLOCK, 1), lens + + +def _check_decode_rows(logits, lens, mask): + """Row i sees columns [0, lens[i]) (decode, next_n 1): the whole mask against the reference.""" + op.compare(mask, op.reference(logits, None, lens, TOPK_BLOCKS, BLOCK, 1)) + + +# --------------------------------------------------------------------------------------------- +# the certified cells (op.CELLS): prefill, decode, ties, -inf blocks, empty rows, strided logits, +# and the packed prefill steps untiled and in the target's row tiles +# --------------------------------------------------------------------------------------------- + + +def _kept_block_scores(logits_row, keep_row): + width = logits_row.numel() + s = F.pad(logits_row.float(), (0, -width % BLOCK), value=-torch.inf).view(-1, BLOCK).amax(-1) + k = F.pad(keep_row, (0, -width % BLOCK), value=False).view(-1, BLOCK).any(-1) + return torch.sort(s[k]).values + + +def _check_kept_multiset(logits, lens, got): + for i, n in enumerate(lens.tolist()): + ref = _ref_select(logits[i : i + 1, :n], n)[0] + assert torch.equal( + _kept_block_scores(logits[i, :n], got[i, :n]), _kept_block_scores(logits[i, :n], ref) + ), f"row {i}: kept block scores differ from the checkpoint's selection" + + +def _build(spec): + """The call's arguments for a non-packed cell, from the spec (seeds fixed per kind).""" + kind = spec["kind"] + i32 = dict(dtype=torch.int32, device=DEV) + if kind == "prefill_one": + seq = spec["seq"] + starts = torch.zeros(seq, **i32) if spec["explicit_starts"] else None + lens = torch.arange(1, seq + 1, **i32) + return _distinct(seq, seq, seed=seq), starts, lens, TOPK_BLOCKS, BLOCK, 1 + if kind == "prefill_chunk": + first, q = spec["first"], spec["q"] + w = first + q + logits = _distinct(q, w + spec["pad"], seed=7)[:, :w] + assert logits.stride(0) == w + spec["pad"] + lens = torch.arange(first + 1, first + q + 1, **i32) + return logits, torch.zeros(q, **i32), lens, TOPK_BLOCKS, BLOCK, 1 + if kind == "prefill_ragged": + keys, queries = spec["keys"], spec["queries"] + w = sum(keys) + spec["spare"] + logits = _distinct(sum(q for _, q in queries), w, seed=11) + starts, ends, base = [], [], 0 + for k, (first, q) in zip(keys, queries): + starts.append(torch.full((q,), base, **i32)) + ends.append(base + torch.arange(first + 1, first + q + 1, **i32)) + base += k + return logits, torch.cat(starts), torch.cat(ends), TOPK_BLOCKS, BLOCK, 1 + if kind == "decode": + lens = torch.tensor(spec["lens"], **i32) + w = int(lens.max()) + 3 + logits = _distinct(lens.numel(), w, seed=3) + # columns past each row's visible range hold NaN / +inf: they must be neither read nor kept + logits[torch.arange(w, device=DEV) >= lens[:, None]] = float("nan") + logits[:, -1] = float("inf") + return logits, None, lens, TOPK_BLOCKS, BLOCK, 1 + if kind == "decode_next_n": + visible, next_n = torch.tensor(spec["visible"], **i32), spec["next_n"] + logits = _distinct(visible.numel() * next_n, int(visible.max()), seed=5) + return logits, None, visible, TOPK_BLOCKS, BLOCK, next_n + if kind == "bf16": + lens = torch.tensor(spec["lens"], **i32) + g = torch.Generator(device=DEV).manual_seed(9) + logits = torch.randn(lens.numel(), int(lens.max()), generator=g, device=DEV) + return logits.to(torch.bfloat16), None, lens, TOPK_BLOCKS, BLOCK, 1 + if kind == "ties": + lens = torch.tensor(spec["lens"], **i32) + w = int(lens.max()) + g = torch.Generator(device=DEV).manual_seed(13) + choices = torch.tensor([-1.0, -0.0, 0.0, 1.0, 2.0], device=DEV) + logits = choices[torch.randint(0, 5, (lens.numel(), w), generator=g, device=DEV)] + # last row: almost every block scores -0.0 or +0.0, which rank equal (the threshold tie) + logits[3] = choices[torch.randint(0, 3, (w,), generator=g, device=DEV)] + return logits, None, lens, TOPK_BLOCKS, BLOCK, 1 + if kind == "minus_inf": + lens = torch.tensor(spec["lens"], **i32) + logits = _distinct(lens.numel(), int(lens.max()), seed=17) + blocks = logits.view(lens.numel(), -1, BLOCK) + # row 0: 1000 all -inf blocks of 2500 -> 1500 finite blocks, fewer than 2048: all kept + blocks[0, 100:1100] = float("-inf") + # row 1: same, and the pinned (newest) block itself all -inf: kept anyway + blocks[1, 100:1100] = float("-inf") + blocks[1, (20000 - 1) // BLOCK] = float("-inf") + # row 2: scattered -inf positions inside finite blocks do not change a block's score + logits[2, ::3] = float("-inf") + return logits, None, lens, TOPK_BLOCKS, BLOCK, 1 + if kind == "empty_rows": + lens = torch.tensor(spec["lens"], **i32) + return _distinct(lens.numel(), spec["width"], seed=19), None, lens, TOPK_BLOCKS, BLOCK, 1 + if kind == "no_rows": + logits = _distinct(3, spec["width"], seed=19)[:0] + return logits, None, torch.zeros(0, **i32), TOPK_BLOCKS, BLOCK, 1 + if kind == "col_strided": + lens = torch.tensor(spec["lens"], **i32) + logits = _distinct(lens.numel(), int(lens.max()), seed=23).t().contiguous().t() + assert logits.stride(1) == lens.numel() + return logits, None, lens, TOPK_BLOCKS, BLOCK, 1 + raise AssertionError(kind) + + +def _drive_packed(spec): + """A packed prefill step called as the target calls it: whole, or one call per row tile on + the step's logits rows [a:b] with the views ks[a:b] / ke[a:b]. Every mask lands on a block + pre-filled with 1s, so every byte of it has to be written.""" + q, kv = list(spec["q_lens"]), list(spec["kv_lens"]) + seed = 8 * (sum(kv) % 100_000) + len(q) if spec["tiled"] else sum(kv) + len(q) + try: + step = _packed_step(q, kv, seed=seed, with_want=False) + T, W = step["logits"].shape + tiles = _tiles(T, W) if spec["tiled"] else [(0, T)] + reused = 0 + for a, b in tiles: + torch.cuda.empty_cache() + junk = torch.ones((b - a) * W, dtype=torch.uint8, device=DEV) + junk_ptr = junk.data_ptr() + del junk + args = (step["logits"][a:b], step["ks"][a:b], step["ke"][a:b], TOPK_BLOCKS, BLOCK, 1) + got = op(*args) + reused += got.data_ptr() == junk_ptr + op.compare(got, op.reference(*args)) + del got + vis = step["ke"] - step["ks"] + print( + f"\n T={T} W={W} k={len(q)} tiles {[b - a for a, b in tiles][:6]}" + f"{'...' if len(tiles) > 6 else ''} ({len(tiles)}): exact; " + f"{int((vis > BINDS_ABOVE).sum())} binding rows, longest {int(vis.max())}; " + f"{reused} of {len(tiles)} masks on a poisoned block" + ) + del step + finally: + _release_big_logits() + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + if spec["kind"] == "packed": + _drive_packed(spec) + return + args = _build(spec) + logits, _, lens, _, _, _ = args + got = op(*args) + assert got.dtype == torch.bool and got.shape == logits.shape and got.is_contiguous() + op.compare(got, op.reference(*args)) + + kind = spec["kind"] + if kind == "prefill_one": + # masking is active exactly on the rows that see more than 2048 blocks + seq = spec["seq"] + reachable = torch.arange(seq, device=DEV) < lens[:, None] + narrowed = (got != reachable).any(dim=1) + binding = lens > BINDS_ABOVE + assert torch.equal(narrowed, binding) + if binding.any(): + kept_blocks = got[binding].view(-1, seq // BLOCK, BLOCK).any(-1).sum(-1) + assert bool((kept_blocks == TOPK_BLOCKS).all()) + elif kind == "prefill_chunk": + assert int((lens > BINDS_ABOVE).sum()) == spec["first"] + spec["q"] - BINDS_ABOVE + elif kind in ("bf16", "ties"): + # ties are frequent: the kept block-score multiset equals the checkpoint's own selection + _check_kept_multiset(logits, lens, got) + elif kind == "minus_inf": + assert int(got[0].sum()) == (2500 - 1000) * BLOCK + assert bool(got[1, 19992:20000].all()) + elif kind == "empty_rows": + assert not got[0].any() and not got[2].any() and bool(got[1].all()) + elif kind == "no_rows": + assert got.shape == (0, spec["width"]) + + +def test_both_launch_configs_agree(): + """Calls of <= 256 rows and of more rows take different launch configurations.""" + rows = 512 + g = torch.Generator(device=DEV).manual_seed(59) + lens = torch.randint(15000, 40000, (rows,), generator=g, device=DEV, dtype=torch.int32) + logits = _distinct(rows, 40000, seed=61) + whole = op(logits, None, lens, TOPK_BLOCKS, BLOCK, 1) + halves = torch.cat( + [ + op(logits[:256], None, lens[:256], TOPK_BLOCKS, BLOCK, 1), + op(logits[256:], None, lens[256:], TOPK_BLOCKS, BLOCK, 1), + ] + ) + assert torch.equal(whole, halves) + _check_decode_rows(logits, lens, whole) + + +def test_is_valid_rejects_mismatched_length_tensors(): + """The kernel would read these wrongly without an error; the entry's guard refuses them.""" + logits = _distinct(8, 64, seed=67) + lens = torch.full((8,), 64, dtype=torch.int32, device=DEV) + zeros = torch.zeros(8, dtype=torch.int32, device=DEV) + rejected = ( + ("row_ends", (logits, None, lens[:4], TOPK_BLOCKS, BLOCK, 1)), + ("row_ends", (logits, None, lens, TOPK_BLOCKS, BLOCK, 6)), + ("row_ends", (logits, None, torch.stack([lens, lens], 1)[:, 0], TOPK_BLOCKS, BLOCK, 1)), + ("row_starts", (logits, zeros[:4], lens, TOPK_BLOCKS, BLOCK, 1)), + ("row_starts", (logits, torch.stack([zeros, zeros], 1)[:, 0], lens, TOPK_BLOCKS, BLOCK, 1)), + ) + for match, args in rejected: + with pytest.raises(AssertionError, match=match): + with validating(op): + op(*args) + # the guard admits the well-formed call + with validating(op): + op(logits, zeros, lens, TOPK_BLOCKS, BLOCK, 1) + + +# --------------------------------------------------------------------------------------------- +# CUDA graph +# --------------------------------------------------------------------------------------------- + + +def test_cuda_graph_capture_and_replay(): + rows, w = 8, 40000 + logits = torch.empty(rows, w, device=DEV) + lens = torch.empty(rows, dtype=torch.int32, device=DEV) + logits.copy_(_distinct(rows, w, seed=29)) + lens.copy_( + torch.tensor([16385, 20000, 40000, 30000, 17000, 100, 16384, 39999], dtype=torch.int32) + ) + op(logits, None, lens, TOPK_BLOCKS, BLOCK, 1) # compile outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(logits, None, lens, TOPK_BLOCKS, BLOCK, 1) + for seed, new_lens in ( + (31, [40000] * 8), + (37, [16400, 16393, 20001, 1, 35000, 16385, 39000, 22222]), + ): + logits.copy_(_distinct(rows, w, seed=seed)) + lens.copy_(torch.tensor(new_lens, dtype=torch.int32)) + graph.replay() + torch.cuda.synchronize() + _check_decode_rows(logits, lens, out) + + +# --------------------------------------------------------------------------------------------- +# deliberately-wrong controls: each matches above 16384 visible positions only if the test +# cannot see the difference; each must match below it (where masking is a no-op) +# --------------------------------------------------------------------------------------------- + +CONTROL_LENS = [ + 4096, + 12000, + 16384, + 16385, + 16400, + 17000, + 18000, + 20000, + 24000, + 28000, + 32768, + 40000, + 50000, + 65536, + 100000, + 131072, +] + + +def _control_masks(logits, lens): + """(entry mask, {control: mask}) per row, all over each row's own visible width.""" + got = op(logits, None, lens, TOPK_BLOCKS, BLOCK, 1) + got16 = op(logits, None, lens, TOPK_BLOCKS, 16, 1) + rows = [] + for i, n in enumerate(lens.tolist()): + row = logits[i : i + 1, :n] + rows.append( + ( + n, + got[i, :n], + { + "no pin of the newest block": _ref_select(row, n, pin=False)[0], + "block size 16": _ref_select(row, n, block_size=16)[0], + "top-k of positions, not blocks": _ref_top_positions(row)[0], + "entry kernel run with block_size 16": got16[i, :n], + }, + ) + ) + return rows + + +def test_controls_differ_above_16384_only(): + lens = torch.tensor(CONTROL_LENS, dtype=torch.int32, device=DEV) + logits = _distinct(lens.numel(), int(lens.max()), seed=41) + consumer = _distinct(lens.numel(), int(lens.max()), seed=43) # a Reindex layer's own scores + rows = _control_masks(logits, lens) + for _, _, controls in rows: # the kernel honours block_size: its 16 run is the 16 reference + assert torch.equal( + controls["entry kernel run with block_size 16"], controls["block size 16"] + ) + wrong_rows = {} + wrong_top512 = {} + for i, (n, entry, controls) in enumerate(rows): + op.compare( + entry[None], + op.reference(logits[i : i + 1, :n], None, lens[i : i + 1], TOPK_BLOCKS, BLOCK, 1), + ) + entry_top = consumer[i, :n].masked_fill(~entry, -torch.inf).topk(512).indices.sort().values + for name, ctl in controls.items(): + differs = not torch.equal(ctl, entry) + if n <= BINDS_ABOVE: + assert not differs, f"control '{name}' differs at visible {n} <= 16384" + continue + wrong_rows[name] = wrong_rows.get(name, 0) + int(differs) + ctl_top = consumer[i, :n].masked_fill(~ctl, -torch.inf).topk(512).indices.sort().values + wrong_top512[name] = wrong_top512.get(name, 0) + int( + not torch.equal(ctl_top, entry_top) + ) + binding_rows = sum(1 for n in CONTROL_LENS if n > BINDS_ABOVE) + print(f"\ncontrols over {binding_rows} rows with more than 16384 visible positions:") + for name in wrong_rows: + print( + f" {name}: mask differs on {wrong_rows[name]} rows, consumer top-512 differs on {wrong_top512[name]} rows" + ) + for name, n_wrong in wrong_rows.items(): + assert n_wrong > 0, f"control '{name}' was not detected above 16384" + + +def test_controls_on_a_20480_token_prefill(): + """The same controls over every query row of a 20480-token prefill (visible 1 .. 20480).""" + seq = 20480 + logits = _distinct(seq, seq, seed=47) + consumer = _distinct(seq, seq, seed=53) + compress_lens = torch.arange(1, seq + 1, device=DEV).unsqueeze(-1) + unreachable = torch.arange(seq, device=DEV) >= compress_lens + masked = logits.masked_fill(unreachable, -torch.inf) + got = op(logits, None, compress_lens.squeeze(-1).int(), TOPK_BLOCKS, BLOCK, 1) + op.compare( + got, op.reference(logits, None, compress_lens.squeeze(-1).int(), TOPK_BLOCKS, BLOCK, 1) + ) + controls = { + "no pin of the newest block": _ref_select(masked, compress_lens, pin=False) & ~unreachable, + "block size 16": _ref_select(masked, compress_lens, block_size=16) & ~unreachable, + "top-k of positions, not blocks": _ref_top_positions(masked) & ~unreachable, + } + binding = compress_lens.squeeze(-1) > BINDS_ABOVE + cons = consumer.masked_fill(unreachable, -torch.inf)[binding] + entry_top = ( + cons.masked_fill(~got[binding], -torch.inf).topk(512, dim=-1).indices.sort(dim=-1).values + ) + print( + f"\ncontrols over the {int(binding.sum())} prefill rows that see more than 16384 positions:" + ) + for name, ctl in controls.items(): + differs = (ctl != got).any(dim=1) + assert not bool(differs[~binding].any()), f"control '{name}' differs at or below 16384" + ctl_top = ( + cons.masked_fill(~ctl[binding], -torch.inf) + .topk(512, dim=-1) + .indices.sort(dim=-1) + .values + ) + top_differs = (ctl_top != entry_top).any(dim=1) + print( + f" {name}: mask differs on {int(differs.sum())} rows, " + f"consumer top-512 differs on {int(top_differs.sum())} rows" + ) + assert int(differs[binding].sum()) > 0, f"control '{name}' was not detected above 16384" + + +# ============================================================================================= +# The DeepSeek-V4.1 core target's first envelope: layer 20 (ratio 1) calls this entry +# +# - prefill, whenever the step's packed context key count S_1 = sum_b kv_len_b exceeds 16384: +# logits = deepgemm_fp8_fp4_mqa_logits(...) fp32 [tc, S_1] in DeepGEMM's full-width layout (row +# stride align(S_1 + 256, 256)), row_starts = ks, row_ends = ke (fresh int32 [tc]), next_n 1. +# Envelope: tc <= 8192 query rows (max_num_tokens), k <= 32 context requests (max_batch_size), +# kv_len_b <= 16384 (max_seq_len), so S_1 <= 32 x 16384 = 524288 and every row sees <= 16384 +# positions (the pre-filter never binds there; the mask is the row's range minus non-pinned +# all -inf blocks). +# - decode, in every step (captured ones included) of an engine whose static width +# max_ctx = max_seq_len + max_draft_tokens exceeds 16384: logits = +# deepgemm_fp8_fp4_paged_mqa_logits(...) fp32 [G * next_n, max_ctx] (row stride ceil256(max_ctx)), +# row_starts None, row_ends = compressed_visible_lens[tc:T] (a row-offset view of a +# [max_num_tokens] buffer), next_n argument 1 (the verify rows carry their own lengths), rows +# G * next_n <= 32 x 6. The runtime dumps show max_ctx 16398 / 16403 (DSpark engines); an engine +# without drafts has max_ctx 16384 and never calls it. +# ============================================================================================= + +ENV_ROWS = 8192 +ENV_REQS = 32 +ENV_KV = 16384 +ENV_W = ENV_REQS * ENV_KV # 524288 +KERNEL_NAME = "_candidate_block_select_kernel" +# max_ctx = (max_seq_len + max_draft_tokens) // 1 at max_seq_len 16384: 0 / 5 / 6 draft tokens, the +# two DSpark widths of the runtime dumps, and a W % 16 == 0 width past them +DECODE_WIDTHS = (16384, 16389, 16390, 16398, 16403, 32768) +CAPTURE_BATCHES = (1, 2, 3, 4, 5, 6, 7, 8, 16, 24, 32) # boot 1..8, the accuracy padded set +_NAN = float("nan") +_ENV_STORE = {} + + +def _align(x: int, a: int) -> int: + return (x + a - 1) // a * a + + +def _ri(g: torch.Generator, lo: int, hi: int) -> int: + """Uniform int in [lo, hi].""" + return lo + int(torch.randint(0, hi - lo + 1, (1,), generator=g)) + + +_DEFAULT_LOGITS_ELEMS = ENV_ROWS * _align(ENV_W + 256, 256) + + +def _prefill_logits(rows: int, width: int) -> torch.Tensor: + """A [rows, width] view in DeepGEMM's full-width prefill layout (row stride + align(width + 256, 256)), carved from one reusable buffer (17.2 GB, grown for a larger step + and shrunk back by ``_release_big_logits``); its whole span is NaN.""" + stride = _align(width + 256, 256) + buf = _ENV_STORE.get("logits") + if buf is None or buf.numel() < rows * stride: + _ENV_STORE.pop("logits", None) + del buf + torch.cuda.empty_cache() + buf = torch.empty(max(_DEFAULT_LOGITS_ELEMS, rows * stride), device=DEV) + _ENV_STORE["logits"] = buf + view = buf[: rows * stride].view(rows, stride) + view.fill_(_NAN) + return view[:, :width] + + +def _release_big_logits() -> None: + """Drop the prefill buffer if a step grew it past its default size.""" + buf = _ENV_STORE.get("logits") + if buf is not None and buf.numel() > _DEFAULT_LOGITS_ELEMS: + _ENV_STORE.pop("logits") + del buf + torch.cuda.empty_cache() + + +def _dead_blocks(g: torch.Generator, rows: int, width: int, rate: float = 1 / 32) -> torch.Tensor: + """Random whole blocks of BLOCK columns (block j = columns [8j, 8j + 8) of the row's range).""" + nblk = -(-width // BLOCK) + return (torch.rand(rows, nblk, generator=g, device=DEV) < rate).repeat_interleave(BLOCK, dim=1)[ + :, :width + ] + + +def _packed_step(q_lens, kv_lens, seed: int, with_want: bool = True) -> dict: + """One packed prefill call as layer 20 makes it. + + Request b holds kv_lens[b] keys at columns [base_b, base_b + kv_lens[b]) (base_b = the keys + of the requests before it) and its last q_lens[b] positions are this step's query rows; the + row at position p sees [base_b, base_b + p + 1). Scores in a row's visible range are distinct + (a random permutation), with whole blocks set to -inf at random: those must be dropped unless + pinned, which makes every row's logits addressing observable even where the pre-filter does + not bind. Every other element (unreachable columns, other requests' columns, the stride + padding) is NaN and must never be read. A request with q_lens[b] = 0 only holds keys (its rows + are not in the call, as for the decoder-replay tail). ``want`` is ``op.reference`` over each + request's own rows and keys, independent of ``ks`` / ``ke`` (skipped with ``with_want=False``). + """ + T, W = sum(q_lens), sum(kv_lens) + logits = _prefill_logits(T, W) + want = torch.zeros(T, W, dtype=torch.bool, device=DEV) if with_want else None + g = torch.Generator(device=DEV).manual_seed(seed) + ks, ke, bases = [], [], [] + r0 = base = 0 + for b, (q, kv) in enumerate(zip(q_lens, kv_lens)): + first = kv - q + if ( + q == 0 + ): # keys only: the request's keys sit in the gather, none of its rows is in the call + bases.append(base) + base += kv + continue + if q * kv <= 1 << 29: + vals = _distinct(q, kv, seed=seed * 64 + b) + vals.masked_fill_(_dead_blocks(g, q, kv), -torch.inf) + lens = torch.arange(first + 1, kv + 1, device=DEV) + if with_want: + req_want, _ = _prefill_reference(vals, first) + want[r0 : r0 + q, base : base + kv] = req_want + logits[r0 : r0 + q, base : base + kv] = vals.masked_fill_( + torch.arange(kv, device=DEV) >= lens[:, None], _NAN + ) + else: # a long request: the same construction in row chunks of at most 2^27 elements + gq = torch.Generator(device=DEV).manual_seed(seed * 64 + b) + step = max(1, (1 << 27) // kv) + for c0 in range(0, q, step): + c1 = min(q, c0 + step) + vals = _distinct_from(gq, c1 - c0, kv) + vals.masked_fill_(_dead_blocks(g, c1 - c0, kv), -torch.inf) + if with_want: + req_want, _ = _prefill_reference(vals, first + c0) + want[r0 + c0 : r0 + c1, base : base + kv] = req_want + del req_want + ctx_lens = torch.arange(first + c0 + 1, first + c1 + 1, device=DEV) + logits[r0 + c0 : r0 + c1, base : base + kv] = vals.masked_fill_( + torch.arange(kv, device=DEV) >= ctx_lens[:, None], _NAN + ) + del vals + lens = torch.arange(first + 1, kv + 1, device=DEV) + ks.append(torch.full((q,), base, dtype=torch.int32, device=DEV)) + ke.append((base + lens).to(torch.int32)) + bases.append(base) + r0 += q + base += kv + return { + "logits": logits, + "ks": torch.cat(ks), + "ke": torch.cat(ke), + "want": want, + "bases": bases, + "q_lens": list(q_lens), + "kv_lens": list(kv_lens), + } + + +def _envelope_partition(g: torch.Generator, k: int): + """Random (q_lens, kv_lens) of a context step inside the envelope that reaches the call: + k >= 2 requests, T = sum q in [k, 8192], q_b <= kv_b <= 16384, S_1 = sum kv in (16384, 524288]. + Fresh prompts (kv = q), cached prefixes / later chunks and requests at max_seq_len are mixed.""" + mode = _ri(g, 0, 3) + T = [ + _ri(g, k, max(k, 64 * k)), + _ri(g, k, ENV_ROWS), + ENV_ROWS - _ri(g, 0, min(7, ENV_ROWS - k)), + _ri(g, k, 300), + ][mode] + cuts = sorted((torch.randperm(T - 1, generator=g)[: k - 1] + 1).tolist()) if k > 1 else [] + q = [b - a for a, b in zip([0] + cuts, cuts + [T])] + kv = [] + for qb in q: + c = _ri(g, 0, 3) + kv.append([qb, ENV_KV, _ri(g, qb, ENV_KV), _ri(g, max(qb, ENV_KV - 64), ENV_KV)][c]) + while sum(kv) <= ENV_KV: + kv[_ri(g, 0, k - 1)] = ENV_KV + return q, kv + + +def _binding_partition(g: torch.Generator, k: int): + """Outside the gates' envelope (a request longer than 16384 keys, the long-context shape): + k in 1..32, T <= 8192, at least one request with more than 16384 keys whose rows bind, and + S_1 <= 524288.""" + T = _ri(g, k, ENV_ROWS) + cuts = sorted((torch.randperm(T - 1, generator=g)[: k - 1] + 1).tolist()) if k > 1 else [] + q = [b - a for a, b in zip([0] + cuts, cuts + [T])] + kv = [_ri(g, qb, ENV_KV) for qb in q] + for b in range(k): + if b == 0 or _ri(g, 0, 3) == 0: + kv[b] = max(kv[b], _ri(g, ENV_KV + 1, 40960)) + while sum(kv) > ENV_W: + b = _ri(g, 1, k - 1) + kv[b] = q[b] + return q, kv + + +def _call_prefill(step, ks=None, ke=None): + return op( + step["logits"], + step["ks"] if ks is None else ks, + step["ke"] if ke is None else ke, + TOPK_BLOCKS, + BLOCK, + 1, + ) + + +def _first_bad_row(got, want): + bad = (got != want).any(dim=1).nonzero() + return int(bad[0]) if bad.numel() else None + + +def _fill_decode( + logits: torch.Tensor, vis: torch.Tensor, next_n: int, seed: int, first_end: Optional[int] = None +) -> torch.Tensor: + """Fill a decode call's logits [G * next_n, W] and visible lengths in place; return the + reference mask. Request b's last row sees e_b and its row j sees e_b - (next_n - 1 - j): e_b + past 16384 (binding), around 16384, anywhere in [1, W], or 0 / 1 (a capture padding row); + ``first_end`` pins e_0 (e.g. W: the far-end column visible). Columns past a row's length are + NaN (never read).""" + rows, W = logits.shape + G = rows // next_n + g = torch.Generator().manual_seed(seed) + ends = [] + for _ in range(G): + c = _ri(g, 0, 3) + past = _ri(g, ENV_KV + 1, W) if W > ENV_KV else _ri(g, 1, W) + ends.append([past, _ri(g, ENV_KV - 40, min(W, ENV_KV + 8)), _ri(g, 1, W), _ri(g, 0, 1)][c]) + if first_end is not None: + ends[0] = first_end + e = torch.tensor(ends, dtype=torch.int32, device=DEV).repeat_interleave(next_n) + j = torch.arange(next_n, dtype=torch.int32, device=DEV).repeat(G) + lens = (e - (next_n - 1 - j)).clamp(min=0).to(torch.int32) + vals = _distinct_any(rows, W, seed=seed) + vals.masked_fill_( + _dead_blocks(torch.Generator(device=DEV).manual_seed(seed), rows, W), -torch.inf + ) + reach = torch.arange(W, device=DEV) < lens[:, None] + want = op.reference(vals, None, lens, TOPK_BLOCKS, BLOCK, 1) + logits.copy_(vals.masked_fill_(~reach, _NAN)) + vis.copy_(lens) + return want + + +def _decode_buffers(rows: int, W: int): + """Static decode inputs in the paged logits layout (row stride ceil256(W), padding +inf) and + a [max_num_tokens] visible-length buffer, as the step metadata holds it.""" + buf = torch.full((rows, _align(W, 256)), float("inf"), device=DEV) + return buf[:, :W], torch.zeros(ENV_ROWS, dtype=torch.int32, device=DEV) + + +def _jit_keys(fn) -> set: + """The compiled-variant keys Triton holds for ``fn`` (Triton >= 3.2: ``device_caches``).""" + dc = getattr(fn, "device_caches", None) + if dc is not None: + return {(str(d), k) for d, v in dc.items() for k in v[0].keys()} + return {(str(d), k) for d, v in fn.cache.items() for k in v.keys()} + + +def _jit_clear(fn) -> None: + dc = getattr(fn, "device_caches", None) + if dc is not None: + for v in dc.values(): + v[0].clear() + else: + for v in fn.cache.values(): + v.clear() + + +def _profiled(call): + """Run ``call`` under the profiler; return (result, names of the CUDA kernels it launched).""" + from torch.profiler import ProfilerActivity, profile + + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + out = call() + torch.cuda.synchronize() + names = [e.name for e in prof.events() if str(e.device_type).endswith("CUDA")] + return out, names + + +# ─── (a) the kernel-variant map over the envelope ───────────────────────────────────────────── +# +# One call = one launch of _candidate_block_select_kernel[(rows,)]. Triton specializes on the +# constexprs (HAS_ROW_STARTS: prefill / decode; CHUNK_BLOCKS and num_warps: rows <= 256 or more), +# integer arguments divisible by 16 (num_cols and the mask row stride are both W; the logits row +# strides are multiples of 256 in both layouts) and 16-byte pointer alignment (the decode lengths +# are a view at row tc of the step's buffer). The map below compiles each class once, checks +# every call against the reference, and pins that nothing else in the class compiles. + +_PREFILL_CLASSES = { + "prefill rows<=256 W%16==0": [ + ([1, 1], [16384, 16384]), + ([100, 156], [16384, 4096]), + ([8] * 32, [ENV_KV] * 32), + ], + "prefill rows<=256 W%16!=0": [ + ([1, 1], [8193, 8192]), + ([200, 56], [16000, 3001]), + ([8] * 32, [ENV_KV] * 31 + [16383]), + ], + "prefill rows>256 W%16==0": [ + ([257, 1], [16384, 16384]), + ([4000, 4192], [16384, 8000]), + ([256] * 32, [ENV_KV] * 32), + ], + "prefill rows>256 W%16!=0": [ + ([200, 57], [8193, 8192]), + ([4096, 4096], [16383, 16384]), + ([256] * 32, [ENV_KV] * 31 + [16383]), + ], +} +# (W, next_n, G, tc): tc is the row offset of the lengths view (captured steps: 0) +_DECODE_CLASSES = { + "decode W%16!=0 lengths 16B-aligned": [ + (16398, 6, 1, 0), + (16403, 6, 32, 0), + (16403, 1, 32, 4), + (16398, 6, 5, 7672), + ], + "decode W%16!=0 lengths misaligned": [ + (16398, 6, 1, 1), + (16403, 6, 32, 2), + (16403, 1, 7, 3), + (16398, 6, 24, 573), + ], + "decode W%16==0 lengths 16B-aligned": [(32768, 1, 8, 0), (32768, 6, 32, 0), (16400, 6, 3, 8)], + "decode W%16==0 lengths misaligned": [(32768, 1, 8, 1), (32768, 6, 32, 7674), (16400, 6, 3, 5)], +} + + +def _decode_call(W, next_n, G, tc, seed): + logits, vis_buf = _decode_buffers(G * next_n, W) + vis = vis_buf[tc : tc + G * next_n] + want = _fill_decode(logits, vis, next_n, seed) + return logits, vis, want + + +def test_envelope_kernel_variant_map(): + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.candidate_block_select import ( # noqa: E501 + _candidate_block_select_kernel as fn, + ) + + variant_of, calls, bad = {}, 0, [] + _jit_clear(fn) # the map starts from no variant, so each class compiles its own + assert not _jit_keys(fn) + + def run(cls, i, desc, call, want): + nonlocal calls + before = _jit_keys(fn) + got, names = _profiled(call) + new = _jit_keys(fn) - before + calls += 1 + assert names == [KERNEL_NAME], f"{cls} {desc}: launched {names}" + if i == 0: + assert len(new) == 1, f"{cls}: its first call compiled {len(new)} variants" + variant_of[next(iter(new))] = f"{cls} (first at {desc})" + else: + assert not new, f"{cls} {desc}: compiled {len(new)} new variant(s)" + if not torch.equal(got, want): + bad.append((cls, desc, _first_bad_row(got, want))) + + for cls, cells in _PREFILL_CLASSES.items(): + for i, (q, kv) in enumerate(cells): + step = _packed_step(q, kv, seed=1000 + 17 * i + len(cls)) + run( + cls, + i, + f"T={sum(q)} W={sum(kv)} k={len(q)}", + lambda: _call_prefill(step), + step["want"], + ) + for cls, cells in _DECODE_CLASSES.items(): + for i, (W, next_n, G, tc) in enumerate(cells): + logits, vis, want = _decode_call(W, next_n, G, tc, seed=2000 + 31 * i + len(cls)) + run( + cls, + i, + f"W={W} next_n={next_n} G={G} tc={tc}", + lambda: op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1), + want, + ) + + # closure: random envelope calls in the target's forms compile nothing further + before = _jit_keys(fn) + g = torch.Generator().manual_seed(4242) + for n in range(24): + if n % 2: + step = _packed_step(*_envelope_partition(g, _ri(g, 2, ENV_REQS)), seed=3000 + n) + got, names = _profiled(lambda: _call_prefill(step)) + want = step["want"] + else: + W, next_n = DECODE_WIDTHS[_ri(g, 0, len(DECODE_WIDTHS) - 1)], (1, 6)[_ri(g, 0, 1)] + G = CAPTURE_BATCHES[_ri(g, 0, len(CAPTURE_BATCHES) - 1)] + logits, vis, want = _decode_call( + W, next_n, G, _ri(g, 0, ENV_ROWS - G * next_n), seed=3000 + n + ) + got, names = _profiled(lambda: op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1)) + calls += 1 + assert names == [KERNEL_NAME] + if not torch.equal(got, want): + bad.append(("closure", n, _first_bad_row(got, want))) + after = _jit_keys(fn) + print( + f"\n variant map: {len(variant_of)} compiled variant(s) over {calls} profiled calls " + f"(one {KERNEL_NAME} launch each):" + ) + for key, cls in sorted(variant_of.items(), key=lambda kv: kv[1]): + print(f" {cls}: {key[1][:220]}") + print( + f" 24 random envelope calls after the map compiled {len(after - before)} further variant(s); " + f"{len(bad)} calls differ from the reference" + ) + assert not bad, bad[:5] + assert after == before + assert len(variant_of) == len(_PREFILL_CLASSES) + len(_DECODE_CLASSES) + + +# ─── (b) >= 200 seeded random partitions of the packed prefill call, plus binding extras ───── + + +def test_envelope_random_packed_partitions(): + g = torch.Generator().manual_seed(2323) + stats = { + "calls": 0, + "rows": 0, + "max_T": 0, + "max_W": 0, + "min_W": 1 << 30, + "k": set(), + "binding_rows": 0, + } + ctl_rows = { + "visible end one short (ke - 1)": 0, + "start of the neighbouring request": 0, + "no -inf block drop": 0, + } + ctl_checked = 0 + bad = [] + cases = [("envelope", 2 + i % (ENV_REQS - 1)) for i in range(200)] + [ + ("binding", 1 + i % ENV_REQS) for i in range(24) + ] + for n, (kind, k) in enumerate(cases): + q, kv = (_envelope_partition if kind == "envelope" else _binding_partition)(g, k) + T, W = sum(q), sum(kv) + if kind == "envelope": + assert k <= ENV_REQS and T <= ENV_ROWS and max(kv) <= ENV_KV and ENV_KV < W <= ENV_W + step = _packed_step(q, kv, seed=5000 + n) + got = _call_prefill(step) + if not torch.equal(got, step["want"]): + bad.append((kind, n, q, kv, _first_bad_row(got, step["want"]))) + stats["calls"] += 1 + stats["rows"] += T + stats["max_T"] = max(stats["max_T"], T) + stats["max_W"] = max(stats["max_W"], W) + stats["min_W"] = min(stats["min_W"], W) + stats["k"].add(k) + stats["binding_rows"] += int((step["ke"] - step["ks"] > BINDS_ABOVE).sum()) + if kind == "envelope": + assert int((step["ke"] - step["ks"]).max()) <= BINDS_ABOVE + if n % 10 == 0 and k > 1: # deliberately-wrong controls on the same data + want = step["want"] + ks, ke = step["ks"], step["ke"] + req = torch.repeat_interleave(torch.arange(k, device=DEV), torch.tensor(q, device=DEV)) + bases = torch.tensor(step["bases"], dtype=torch.int32, device=DEV) + col = torch.arange(W, device=DEV) + ctls = { + "visible end one short (ke - 1)": _call_prefill(step, ke=ke - 1), + "start of the neighbouring request": _call_prefill( + step, ks=bases[(req - 1) % k].contiguous() + ), + "no -inf block drop": (col >= ks[:, None]) & (col < ke[:, None]), + } + for name, ctl in ctls.items(): + ctl_rows[name] += int((ctl != want).any(dim=1).sum()) + ctl_checked += T + print( + f"\n {stats['calls']} packed calls (200 inside the envelope, 24 binding): {stats['rows']} rows, " + f"T <= {stats['max_T']}, S_1 {stats['min_W']}..{stats['max_W']}, k in {sorted(stats['k'])}, " + f"{stats['binding_rows']} binding rows; {len(bad)} calls differ from the reference" + ) + print(f" controls over {ctl_checked} rows of every 10th call:") + for name, rows in ctl_rows.items(): + print(f" {name}: {rows} rows differ") + assert not bad, bad[:3] + assert stats["k"] == set(range(1, ENV_REQS + 1)) + for name, rows in ctl_rows.items(): + assert rows > 0, f"control '{name}' was not detected" + + +# ─── the maximal cell's far end (the cells themselves are in op.CELLS) ────────────────────── + + +def test_envelope_maximal_cell_far_end(): + """8192 x 524288: logits rows 8189..8191 start past element 2^32 (rows 4095.. past 2^31) and + the mask spans 2^32 bytes. Controls: the same kernel run on the rows a 32-bit (signed or + unsigned) row offset would read.""" + q, kv = [256] * 32, [ENV_KV] * 32 + step = _packed_step(q, kv, seed=777) + logits, want = step["logits"], step["want"] + T, W = logits.shape + stride = logits.stride(0) + assert (T - 1) * stride > 2**32 and T * W == 2**32 + torch.cuda.empty_cache() + junk = torch.ones(T * W, dtype=torch.uint8, device=DEV) + junk_ptr = junk.data_ptr() + del junk + got = _call_prefill(step) + assert got.data_ptr() == junk_ptr, "the mask did not land on the poisoned block" + assert torch.equal(got, want) and int(got.view(torch.uint8).max()) <= 1 + del got + flat = logits.as_strided((T * stride,), (1,), logits.storage_offset()) + for bits in (31, 32): + wrapped = torch.empty(T, W, device=DEV) + moved = 0 + for i in range(T): + off = (i * stride) % (1 << bits) + moved += off != i * stride + wrapped[i] = flat[off : off + W] + ctl = op(wrapped, step["ks"], step["ke"], TOPK_BLOCKS, BLOCK, 1) + differ = int((ctl != want).any(dim=1).sum()) + print(f"\n row offset mod 2^{bits}: {moved} rows read elsewhere, {differ} rows differ") + assert differ == moved > 0 + del wrapped, ctl + + +# ─── (c) the row axis: every prefix 1..8192 bitwise ──────────────────────────────────────────── + + +@pytest.mark.parametrize( + "q,kv", + [([256] * 32, [ENV_KV] * 32), ([4096, 4096], [24576, 20480]), ([8192, 8192], [24576, 20480])], + ids=["maximal 8192 x 524288", "binding 8192 x 45056", "binding 16384 x 45056"], +) +def test_envelope_rows_prefix_bitwise(q, kv): + """Each row is one program: the r-row call (rows <= 256 and more take different launch + configurations) must equal the first r rows of the full call, bitwise, for every r.""" + step = _packed_step(q, kv, seed=sum(kv)) + full = _call_prefill(step) + assert torch.equal(full, step["want"]) + logits, ks, ke = step["logits"], step["ks"], step["ke"] + T, W = logits.shape + bad, reused = [], 0 + for r in range(T, 0, -1): # largest first: every smaller mask reuses the cached block + junk = torch.ones(r * W, dtype=torch.uint8, device=DEV) + ptr = junk.data_ptr() + del junk + part = op(logits[:r], ks[:r], ke[:r], TOPK_BLOCKS, BLOCK, 1) + reused += part.data_ptr() == ptr + if not torch.equal(part, full[:r]) or int(part.view(torch.uint8).max()) > 1: + bad.append(r) + del part + print( + f"\n every r in 1..{T} at W {W}: {len(bad)} prefixes differ; {reused} of {T} masks on a poisoned block" + ) + assert not bad, bad[:10] + + +# ─── decode rows (next_n 1 and 6) under CUDA-graph capture at the capture set ──────────────── + + +@pytest.mark.parametrize("W", DECODE_WIDTHS) +def test_envelope_decode_rows_under_capture(W): + bad, graphs = [], 0 + for next_n in (1, 6): + for G in CAPTURE_BATCHES: + rows = G * next_n + logits, vis_buf = _decode_buffers(rows, W) + vis = vis_buf[:rows] + want = _fill_decode(logits, vis, next_n, seed=W + 97 * G + next_n) + eager = op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1) + if not torch.equal(eager, want): + bad.append(("eager", next_n, G, _first_bad_row(eager, want))) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1) + graphs += 1 + for rep in (1, 2): + want = _fill_decode(logits, vis, next_n, seed=W + 97 * G + next_n + 1000 * rep) + graph.replay() + torch.cuda.synchronize() + ref_eager = op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1) + if not (torch.equal(out, ref_eager) and torch.equal(out, want)): + bad.append(("replay", next_n, G, rep, _first_bad_row(out, want))) + del graph + print( + f"\n W {W}: {graphs} graphs (next_n 1 and 6 x batches {CAPTURE_BATCHES}), 2 replays each; " + f"{len(bad)} mismatches" + ) + assert not bad, bad[:5] + + +def test_envelope_decode_rows_at_row_offsets(): + """Mixed steps: the lengths are the view [tc:T] of the step's buffer at any tc.""" + bad, calls = [], 0 + for W in (16389, 16398, 16403): + for next_n, G in ((6, 1), (6, 32), (1, 32), (6, 7)): + rows = G * next_n + for tc in (1, 2, 3, 4, 5, 573, 7674, ENV_ROWS - rows): + logits, vis_buf = _decode_buffers(rows, W) + vis = vis_buf[tc : tc + rows] + want = _fill_decode(logits, vis, next_n, seed=W + tc + rows) + vis_buf[:tc] = -7 + vis_buf[tc + rows :] = 1 << 20 + got = op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1) + same = op(logits, None, vis.clone(), TOPK_BLOCKS, BLOCK, 1) + calls += 1 + if not (torch.equal(got, want) and torch.equal(got, same)): + bad.append((W, next_n, G, tc, _first_bad_row(got, want))) + print(f"\n {calls} decode calls at row offsets: {len(bad)} differ") + assert not bad, bad[:5] + + +# ============================================================================================= +# The long-context envelope: max_num_tokens <= 16384 per rank, max_batch_size <= 256, max_seq_len L +# up to 1,048,576, CUDA graphs at every G in GSET +# +# - prefill: layer 20 runs the step's context rows in row tiles (`_index_row_tiles` in the target): +# ranges [a, b) whose fp32 logits (b - a) x align(S_1 + 256, 256) stay within 2^28 elements, one +# range when the whole pass fits, otherwise a multiple of 64 rows per range where that is at +# least 64. Each tile is one call on a fresh [b - a, S_1] logits tile with ks[a:b] / ke[a:b] +# (views at row a). A request holds up to Lm keys (Lm = L + 14 with DSpark, the cache manager's +# length; L without), so rows see up to Lm positions and the pre-filter binds. S_1 = sum_b kv_b +# is the step's gathered key count: Lm or less for one request, more when several long requests +# share a step; past S_1 = 4,194,048 a tile holds fewer than 64 rows and ks[a:b] is no longer +# 16-byte aligned. +# - tail (decoder bounded replay): the same call on 128 rows per tail request, tiled the same way, +# ks / ke at the tail requests' own key ranges inside the step's gather. +# - decode: G * next_n rows, G <= 256 (any G eager, G in GSET captured), next_n 1 or 6 (DSpark +# verify), so up to 1536 rows; calls above 256 rows take the many-rows launch configuration. +# The width is md.max_compressed_kv_lens[1] = Lm + max_draft_tokens: L + 19 under the graph +# metadata with DSpark, L + 14 on eager steps (the dumps show both, e.g. 1048595 / 1048590), and +# L without DSpark (an engine at L 16384 without drafts never calls). The lengths are the view +# [tc:T] of the step's [max_num_tokens] (<= 16384) visible-length buffer. +# ============================================================================================= + +GSET = tuple(range(1, 33)) + (40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256) +LONG_ROWS = 16384 +LONG_REQS = 256 +LONG_L = (16384, 32768, 65536, 131072, 262144, 1048576) +LONG_LM = LONG_L[-1] + 14 # the longest request with DSpark: 1,048,590 keys +LONG_WIDTHS = ( + tuple(L + 14 for L in LONG_L) # DSpark, eager steps: Lm + + tuple(L + 19 for L in LONG_L) # DSpark, graph metadata: Lm + 5 draft tokens + + tuple(L for L in LONG_L if L > BINDS_ABOVE) # DSpark off: L +) +LONG_TILE_ELEMS = 1 << 28 +LONG_BUDGET = (1 << 32) - ( + 1 << 24 +) # elements of a random step's whole fp32 [T, S_1] logits (< 16 GiB) + + +def _tiles(rows: int, s_len: int) -> list: + """The target's prefill row tiles, restated: [a, b) ranges whose logits rows x + align(s_len + 256, 256) fit 2^28 elements; one range when all rows fit, else ranges of a + multiple of 64 rows where that is at least 64.""" + width = _align(s_len + 256, 256) + per = max(1, LONG_TILE_ELEMS // width) + if per >= rows: + return [(0, rows)] + if per >= 64: + per -= per % 64 + return [(a, min(rows, a + per)) for a in range(0, rows, per)] + + +def _tile_masks(step): + """The target's tiled call: one launch per tile on the step's logits rows [a:b] (DeepGEMM's + row stride, a 16-byte-aligned tile start) with the views ks[a:b] / ke[a:b]. Returns the tile + masks and the tiles.""" + T, W = step["logits"].shape + tiles = _tiles(T, W) + masks = [ + op(step["logits"][a:b], step["ks"][a:b], step["ke"][a:b], TOPK_BLOCKS, BLOCK, 1) + for a, b in tiles + ] + return masks, tiles + + +def _call_tiled(step): + """``_tile_masks``, the tile masks joined into the step's [T, S_1] mask.""" + masks, tiles = _tile_masks(step) + return (masks[0] if len(masks) == 1 else torch.cat(masks)), tiles + + +def _long_decode_call(W, next_n, G, tc, seed, first_end=None): + """A decode call in the paged layout (row stride ceil256(W), padding +inf) whose lengths are the + view [tc:tc + G * next_n] of a [16384] buffer.""" + rows = G * next_n + buf = torch.full((rows, _align(W, 256)), float("inf"), device=DEV) + logits, vis_buf = buf[:, :W], torch.zeros(LONG_ROWS, dtype=torch.int32, device=DEV) + vis = vis_buf[tc : tc + rows] + want = _fill_decode(logits, vis, next_n, seed, first_end=first_end) + return logits, vis, vis_buf, want + + +def _long_partition(g: torch.Generator, k: int, budget: int = LONG_BUDGET): + """Random (q_lens, kv_lens) of a context step of k requests that reaches the call: T = sum q in + [k, 16384], q_b <= kv_b <= Lm, S_1 = sum kv > 16384, T * S_1 <= budget (the test builds the + whole step's logits to take its tiles from). Request lengths mix fresh prompts (kv = q), long + cached prefixes up to Lm, requests at Lm, and short ones; long requests are cut back at random + until the step fits the budget.""" + mode = _ri(g, 0, 3) + T = [ + _ri(g, k, max(k, 64 * k)), + _ri(g, k, LONG_ROWS), + LONG_ROWS - _ri(g, 0, min(7, LONG_ROWS - k)), + _ri(g, k, max(k, 300)), + ][mode] + T = min(T, LONG_ROWS) + cuts = sorted((torch.randperm(T - 1, generator=g)[: k - 1] + 1).tolist()) if k > 1 else [] + q = [b - a for a, b in zip([0] + cuts, cuts + [T])] + kv = [] + for qb in q: + c = _ri(g, 0, 4) + kv.append( + [ + qb, + _ri(g, qb, LONG_LM), + _ri(g, max(qb, LONG_LM - 64), LONG_LM), + _ri(g, qb, max(qb, 32768)), + _ri(g, qb, max(qb, 4096)), + ][c] + ) + while T * sum(kv) > budget: + b = _ri(g, 0, k - 1) + kv[b] = q[b] + (kv[b] - q[b]) // 4 + while sum(kv) <= BINDS_ABOVE: + b = _ri(g, 0, k - 1) + kv[b] = min(LONG_LM, kv[b] + BINDS_ABOVE + 1 - sum(kv) + _ri(g, 0, 4096)) + return q, kv + + +# ─── (a) the kernel-variant map over the long-context envelope ──────────────────────────────── +# +# Fourteen classes. Prefill tiles: (rows <= 256 or more) x (W % 16), with ks[a:b] 16-byte aligned, +# plus the tiles of fewer than 64 rows past S_1 = 4,194,048 (rows <= 256, ks misaligned) x (W % 16). +# Decode: (rows <= 256 or more) x (W % 16) x (lengths view 16-byte aligned or not: tc % 4). The +# "decode rows > 256" classes are new at this envelope (decode had at most 32 x 6 rows); the two +# at W % 16 == 0 are not reachable from the target (without DSpark next_n is 1, so rows = G <= 256, +# and every DSpark width L + 14 / L + 19 has W % 16 != 0) but run in the capture sweep below, which +# captures 6G at the DSpark-off widths as well. Each prefill cell is a step whose every tile lands +# in the class (a misaligned cell's first tile starts at row 0, in an aligned class mapped before). + +_LONG_PREFILL_CLASSES = { + "prefill rows<=256 W%16==0": [ + ([1], [32768]), + ([128, 64], [1048576, 4096]), + ([1] * 192, [4096] * 192), + ], + "prefill rows<=256 W%16!=0": [ + ([1], [LONG_LM]), + ([128, 128], [16393, 128]), + ([1] * 192, [4097] + [4096] * 191), + ], + "prefill rows>256 W%16==0": [ + ([15872], [16400]), + ([256, 192], [524288, 2048]), + ([4] * 256, [512] * 256), + ], + "prefill rows>256 W%16!=0": [ + ([10441, 5851], [10441, 6184]), + ([300], [262163]), + ([10451, 5927], [10451, 6244]), + ], + "prefill tiles<64 rows W%16==0 ks misaligned": [ + ([32] * 4, [1048576] * 4), + ([16] * 5, [1048576] * 5), + ([8] * 8, [1048576] * 8), + ], + "prefill tiles<64 rows W%16!=0 ks misaligned": [ + ([32] * 4, [LONG_LM] * 4), + ([16] * 5, [LONG_LM] * 5), + ([8] * 8, [LONG_LM - 1] * 8), + ], +} +# (W, next_n, G, tc) +_LONG_DECODE_CLASSES = { + "decode rows<=256 W%16!=0 lengths 16B-aligned": [ + (1048595, 6, 1, 0), + (32787, 6, 42, 0), + (65550, 1, 256, 4), + ], + "decode rows<=256 W%16!=0 lengths misaligned": [ + (16403, 6, 5, 1), + (131091, 1, 200, 2), + (1048590, 6, 3, 16365), + ], + "decode rows<=256 W%16==0 lengths 16B-aligned": [ + (32768, 1, 256, 0), + (1048576, 1, 1, 0), + (262144, 1, 17, 8192), + ], + "decode rows<=256 W%16==0 lengths misaligned": [ + (65536, 1, 3, 1), + (1048576, 1, 256, 16127), + (131072, 1, 40, 5), + ], + "decode rows>256 W%16!=0 lengths 16B-aligned": [ + (16403, 6, 43, 0), + (1048595, 6, 256, 0), + (32782, 6, 100, 8), + ], + "decode rows>256 W%16!=0 lengths misaligned": [ + (32787, 6, 48, 1), + (1048590, 6, 256, 14847), + (262163, 6, 200, 2), + ], + "decode rows>256 W%16==0 lengths 16B-aligned": [ + (32768, 6, 43, 0), + (1048576, 6, 256, 0), + (131072, 6, 64, 4096), + ], + "decode rows>256 W%16==0 lengths misaligned": [ + (65536, 6, 48, 3), + (1048576, 6, 256, 14845), + (262144, 6, 100, 7), + ], +} + + +def test_long_context_kernel_variant_map(): + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.candidate_block_select import ( # noqa: E501 + _candidate_block_select_kernel as fn, + ) + + variant_of, calls, bad = {}, 0, [] + _jit_clear(fn) + assert not _jit_keys(fn) + + def run(cls, i, desc, call, want, launches=1): + nonlocal calls + before = _jit_keys(fn) + got, names = _profiled(call) # the tile masks are joined outside the profiled region + got = torch.cat(got) if isinstance(got, list) else got + new = _jit_keys(fn) - before + calls += launches + assert names == [KERNEL_NAME] * launches, f"{cls} {desc}: launched {names}" + if i == 0: + assert len(new) == 1, f"{cls}: its first cell compiled {len(new)} variants" + variant_of[next(iter(new))] = f"{cls} (first at {desc})" + else: + assert not new, f"{cls} {desc}: compiled {len(new)} new variant(s)" + if not torch.equal(got, want): + bad.append((cls, desc, _first_bad_row(got, want))) + + for cls, cells in _LONG_PREFILL_CLASSES.items(): + for i, (q, kv) in enumerate(cells): + step = _packed_step(q, kv, seed=7000 + 16 * i + len(cls)) + tiles = _tiles(sum(q), sum(kv)) + run( + cls, + i, + f"T={sum(q)} W={sum(kv)} k={len(q)} tiles={[b - a for a, b in tiles]}", + lambda: _tile_masks(step)[0], + step["want"], + launches=len(tiles), + ) + del step + for cls, cells in _LONG_DECODE_CLASSES.items(): + for i, (W, next_n, G, tc) in enumerate(cells): + logits, vis, _, want = _long_decode_call( + W, next_n, G, tc, seed=8000 + 31 * i + len(cls) + ) + run( + cls, + i, + f"W={W} next_n={next_n} G={G} tc={tc}", + lambda: op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1), + want, + ) + del logits, vis, want + + # closure: random calls in the target's forms over the new envelope compile nothing further + before = _jit_keys(fn) + g = torch.Generator().manual_seed(4343) + for n in range(32): + if n % 2: + q, kv = _long_partition(g, _ri(g, 1, LONG_REQS), budget=1 << 30) + step = _packed_step(q, kv, seed=9000 + 8 * n) + (masks, tiles), names = _profiled(lambda: _tile_masks(step)) + got, want, launches = torch.cat(masks), step["want"], len(tiles) + else: + W, next_n = LONG_WIDTHS[_ri(g, 0, len(LONG_WIDTHS) - 1)], (1, 6)[_ri(g, 0, 1)] + G = _ri(g, 1, LONG_REQS) + logits, vis, _, want = _long_decode_call( + W, next_n, G, _ri(g, 0, LONG_ROWS - G * next_n), seed=9000 + n + ) + got, names = _profiled(lambda: op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1)) + launches = 1 + calls += launches + assert names == [KERNEL_NAME] * launches + if not torch.equal(got, want): + bad.append(("closure", n, _first_bad_row(got, want))) + after = _jit_keys(fn) + print( + f"\n long-context variant map: {len(variant_of)} compiled variant(s) over {calls} profiled launches " + f"(one {KERNEL_NAME} launch per call / tile):" + ) + for key, cls in sorted(variant_of.items(), key=lambda kv: kv[1]): + print(f" {cls}: {key[1][:220]}") + print( + f" 32 random calls after the map compiled {len(after - before)} further variant(s); " + f"{len(bad)} calls differ from the reference" + ) + assert not bad, bad[:5] + assert after == before + assert len(variant_of) == len(_LONG_PREFILL_CLASSES) + len(_LONG_DECODE_CLASSES) + + +# ─── (b) >= 200 seeded random packed prefill steps over the new envelope ────────── + + +def test_long_context_random_packed_partitions(): + """Each step is called as the target calls it (one call per row tile) and, as a superset, in + one untiled call; both are compared with the reference on every row.""" + g = torch.Generator().manual_seed(2424) + fixed_k = [1, 2, 3, 128, 255, 256] + ks_list = fixed_k + [_ri(g, 1, LONG_REQS) for _ in range(200 - len(fixed_k))] + stats = { + "rows": 0, + "max_T": 0, + "max_W": 0, + "min_W": 1 << 40, + "k": set(), + "binding_rows": 0, + "max_vis": 0, + "vis_over_512k": 0, + "T_16384": 0, + "tiles": 0, + "multi_tile": 0, + "min_tile": 1 << 30, + "max_tile": 0, + "small_tiles": 0, + } + ctl_rows = { + "visible end one short (ke - 1)": 0, + "start of the neighbouring request": 0, + "no -inf block drop": 0, + } + ctl_checked, bad = 0, [] + for n, k in enumerate(ks_list): + q, kv = _long_partition(g, k) + T, W = sum(q), sum(kv) + assert ( + len(q) == k <= LONG_REQS + and k <= T <= LONG_ROWS + and max(kv) <= LONG_LM + and BINDS_ABOVE < W + ) + assert T * W <= LONG_BUDGET + step = _packed_step(q, kv, seed=1_000_000 + 8 * n) + tiled, tiles = _call_tiled(step) + whole = _call_prefill(step) + if not (torch.equal(tiled, step["want"]) and torch.equal(whole, step["want"])): + bad.append( + ( + n, + k, + T, + W, + _first_bad_row(tiled, step["want"]), + _first_bad_row(whole, step["want"]), + ) + ) + del tiled, whole + vis = step["ke"] - step["ks"] + stats["rows"] += T + stats["max_T"] = max(stats["max_T"], T) + stats["max_W"] = max(stats["max_W"], W) + stats["min_W"] = min(stats["min_W"], W) + stats["k"].add(k) + stats["binding_rows"] += int((vis > BINDS_ABOVE).sum()) + stats["max_vis"] = max(stats["max_vis"], int(vis.max())) + stats["vis_over_512k"] += int(int(vis.max()) > 1 << 19) + stats["T_16384"] += int(T == LONG_ROWS) + stats["tiles"] += len(tiles) + stats["multi_tile"] += int(len(tiles) > 1) + stats["min_tile"] = min(stats["min_tile"], min(b - a for a, b in tiles)) + stats["max_tile"] = max(stats["max_tile"], max(b - a for a, b in tiles)) + stats["small_tiles"] += int(len(tiles) > 1 and tiles[0][1] < 64) + if n % 10 == 0 and k > 1: # deliberately-wrong controls on the same data + want, ks, ke = step["want"], step["ks"], step["ke"] + req = torch.repeat_interleave(torch.arange(k, device=DEV), torch.tensor(q, device=DEV)) + bases = torch.tensor(step["bases"], dtype=torch.int32, device=DEV) + col = torch.arange(W, device=DEV) + ctls = { + "visible end one short (ke - 1)": _call_prefill(step, ke=ke - 1), + "start of the neighbouring request": _call_prefill( + step, ks=bases[(req - 1) % k].contiguous() + ), + "no -inf block drop": (col >= ks[:, None]) & (col < ke[:, None]), + } + for name, ctl in ctls.items(): + ctl_rows[name] += int((ctl != want).any(dim=1).sum()) + ctl_checked += T + del ctls + del step + print( + f"\n {len(ks_list)} packed steps: {stats['rows']} rows, T <= {stats['max_T']} ({stats['T_16384']} at 16384), " + f"S_1 {stats['min_W']}..{stats['max_W']}, k {min(stats['k'])}..{max(stats['k'])} ({len(stats['k'])} distinct), " + f"{stats['binding_rows']} binding rows, longest row {stats['max_vis']} positions " + f"({stats['vis_over_512k']} steps with a row past 524288)" + ) + print( + f" {stats['tiles']} tiles ({stats['multi_tile']} steps tiled, {stats['small_tiles']} with tiles under 64 " + "rows), " + f"tile rows {stats['min_tile']}..{stats['max_tile']}; {len(bad)} steps differ from the reference " + f"(tiled or whole)" + ) + print(f" controls over {ctl_checked} rows of every 10th step:") + for name, rows in ctl_rows.items(): + print(f" {name}: {rows} rows differ") + assert not bad, bad[:3] + assert set(fixed_k) <= stats["k"] + for name, rows in ctl_rows.items(): + assert rows > 0, f"control '{name}' was not detected" + + +# ─── the tail's premise: a row's mask depends only on its own scores over [ks, ke) ──────────── + + +def _deepgemm_layout(rows: torch.Tensor) -> torch.Tensor: + """A copy of ``rows`` [n, W] in DeepGEMM's prefill layout (row stride align(W + 256, 256)), + the layout (and the compiled variant) the target's calls use.""" + n, w = rows.shape + out = torch.full((n, _align(w + 256, 256)), _NAN, device=DEV)[:, :w] + out.copy_(rows) + return out + + +def test_long_context_prefill_cells_row_independence_and_key_offset(): + """Decoder replay scores the tail rows' masks against the step's gathered keys, and the target + tiles rows: so a row must get the same mask alone, at any position inside a larger call, and + with its request's keys at any offset of the gather. Two tail requests of 128 rows with other + requests' keys around them (S_1 1,048,219: one 256-row tile).""" + q, kv = [0, 128, 0, 128], [300001, 400000, 48218, 300000] + step = _packed_step(q, kv, seed=5_000_000) + T, W = step["logits"].shape + assert W == 1048219 and _tiles(T, W) == [(0, T)] + logits = _deepgemm_layout( + step["logits"] + ) # our own copy: the next _packed_step reuses the buffer + ks, ke, want, bases = step["ks"], step["ke"], step["want"], step["bases"] + full = op(logits, ks, ke, TOPK_BLOCKS, BLOCK, 1) + assert torch.equal(full, want) + assert bool((ke - ks > BINDS_ABOVE).all()) + # (1) every row alone + alone_bad = [ + i + for i in range(T) + if not torch.equal( + op(logits[i : i + 1], ks[i : i + 1], ke[i : i + 1], TOPK_BLOCKS, BLOCK, 1)[0], + full[i], + ) + ] + # (2) the rows in reverse order + perm = torch.arange(T - 1, -1, -1, device=DEV) + rev = op( + _deepgemm_layout(logits[perm]), + ks[perm].contiguous(), + ke[perm].contiguous(), + TOPK_BLOCKS, + BLOCK, + 1, + ) + rev_ok = torch.equal(rev, full[perm]) + del rev + # (3) inside a larger call: 100 rows of another step before ours and 60 after (416 rows, the + # many-rows launch configuration where ours alone took the few-rows one) + other = _packed_step([160], [W], seed=5_000_001) + big = _deepgemm_layout(torch.cat([other["logits"][:100], logits, other["logits"][100:]])) + big_ks = torch.cat([other["ks"][:100], ks, other["ks"][100:]]) + big_ke = torch.cat([other["ke"][:100], ke, other["ke"][100:]]) + in_big = op(big, big_ks, big_ke, TOPK_BLOCKS, BLOCK, 1) + big_ok = ( + torch.equal(in_big[100 : 100 + T], full) + and torch.equal(in_big[:100], other["want"][:100]) + and torch.equal(in_big[100 + T :], other["want"][100:]) + ) + del other, big, in_big + # (4) each tail request's keys moved to other offsets of a gather: the mask moves with them + shift_bad, shifts = [], 0 + r0 = 0 + for b, (qb, kvb) in enumerate(zip(q, kv)): + if qb == 0: + continue + rows, base = slice(r0, r0 + qb), bases[b] + mine = logits[rows, base : base + kvb] + for d in (0, 1, 7, 8, 13, 4097, 524289, base + 3): + buf = torch.full((qb, _align(d + kvb + 5 + 256, 256)), _NAN, device=DEV) + moved = buf[:, : d + kvb + 5] + moved[:, d : d + kvb] = mine + got = op( + moved, + (ks[rows] - base + d).contiguous(), + (ke[rows] - base + d).contiguous(), + TOPK_BLOCKS, + BLOCK, + 1, + ) + shifts += 1 + if not ( + torch.equal(got[:, d : d + kvb], full[rows, base : base + kvb]) + and not got[:, :d].any() + and not got[:, d + kvb :].any() + ): + shift_bad.append((b, d)) + del buf, moved, got + r0 += qb + print( + f"\n {T} tail rows at S_1 {W}: {len(alone_bad)} rows differ alone; reversed order equal: {rev_ok}; " + f"inside a 416-row call equal: {big_ok}; {len(shift_bad)} of {shifts} key-offset moves differ" + ) + assert not alone_bad and rev_ok and big_ok and not shift_bad, (alone_bad[:5], shift_bad[:5]) + + +# ─── the row axis at the widest tiles: every r in 1..tile rows, bitwise ────────────────────── + + +@pytest.mark.parametrize( + "q,kv", + [([256], [1048219]), ([192], [LONG_LM]), ([0, 256], [48219, 1000000])], + ids=[ + "256-row tile at S_1 1048219", + "192-row tile at Lm 1048590", + "256 tail rows past other keys", + ], +) +def test_long_context_maximal_tile_rows_prefix_bitwise(q, kv): + """The widest tiles hold 256 rows (S_1 up to 1,048,320) or 192 (up to Lm); every row binds. + Every r-row call equals the first r rows of the whole tile, bitwise, on a poisoned block.""" + step = _packed_step(q, kv, seed=6_000_000 + sum(kv) % 1000) + logits, ks, ke = step["logits"], step["ks"], step["ke"] + T, W = logits.shape + assert _tiles(T, W) == [(0, T)] and bool((ke - ks > BINDS_ABOVE).all()) + start, stop = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + full = _call_prefill(step) + stop.record() + torch.cuda.synchronize() + assert torch.equal(full, step["want"]) + bad, reused = [], 0 + for r in range(T, 0, -1): + junk = torch.ones(r * W, dtype=torch.uint8, device=DEV) + ptr = junk.data_ptr() + del junk + part = op(logits[:r], ks[:r], ke[:r], TOPK_BLOCKS, BLOCK, 1) + reused += part.data_ptr() == ptr + if not torch.equal(part, full[:r]) or int(part.view(torch.uint8).max()) > 1: + bad.append(r) + del part + print( + f"\n {T} x {W} (one call {start.elapsed_time(stop):.2f} ms): every r in 1..{T}: {len(bad)} prefixes differ; " + f"{reused} of {T} masks on a poisoned block" + ) + assert not bad, bad[:10] + + +# ─── decode: capture at every G in GSET, rows G and 6G, at every ratio-1 width ──────────────── + + +@pytest.mark.parametrize("W", LONG_WIDTHS) +def test_long_context_decode_capture_at_gset(W): + """Each (next_n, G) is captured after an eager call at its shape, then replayed twice on new + scores and lengths. Request 0 sees the whole width (column W - 1) at the eager call and the + second replay, and a shorter binding length at the first, so every replay's expected mask + differs from the previous step's and a stale replay is seen.""" + bad, graphs, stale_min = [], 0, None + mid = BINDS_ABOVE + 1 + (W - BINDS_ABOVE - 1) // 2 + for next_n in (1, 6): + for G in GSET: + seed = 10 * W + 977 * G + next_n + logits, vis, _, want = _long_decode_call(W, next_n, G, 0, seed, first_end=W) + eager = op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1) + if not torch.equal(eager, want): + bad.append(("eager", next_n, G, _first_bad_row(eager, want))) + del eager + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1) + graphs += 1 + prev = want + for rep, end in ((1, mid), (2, W)): + want = _fill_decode(logits, vis, next_n, seed=seed + 100_000 * rep, first_end=end) + changed = int((want != prev).any(dim=1).sum()) + stale_min = changed if stale_min is None else min(stale_min, changed) + graph.replay() + torch.cuda.synchronize() + ref_eager = op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1) + if not (torch.equal(out, ref_eager) and torch.equal(out, want)): + bad.append(("replay", next_n, G, rep, _first_bad_row(out, want))) + del ref_eager + prev = want + del graph, out, logits, vis, want, prev + print( + f"\n W {W}: {graphs} graphs (next_n 1 and 6 x {len(GSET)} batches, rows up to {6 * GSET[-1]}), " + f"2 replays each; {len(bad)} mismatches; every replay's mask differs from the previous step's " + f"on >= {stale_min} row(s)" + ) + assert not bad, bad[:5] + assert stale_min >= 1 + + +# ─── decode rows on the row axis: every r in 1..1536, bitwise ───────────────────────────────── + + +@pytest.mark.parametrize("W", (32782, 1048595)) +def test_long_context_decode_rows_prefix_bitwise(W): + """Eager steps hold any G <= 256 at next_n 1 or 6, so a decode call has any row count up to + 1536. The 1536-row call (256 DSpark requests, request 0 at the far end) against the reference, + then every r in 1..1536 bitwise equal to its first r rows, across the 256-row launch switch.""" + rows = 6 * LONG_REQS + logits, vis, _, want = _long_decode_call(W, 6, LONG_REQS, 0, seed=W + 11, first_end=W) + start, stop = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + full = op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1) + torch.cuda.synchronize() + start.record() + for _ in range(3): + op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1) + stop.record() + torch.cuda.synchronize() + assert torch.equal(full, want) + print( + f"\n {rows} x {W}: exact; {int((vis > BINDS_ABOVE).sum())} binding rows; " + f"eager call {start.elapsed_time(stop) / 3:.2f} ms" + ) + bad, reused = [], 0 + for r in range(rows, 0, -1): + junk = torch.ones(r * W, dtype=torch.uint8, device=DEV) + ptr = junk.data_ptr() + del junk + part = op(logits[:r], None, vis[:r], TOPK_BLOCKS, BLOCK, 1) + reused += part.data_ptr() == ptr + if not torch.equal(part, full[:r]) or int(part.view(torch.uint8).max()) > 1: + bad.append(r) + del part + print( + f" every r in 1..{rows} at W {W}: {len(bad)} prefixes differ; {reused} of {rows} masks on a poisoned block" + ) + assert not bad, bad[:10] + + +# ─── decode in mixed (eager) steps: the lengths view at row offsets of the [16384] buffer ───── + + +def test_long_context_decode_rows_at_row_offsets(): + bad, calls = [], 0 + for W in (16398, 32782, 1048590, 1048576): + for next_n, G in ((6, 1), (1, 7), (6, 43), (1, 256), (6, 256)): + rows = G * next_n + for tc in (1, 2, 3, 5, 8193, LONG_ROWS - rows): + logits, vis, vis_buf, want = _long_decode_call( + W, next_n, G, tc, seed=W + tc + rows, first_end=W + ) + vis_buf[:tc] = -7 + vis_buf[tc + rows :] = 1 << 22 + got = op(logits, None, vis, TOPK_BLOCKS, BLOCK, 1) + same = op(logits, None, vis.clone(), TOPK_BLOCKS, BLOCK, 1) + calls += 1 + if not (torch.equal(got, want) and torch.equal(got, same)): + bad.append((W, next_n, G, tc, _first_bad_row(got, want))) + del logits, vis, vis_buf, want, got, same + print(f"\n {calls} decode calls at row offsets of the 16384 buffer: {len(bad)} differ") + assert not bad, bad[:5] diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_compressor_paged_kv_compress.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_compressor_paged_kv_compress.py new file mode 100644 index 000000000000..ed9ac2edc8d9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_compressor_paged_kv_compress.py @@ -0,0 +1,4042 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the compressor_paged_kv_compress catalog entry. + +The entry owns what is certified: `CELLS` is the list, `reference` and `compare` are the gate, and +`test_certified_cells` drives every cell through them, each step of each cell gated against +`op.reference` taken from the pools as they were before the call. The cells are the DeepSeek-V4.1 +ratio-2 compressor of the KV-source layers 2/8/14 at generation time -- compress_ratio 2, head_dim +512, fp32 kv_score rows ``[kv || score]``, fp32 paged state, zero APE (V4.1 has none), fp32 or bf16 +output -- for plain decode (next_n 1) and DSpark verify (next_n 6), next_n 1..8 in all. + +This file keeps the checks a cell cannot carry: DSpark verify with per-request acceptance rewinding +the next step's start, the prefill op's parked token consumed in one shared buffer, CUDA-graph +replay, the real cache manager's pools and metadata slices (generation-only and mixed steps), the +op's own rejections, the wrong-variant controls, the silent-misuse guards, and the engine-size +coverage (every batch size at every table width, every context lead, the capture set, offset +hazards, pool-size invariance, the multi-request envelope, every dumped step, the DSpark tail). + +Those engine-size checks gate on ``_pool_windows``: the checkpoint's own decode rule (its +inference/model.py, Compressor.forward, start_pos > 0 branch -- token ``p`` fills slot ``p % 2`` and, +when ``(p + 1) % 2 == 0``, the pair is pooled per channel as ``sum(kv * softmax_over_the_pair(score))``) +over a per-position array holding what each position last received. It is independent of the pool +layout, which is what lets it run on the engine's pools (up to 65 GB) that `op.reference`, which +returns whole expected pools, cannot copy. + +State geometry. DeepseekV4CacheManager in V4.1 mode puts COMPRESSOR_KV and +COMPRESSOR_SCORE of every owner layer in ONE fp32 pool ``[pages, 128, 512]`` +(same base pointer; the kv and score tables name different pages of it), so +the synthetic tests share one tensor between the two arguments by default. +Synthetic pools have POOL_PAGES pages -- chosen plainly larger than any batch +here -- with a random page permutation, distinct pages for kv and score, and a +trap page behind every unused table column. test_v41_cache_manager_pools runs +the op on the real manager's pools, tables and per-step metadata instead. +""" + +import math +import zlib +from types import SimpleNamespace +from typing import List, Sequence, Tuple + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.compressor_paged_kv_compress import ( # noqa: E501 + compressor_paged_kv_compress as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "compressor_paged_kv_compress requires a CUDA device" + +DEV = torch.device("cuda") +HEAD_DIM = 512 +RATIO = 2 +POOL_PAGES = 1024 +TRAP_PAGE = 0 +SENTINEL = 7777.0 + + +def _gen(seed: int) -> torch.Generator: + g = torch.Generator(device=DEV) + g.manual_seed(seed) + return g + + +def _rand_rows(n: int, gen: torch.Generator) -> Tuple[torch.Tensor, torch.Tensor]: + """fp32 kv and score rows. Channel 0 is an exact tie in every pair and channel + 1 a gap of 240 between even and odd positions (exp underflows to exactly 0).""" + kv = torch.randn(n, HEAD_DIM, generator=gen, device=DEV) * 2.0 + score = torch.randn(n, HEAD_DIM, generator=gen, device=DEV) * 4.0 + score[:, 0] = 1.5 + score[:, 1] = 120.0 * (1 - 2 * (torch.randint(0, 2, (n,), generator=gen, device=DEV).float())) + return kv, score + + +def _pool_windows( + kv: torch.Tensor, score: torch.Tensor, ape: torch.Tensor, w0: int, w1: int +) -> torch.Tensor: + """model.py's ratio-2 pooling of windows [w0, w1) of a per-position array, fp32.""" + p = torch.arange(w0 * RATIO, w1 * RATIO, device=DEV) + k = kv[p].unflatten(0, (w1 - w0, RATIO)) + s = (score[p] + ape[p % RATIO]).unflatten(0, (w1 - w0, RATIO)) + return (k * s.softmax(dim=1)).sum(dim=1) + + +def _zero_ape() -> torch.Tensor: + return torch.zeros(RATIO, HEAD_DIM, device=DEV) + + +def _i32(values) -> torch.Tensor: + return torch.tensor(list(values), device=DEV, dtype=torch.int32) + + +class _Rows: + """Per-position rows of one request over absolute positions [lo, hi), stored from lo, so a + request whose history reaches position 1,048,594 costs only its window. Indexed with + absolute positions (int, tensor or slice), like a [hi, HEAD_DIM] tensor would be.""" + + def __init__(self, lo: int, hi: int): + self.lo, self.hi = lo, hi + self.t = torch.zeros(hi - lo, HEAD_DIM, device=DEV) + + def _k(self, k): + if isinstance(k, slice): + start = self.lo if k.start is None else k.start + stop = self.hi if k.stop is None else k.stop + assert self.lo <= start <= stop <= self.hi, (start, stop, self.lo, self.hi) + return slice(start - self.lo, stop - self.lo) + if isinstance(k, torch.Tensor): + return k - self.lo + assert self.lo <= k < self.hi, (k, self.lo, self.hi) + return k - self.lo + + def __getitem__(self, k): + return self.t[self._k(k)] + + def __setitem__(self, k, v): + self.t[self._k(k)] = v + + def clone(self) -> "_Rows": + c = _Rows.__new__(_Rows) + c.lo, c.hi, c.t = self.lo, self.hi, self.t.clone() + return c + + +class _Pool: + """Paged fp32 compressor state and per-sequence kv / score block tables. + + ``shared=True`` is the V4.1 manager's arrangement: kv and score live in one + tensor, on different pages.""" + + def __init__(self, max_lens: Sequence[int], page_size: int, seed: int, shared: bool = True): + self.page_size = page_size + self.kv = torch.full((POOL_PAGES, page_size, HEAD_DIM), SENTINEL, device=DEV) + self.score = self.kv if shared else torch.full_like(self.kv, SENTINEL) + n_log = [max(math.ceil(L / page_size), 1) for L in max_lens] + width = max(n_log) + 1 + cpu = torch.Generator().manual_seed(seed) + pages = (torch.randperm(POOL_PAGES - 1, generator=cpu) + 1).tolist() + assert 2 * sum(n_log) <= len(pages), "raise POOL_PAGES" + self.bt_kv = torch.full((len(max_lens), width), TRAP_PAGE, dtype=torch.int32) + self.bt_score = torch.full_like(self.bt_kv, TRAP_PAGE) + for b, n in enumerate(n_log): + self.bt_kv[b, :n] = torch.tensor(pages[:n], dtype=torch.int32) + del pages[:n] + self.bt_score[b, :n] = torch.tensor(pages[:n], dtype=torch.int32) + del pages[:n] + self.bt_kv = self.bt_kv.to(DEV) + self.bt_score = self.bt_score.to(DEV) + + def slots(self, table: torch.Tensor, b: int, pos: torch.Tensor): + return table[b].long()[pos // self.page_size], pos % self.page_size + + def seed( + self, + b: int, + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + upto: int, + poison: bool, + ): + """The state a prefill call leaves for positions [0, upto): kv and score + ape. + With ``poison`` every position but ``upto - 1`` holds NaN instead.""" + if upto == 0: + return + pos = torch.arange(upto, device=DEV) + kv_v = kv[:upto].clone() + sc_v = score[:upto] + ape[pos % RATIO] + if poison: + kv_v[:-1] = float("nan") + sc_v[:-1] = float("nan") + self.kv[self.slots(self.bt_kv, b, pos)] = kv_v + self.score[self.slots(self.bt_score, b, pos)] = sc_v + + def check(self, kvs, scores, lens, ape: torch.Tensor, first: Sequence[int]): + """Positions [first[b], lens[b]) of sequence b hold its kv and score + ape bit + for bit, and no slot outside the sequences' positions [0, lens[b]) changed.""" + touched = torch.zeros(POOL_PAGES, self.page_size, dtype=torch.bool, device=DEV) + touched_sc = torch.zeros_like(touched) + for b, (kv, sc, n, f) in enumerate(zip(kvs, scores, lens, first)): + pos = torch.arange(n, device=DEV) + pk, ps = self.slots(self.bt_kv, b, pos), self.slots(self.bt_score, b, pos) + touched[pk] = True + touched_sc[ps] = True + live = torch.arange(f, n, device=DEV) + assert torch.equal(self.kv[self.slots(self.bt_kv, b, live)], kv[f:n]), ( + f"seq {b}: paged kv state" + ) + assert torch.equal( + self.score[self.slots(self.bt_score, b, live)], sc[f:n] + ape[live % RATIO] + ), f"seq {b}: paged score state" + if self.score is self.kv: + touched |= touched_sc + assert (self.kv[~touched] == SENTINEL).all(), ( + "pool written outside the sequences' slots" + ) + else: + assert (self.kv[~touched] == SENTINEL).all(), ( + "kv pool written outside the sequences' slots" + ) + assert (self.score[~touched_sc] == SENTINEL).all(), ( + "score pool written outside the sequences' slots" + ) + + +def _step( + pool: _Pool, + kvs, + scores, + starts: Sequence[int], + next_n: int, + ape, + out_dtype=torch.float32, + lead_rows: int = 0, + lead_out: int = 0, + certify: bool = False, +): + """One generation step: request b feeds positions [starts[b], starts[b] + next_n). + + ``lead_rows`` / ``lead_out`` put NaN rows ahead of the generation rows in + kv_score and ahead of the outputs, where a mixed batch's context rows sit: + the op must neither read the former nor write the latter. The output buffer + has the runtime's reserved size (ceil(next_n / 2) rows per request) plus one + guard row, and everything past the packed outputs must stay NaN. + + ``certify`` gates the call on the entry's own ``reference`` (taken from the pools before the + call) and ``compare``, and checks the op returns the three tensors it wrote.""" + B = len(starts) + rows = [ + torch.cat([kv[s : s + next_n], sc[s : s + next_n]], dim=1) + for kv, sc, s in zip(kvs, scores, starts) + ] + lead = torch.full((lead_rows, 2 * HEAD_DIM), float("nan"), device=DEV) + kv_score = torch.cat([lead] + rows, dim=0).contiguous() + n_out = [(s + next_n) // RATIO - s // RATIO for s in starts] + cu_out = [lead_out] + for o in n_out: + cu_out.append(cu_out[-1] + o) + reserved = B * math.ceil(next_n / RATIO) + buf = torch.full((lead_out + reserved + 1, HEAD_DIM), float("nan"), device=DEV, dtype=out_dtype) + output = buf[: lead_out + reserved] + args = ( + kv_score, + ape, + pool.kv, + pool.score, + pool.bt_kv, + pool.bt_score, + output, + _i32(s + next_n for s in starts), + _i32(lead_rows + b * next_n for b in range(B + 1)), + _i32(cu_out), + B, + pool.page_size, + HEAD_DIM, + RATIO, + next_n, + ) + expected = op.reference(*args) if certify else None + result = op(*args) + torch.cuda.synchronize() + if certify: + assert result[0] is output and result[1] is pool.kv and result[2] is pool.score, ( + "the op writes in place and returns (output, paged_kv, paged_score)" + ) + op.compare(result, expected) + assert torch.isnan(buf[:lead_out]).all(), "op wrote into the rows ahead of cu_kv_comp[0]" + assert torch.isnan(buf[cu_out[-1] :]).all(), "op wrote past the last packed output row" + return [output[cu_out[b] : cu_out[b + 1]] for b in range(B)] + + +def _decode_run( + cached: List[int], + steps: int, + next_n: int, + page_size: int, + ape, + out_dtype, + seed: int, + shared: bool = True, + poison_history: bool = False, + lead: Tuple[int, int] = (0, 0), + certify: bool = False, +): + """``steps`` all-accepted generation steps after a seeded history; returns per-step + outputs and the per-position arrays the reference pools from.""" + gen = _gen(seed) + lens = [c + steps * next_n for c in cached] + seqs = [_rand_rows(L, gen) for L in lens] + kvs, scores = [s[0] for s in seqs], [s[1] for s in seqs] + pool = _Pool(lens, page_size, seed, shared) + for b, c in enumerate(cached): + pool.seed(b, kvs[b], scores[b], ape, c, poison_history) + outs = [] + starts = list(cached) + for _ in range(steps): + got = _step(pool, kvs, scores, starts, next_n, ape, out_dtype, *lead, certify=certify) + outs.append((list(starts), got)) + starts = [s + next_n for s in starts] + return pool, kvs, scores, lens, outs + + +def _check_outputs(kvs, scores, ape, outs, next_n): + for starts, got in outs: + for b, (s, g) in enumerate(zip(starts, got)): + ref = _pool_windows(kvs[b], scores[b], ape, s // RATIO, (s + next_n) // RATIO) + assert g.shape == ref.shape, ( + f"seq {b} start {s}: {tuple(g.shape)} vs {tuple(ref.shape)}" + ) + # Default fp32 tolerances (rtol 1.3e-6, atol 1e-5): the kernel is built with + # --use_fast_math and runs an online softmax that divides once at the end, so it + # is not bit-identical to torch's softmax-then-sum, but stays inside the dtype's + # own tolerance. + torch.testing.assert_close(g, ref, msg=lambda m: f"seq {b} start {s}: {m}") + + +# Histories that put the first new token at both parities and on both sides of a +# 128- and a 256-token page boundary. +CACHED = [1, 2, 127, 128, 255, 256, 1001, 8191] + + +# --------------------------------------------------------------------------- +# certified surface +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Every step of every cell gated on the entry's reference and compare: the step's rows, the + whole pool(s) bitwise, and every output row the op must leave alone. A step after the first + reads the token the op itself parked. At the end every position holds what it last received + (with a poisoned history: from start - 1 on) and no other slot of the pool changed.""" + certifying(op) + spec = cell.spec + seed = zlib.crc32(cell.why.encode()) % 2**31 + if spec["ape"] == "zero": + ape = _zero_ape() + else: + ape = torch.randn(RATIO, HEAD_DIM, generator=_gen(seed + 1), device=DEV) * 2.0 + cached = list(spec["cached"]) + pool, kvs, scores, lens, _ = _decode_run( + cached, + spec["steps"], + spec["next_n"], + spec["page_size"], + ape, + spec["out_dtype"], + seed=seed, + shared=spec["shared_pool"], + poison_history=spec["poison_history"], + lead=(spec["lead_rows"], spec["lead_out"]), + certify=True, + ) + first = [c - 1 if c else 0 for c in cached] if spec["poison_history"] else [0] * len(lens) + pool.check(kvs, scores, lens, ape, first=first) + + +def test_dspark_verify_with_rejection(): + """next_n = 6 (DSpark: 5 drafts + 1). After every step each request accepts a + random 0..5 drafts, so the next step starts at start + 1 + accepted, at either + parity, and rewrites the rejected positions with new values. Each step must + pool from what every position last received -- the accepted token at start - 1 + included -- which is exactly the per-token decode rule applied in order.""" + ape = _zero_ape() + next_n, steps, page_size = 6, 10, 128 + cached = [1, 34, 127, 128, 129, 294, 1094, 4095] + gen = _gen(401) + cpu = torch.Generator().manual_seed(401) + cap = [c + steps * next_n for c in cached] + kvs = [torch.zeros(L, HEAD_DIM, device=DEV) for L in cap] + scores = [torch.zeros(L, HEAD_DIM, device=DEV) for L in cap] + for b, c in enumerate(cached): + kvs[b][:c], scores[b][:c] = _rand_rows(c, gen) + pool = _Pool(cap, page_size, 401) + for b, c in enumerate(cached): + pool.seed(b, kvs[b], scores[b], ape, c, poison=False) + starts = list(cached) + parities = set() + for _ in range(steps): + for b, s in enumerate(starts): + kvs[b][s : s + next_n], scores[b][s : s + next_n] = _rand_rows(next_n, gen) + parities.add(s % RATIO) + got = _step(pool, kvs, scores, starts, next_n, ape) + _check_outputs(kvs, scores, ape, [(list(starts), got)], next_n) + accepted = torch.randint(0, next_n, (len(starts),), generator=cpu).tolist() + starts = [s + 1 + a for s, a in zip(starts, accepted)] + assert parities == {0, 1} + # Every position up to the last one written holds what it last received. + pool.check( + kvs, + scores, + [s - 1 + next_n - a for s, a in zip(starts, accepted)], + ape, + first=[0] * len(starts), + ) + + +def test_prefill_state_then_decode_in_one_buffer(): + """The engine's mixed step: the prefill op pools the contexts into the first + rows of one output buffer and parks their odd tail token in the pool; the + decode op runs the generation requests behind them in the same kv_score and + output. The next steps decode everyone, the former contexts included, so the + tail token the prefill op parked is the pending token decode reads back.""" + ape = _zero_ape() + next_n, page_size = 6, 128 + ctx_lens = [301, 130, 7] + gen_cached = [5, 128, 999] + steps = 3 + gen = _gen(601) + cap = [L + steps * next_n for L in ctx_lens] + [c + (steps + 1) * next_n for c in gen_cached] + seqs = [_rand_rows(L, gen) for L in cap] + kvs, scores = [s[0] for s in seqs], [s[1] for s in seqs] + pool = _Pool(cap, page_size, 601) + nctx = len(ctx_lens) + for j, c in enumerate(gen_cached): + pool.seed(nctx + j, kvs[nctx + j], scores[nctx + j], ape, c, poison=False) + + # Step 1: contexts through the prefill op, generation requests through this op. + ctx_rows = [torch.cat([kvs[b][:L], scores[b][:L]], dim=1) for b, L in enumerate(ctx_lens)] + gen_rows = [ + torch.cat([kvs[nctx + j][c : c + next_n], scores[nctx + j][c : c + next_n]], dim=1) + for j, c in enumerate(gen_cached) + ] + kv_score = torch.cat(ctx_rows + gen_rows, dim=0).contiguous() + ctx_out = [L // RATIO for L in ctx_lens] + gen_out = [(c + next_n) // RATIO - c // RATIO for c in gen_cached] + cu_seq = [0] + for n in ctx_lens + [next_n] * len(gen_cached): + cu_seq.append(cu_seq[-1] + n) + cu_out = [0] + for o in ctx_out + gen_out: + cu_out.append(cu_out[-1] + o) + total = sum(ctx_out) + len(gen_cached) * math.ceil(next_n / RATIO) + buf = torch.full((total + 1, HEAD_DIM), float("nan"), device=DEV) + output = buf[:total] + torch.ops.trtllm.compressor_prefill_reduction( + kv_score[: cu_seq[nctx]], + ape, + pool.kv, + pool.score, + pool.bt_kv[:nctx], + pool.bt_score[:nctx], + output, + _i32(ctx_lens), + _i32([0] * nctx), + _i32(cu_seq[: nctx + 1]), + _i32(cu_out[: nctx + 1]), + nctx, + page_size, + HEAD_DIM, + RATIO, + max(ctx_out), + ) + op( + kv_score, + ape, + pool.kv, + pool.score, + pool.bt_kv[nctx:], + pool.bt_score[nctx:], + output, + _i32(c + next_n for c in gen_cached), + _i32(cu_seq[nctx:]), + _i32(cu_out[nctx:]), + len(gen_cached), + page_size, + HEAD_DIM, + RATIO, + next_n, + ) + torch.cuda.synchronize() + assert torch.isnan(buf[cu_out[-1] :]).all(), "write past the packed outputs" + for b, L in enumerate(ctx_lens): + torch.testing.assert_close( + output[cu_out[b] : cu_out[b + 1]], _pool_windows(kvs[b], scores[b], ape, 0, L // 2) + ) + for j, c in enumerate(gen_cached): + b = nctx + j + torch.testing.assert_close( + output[cu_out[b] : cu_out[b + 1]], + _pool_windows(kvs[b], scores[b], ape, c // 2, (c + next_n) // 2), + ) + + # Steps 2..: everyone decodes. + starts = list(ctx_lens) + [c + next_n for c in gen_cached] + assert any(s % RATIO for s in starts[:nctx]), "a context must leave a pending token" + outs = [] + for _ in range(steps): + outs.append((list(starts), _step(pool, kvs, scores, starts, next_n, ape))) + starts = [s + next_n for s in starts] + _check_outputs(kvs, scores, ape, outs, next_n) + pool.check(kvs, scores, starts, ape, first=[0] * len(starts)) + + +def test_bf16_output_against_a_bf16_operand_control(): + """bf16 output = round-to-nearest-even of the fp32 pooled value, i.e. the + reference's ``kv.to(bf16)`` before its RMSNorm. The bf16 cells gate it call by call; + this pools next_n 1 and 6 into one population, the size the 0.999 gate was set on, and + shows a bf16-operand control misses that gate.""" + ape = _zero_ape() + got, ref, wrong = [], [], [] + for next_n in (1, 6): + _, kvs, scores, _, outs = _decode_run( + CACHED, 6, next_n, 128, ape, torch.bfloat16, seed=700 + next_n + ) + for starts, rows in outs: + for b, (s, g) in enumerate(zip(starts, rows)): + w0, w1 = s // 2, (s + next_n) // 2 + got.append(g) + ref.append(_pool_windows(kvs[b], scores[b], ape, w0, w1).to(torch.bfloat16)) + wrong.append( + _pool_windows( + kvs[b].bfloat16().float(), scores[b].bfloat16().float(), ape, w0, w1 + ).to(torch.bfloat16) + ) + got, ref, wrong = torch.cat(got), torch.cat(ref), torch.cat(wrong) + torch.testing.assert_close(got, ref) + exact = (got == ref).float().mean().item() + print(f"bf16 output: {exact:.6f} of {got.numel()} elements bit-identical to the reference") + # The fp32 value before the cast is within a few fp32 ulp of the reference, so a + # bf16 element can differ only where it straddles a bf16 rounding midpoint + # (~1e-4 expected). 0.999 is ~10x headroom; the bf16-operand control lands far below. + assert exact >= 0.999, f"bf16 output bit-exact fraction {exact}" + frac = (got == wrong).float().mean().item() + print(f"control bf16_operands vs bf16 output: bit-exact fraction {frac:.6f}") + assert frac < 0.999 + + +def test_cuda_graph_replay(): + """Captured once per next_n, replayed over steps whose kv_score, kv_lens and + cu_kv_comp change in place (batch, next_n and the tables stay fixed).""" + ape = _zero_ape() + for next_n in (1, 6): + steps = 6 + cached = [1, 128, 255, 1001] + B = len(cached) + gen = _gen(901 + next_n) + lens = [c + (steps + 1) * next_n for c in cached] + seqs = [_rand_rows(L, gen) for L in lens] + kvs, scores = [s[0] for s in seqs], [s[1] for s in seqs] + pool = _Pool(lens, 128, 901) + reserved = B * math.ceil(next_n / RATIO) + s_kv_score = torch.zeros(B * next_n, 2 * HEAD_DIM, device=DEV) + s_kv_lens = torch.zeros(B, dtype=torch.int32, device=DEV) + s_cu_seq = _i32(b * next_n for b in range(B + 1)) + s_cu_out = torch.zeros(B + 1, dtype=torch.int32, device=DEV) + s_out = torch.zeros(reserved, HEAD_DIM, device=DEV) + + def load(starts): + s_kv_score.copy_( + torch.cat( + [ + torch.cat([kv[s : s + next_n], sc[s : s + next_n]], 1) + for kv, sc, s in zip(kvs, scores, starts) + ] + ) + ) + s_kv_lens.copy_(_i32(s + next_n for s in starts)) + n_out = [(s + next_n) // 2 - s // 2 for s in starts] + s_cu_out.copy_(_i32([0] + torch.tensor(n_out).cumsum(0).tolist())) + return n_out + + def call(): + op( + s_kv_score, + ape, + pool.kv, + pool.score, + pool.bt_kv, + pool.bt_score, + s_out, + s_kv_lens, + s_cu_seq, + s_cu_out, + B, + 128, + HEAD_DIM, + RATIO, + next_n, + ) + + # Warm up and capture on step "-1" (positions [c - next_n, c)); the replayed + # steps then start from a pool whose history is fully written. + for b, c in enumerate(cached): + pool.seed(b, kvs[b], scores[b], ape, c, poison=False) + load(cached) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + call() + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + # Replays advance from the (already correct) state after the eager warm-up + # step; the captured run itself wrote nothing (capture does not execute). + starts = [c + next_n for c in cached] + for _ in range(steps): + n_out = load(starts) + s_out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + off = 0 + for b, s in enumerate(starts): + ref = _pool_windows(kvs[b], scores[b], ape, s // 2, (s + next_n) // 2) + torch.testing.assert_close(s_out[off : off + n_out[b]], ref) + off += n_out[b] + assert torch.isnan(s_out[off:]).all() + starts = [s + next_n for s in starts] + + +def test_v41_cache_manager_pools(): + """The op on DeepseekV4CacheManager's own V4.1 compressor-state pools: kv and + score pools of owners 2/8/14 from ``get_buffers``, tables from the metadata's + ``sliding_block_tables[layer_offsets[owner], COMPRESSOR_KV / COMPRESSOR_SCORE]``, + and ``kv_lens_cuda_runtime``, ``cu_seq_lens_cuda``, ``cu_new_comp_kv_cuda[2]``, + ``num_gen_tokens_per_seq``, ``num_total_compressed_tokens[2]`` and + ``tokens_per_block`` of each step's prepared metadata. Plain decode (no + speculation, window 2) and DSpark verify (max_draft_len 5, window 7, next_n 6). + Every page the tables map for these requests is NaN except the pending token + at start - 1, and no other byte of the shared pool may change.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + DeepseekV4CacheManager, + DeepseekV4TrtllmAttentionMetadata, + ) + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + from tensorrt_llm._torch.metadata import KVCacheParams + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from tensorrt_llm.bindings import DataType, SamplingConfig + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, + KvCacheConfig, + ) + from tensorrt_llm.mapping import Mapping + + # The checkpoint's text_config fields the V4.1 layout reads. + text = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=128, + head_dim=HEAD_DIM, + index_head_dim=128, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, + ) + config = SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **{k: v for k, v in text.items() if k != "model_type"}, + ) + ape = _zero_ape() + steps = 3 + for max_draft_len, cached in ((0, [1, 127, 128, 299, 2048]), (5, [34, 127, 128, 129, 1094])): + next_n = max_draft_len + 1 + sparse = DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ) + manager = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, max_tokens=4096 * 8, event_buffer_max_size=0 + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=40, + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=128, + max_seq_len=4096, + max_batch_size=8, + max_input_len=4096, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=8192, + sparse_attn_config=sparse, + pretrained_config=config, + spec_config=DSparkDecodingConfig(max_draft_len=max_draft_len, block_size=max_draft_len) + if max_draft_len + else None, + ) + requests = [] + try: + assert manager.v41_layout.window(T.COMPRESSOR_KV, max_draft_len) == 2 + max_draft_len + B = len(cached) + cap = [c + steps * next_n for c in cached] + for i, L in enumerate(cap): + req = LlmRequest( + request_id=i, + max_new_tokens=64, + input_tokens=list(range(L)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + assert manager.prepare_context(req) + assert manager.resize_context(req, req.context_chunk_size) + requests.append(req) + gen = _gen(1001 + max_draft_len) + seqs = [_rand_rows(L, gen) for L in cap] + kvs, scores = [s[0] for s in seqs], [s[1] for s in seqs] + for owner in (2, 8, 14): + pool_kv = manager.get_buffers(owner, T.COMPRESSOR_KV) + pool_sc = manager.get_buffers(owner, T.COMPRESSOR_SCORE) + assert pool_kv.dtype == torch.float32 and tuple(pool_kv.shape[1:]) == ( + 128, + HEAD_DIM, + ) + assert pool_kv.data_ptr() == pool_sc.data_ptr(), ( + "V4.1 keeps kv and score in one pool" + ) + assert pool_kv.shape == pool_sc.shape + starts = list(cached) + for step in range(steps): + md = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor([next_n] * B, dtype=torch.int), + request_ids=list(range(B)), + max_num_requests=manager.max_batch_size, + num_contexts=0, + prompt_lens=[s + next_n for s in starts], + max_num_tokens=8192, + kv_cache_manager=manager, + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=list(starts) + ), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + sparse_attention_config=sparse, + ) + md.prepare() + torch.cuda.synchronize() + assert md.num_gen_tokens_per_seq == next_n + lo = manager.layer_offsets[owner] + bt_kv = md.sliding_block_tables[lo, T.COMPRESSOR_KV.value][:B] + bt_sc = md.sliding_block_tables[lo, T.COMPRESSOR_SCORE.value][:B] + page = manager.tokens_per_block + if step == 0: + # Poison every mapped page of these rows, then park the pending token. + for b, s in enumerate(starts): + for t in (bt_kv[b], bt_sc[b]): + mapped = t[t >= 0].long() + pool_kv[mapped] = float("nan") + if s % RATIO: + p = s - 1 + pool_kv[bt_kv[b, p // page].long(), p % page] = kvs[b][p] + pool_sc[bt_sc[b, p // page].long(), p % page] = ( + scores[b][p] + ape[p % RATIO] + ) + for b, s in enumerate(starts): + for p in range(s - (s % RATIO), s + next_n): + assert ( + int(bt_kv[b, p // page]) >= 0 and int(bt_sc[b, p // page]) >= 0 + ), (b, p) + before = pool_kv.clone() + kv_lens = md.kv_lens_cuda_runtime[:B] + assert kv_lens.tolist() == [s + next_n for s in starts] + cu_seq = md.cu_seq_lens_cuda[: B + 1] + cu_out = md.cu_new_comp_kv_cuda[RATIO][: B + 1] + total = md.num_total_compressed_tokens[RATIO] + buf = torch.full((total + 1, HEAD_DIM), float("nan"), device=DEV) + kv_score = torch.cat( + [ + torch.cat([kvs[b][s : s + next_n], scores[b][s : s + next_n]], 1) + for b, s in enumerate(starts) + ] + ).contiguous() + op( + kv_score, + ape, + pool_kv, + pool_sc, + bt_kv, + bt_sc, + buf[:total], + kv_lens, + cu_seq, + cu_out, + B, + page, + HEAD_DIM, + RATIO, + md.num_gen_tokens_per_seq, + ) + torch.cuda.synchronize() + cu = cu_out.tolist() + for b, s in enumerate(starts): + ref = _pool_windows(kvs[b], scores[b], ape, s // 2, (s + next_n) // 2) + torch.testing.assert_close( + buf[cu[b] : cu[b + 1]], ref, msg=lambda m: f"owner {owner} b {b}: {m}" + ) + assert torch.isnan(buf[cu[-1] :]).all() + # Exactly the new positions' kv and score slots changed. + expect = before.clone() + for b, s in enumerate(starts): + pos = torch.arange(s, s + next_n, device=DEV) + expect[bt_kv[b].long()[pos // page], pos % page] = kvs[b][pos] + expect[bt_sc[b].long()[pos // page], pos % page] = ( + scores[b][pos] + ape[pos % RATIO] + ) + assert torch.equal(pool_kv.view(torch.int32), expect.view(torch.int32)), ( + f"owner {owner}: pool" + ) + starts = [s + next_n for s in starts] + finally: + for req in requests: + manager.free_resources(req) + manager.shutdown() + + +# Mixed steps on the real manager: (cached, new) of the context requests, which come +# first in the batch, then the generation requests' histories, then max_draft_len. +# ctx3_gen32 is the target's widest decode graph batch (32 generation rows) behind +# two fresh prompts (even 300 and odd 257 tokens) and a chunked context whose chunk +# starts mid-page (513) with an odd end; ctx1tok_gen1 has a context that emits no +# compressed row, so the generation rows start at output row 0 while their kv_score +# rows start at 1; the chunked case has an odd cached count (a pair straddling the +# chunk boundary) and a history ending at max_seq_len - 95; the verify case puts +# DSpark rows (next_n 6) behind contexts. +_MIXED_GEN32 = torch.randint( + 1, 400, (28,), generator=torch.Generator().manual_seed(2024) +).tolist() + [127, 128, 2047, 2048] +_MIXED_STEPS = { + "ctx3_gen32": ([(0, 300), (0, 257), (513, 388)], _MIXED_GEN32, 0), + "ctx1tok_gen1": ([(0, 1)], [2047], 0), + "ctx2_gen5_chunked": ([(255, 301), (1024, 476)], [1, 2, 129, 1000, 4000], 0), + "ctx2_verify4": ([(0, 130), (513, 300)], [34, 127, 128, 1094], 5), +} +_MIXED_MAX_BATCH = 40 +_MIXED_MAX_TOKENS = 4096 * _MIXED_MAX_BATCH * 2 + + +@pytest.mark.parametrize("scenario", list(_MIXED_STEPS)) +def test_v41_cache_manager_mixed_step_behind_contexts(scenario): + """The target's call on a real mixed step: DeepseekV4CacheManager (V4.1 mode, + tokens_per_block 128, max_batch_size 40, max_seq_len 4096, max_tokens 327680; + the shared compressor-state pool it allocates is fp32 [1062, 128, 512]) and a + DeepseekV4TrtllmAttentionMetadata prepared with the context requests ahead of the + generation rows. The op gets exactly the target's arguments: the step's whole + kv_score [T, 1024] (context rows NaN -- they belong to the prefill op and must not + be read), the generation slices ``[nc:B]`` of the owner's COMPRESSOR_KV / SCORE + sliding tables, ``kv_lens_cuda_runtime``, ``cu_seq_lens_cuda`` (B + 1 entries) and + ``cu_new_comp_kv_cuda[2]``, and an output buffer of + ``num_total_compressed_tokens[2]`` rows whose context rows hold what the prefill + op wrote (a sentinel here). fp32 and bf16 (the target's) output, owners 2/8/14. + The whole shared pool is NaN except each odd-start request's pending token.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, + KvCacheConfig, + ) + from tensorrt_llm.mapping import Mapping + + ctx, gen_cached, max_draft_len = _MIXED_STEPS[scenario] + text = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=128, + head_dim=HEAD_DIM, + index_head_dim=128, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, + ) + config = SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **{k: v for k, v in text.items() if k != "model_type"}, + ) + sparse = DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ) + manager = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, max_tokens=_MIXED_MAX_TOKENS, event_buffer_max_size=0 + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=40, + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=128, + max_seq_len=4096, + max_batch_size=_MIXED_MAX_BATCH, + max_input_len=4096, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=8192, + sparse_attn_config=sparse, + pretrained_config=config, + spec_config=DSparkDecodingConfig(max_draft_len=max_draft_len, block_size=max_draft_len) + if max_draft_len + else None, + ) + try: + _mixed_step_on_manager(manager, scenario, ctx, gen_cached, max_draft_len) + finally: + manager.shutdown() + + +def _mixed_step_on_manager(manager, scenario, ctx, gen_cached, max_draft_len, max_num_tokens=8192): + """One mixed step on ``manager``: contexts ``ctx`` = [(cached, new)] ahead of + generation requests with histories ``gen_cached`` at next_n = max_draft_len + 1, + prepared by DeepseekV4TrtllmAttentionMetadata and handed to the op exactly as the + target hands it (see test_v41_cache_manager_mixed_step_behind_contexts).""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + DeepseekV4TrtllmAttentionMetadata, + ) + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + from tensorrt_llm._torch.metadata import KVCacheParams + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from tensorrt_llm.bindings import SamplingConfig + from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig + from tensorrt_llm.mapping import Mapping + + next_n = max_draft_len + 1 + sparse = DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ) + requests = [] + try: + cached = [c for c, _ in ctx] + list(gen_cached) + new = [n for _, n in ctx] + [next_n] * len(gen_cached) + nc, B = len(ctx), len(cached) + G = B - nc + assert {c % RATIO for c in gen_cached} == {0, 1} or G == 1 + for i, (c, n) in enumerate(zip(cached, new)): + req = LlmRequest( + request_id=i, + max_new_tokens=64, + input_tokens=list(range(c + n)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + assert manager.prepare_context(req), f"manager could not hold request {i}" + assert manager.resize_context(req, req.context_chunk_size) + requests.append(req) + md = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor(new, dtype=torch.int), + request_ids=list(range(B)), + max_num_requests=manager.max_batch_size, + num_contexts=nc, + prompt_lens=[c + n for c, n in zip(cached, new)], + max_num_tokens=max_num_tokens, + kv_cache_manager=manager, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + sparse_attention_config=sparse, + ) + md.prepare() + torch.cuda.synchronize() + + # The offsets the op consumes start past the context part of the step. + ctx_tokens = sum(n for _, n in ctx) + ctx_rows = sum((c + n) // RATIO - c // RATIO for c, n in ctx) + num_tokens = sum(new) + assert (md.num_contexts, md.num_seqs, md.num_ctx_tokens) == (nc, B, ctx_tokens) + assert md.num_gen_tokens_per_seq == next_n + kv_lens = md.kv_lens_cuda_runtime[:B] + cu_seq = md.cu_seq_lens_cuda[: B + 1] + cu_out = md.cu_new_comp_kv_cuda[RATIO][: B + 1] + total = md.num_total_compressed_tokens[RATIO] + assert kv_lens.tolist() == [c + n for c, n in zip(cached, new)] + assert cu_seq.tolist() == [0] + torch.tensor(new).cumsum(0).tolist() + cu = cu_out.tolist() + assert cu[nc] == ctx_rows and total == ctx_rows + G * math.ceil(next_n / RATIO) + assert cu[B] <= total + + gen = _gen(1300 + len(scenario)) + kvs, scores = _truth_rows(gen_cached, next_n, gen) + kv_score = torch.full((num_tokens, 2 * HEAD_DIM), float("nan"), device=DEV) + for j, c in enumerate(gen_cached): + r0 = ctx_tokens + j * next_n + kv_score[r0 : r0 + next_n] = torch.cat( + [kvs[j][c : c + next_n], scores[j][c : c + next_n]], 1 + ) + ape = _zero_ape() + page = manager.tokens_per_block + pool_kv = manager.get_buffers(2, T.COMPRESSOR_KV) + for owner in (2, 8, 14): + p_kv = manager.get_buffers(owner, T.COMPRESSOR_KV) + p_sc = manager.get_buffers(owner, T.COMPRESSOR_SCORE) + assert ( + p_kv.data_ptr() == p_sc.data_ptr() == pool_kv.data_ptr() + and p_kv.shape == pool_kv.shape + ) + assert p_kv.dtype == torch.float32 and tuple(p_kv.shape[1:]) == (128, HEAD_DIM) + lo = manager.layer_offsets[owner] + bt_kv = md.sliding_block_tables[lo, T.COMPRESSOR_KV.value] + bt_sc = md.sliding_block_tables[lo, T.COMPRESSOR_SCORE.value] + p_kv.fill_(float("nan")) + for j, c in enumerate(gen_cached): + b = nc + j + for p in range(c - (c % RATIO), c + next_n): + assert int(bt_kv[b, p // page]) >= 0 and int(bt_sc[b, p // page]) >= 0, ( + owner, + b, + p, + ) + if c % RATIO: + p = c - 1 + p_kv[bt_kv[b, p // page].long(), p % page] = kvs[j][p] + p_sc[bt_sc[b, p // page].long(), p % page] = scores[j][p] + ape[p % RATIO] + for out_dtype in (torch.float32, torch.bfloat16): + buf = torch.full((total + 1, HEAD_DIM), float("nan"), device=DEV, dtype=out_dtype) + buf[:ctx_rows] = SENTINEL + ctx_before = buf[:ctx_rows].clone() + before = p_kv.clone() + op( + kv_score, + ape, + p_kv, + p_sc, + bt_kv[nc:B], + bt_sc[nc:B], + buf[:total], + kv_lens[nc:B], + cu_seq[nc : B + 1], + cu_out[nc : B + 1], + G, + page, + HEAD_DIM, + RATIO, + md.num_gen_tokens_per_seq, + ) + torch.cuda.synchronize() + assert torch.equal(buf[:ctx_rows], ctx_before), ( + f"{scenario} owner {owner}: context output rows written" + ) + assert torch.isnan(buf[cu[B] :]).all(), ( + f"{scenario} owner {owner}: rows past the packed outputs" + ) + got, ref = [], [] + for j, c in enumerate(gen_cached): + b = nc + j + got.append(buf[cu[b] : cu[b + 1]].float()) + ref.append( + _pool_windows(kvs[j], scores[j], ape, c // RATIO, (c + next_n) // RATIO) + ) + assert got[-1].shape == ref[-1].shape, (scenario, owner, b) + got, ref = torch.cat(got), torch.cat(ref) + assert got.shape[0] == cu[B] - cu[nc] and got.shape[0] > 0 + if out_dtype == torch.float32: + # Default fp32 tolerances, as in _check_outputs. + torch.testing.assert_close( + got, ref, msg=lambda m: f"{scenario} owner {owner}: {m}" + ) + else: + torch.testing.assert_close( + got.bfloat16(), + ref.bfloat16(), + msg=lambda m: f"{scenario} owner {owner}: {m}", + ) + # Same gate as test_bf16_output: RNE of an fp32 value a few ulp from the reference. + exact = (got.bfloat16() == ref.bfloat16()).float().mean().item() + assert exact >= 0.999, ( + f"{scenario} owner {owner}: bf16 bit-exact fraction {exact}" + ) + # Exactly the generation requests' new kv and score slots changed. + expect = before.clone() + for j, c in enumerate(gen_cached): + b = nc + j + pos = torch.arange(c, c + next_n, device=DEV) + expect[bt_kv[b].long()[pos // page], pos % page] = kvs[j][pos] + expect[bt_sc[b].long()[pos // page], pos % page] = ( + scores[j][pos] + ape[pos % RATIO] + ) + assert torch.equal(p_kv.view(torch.int32), expect.view(torch.int32)), ( + f"{scenario} owner {owner}: pool" + ) + print( + f"[mixed {scenario} owner {owner}] nc={nc} G={G} next_n={next_n} ctx_tokens={ctx_tokens} " + f"ctx_rows={ctx_rows} gen_rows={cu[B] - cu[nc]} total={total} pool_pages={p_kv.shape[0]} " + f"table_width={bt_kv.shape[-1]} max_kv_len={max(c + n for c, n in zip(cached, new))}" + ) + + # Control: the same call with the generation slices taken from row 0, as a caller that + # forgot the contexts ahead would pass them. The checks above must see it. + if nc and ctx_rows: + lo = manager.layer_offsets[2] + bt_kv = md.sliding_block_tables[lo, T.COMPRESSOR_KV.value] + bt_sc = md.sliding_block_tables[lo, T.COMPRESSOR_SCORE.value] + buf = torch.full((total + 1, HEAD_DIM), float("nan"), device=DEV) + buf[:ctx_rows] = SENTINEL + op( + kv_score, + ape, + pool_kv, + pool_kv, + bt_kv[:G], + bt_sc[:G], + buf[:total], + kv_lens[:G], + cu_seq[: G + 1], + cu_out[: G + 1], + G, + page, + HEAD_DIM, + RATIO, + next_n, + ) + torch.cuda.synchronize() + ctx_hit = int((buf[:ctx_rows] != SENTINEL).any(-1).sum()) + gen_rows = buf[cu[nc] : cu[B]] + bad_gen = int((~torch.isfinite(gen_rows)).any(-1).sum()) + print( + f"[mixed {scenario} control: slices from row 0] context rows overwritten {ctx_hit}/{ctx_rows}, " + f"generation rows not produced {bad_gen}/{cu[B] - cu[nc]}" + ) + assert ctx_hit > 0 or bad_gen > 0, ( + "the mixed-step checks cannot see un-offset generation slices" + ) + finally: + for req in requests: + manager.free_resources(req) + + +def test_rejected_by_the_op(): + """Ratio 1 (V4.1's layer 20, which does no pooling), next_n 0 and a head_dim + without an instantiation raise from the launcher.""" + pk = torch.zeros(4, 128, HEAD_DIM, device=DEV) + bt = torch.arange(4, device=DEV, dtype=torch.int32).view(1, 4) + + def call(ratio, head_dim, next_n): + op( + torch.zeros(max(next_n, 1), 2 * HEAD_DIM, device=DEV), + torch.zeros(RATIO, HEAD_DIM, device=DEV), + pk, + pk, + bt, + bt, + torch.zeros(4, head_dim, device=DEV), + _i32([4]), + _i32([0, max(next_n, 1)]), + _i32([0, 4]), + 1, + 128, + head_dim, + ratio, + next_n, + ) + + with pytest.raises( + RuntimeError, match=r"only supports compress_ratio 2 \(V4\.1\), 4 or 128, got 1" + ): + call(1, HEAD_DIM, 1) + with pytest.raises(RuntimeError, match=r"next_n must be in \[1, "): + call(RATIO, HEAD_DIM, 0) + with pytest.raises(RuntimeError, match="no matching instantiation for HD=256"): + call(RATIO, 256, 1) + + +# --------------------------------------------------------------------------- +# discrimination: wrong variants must fail the same gate +# --------------------------------------------------------------------------- + + +def test_wrong_variants_fail_the_gate(): + ape = _zero_ape() + got, ref, wrong = ( + [], + [], + { + k: [] + for k in ( + "wrong_parity", + "softmax_over_channels", + "uniform_mean", + "pending_token_dropped", + "bf16_operands", + ) + }, + ) + for next_n in (1, 6): + _, kvs, scores, _, outs = _decode_run( + CACHED, 6, next_n, 128, ape, torch.float32, seed=1100 + next_n + ) + for starts, rows in outs: + for b, (s, g) in enumerate(zip(starts, rows)): + w0, w1 = s // 2, (s + next_n) // 2 + if w1 == w0: + continue + kv, sc = kvs[b], scores[b] + got.append(g) + ref.append(_pool_windows(kv, sc, ape, w0, w1)) + p = torch.arange(w0 * 2, w1 * 2, device=DEV) + k, sg = kv[p].unflatten(0, (-1, 2)), sc[p].unflatten(0, (-1, 2)) + # Emitting when p % 2 == 0 instead: each row pools positions (2k - 1, 2k). + q = (p - 1).clamp(min=0) + wrong["wrong_parity"].append( + (kv[q].unflatten(0, (-1, 2)) * sc[q].unflatten(0, (-1, 2)).softmax(1)).sum(1) + ) + wrong["softmax_over_channels"].append((k * sg.softmax(dim=-1)).sum(1)) + wrong["uniform_mean"].append(k.mean(1)) + # A kernel that ignores the pool pools a straddling pair from its new token alone. + r = _pool_windows(kv, sc, ape, w0, w1).clone() + if s % 2: + r[0] = kv[s] + wrong["pending_token_dropped"].append(r) + wrong["bf16_operands"].append( + _pool_windows(kv.bfloat16().float(), sc.bfloat16().float(), ape, w0, w1) + ) + got, ref = torch.cat(got), torch.cat(ref) + torch.testing.assert_close(got, ref) + print( + f"correct reference: max abs err {(got - ref).abs().max().item():.3e} over {got.shape[0]} rows" + ) + for name, rows in wrong.items(): + w = torch.cat(rows) + print(f"control {name}: max abs diff {(got - w).abs().max().item():.3e}") + with pytest.raises(AssertionError): + torch.testing.assert_close(got, w) + + +# --------------------------------------------------------------------------- +# misuse: what the op accepts silently, and what the wrapper therefore rejects +# --------------------------------------------------------------------------- + + +def test_silent_misuse_is_rejected_by_the_wrapper(): + ape = _zero_ape() + next_n = 6 + gen = _gen(1201) + starts = [127, 128, 1001] + lens = [s + next_n for s in starts] + seqs = [_rand_rows(L, gen) for L in lens] + kvs, scores = [s[0] for s in seqs], [s[1] for s in seqs] + ref = torch.cat( + [ + _pool_windows(kv, sc, ape, s // 2, (s + next_n) // 2) + for kv, sc, s in zip(kvs, scores, starts) + ] + ) + n_rows = ref.shape[0] + kv_score = torch.cat( + [ + torch.cat([kv[s : s + next_n], sc[s : s + next_n]], 1) + for kv, sc, s in zip(kvs, scores, starts) + ] + ).contiguous() + + def args(ks, out): + pool = _Pool(lens, 128, 1201) + for b, s in enumerate(starts): + pool.seed(b, kvs[b], scores[b], ape, s, poison=False) + n_out = [(s + next_n) // 2 - s // 2 for s in starts] + return ( + ks, + ape, + pool.kv, + pool.score, + pool.bt_kv, + pool.bt_score, + out, + _i32(lens), + _i32(b * next_n for b in range(len(starts) + 1)), + _i32([0] + torch.tensor(n_out).cumsum(0).tolist()), + len(starts), + 128, + HEAD_DIM, + RATIO, + next_n, + ) + + wide_in = torch.zeros(kv_score.shape[0], 3 * HEAD_DIM, device=DEV) + wide_in[:, : 2 * HEAD_DIM] = kv_score + wide_out = torch.zeros(n_rows, 2 * HEAD_DIM, device=DEV) + cases = { + "fp16_kv_score": (kv_score.half(), lambda: torch.zeros(n_rows, HEAD_DIM, device=DEV)), + "fp16_output": ( + kv_score, + lambda: torch.zeros(n_rows, HEAD_DIM, device=DEV, dtype=torch.float16), + ), + "strided_kv_score": ( + wide_in[:, : 2 * HEAD_DIM], + lambda: torch.zeros(n_rows, HEAD_DIM, device=DEV), + ), + "strided_output": (kv_score, lambda: wide_out[:, :HEAD_DIM]), + } + for name, (ks, make_out) in cases.items(): + out = make_out() + op.raw_call(*args(ks, out)) # outside validating: accepted, no error + torch.cuda.synchronize() + diff = (out.float() - ref).abs().max().item() + print(f"raw op, {name}: returned normally, max abs diff vs reference {diff:.3e}") + assert diff > 1.0, f"{name}: expected a silently wrong result from the bare op" + match = "float32 or bfloat16" if name.startswith("fp16") else "must be contiguous" + with pytest.raises(AssertionError, match=match): + with validating(op): + op(*args(ks, make_out())) + + +# --------------------------------------------------------------------------- +# The engine's own pool, every batch size, every context lead, the gates' capture +# sets +# --------------------------------------------------------------------------- +# +# The engine is the DeepSeek-V4.1 core target's accuracy gate: tep4, +# max_seq_len 16384, max_batch_size 32, max_num_tokens 8192, tokens_per_block +# 128, block reuse on, a per-rank KV quota of 72.68959879875183 GiB (the +# gate's log line "KV cache manager v2 device quota set to ..."). The manager +# built below at that quota reproduces the pool shapes the gate's runtime-value +# dump names (asserted in _engine_manager), so its compressor-state pool is the +# engine's pool, and every test in this section runs on it. + +_ENGINE_QUOTA_BYTES = int(round(72.68959879875183 * (1 << 30))) +_ENGINE_MAX_SEQ = 16384 +# DSpark (max_draft_len 5, overlap on) grows the manager's max_seq_len by 5 + 4 + 5. +_ENGINE_MAX_SEQ_DSPARK = 16398 +# Pool shapes of the gate's runtime-value dump (rank 0): the SWA pool and +# the ratio-2 / ratio-1 COMPRESS views the fake-quant ops were handed. +_ENGINE_SWA_PAGES = 137440 +_ENGINE_COMPRESS_R2_PAGES = {453118, 453119, 453120} +_ENGINE_COMPRESS_R1_PAGES = 151040 +_ENGINE_MANAGERS = {} + +# The largest compressor-state pool any engine at the gate settings can be +# handed. kv_cache_fraction 0.5 takes half of the FREE device memory after the weights +# load, so half of the TOTAL device memory (GB300: 297,020,948,480 B) bounds every quota +# an engine can pick. At that quota the manager's shared fp32 state pool has these +# pages (DSpark off / on x the accuracy gate's max_batch_size 32 and the boot gate's 8, at +# max_num_tokens 8192), with bf16 KV containers and with the native uint8 KV storage +# (TRTLLM_DSV41_KV_STORAGE=native), whose smaller SWA / main-KV entries +# leave the state pool more of the quota. The largest over the gate configurations (native, +# max_num_tokens 8192, batch 8, DSpark off) is 133248 pages (34.9 GB, past element offset 2^33); +# 94464 is the native gates' largest. The engine envelope's B_max (172416) is below. +_ENVELOPE_QUOTA_BYTES = 148_510_474_240 +_ENVELOPE_STATE_PAGES = { + ("bf16", 0, 32): 43518, + ("bf16", 5, 32): 40446, + ("bf16", 0, 8): 78078, + ("bf16", 5, 8): 75516, + ("native", 0, 32): 94464, + ("native", 5, 32): 85884, + ("native", 0, 8): 133248, + ("native", 5, 8): 125952, +} +_ENVELOPE_B_MAX_GATES = max(_ENVELOPE_STATE_PAGES.values()) + +# The engine envelope (per rank): max_num_tokens up to 16384, max_batch_size up to 256, +# max_seq_len L in {16384, 32768, 65536, 131072, 262144, 1048576}, DSpark on (max_draft_len 5: +# the manager's length is Lm = L + 14 and the CUDA-graph metadata adds 5 more) and off. The +# compressor-state tables are the manager's max_blocks_per_seq columns wide: 132 / 260 / 516 / +# 1028 / 2052 / 8196 for those L, with and without DSpark (measured on managers of the envelope). +# Shared compressor-state pools at the same 0.5 x total-memory quota, measured maxima over that +# envelope (max_batch_size sampled, not swept): +# - one unsplit manager (960 tuples per storage mode): +# native [172416, 128, 512] (max_num_tokens 16384, max_batch_size 2, DSpark off, L 16384), bf16 +# 120192 (16384, 1, DSpark); +# - a PP2 stage manager built alone: stage 0 +# holds layers 0..19, so the compressor state of owners 2 / 8 / 14 gets a larger share of the quota +# -- native [247806, 128, 512] (16384, 2, DSpark off, L 16384) and bf16 [188796, 128, 512] (16384, +# 1, DSpark, L 16384). Stage 1 (layers 20..39) holds no compressor state. 247806 is B_max. +_ENVELOPE_TOKENS = 16384 +_ENVELOPE_REQS = 256 +_ENVELOPE_L = (16384, 32768, 65536, 131072, 262144, 1048576) +_ENVELOPE_WIDTHS = (132, 260, 516, 1028, 2052, 8196) +_ENVELOPE_WIDTH_OF_L = dict(zip(_ENVELOPE_L, _ENVELOPE_WIDTHS)) +_GSET = tuple(range(1, 33)) + (40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256) +assert len(_GSET) == 46 +_ENVELOPE_STATE_PAGES[("native", 0, 2, 16384)] = 172416 +_ENVELOPE_STATE_PAGES[("bf16", 5, 1, 16384)] = 120192 +_ENVELOPE_STATE_PAGES[("native", 0, 2, 16384, 0)] = 247806 # PP2 stage 0 +_ENVELOPE_STATE_PAGES[("bf16", 5, 1, 16384, 0)] = 188796 # PP2 stage 0 +_UNSPLIT_B_MAX = 172416 +_ENVELOPE_B_MAX = 247806 +_ENVELOPE_BMAX_MANAGER = dict( + max_draft_len=0, + max_batch_size=2, + quota=_ENVELOPE_QUOTA_BYTES, + storage="native", + max_num_tokens=16384, + pp_stage=0, +) + + +def _envelope_kv_max(L: int, next_n: int) -> int: + """The largest kv length after a generation step at max_seq_len L: L without DSpark (next_n 1); + with DSpark (next_n 6) the manager's Lm = L + 14 plus the CUDA-graph metadata's 5 draft tokens.""" + return L + 14 + 5 if next_n > 1 else L + + +def _v41_pretrained_config(): + text = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=128, + head_dim=HEAD_DIM, + index_head_dim=128, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, + ) + return SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **{k: v for k, v in text.items() if k != "model_type"}, + ) + + +def _engine_manager( + max_draft_len: int, + max_batch_size: int = 32, + quota: int = _ENGINE_QUOTA_BYTES, + storage: str = "bf16", + max_num_tokens: int = 8192, + max_seq_len: int = None, + pp_stage: int = None, +): + """The engine-size V4.1 DeepseekV4CacheManager (one alive at a time; released at module end). + + ``max_batch_size`` 32 is the accuracy gate; 8 is the boot gate's batch size at the + same quota, where the compressor-state pool is larger. ``quota`` defaults to the + accuracy gate's logged quota; ``_ENVELOPE_QUOTA_BYTES`` builds the largest pools an + engine at the gate settings can get. ``storage`` "native" builds it with + TRTLLM_DSV41_KV_STORAGE=native (set for the construction only). ``max_seq_len`` is the + manager's length (default 16384, or 16398 with DSpark). ``pp_stage`` 0 / 1 builds the + stage manager of PP2 (``Mapping(world_size=4, rank=2 * stage, tp_size=2, pp_size=2)``; + the pools do not depend on tp), built alone in this process, so + the manager's cross-rank allreduce(MIN) of its token capacity runs over this rank only -- + the stage-alone pools, an upper bound on what a PP2 engine's stage allocates. Host tier off + (``host_cache_size=0``): it only adds a host copy for suspend/resume and does not + change the device pools.""" + import os + + if max_seq_len is None: + max_seq_len = _ENGINE_MAX_SEQ_DSPARK if max_draft_len else _ENGINE_MAX_SEQ + key = (max_draft_len, max_batch_size, quota, storage, max_num_tokens, max_seq_len, pp_stage) + if key in _ENGINE_MANAGERS: + return _ENGINE_MANAGERS[key] + _release_engine_managers() + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, + KvCacheConfig, + ) + from tensorrt_llm.mapping import Mapping + + env_before = os.environ.get("TRTLLM_DSV41_KV_STORAGE") + os.environ["TRTLLM_DSV41_KV_STORAGE"] = storage + try: + mgr = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + max_gpu_total_bytes=quota, + host_cache_size=0, + event_buffer_max_size=0, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=40, + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=128, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + if pp_stage is None + else Mapping(world_size=4, rank=2 * pp_stage, tp_size=2, pp_size=2, gpus_per_node=4), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=max_num_tokens, + sparse_attn_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ), + pretrained_config=_v41_pretrained_config(), + spec_config=DSparkDecodingConfig(max_draft_len=max_draft_len, block_size=max_draft_len) + if max_draft_len + else None, + max_cuda_graph_batch_size=max_batch_size, + disable_overlap_scheduler=False, + ) + finally: + if env_before is None: + os.environ.pop("TRTLLM_DSV41_KV_STORAGE") + else: + os.environ["TRTLLM_DSV41_KV_STORAGE"] = env_before + if pp_stage is not None: + assert pp_stage == 0, "only PP2 stage 0 holds the compressor state (owners 2 / 8 / 14)" + assert list(mgr.pp_layers) == list(range(20)), list(mgr.pp_layers) + pool = mgr.get_buffers(2, T.COMPRESSOR_KV) + shapes = { + "swa": tuple(mgr.get_buffers(0, T.SWA).shape), + "compress_r2": sorted({mgr.get_buffers(o, T.COMPRESS).shape[0] for o in (2, 8, 14)}), + "compress_r1": tuple(mgr.get_buffers(20, T.COMPRESS).shape) if pp_stage is None else None, + "compressor_state": tuple(pool.shape), + "swa_dtype": str(mgr.get_buffers(0, T.SWA).dtype), + } + print( + f"[engine manager max_draft_len={max_draft_len} max_batch_size={max_batch_size} quota={quota} " + f"storage={storage} max_num_tokens={max_num_tokens} max_seq_len={max_seq_len} pp_stage={pp_stage} " + f"max_blocks_per_seq={mgr.max_blocks_per_seq}] {shapes}" + ) + assert (shapes["swa_dtype"] == "torch.uint8") == (storage == "native"), shapes + if key == (0, 32, _ENGINE_QUOTA_BYTES, "bf16", 8192, _ENGINE_MAX_SEQ, None): + # The gate's dump: this manager is the engine's. + assert shapes["swa"][0] == _ENGINE_SWA_PAGES, shapes + assert set(shapes["compress_r2"]) == _ENGINE_COMPRESS_R2_PAGES, shapes + assert shapes["compress_r1"][0] == _ENGINE_COMPRESS_R1_PAGES, shapes + if quota == _ENVELOPE_QUOTA_BYTES and max_seq_len in (_ENGINE_MAX_SEQ, _ENGINE_MAX_SEQ_DSPARK): + assert quota <= torch.cuda.mem_get_info()[1], ( + "the envelope quota exceeds this device's memory" + ) + if pp_stage is not None: + want = _ENVELOPE_STATE_PAGES.get( + (storage, max_draft_len, max_batch_size, max_num_tokens, pp_stage) + ) + elif max_num_tokens != 8192: + want = _ENVELOPE_STATE_PAGES.get( + (storage, max_draft_len, max_batch_size, max_num_tokens) + ) + else: + want = _ENVELOPE_STATE_PAGES[(storage, max_draft_len, max_batch_size)] + if want is not None: + assert shapes["compressor_state"][0] == want, shapes + if max_seq_len - (14 if max_draft_len else 0) in _ENVELOPE_WIDTH_OF_L: + assert ( + mgr.max_blocks_per_seq + == _ENVELOPE_WIDTH_OF_L[max_seq_len - (14 if max_draft_len else 0)] + ), mgr.max_blocks_per_seq + for owner in (2, 8, 14): + for role in (T.COMPRESSOR_KV, T.COMPRESSOR_SCORE): + b = mgr.get_buffers(owner, role) + assert ( + b.data_ptr() == pool.data_ptr() + and b.shape == pool.shape + and b.dtype == torch.float32 + ) + assert pool.is_contiguous() and tuple(pool.shape[1:]) == (128, HEAD_DIM) + _ENGINE_MANAGERS[key] = mgr + return mgr + + +def _release_engine_managers(): + for mgr in _ENGINE_MANAGERS.values(): + mgr.shutdown() + _ENGINE_MANAGERS.clear() + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +@pytest.fixture(scope="module", autouse=True) +def _engine_managers_released_at_module_end(): + yield + _release_engine_managers() + + +def _engine_state_pool(): + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + + return _engine_manager(0).get_buffers(2, T.COMPRESSOR_KV) + + +class _FarEndPool: + """Synthetic kv / score tables over the ENGINE's shared compressor-state pool. + Only the pages covering ``spans[b] = (lo, hi)`` are mapped, to distinct kv and + score pages taken (shuffled) from ``pages`` -- by default the pool's LAST pages; + every other column points at a trap page the op must never touch (by default the + highest page not in ``pages``). The mapped pages and the trap page are filled with + SENTINEL; the rest of the pool is left alone and covered by a checksum compare in + ``check``.""" + + def __init__( + self, + pool: torch.Tensor, + spans: Sequence[Tuple[int, int]], + seed: int, + pages=None, + trap=None, + width=None, + ): + self.kv = self.score = pool + self.page_size = pool.shape[1] + ps = self.page_size + num_pages = pool.shape[0] + per = [list(range(lo // ps, (hi - 1) // ps + 1)) for lo, hi in spans] + need = 2 * sum(len(p) for p in per) + if pages is None: + pages = list(range(num_pages - need, num_pages)) + assert len(set(pages)) >= need and max(pages) < num_pages, "not enough pages" + pages = list(pages)[:need] + if trap is None: + trap = next(t for t in range(num_pages - 1, -1, -1) if t not in set(pages)) + assert trap not in set(pages) and 0 <= trap < num_pages + self.trap = trap + cpu = torch.Generator().manual_seed(seed) + far = [pages[i] for i in torch.randperm(need, generator=cpu).tolist()] + width = max(p[-1] for p in per) + 2 if width is None else width + assert width > max(p[-1] for p in per) + bt_kv = torch.full((len(spans), width), self.trap, dtype=torch.int32) + bt_sc = torch.full_like(bt_kv, self.trap) + for b, logical in enumerate(per): + for q in logical: + bt_kv[b, q] = far.pop() + bt_sc[b, q] = far.pop() + self.bt_kv, self.bt_score = bt_kv.to(DEV), bt_sc.to(DEV) + self.bt_kv_all, self.bt_score_all = self.bt_kv, self.bt_score + self.mapped = torch.tensor(sorted(set(pages) | {self.trap}), device=DEV) + self.lo_page = int(self.mapped[0]) + pool[self.mapped] = SENTINEL + self.outside = self._outside_sum() + + def _outside_sum(self) -> int: + # Summed 8192 pages at a time: one int64 reduction over a B_max pool would need a + # temporary as large as the pool itself. + whole = sum( + int(self.kv[i : i + 8192].view(torch.int32).sum(dtype=torch.int64)) + for i in range(0, self.kv.shape[0], 8192) + ) + return whole - int(self.kv[self.mapped].view(torch.int32).sum(dtype=torch.int64)) + + def slots(self, table: torch.Tensor, b: int, pos: torch.Tensor): + return table[b].long()[pos // self.page_size], pos % self.page_size + + def put(self, b: int, pos: int, kv_row: torch.Tensor, score_row: torch.Tensor) -> None: + """Park one token's state, as a previous call leaves it (V4.1: zero ape).""" + ps = self.page_size + self.kv[int(self.bt_kv[b, pos // ps]), pos % ps] = kv_row + self.score[int(self.bt_score[b, pos // ps]), pos % ps] = score_row + + def written(self) -> torch.Tensor: + """Number of mapped-page rows (slots) that no longer hold SENTINEL, on device.""" + return (self.kv[self.mapped] != SENTINEL).any(-1).sum() + + def check(self, truth: Sequence[Tuple[int, torch.Tensor, torch.Tensor, Sequence[int]]]) -> None: + """truth: (b, kv, score, positions) -- every listed position holds kv / score bit for + bit; no other slot of the mapped pages, the trap page included, and no byte of + the rest of the pool changed.""" + n = 0 + for b, kv, sc, positions in truth: + pos = torch.as_tensor(list(positions), device=DEV) + assert torch.equal(self.kv[self.slots(self.bt_kv, b, pos)], kv[pos]), ( + f"seq {b}: kv state" + ) + assert torch.equal(self.score[self.slots(self.bt_score, b, pos)], sc[pos]), ( + f"seq {b}: score state" + ) + n += 2 * pos.numel() + assert int(self.written()) == n, "the op wrote a slot outside the named positions" + assert (self.kv[self.trap] == SENTINEL).all(), "the trap page was written" + assert self._outside_sum() == self.outside, "pool written outside the mapped pages" + + +def _hazard_pages(num_pages: int) -> List[int]: + """Pages on both sides of every 32-bit offset hazard of an fp32 [pages, 128, 512] + pool (65536 elements, 256 KiB per page): byte offsets 2^31 (page 8192) and 2^32 + (16384), element offset 2^31 = byte offset 2^33 (32768), element offset 2^32 = byte + offset 2^34 (65536), element offset 2^33 = byte offset 2^35 (131072); those the pool + reaches.""" + return sorted( + c + for c in (8191, 8192, 16383, 16384, 32767, 32768, 65535, 65536, 131071, 131072) + if c < num_pages + ) + + +# Histories of the 32 generation requests: both parities, both sides of a +# 128-token page edge, and up to the engine's max_seq_len. +_GEN_HISTORIES = [ + 1, + 2, + 3, + 126, + 127, + 128, + 129, + 255, + 256, + 257, + 511, + 1000, + 1001, + 2047, + 2048, + 4095, + 4096, + 5003, + 6000, + 7777, + 8191, + 8192, + 9001, + 10240, + 11111, + 12287, + 13000, + 14335, + 15000, + 16001, + 16376, + 16377, +] + + +def _truth_rows(cached: Sequence[int], extra: int, gen: torch.Generator): + """Per-request per-position kv / score arrays over [start - 1, start + extra) (``_Rows``, + absolute positions): only positions from start - 1 on are ever read.""" + kvs, scores = [], [] + for c in cached: + lo = max(c - 1, 0) + kv, sc = _Rows(lo, c + extra), _Rows(lo, c + extra) + kv.t, sc.t = _rand_rows(c + extra - lo, gen) + kvs.append(kv) + scores.append(sc) + return kvs, scores + + +def _envelope_histories(n: int, W: int, next_n: int, seed: int) -> List[int]: + """n generation histories for a table W columns wide (max_seq_len L = the envelope's L of W): + request 0 ends on the table's last slot (W * 128 - 1, past any position the engine writes), + request 1 at the envelope's largest kv length (_envelope_kv_max), request 2 at L, then page and + parity edges, and the rest uniform up to the envelope's largest start -- both parities.""" + L = _ENVELOPE_L[_ENVELOPE_WIDTHS.index(W)] + top = _envelope_kv_max(L, next_n) - next_n + fixed = [ + W * 128 - next_n, + top, + L - next_n, + 1, + 2, + 3, + 126, + 127, + 128, + 129, + 255, + 256, + 257, + 16376, + 16377, + top - 1, + top - 2, + L - next_n - 1, + ] + cpu = torch.Generator().manual_seed(seed) + rand = torch.randint(1, top + 1, (max(n - len(fixed), 0),), generator=cpu).tolist() + return (fixed + rand)[:n] + + +@pytest.mark.parametrize("next_n", [1, 6]) +@pytest.mark.parametrize("W", _ENVELOPE_WIDTHS) +def test_every_batch_size_1_to_256_at_every_table_width(W: int, next_n: int): + """On the batch axis: G generation requests size the launch grid, and the table + width W is the row stride of both tables. At every width the manager builds over the + envelope (132 / 260 / 516 / 1028 / 2052 / 8196 columns: L 16384 .. 1048576), calls at every + G in 1..256 over the first G of 256 requests -- histories up to the envelope's largest kv + length (_envelope_kv_max) and request 0 on the table's last column -- on the engine's pool with + every request's pages at its far end, return for each request exactly the rows of the + G = 256 call bit for bit (re-running a step stores the same state, so the calls are + idempotent), and the G = 256 rows match the reference.""" + pool_t = _engine_state_pool() + ape = _zero_ape() + cached = _envelope_histories(_ENVELOPE_REQS, W, next_n, seed=1400 + W + next_n) + assert len(cached) == _ENVELOPE_REQS and max(cached) + next_n == W * 128 + gen = _gen(1400 + W + next_n) + kvs, scores = _truth_rows(cached, next_n, gen) + pool = _FarEndPool( + pool_t, [(max(c - 1, 0), c + next_n) for c in cached], seed=1400 + next_n, width=W + ) + assert pool.bt_kv.shape == (_ENVELOPE_REQS, W) + pending = [(b, c - 1) for b, c in enumerate(cached) if c % RATIO] + for b, p in pending: + pool.put(b, p, kvs[b][p], scores[b][p]) + full = _step(pool, kvs, scores, cached, next_n, ape) + _check_outputs(kvs, scores, ape, [(list(cached), full)], next_n) + bad = torch.zeros((), dtype=torch.int64, device=DEV) + for G in range(1, _ENVELOPE_REQS + 1): + pool.bt_kv, pool.bt_score = pool.bt_kv_all[:G], pool.bt_score_all[:G] + rows = _step(pool, kvs[:G], scores[:G], cached[:G], next_n, ape) + for b in range(G): + bad += (rows[b].view(torch.int32) != full[b].view(torch.int32)).sum() + # Idempotent state: the same new slots as the G = 256 call, plus the parked tokens. + bad += (pool.written() - 2 * (len(pending) + _ENVELOPE_REQS * next_n)).abs() + pool.bt_kv, pool.bt_score = pool.bt_kv_all, pool.bt_score_all + assert int(bad) == 0, f"{int(bad)} elements / slots differ from the G = {_ENVELOPE_REQS} call" + pool.check( + [ + (b, kvs[b], scores[b], range(max(c - 1, 0) if c % RATIO else c, c + next_n)) + for b, c in enumerate(cached) + ] + ) + print( + f"[G sweep W={W} next_n={next_n}] engine pool {tuple(pool_t.shape)}, far-end pages " + f"{pool.lo_page + 1}..{pool_t.shape[0] - 1}, positions up to {max(cached) + next_n - 1}: " + f"G 1..{_ENVELOPE_REQS} bitwise equal to the G = {_ENVELOPE_REQS} rows" + ) + + +@pytest.mark.parametrize("next_n", [1, 6]) +def test_every_context_lead_on_the_engine_pool(next_n: int): + """On the mixed-step offsets: the target hands the op the step's whole + kv_score and output buffer, so the generation rows start at the context token count + L (cu_seq_lens[0]) and the outputs at the contexts' compressed-row count O + (cu_kv_comp[0]). For every L in 0..16384 - 8 * next_n (O = L // 2; max_num_tokens 16384), + ten generation requests at both parities, two of them at the 1,048,576-token envelope's + far end on 8196-column tables, return exactly the rows of the L = 0 call, bit for bit, + and write nothing else into the output buffer; the rows ahead of them are NaN and + never read. The dumps' leads (1016 / 508, 1934 / 967, and the batch-256 engines' up to + 16290 / 8145) are inside the sweep.""" + pool_t = _engine_state_pool() + ape = _zero_ape() + top = _envelope_kv_max(_ENVELOPE_L[-1], next_n) - next_n + cached = [1, 2, 127, 128, 1001, 2048, 8191, 16376, top - 1, top] + G = len(cached) + gen = _gen(1500 + next_n) + kvs, scores = _truth_rows(cached, next_n, gen) + pool = _FarEndPool( + pool_t, + [(max(c - 1, 0), c + next_n) for c in cached], + seed=1500 + next_n, + width=_ENVELOPE_WIDTHS[-1], + ) + for b, c in enumerate(cached): + if c % RATIO: + pool.put(b, c - 1, kvs[b][c - 1], scores[b][c - 1]) + base = torch.cat(_step(pool, kvs, scores, cached, next_n, ape)) + n = base.shape[0] + rows = torch.cat( + [ + torch.cat([kv[c : c + next_n], sc[c : c + next_n]], 1) + for kv, sc, c in zip(kvs, scores, cached) + ] + ) + kv_lens = _i32(c + next_n for c in cached) + step_cu_out = _i32( + [0] + torch.tensor([(c + next_n) // 2 - c // 2 for c in cached]).cumsum(0).tolist() + ) + step_cu_seq = torch.arange(G + 1, device=DEV, dtype=torch.int32) * next_n + max_lead = _ENVELOPE_TOKENS - G * next_n + kv_score = torch.full((_ENVELOPE_TOKENS, 2 * HEAD_DIM), float("nan"), device=DEV) + out = torch.full((max_lead // 2 + n + 1, HEAD_DIM), float("nan"), device=DEV) + bad = torch.zeros((), dtype=torch.int64, device=DEV) + for L in range(max_lead + 1): + out_lead = L // 2 + kv_score[L : L + G * next_n] = rows + op( + kv_score, + ape, + pool.kv, + pool.score, + pool.bt_kv, + pool.bt_score, + out, + kv_lens, + step_cu_seq + L, + step_cu_out + out_lead, + G, + pool.page_size, + HEAD_DIM, + RATIO, + next_n, + ) + bad += (out[out_lead : out_lead + n].view(torch.int32) != base.view(torch.int32)).sum() + bad += ((~torch.isnan(out)).sum() - n * HEAD_DIM).abs() + out[out_lead : out_lead + n] = float("nan") + kv_score[L : L + G * next_n] = float("nan") + torch.cuda.synchronize() + assert int(bad) == 0, ( + f"{int(bad)} output elements differ from the lead-0 call or were written elsewhere" + ) + pool.check( + [ + (b, kvs[b], scores[b], range(max(c - 1, 0) if c % RATIO else c, c + next_n)) + for b, c in enumerate(cached) + ] + ) + print( + f"[lead sweep next_n={next_n}] leads L 0..{max_lead} (O = L // 2): {n} rows bitwise equal to lead 0" + ) + + +@pytest.mark.parametrize("next_n", [1, 6]) +@pytest.mark.parametrize("W", _ENVELOPE_WIDTHS) +def test_cuda_graph_capture_set_on_the_engine_pool(W: int, next_n: int): + """Captured at every graph batch G of Gset (every integer 1..32 plus the padded + default list for cuda_graph_config {enable_padding: true, max_batch_size: 256}: 40..128 step + 8, 192, 256 -- 46 values; the gates' earlier sets 1..8 and {1, 2, 4, 8, 16, 24, 32} are + subsets), at next_n 1 (decode: n = G output rows) and 6 (DSpark verify: n = 3G, with a + random 0..5 accepted drafts per request between steps, so replays rewind to either parity + and overwrite rejected positions), on tables of every width W the graph metadata uses over + the envelope (132 .. 8196 columns) with histories up to the envelope's largest kv length. + Warm-up and capture on step 0, then 3 replays on new kv_score / kv_lens / cu_kv_comp + written in place. Every replay matches the reference built from what each position + last received, an eager call on the same buffers returns the replay's rows bit for + bit, and the pool holds exactly the new state (far-end pages, rest untouched).""" + pool_t = _engine_state_pool() + ape = _zero_ape() + steps = 3 + L = _ENVELOPE_L[_ENVELOPE_WIDTHS.index(W)] + top = ( + _envelope_kv_max(L, next_n) - (steps + 1) * next_n + ) # every replay's kv length stays <= _envelope_kv_max + cpu = torch.Generator().manual_seed(1600 + W + next_n) + parities = set() + for G in _GSET: + cached = sorted(torch.randint(1, top, (G,), generator=cpu).tolist()) + cached[0] |= 1 # at least one odd start (a pending token read back) + if G > 1: + cached[-1] = top - top % 2 # an even start at the envelope's far end + extra = (steps + 1) * next_n + gen = _gen(1600 + 10 * G + next_n) + kvs, scores = _truth_rows(cached, extra, gen) + pool = _FarEndPool( + pool_t, [(max(c - 1, 0), c + extra) for c in cached], seed=1600 + G, width=W + ) + for b, c in enumerate(cached): + if c % RATIO: + pool.put(b, c - 1, kvs[b][c - 1], scores[b][c - 1]) + reserved = G * math.ceil(next_n / RATIO) + s_kv_score = torch.zeros(G * next_n, 2 * HEAD_DIM, device=DEV) + s_kv_lens = torch.zeros(G, dtype=torch.int32, device=DEV) + s_cu_seq = torch.arange(G + 1, device=DEV, dtype=torch.int32) * next_n + s_cu_out = torch.zeros(G + 1, dtype=torch.int32, device=DEV) + s_out = torch.zeros(reserved, HEAD_DIM, device=DEV) + e_out = torch.zeros_like(s_out) + + def load(starts): + s_kv_score.copy_( + torch.cat( + [ + torch.cat([kv[s : s + next_n], sc[s : s + next_n]], 1) + for kv, sc, s in zip(kvs, scores, starts) + ] + ) + ) + s_kv_lens.copy_(_i32(s + next_n for s in starts)) + n_out = [(s + next_n) // 2 - s // 2 for s in starts] + s_cu_out.copy_(_i32([0] + torch.tensor(n_out).cumsum(0).tolist())) + return n_out + + def call(out): + op( + s_kv_score, + ape, + pool.kv, + pool.score, + pool.bt_kv, + pool.bt_score, + out, + s_kv_lens, + s_cu_seq, + s_cu_out, + G, + pool.page_size, + HEAD_DIM, + RATIO, + next_n, + ) + + def check(starts, n_out, got): + off = 0 + for b, s in enumerate(starts): + ref = _pool_windows(kvs[b], scores[b], ape, s // 2, (s + next_n) // 2) + torch.testing.assert_close( + got[off : off + n_out[b]], ref, msg=lambda m: f"G {G} b {b}: {m}" + ) + off += n_out[b] + assert torch.isnan(got[off:]).all(), f"G {G}: rows past the packed outputs written" + + starts = list(cached) + n_out = load(starts) + s_out.fill_(float("nan")) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + call(s_out) # the engine's warm-up at this batch + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + check(starts, n_out, s_out) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call(s_out) + for _ in range(steps): + accepted = ( + torch.randint(0, next_n, (G,), generator=cpu).tolist() if next_n > 1 else [0] * G + ) + starts = [s + 1 + a for s, a in zip(starts, accepted)] + for b, s in enumerate(starts): + kvs[b][s : s + next_n], scores[b][s : s + next_n] = _rand_rows(next_n, gen) + parities.add(s % RATIO) + n_out = load(starts) + s_out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + check(starts, n_out, s_out) + e_out.fill_(float("nan")) + call(e_out) + torch.cuda.synchronize() + assert torch.equal(s_out.view(torch.int32), e_out.view(torch.int32)), ( + f"G {G}: replay != eager" + ) + # The pool holds what each position last received: the parked token and every + # position a step wrote (rejected drafts included, as later steps overwrote them). + last = [s + next_n for s in starts] + pool.check( + [ + (b, kvs[b], scores[b], range(max(c - 1, 0) if c % RATIO else c, last[b])) + for b, c in enumerate(cached) + ] + ) + del graph + assert parities == {0, 1}, f"replays reached only start parities {parities}" + print( + f"[capture W={W} next_n={next_n}] {len(_GSET)} batches G {_GSET[0]}..{_GSET[-1]} (Gset), output rows " + f"{math.ceil(next_n / RATIO)}G, histories up to {top}: every replay == reference and == eager bitwise" + ) + + +def test_graph_replay_equal_to_eager_cannot_see_a_stale_pending_token(): + """Control for the capture test: a replay and an eager call over the same pool read + the same state, so replay == eager holds even when that state is wrong. A parked + token overwritten by a rejected draft's value (a missed rewind) leaves replay and + eager bitwise equal, and only the reference check rejects it -- which is why the + capture test gates every replay on the reference, not on eager alone.""" + pool_t = _engine_state_pool() + ape = _zero_ape() + next_n, G = 6, 4 + cached = [33, 128, 1001, 4094] + gen = _gen(1700) + kvs, scores = _truth_rows(cached, next_n, gen) + pool = _FarEndPool(pool_t, [(max(c - 1, 0), c + next_n) for c in cached], seed=1700) + stale_kv, stale_sc = _rand_rows(1, gen) + for b, c in enumerate(cached): + if c % RATIO: + pool.put(b, c - 1, stale_kv[0], stale_sc[0]) # not what position c - 1 last received + s_kv_score = torch.cat( + [ + torch.cat([kv[c : c + next_n], sc[c : c + next_n]], 1) + for kv, sc, c in zip(kvs, scores, cached) + ] + ) + n_out = [(c + next_n) // 2 - c // 2 for c in cached] + cu_out = _i32([0] + torch.tensor(n_out).cumsum(0).tolist()) + args = (s_kv_score, ape, pool.kv, pool.score, pool.bt_kv, pool.bt_score) + tail = ( + _i32(c + next_n for c in cached), + _i32(b * next_n for b in range(G + 1)), + cu_out, + G, + pool.page_size, + HEAD_DIM, + RATIO, + next_n, + ) + s_out = torch.zeros(sum(n_out), HEAD_DIM, device=DEV) + op(*args, s_out, *tail) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(*args, s_out, *tail) + s_out.fill_(float("nan")) + graph.replay() + e_out = torch.zeros_like(s_out) + op(*args, e_out, *tail) + torch.cuda.synchronize() + assert torch.equal(s_out.view(torch.int32), e_out.view(torch.int32)) + ref = torch.cat( + [ + _pool_windows(kvs[b], scores[b], ape, c // 2, (c + next_n) // 2) + for b, c in enumerate(cached) + ] + ) + diff = (s_out - ref).abs().max().item() + print( + f"[control stale pending token] replay == eager bitwise; max abs diff vs reference {diff:.3e}" + ) + with pytest.raises(AssertionError): + torch.testing.assert_close(s_out, ref) + + +# Mixed steps on the engine-size manager at the runtime-value dump's leads: a +# 1016-token context ahead of 2 decode rows ((G 2, next_n 1, n 510) in the dump; +# the prefill op's (T 1018, nc 1, max 508)), contexts of 1920 + 14 tokens ahead of +# one decode row ((G 1, n 968); prefill (T 1935, nc 2, max 960)), G = 7 behind a +# chunked context with an odd cached count, a full 8192-token step of 32 requests +# (8161 context tokens ahead of 31 decode rows), and the first two again with +# DSpark verify rows. +_ENGINE_MIXED = { + "dump_ctx1016_gen2": ([(0, 1016)], [2047, 16000], 0), + "dump_ctx1934_gen1": ([(0, 1920), (0, 14)], [5003], 0), + "ctx_chunked_gen7": ([(1023, 1017)], [1, 2, 127, 128, 999, 8192, 16383], 0), + "lead8161_gen31": ([(0, 8161)], _GEN_HISTORIES[:29] + [16382, 16383], 0), + "dump_ctx1016_verify2": ([(0, 1016)], [2047, 16000], 5), + "dump_ctx1934_verify1": ([(0, 1920), (0, 14)], [5003], 5), + # Many contexts ahead (nc 12 / 31 / 24, B = 32), odd and even cached prefixes. + "ctx12_gen20": ( + [ + (0, 700), + (129, 611), + (0, 1), + (2048, 1000), + (4095, 333), + (0, 64), + (777, 777), + (16000, 300), + (1, 2), + (0, 1500), + (8191, 900), + (300, 1), + ], + _GEN_HISTORIES[:20], + 0, + ), + "ctx31_gen1": ([((37 * i) % 500, 200 + 3 * i) for i in range(31)], [5003], 0), + "ctx24_verify8": ( + [((101 * i) % 3000, 150 + (7 * i) % 100) for i in range(24)], + [34, 127, 128, 1094, 2047, 4096, 9001, 16001], + 5, + ), +} + + +@pytest.mark.parametrize("scenario", list(_ENGINE_MIXED)) +def test_engine_size_manager_mixed_step(scenario): + """The target's exact call on a mixed step prepared by the engine-size manager + (see _engine_manager): the same checks as test_v41_cache_manager_mixed_step_behind_contexts.""" + ctx, gen_cached, max_draft_len = _ENGINE_MIXED[scenario] + _mixed_step_on_manager(_engine_manager(max_draft_len), scenario, ctx, gen_cached, max_draft_len) + + +# Mixed steps on managers of the engine envelope: bf16 storage at 0.5 x total memory, +# max_num_tokens 16384, max_batch_size 256, the manager's length Lm (L + 14 with DSpark), so the +# metadata's tables are 260 / 8196 columns wide. (contexts [(cached, new)], generation histories, +# max_draft_len, L). The requests are allocated at their full length in one resize_context, so +# they stay short: a 1,048,576-token request allocated that way does not fit +# (resize_context returned False), and at L 1,048,576 with max_batch_size 256 the manager +# (compressor-state pool 1662 pages, SWA pool 11000) did not hold 256 short requests +# (prepare_context returned False for the 248th behind a 2000-token context, and for the 255th +# with DSpark behind two short contexts). Positions near 1M on 8196-column tables are +# the synthetic tests' (test_every_batch_size_1_to_256_at_every_table_width, the envelope steps). 39 +# contexts ahead of 217 verify rows at L 32768 (256 requests, the batch-256 engines' nc 39); at +# L 1,048,576 a 2000-token context ahead of 239 decode rows, and two contexts ahead of 238 +# verify rows (240 requests). +_ENVELOPE_RNG = torch.Generator().manual_seed(2026) +_ENVELOPE_CTX39 = [ + (int(c), 150 + (97 * i) % 300) + for i, c in enumerate(torch.randint(0, 400, (39,), generator=_ENVELOPE_RNG).tolist()) +] +_ENVELOPE_SHORT = torch.randint(1, 400, (256,), generator=_ENVELOPE_RNG).tolist() +_ENVELOPE_MIXED_STEPS = { + "gc_L32k_ctx39_verify217": (_ENVELOPE_CTX39, _ENVELOPE_SHORT[:217], 5, 32768), + "gc_L1m_nods_ctx1_gen239": ([(0, 2000)], [1 + j % 120 for j in range(239)], 0, 1_048_576), + "gc_L1m_ctx2_verify238": ( + [(0, 700), (129, 301)], + [1 + j % 110 for j in range(238)], + 5, + 1_048_576, + ), +} + + +@pytest.mark.parametrize("scenario", list(_ENVELOPE_MIXED_STEPS)) +def test_envelope_manager_mixed_step(scenario): + """The target's exact call on a mixed step prepared by a manager of the engine envelope + (max_batch_size 256, max_num_tokens 16384, max_seq_len 32768 + 14 and 1,048,576 (+ 14); the + metadata's tables 260 / 8196 columns wide, 256 requests in the step at L 32768 and 240 at L + 1,048,576): the same checks as test_v41_cache_manager_mixed_step_behind_contexts.""" + ctx, gen_cached, max_draft_len, L = _ENVELOPE_MIXED_STEPS[scenario] + lm = L + (14 if max_draft_len else 0) + assert sum(n for _, n in ctx) + len(gen_cached) * (max_draft_len + 1) <= _ENVELOPE_TOKENS + assert len(ctx) + len(gen_cached) <= _ENVELOPE_REQS + assert all(c + n <= lm for c, n in ctx) and all(c + max_draft_len + 1 <= lm for c in gen_cached) + mgr = _engine_manager( + max_draft_len, + _ENVELOPE_REQS, + _ENVELOPE_QUOTA_BYTES, + "bf16", + max_num_tokens=_ENVELOPE_TOKENS, + max_seq_len=lm, + ) + assert mgr.max_blocks_per_seq == _ENVELOPE_WIDTH_OF_L[L] + _mixed_step_on_manager( + mgr, scenario, ctx, gen_cached, max_draft_len, max_num_tokens=_ENVELOPE_TOKENS + ) + + +# The four engine-size pools: DSpark off / on (max_draft_len 5) at the accuracy gate's +# max_batch_size 32 and at the boot gate's 8, one quota. Their compressor-state pools +# are [22908 | 41088 | 21246 | 39678, 128, 512] fp32; the two at batch 8 pass element +# offset 2^31. Then three pools at the envelope quota: [43518 | 78078 | 75516, 128, +# 512] (DSpark off at 32 and 8, on at 8); the two at batch 8 pass element offset 2^32. +# Then two with the native KV storage at that quota: [94464, 128, 512] (batch 32, the +# size of the native gates' largest pool) and [133248, 128, 512] (batch 8), which +# passes element offset 2^33. Then the engine envelope's largest unsplit pools: the bf16 +# storage's [120192, 128, 512] (max_num_tokens 16384, batch 1, DSpark) and the native +# [172416, 128, 512] (max_num_tokens 16384, batch 2, DSpark off; 45.2 GB). Then the PP2 stage-0 +# managers at the same tuples: bf16 [188796, 128, 512] and B_max, native [247806, 128, 512] +# (65.0 GB). +_ENGINE_POOLS = [ + (0, 32, _ENGINE_QUOTA_BYTES, "bf16", 8192, None), + (0, 8, _ENGINE_QUOTA_BYTES, "bf16", 8192, None), + (5, 32, _ENGINE_QUOTA_BYTES, "bf16", 8192, None), + (5, 8, _ENGINE_QUOTA_BYTES, "bf16", 8192, None), + (0, 32, _ENVELOPE_QUOTA_BYTES, "bf16", 8192, None), + (0, 8, _ENVELOPE_QUOTA_BYTES, "bf16", 8192, None), + (5, 8, _ENVELOPE_QUOTA_BYTES, "bf16", 8192, None), + (0, 32, _ENVELOPE_QUOTA_BYTES, "native", 8192, None), + (0, 8, _ENVELOPE_QUOTA_BYTES, "native", 8192, None), + (5, 1, _ENVELOPE_QUOTA_BYTES, "bf16", 16384, None), + (0, 2, _ENVELOPE_QUOTA_BYTES, "native", 16384, None), + (5, 1, _ENVELOPE_QUOTA_BYTES, "bf16", 16384, 0), + (0, 2, _ENVELOPE_QUOTA_BYTES, "native", 16384, 0), +] + + +@pytest.mark.parametrize( + "max_draft_len,max_batch_size,quota,storage,max_num_tokens,pp_stage", + _ENGINE_POOLS, + ids=[ + f"D{d}-bs{b}-{'envelope' if q == _ENVELOPE_QUOTA_BYTES else 'gate'}" + + ("-native" if st == "native" else "") + + (f"-mnt{t}" if t != 8192 else "") + + (f"-pp2stage{pp}" if pp is not None else "") + for d, b, q, st, t, pp in _ENGINE_POOLS + ], +) +def test_offset_hazard_pages_of_the_engine_pools( + max_draft_len: int, max_batch_size: int, quota: int, storage: str, max_num_tokens: int, pp_stage +): + """At the far end, run: 32 generation requests at next_n = max_draft_len + 1 + whose kv / score pages are the pages on both sides of every 32-bit offset hazard the + pool reaches (_hazard_pages) and otherwise its last pages. Outputs match the + reference, the pool holds exactly the new state, and no other byte of the pool + changed. Control: the same step with every page id at or past the pool's highest + hazard wrapped below it (as a 32-bit offset would wrap) fails that check.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + + pool_t = _engine_manager( + max_draft_len, + max_batch_size, + quota, + storage, + max_num_tokens=max_num_tokens, + pp_stage=pp_stage, + ).get_buffers(2, T.COMPRESSOR_KV) + num_pages = pool_t.shape[0] + next_n = max_draft_len + 1 + ape = _zero_ape() + cached = _GEN_HISTORIES + spans = [(max(c - 1, 0), c + next_n) for c in cached] + need = 2 * sum((hi - 1) // 128 - lo // 128 + 1 for lo, hi in spans) + hz = _hazard_pages(num_pages) + pages = hz + [p for p in range(num_pages - need, num_pages) if p not in hz][: need - len(hz)] + gen = _gen(1800 + num_pages) + kvs, scores = _truth_rows(cached, next_n, gen) + + def run(wrap: int): + pool = _FarEndPool(pool_t, spans, seed=1800 + num_pages, pages=pages) + for b, c in enumerate(cached): + if c % RATIO: + pool.put(b, c - 1, kvs[b][c - 1], scores[b][c - 1]) + truth_bt = (pool.bt_kv, pool.bt_score) + if wrap: + pool.bt_kv = torch.where(pool.bt_kv >= wrap, pool.bt_kv % wrap, pool.bt_kv) + pool.bt_score = torch.where(pool.bt_score >= wrap, pool.bt_score % wrap, pool.bt_score) + rows = _step(pool, kvs, scores, cached, next_n, ape) + pool.bt_kv, pool.bt_score = truth_bt + return pool, rows + + pool, rows = run(0) + _check_outputs(kvs, scores, ape, [(list(cached), rows)], next_n) + truth = [ + (b, kvs[b], scores[b], range(max(c - 1, 0) if c % RATIO else c, c + next_n)) + for b, c in enumerate(cached) + ] + pool.check(truth) + wrap = max(hz) + used = sorted( + {int(p) for p in pool.bt_kv.flatten().tolist() + pool.bt_score.flatten().tolist()} + - {pool.trap} + ) + print( + f"[hazard pages D={max_draft_len} bs={max_batch_size} {storage} pp_stage={pp_stage}] pool " + f"{tuple(pool_t.shape)} " + f"({pool_t.numel() / 2**31:.2f} x 2^31 elements): pages {used[: len(hz)]} ... {used[-1]}; " + "outputs and state exact" + ) + pool, _ = run(wrap) + with pytest.raises(AssertionError): + pool.check(truth) + print(f"[hazard pages control] page ids >= {wrap} wrapped mod {wrap}: the check fails") + + +# --------------------------------------------------------------------------- +# Pool-size invariance up to the largest pool an engine can get +# --------------------------------------------------------------------------- + + +def _invariance_schedule(next_n: int, steps: int, seed: int): + """32 generation requests (histories of both parities, up to max_seq_len 16384, and two + at the 1,048,576-token envelope's far end) and ``steps`` generation steps after them; at next_n 6 each later + step starts at s + 1 + a random 0..5 accepted drafts, so it rewinds over rejected + positions and reaches both start parities. Returns the histories, the per-request + kv / score arrays holding the step-0 rows, and for every later step its starts + and the new rows it feeds.""" + cpu = torch.Generator().manual_seed(seed) + extra = (steps + 1) * next_n + far = ( + _envelope_kv_max(_ENVELOPE_L[-1], next_n) - extra + ) # the 1,048,576-token envelope's far end + cached = [1, 2, 127, 128, 129, 16384 - extra, 16383 - extra, far, far - 1] + cached += torch.randint(1, 16384 - extra, (32 - len(cached),), generator=cpu).tolist() + gen = _gen(seed) + kvs, scores = _truth_rows(cached, extra, gen) + later, starts = [], list(cached) + for _ in range(steps - 1): + accepted = ( + torch.randint(0, next_n, (32,), generator=cpu).tolist() if next_n > 1 else [0] * 32 + ) + starts = ( + [s + 1 + a for s, a in zip(starts, accepted)] if next_n > 1 else [s + 1 for s in starts] + ) + later.append((list(starts), [_rand_rows(next_n, gen) for _ in starts])) + return cached, kvs, scores, later + + +def _run_invariance_schedule( + pool_t, cached, kvs0, scores0, later, next_n, seed, pages=None, trap=None +): + """The schedule on ``pool_t`` (tables into ``pages``, or the pool's last pages): every + step's rows match the reference built from what each position last received, and at + the end the pool holds exactly each position's last state and nothing else changed. + Returns the steps' packed outputs and the final contents of the mapped pages.""" + ape = _zero_ape() + extra = (len(later) + 2) * next_n # _invariance_schedule's (steps + 1) * next_n + kvs = [k.clone() for k in kvs0] + scores = [sc.clone() for sc in scores0] + pool = _FarEndPool( + pool_t, [(max(c - 1, 0), c + extra) for c in cached], seed=seed, pages=pages, trap=trap + ) + for b, c in enumerate(cached): + if c % RATIO: + pool.put(b, c - 1, kvs[b][c - 1], scores[b][c - 1]) + outs = [] + plan = [(list(cached), None)] + later + for starts, rows in plan: + if rows is not None: + for b, (s, (kv, sc)) in enumerate(zip(starts, rows)): + kvs[b][s : s + next_n], scores[b][s : s + next_n] = kv, sc + got = _step(pool, kvs, scores, starts, next_n, ape) + _check_outputs(kvs, scores, ape, [(starts, got)], next_n) + outs.append(torch.cat(got)) + last = [s + next_n for s in plan[-1][0]] + pool.check( + [ + (b, kvs[b], scores[b], range(max(c - 1, 0) if c % RATIO else c, last[b])) + for b, c in enumerate(cached) + ] + ) + return outs, pool_t[pool.mapped].clone(), pool + + +@pytest.mark.parametrize("next_n", [1, 6]) +def test_pool_size_invariance_up_to_the_envelope_pool(next_n: int): + """One schedule of 32 generation requests (a decode run at next_n 1; a + DSpark verify run with random rewinds at next_n 6; odd starts read a parked token; two + requests at the 1,048,576-token envelope's far end) on the same logical rows at the same + page indices of pools of B pages: views [:B] of the largest compressor-state pool the + engine envelope gives (B_max = 247806: the live PP2 stage-0 manager with native KV + storage, max_num_tokens 16384, batch 2, DSpark off, at 0.5 x total memory) for B = the + smallest pool that holds those pages, the accuracy gate's 22908, the bf16 gate envelope's + 43518 (batch 32) and 78078 (batch 8), a mid size, the native gate envelope's 94464 (batch + 32) and 133248 (batch 8, the gates' old bound), the bf16 storage's largest unsplit pool + 120192, the native unsplit largest 172416, the bf16 stage-0 largest 188796, and B_max, plus + a separate allocation of exactly the smallest size. Every run matches the reference, and + all runs are bitwise identical -- outputs of every step and the final contents of the + mapped pages. Then the far end: the same schedule with every request's pages among B_max's + last pages (the last page, 247805, owned) returns the same outputs bit for bit, and no + other byte of the 65.0 GB pool changed.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + + pool_max = _engine_manager(**_ENVELOPE_BMAX_MANAGER).get_buffers(2, T.COMPRESSOR_KV) + b_max = pool_max.shape[0] + assert b_max == _ENVELOPE_B_MAX, pool_max.shape + seed = 1900 + next_n + cached, kvs0, scores0, later = _invariance_schedule(next_n, steps=4, seed=seed) + extra = (len(later) + 2) * next_n + need = 2 * sum((c + extra - 1) // 128 - max(c - 1, 0) // 128 + 1 for c in cached) + cpu = torch.Generator().manual_seed(seed) + low = (torch.randperm(2 * need, generator=cpu) + 1)[:need].tolist() # pages 1..2*need, shuffled + trap = 0 + b_small = max(low) + 1 + sizes = [ + b_small, + 22908, + _ENVELOPE_STATE_PAGES[("bf16", 0, 32)], + (b_small + b_max) // 2, + _ENVELOPE_STATE_PAGES[("bf16", 0, 8)], + _ENVELOPE_STATE_PAGES[("native", 0, 32)], + _ENVELOPE_STATE_PAGES[("bf16", 5, 1, 16384)], + _ENVELOPE_B_MAX_GATES, + _UNSPLIT_B_MAX, + _ENVELOPE_STATE_PAGES[("bf16", 5, 1, 16384, 0)], + b_max, + ] + runs = {} + for B in sizes: + runs[B] = _run_invariance_schedule( + pool_max[:B], cached, kvs0, scores0, later, next_n, seed, low, trap + )[:2] + small = torch.full((b_small, 128, HEAD_DIM), SENTINEL, device=DEV) + runs["own"] = _run_invariance_schedule( + small, cached, kvs0, scores0, later, next_n, seed, low, trap + )[:2] + del small + ref_outs, ref_pages = runs[b_small] + for key, (outs, pages_now) in runs.items(): + assert len(outs) == len(ref_outs) + for i, (o, r) in enumerate(zip(outs, ref_outs)): + assert torch.equal(o.view(torch.int32), r.view(torch.int32)), ( + f"B {key} step {i}: outputs differ" + ) + assert torch.equal(pages_now.view(torch.int32), ref_pages.view(torch.int32)), ( + f"B {key}: state differs" + ) + far_outs, _, far = _run_invariance_schedule( + pool_max, cached, kvs0, scores0, later, next_n, seed + ) + used = set(far.bt_kv.flatten().tolist()) | set(far.bt_score.flatten().tolist()) + assert b_max - 1 in used, "the last page is not owned" + for i, (o, r) in enumerate(zip(far_outs, ref_outs)): + assert torch.equal(o.view(torch.int32), r.view(torch.int32)), ( + f"far end step {i}: outputs differ" + ) + n_rows = sum(o.shape[0] for o in ref_outs) + print( + f"[pool invariance next_n={next_n}] {len(ref_outs)} steps, {n_rows} rows, pages {min(low)}..{max(low)}: " + f"B in {sizes} and an own [{b_small}, 128, 512] allocation bitwise identical (outputs + state) and " + f"== reference; far end pages {far.lo_page}..{b_max - 1} of [{b_max}, 128, 512] bitwise the same rows" + ) + + +# --------------------------------------------------------------------------- +# The multi-request envelope of the target's mixed steps +# --------------------------------------------------------------------------- +# A step of the core target over the engine envelope carries T <= 16384 tokens +# (max_num_tokens) of <= 256 requests (max_batch_size), each at most max_seq_len L long (L up +# to 1,048,576; the manager's Lm = L + 14 with DSpark, and generation kv lengths up to Lm + 5 +# under the CUDA-graph metadata, _envelope_kv_max): nc context requests first, then G generation +# requests at next_n 1 (decode, DSpark off) or 6 (DSpark verify). This op sees only the +# generation part: G, next_n, each request's kv_len, the context lead L = cu_seq_lens[0] (the +# contexts' token count), O = cu_kv_comp[0] (their compressed-row count) and the tables' width +# (the manager's max_blocks_per_seq for that L: 132 .. 8196). The steps below draw all of those +# at random over the envelope; the parked token at an odd start comes either from a real +# earlier op call (a decode step, or a verify step whose rejected drafts this step then +# overwrites -- a rewind) or is planted, with rejected-draft garbage planted behind it. + +_ENV_TOKENS = _ENVELOPE_TOKENS +_ENV_REQS = _ENVELOPE_REQS +_ENV_PAGES = 2048 # plainly more than a step's pages; pool size is the invariance test's axis +_ENV_CORNERS = [ + (1, 0), + (1, 0), + (256, 0), + (256, 0), + (1, -1), + (1, -1), + (128, -1), + (128, -1), + (255, 1), + (255, 1), + (2, -1), + (2, -1), + (1, -2), + (1, -2), +] +_ENV_PARTITIONS = ( + len(_ENV_CORNERS) + 2 * _ENVELOPE_REQS +) # the corners, then every G 1..256 at both next_n +_ENV_GROUPS = 8 + +# Every (G, next_n, num_total_compressed_tokens[2]) the core target's runtime-value dump +# recorded for this op (the boot gates' logs and the TEP4 / DEP4 / TP4 accuracy-gate previews' +# logs, a line kept only if it is exactly one 4-tuple; 12 garbled lines dropped), as +# {G: (n, ...)}; next_n is 1 in all 622. The contexts ahead own n - G of the output rows. +_DUMP_DECODE_CALLS = { + 1: ( + 1, + 968, + 4057, + 4069, + ), + 2: ( + 2, + 510, + ), + 3: (3,), + 4: (4,), + 5: (5,), + 6: (6,), + 7: ( + 7, + 4066, + 4073, + ), + 8: (8,), + 9: (4054,), + 10: (4059,), + 16: ( + 16, + 4050, + ), + 17: (4093,), + 18: (4097,), + 19: (4093,), + 24: (24,), + 27: ( + 1583, + 1607, + 1666, + ), + 28: ( + 1054, + 1631, + 1651, + 1671, + 1827, + ), + 29: ( + 1042, + 1071, + 1106, + 1135, + 1144, + 1185, + 1195, + 1212, + 1213, + 1229, + 1233, + 1238, + 1245, + 1257, + 1261, + 1262, + 1268, + 1273, + 1278, + 1282, + 1283, + 1290, + 1299, + 1306, + 1331, + 1341, + 1342, + 1349, + 1350, + 1355, + 1372, + 1375, + 1377, + 1385, + 1386, + 1389, + 1393, + 1403, + 1413, + 1423, + 1435, + 1437, + 1454, + 1460, + 1461, + 1483, + 1512, + 1696, + ), + 30: ( + 480, + 650, + 667, + 672, + 686, + 696, + 702, + 707, + 717, + 723, + 725, + 727, + 728, + 731, + 734, + 739, + 741, + 742, + 746, + 749, + 754, + 761, + 762, + 763, + 766, + 768, + 770, + 771, + 776, + 778, + 782, + 783, + 785, + 786, + 787, + 789, + 790, + 791, + 792, + 793, + 796, + 797, + 798, + 799, + 800, + 801, + 802, + 803, + 804, + 805, + 806, + 807, + 810, + 811, + 813, + 814, + 815, + 816, + 817, + 819, + 820, + 822, + 825, + 826, + 827, + 828, + 830, + 832, + 833, + 834, + 837, + 838, + 840, + 842, + 843, + 844, + 845, + 847, + 848, + 849, + 850, + 851, + 852, + 853, + 854, + 857, + 860, + 861, + 862, + 863, + 864, + 865, + 866, + 867, + 868, + 869, + 870, + 871, + 872, + 873, + 874, + 875, + 876, + 877, + 878, + 879, + 880, + 881, + 882, + 883, + 884, + 885, + 886, + 887, + 888, + 889, + 890, + 891, + 892, + 893, + 895, + 896, + 897, + 898, + 899, + 900, + 901, + 902, + 903, + 904, + 905, + 906, + 907, + 908, + 909, + 910, + 912, + 914, + 915, + 916, + 917, + 919, + 920, + 921, + 922, + 925, + 927, + 928, + 929, + 930, + 931, + 932, + 933, + 934, + 936, + 938, + 939, + 940, + 941, + 942, + 943, + 944, + 945, + 948, + 951, + 954, + 956, + 957, + 958, + 959, + 960, + 961, + 962, + 964, + 965, + 966, + 967, + 968, + 971, + 972, + 973, + 975, + 976, + 977, + 980, + 984, + 985, + 986, + 989, + 990, + 992, + 995, + 997, + 998, + 999, + 1003, + 1006, + 1007, + 1012, + 1013, + 1014, + 1016, + 1017, + 1018, + 1019, + 1020, + 1021, + 1022, + 1023, + 1026, + 1028, + 1031, + 1033, + 1034, + 1036, + 1038, + 1039, + 1043, + 1047, + 1052, + 1055, + 1056, + 1066, + 1068, + 1071, + 1076, + 1082, + 1083, + 1088, + 1090, + 1097, + 1100, + 1103, + 1109, + 1111, + 1112, + 1119, + 1126, + 1143, + 1159, + 1204, + 1236, + ), + 31: ( + 277, + 278, + 283, + 296, + 305, + 306, + 308, + 309, + 312, + 316, + 317, + 320, + 321, + 324, + 326, + 327, + 328, + 329, + 330, + 332, + 334, + 335, + 337, + 338, + 339, + 340, + 342, + 343, + 344, + 345, + 346, + 347, + 348, + 349, + 351, + 352, + 353, + 354, + 355, + 356, + 357, + 358, + 359, + 361, + 362, + 363, + 364, + 365, + 366, + 367, + 368, + 369, + 370, + 371, + 372, + 373, + 374, + 375, + 376, + 377, + 378, + 379, + 380, + 381, + 382, + 383, + 384, + 385, + 386, + 387, + 388, + 390, + 391, + 392, + 393, + 394, + 395, + 396, + 397, + 398, + 399, + 400, + 401, + 402, + 403, + 404, + 405, + 406, + 407, + 408, + 409, + 410, + 411, + 412, + 413, + 414, + 415, + 416, + 417, + 418, + 419, + 420, + 421, + 422, + 423, + 424, + 425, + 426, + 427, + 428, + 429, + 430, + 431, + 432, + 433, + 434, + 435, + 436, + 437, + 438, + 439, + 440, + 441, + 442, + 443, + 444, + 445, + 446, + 447, + 448, + 449, + 450, + 451, + 452, + 453, + 454, + 455, + 456, + 457, + 458, + 459, + 460, + 461, + 462, + 463, + 464, + 465, + 466, + 467, + 468, + 469, + 470, + 471, + 472, + 473, + 474, + 475, + 476, + 477, + 478, + 479, + 480, + 481, + 482, + 483, + 484, + 485, + 486, + 487, + 488, + 489, + 490, + 491, + 492, + 493, + 494, + 495, + 496, + 497, + 498, + 499, + 500, + 501, + 502, + 503, + 504, + 505, + 506, + 507, + 508, + 509, + 510, + 511, + 512, + 513, + 514, + 515, + 516, + 517, + 518, + 519, + 520, + 521, + 522, + 523, + 524, + 525, + 526, + 527, + 528, + 529, + 530, + 531, + 532, + 533, + 534, + 535, + 536, + 537, + 538, + 539, + 540, + 541, + 542, + 543, + 544, + 545, + 546, + 547, + 548, + 549, + 550, + 551, + 552, + 553, + 554, + 555, + 556, + 558, + 559, + 560, + 562, + 563, + 564, + 565, + 566, + 567, + 569, + 570, + 571, + 572, + 573, + 574, + 575, + 576, + 577, + 578, + 580, + 581, + 583, + 584, + 585, + 586, + 588, + 590, + 591, + 592, + 593, + 594, + 595, + 596, + 598, + 599, + 600, + 601, + 602, + 603, + 604, + 605, + 606, + 608, + 609, + 611, + 616, + 617, + 619, + 620, + 622, + 623, + 630, + 631, + 633, + 636, + 639, + 644, + 649, + 661, + 667, + 680, + 681, + 781, + ), + 32: (32,), +} + + +# Every (G, next_n, num_total_compressed_tokens[2], max_seq_len) the runtime-value dumps of the +# envelope engines recorded for this op: four engines at max_batch_size 256, max_num_tokens +# 16384, L 32768 with DSpark (CUDA graph, eager, TP4, DEP4); one at L 131072 without DSpark; one +# at L 1,048,576 with DSpark. The DEP4 engine's ranks see different steps, so a value recorded by +# one rank only is kept. The contexts ahead own n - G * ceil(next_n / 2) of the output rows. +_ENVELOPE_DUMP_DECODE_CALLS = [ + (1, 1, 1, 131072), + (1, 1, 15, 131072), + (1, 6, 3, 32768), + (1, 6, 3, 1048576), + (1, 6, 17, 32768), + (1, 6, 17, 1048576), + (1, 6, 2279, 32768), + (1, 6, 5217, 32768), + (1, 6, 7891, 32768), + (1, 6, 7893, 32768), + (1, 6, 8148, 32768), + (2, 1, 2, 131072), + (2, 6, 6, 32768), + (2, 6, 6, 1048576), + (3, 6, 9, 32768), + (4, 1, 4, 131072), + (4, 6, 12, 32768), + (4, 6, 12, 1048576), + (8, 1, 8, 131072), + (8, 6, 24, 32768), + (8, 6, 24, 1048576), + (10, 6, 30, 32768), + (16, 1, 16, 131072), + (16, 6, 48, 32768), + (18, 6, 54, 32768), + (22, 6, 4617, 32768), + (22, 6, 4759, 32768), + (23, 6, 4594, 32768), + (24, 1, 24, 131072), + (24, 6, 72, 32768), + (25, 6, 4233, 32768), + (28, 6, 84, 32768), + (32, 1, 32, 131072), + (32, 6, 96, 32768), + (38, 6, 114, 32768), + (40, 6, 120, 32768), + (40, 6, 8144, 32768), + (48, 6, 144, 32768), + (55, 6, 165, 32768), + (56, 6, 168, 32768), + (64, 6, 192, 32768), + (72, 6, 216, 32768), + (74, 6, 8159, 32768), + (77, 6, 231, 32768), + (80, 6, 240, 32768), + (88, 6, 264, 32768), + (91, 6, 273, 32768), + (96, 6, 288, 32768), + (104, 6, 312, 32768), + (104, 6, 8153, 32768), + (112, 6, 336, 32768), + (116, 6, 8152, 32768), + (117, 6, 8155, 32768), + (120, 6, 360, 32768), + (128, 6, 384, 32768), + (130, 6, 390, 32768), + (133, 6, 8158, 32768), + (142, 6, 426, 32768), + (143, 6, 4677, 32768), + (144, 6, 4680, 32768), + (192, 6, 576, 32768), + (256, 6, 768, 32768), +] + + +class _GenStep: + """One mixed step as this op sees it. ``ctx``: [(cached, new)] of the contexts ahead + (they only set L and O); ``starts``: each generation request's history sp; ``prior``: + per request the start of a real earlier op call feeding the same request (None: the + parked token and rejected-draft garbage are planted instead); ``seq``: the engine's + max_seq_len L (16384 .. 1048576), which sets the tables' width ``W`` and the largest + generation kv length (_envelope_kv_max).""" + + def __init__(self, ctx, starts, next_n, prior, seed, seq=16384): + self.ctx, self.starts, self.next_n, self.prior, self.seed = ctx, starts, next_n, prior, seed + self.seq = seq + self.W = _ENVELOPE_WIDTH_OF_L[seq] + self.G = len(starts) + self.L = sum(n for _, n in ctx) + self.O = sum((c + n) // RATIO - c // RATIO for c, n in ctx) + lm = seq + (14 if next_n > 1 else 0) + assert self.L + self.G * next_n <= _ENV_TOKENS and len(ctx) + self.G <= _ENV_REQS + assert all(c + n <= lm for c, n in ctx) + assert all(1 <= s and s + next_n <= _envelope_kv_max(seq, next_n) for s in starts) + + +def _rand_contexts(k: int, L: int, cpu, seq_max: int = 16384) -> List[Tuple[int, int]]: + """k contexts of L tokens in all (uniform cuts), random cached prefixes (0, odd, even) + with each context ending by ``seq_max``.""" + cuts = sorted(torch.randperm(L - 1, generator=cpu)[: k - 1].add(1).tolist()) if k > 1 else [] + lens = [b - a for a, b in zip([0] + cuts, cuts + [L])] + ctx = [] + for n in lens: + u = float(torch.rand(1, generator=cpu)) + c = 0 if u < 0.25 else int(torch.randint(0, seq_max - n + 1, (1,), generator=cpu)) + ctx.append((c, n)) + return ctx + + +def _rand_histories(G: int, next_n: int, cpu, seq: int = 16384) -> List[int]: + """G histories at max_seq_len ``seq``: page-edge and envelope-edge values (the largest + start _envelope_kv_max - next_n and its tail) mixed with uniform ones, both parities.""" + top = _envelope_kv_max(seq, next_n) - next_n + edges = [1, 2, 126, 127, 128, 129, 255, 256, seq - next_n - 1, seq - next_n, top - 1, top] + out = [] + for _ in range(G): + u = float(torch.rand(1, generator=cpu)) + if u < 0.25: + out.append(edges[int(torch.randint(0, len(edges), (1,), generator=cpu))]) + elif u < 0.4: + out.append(int(torch.randint(max(top - 256, 1), top + 1, (1,), generator=cpu))) + else: + out.append(int(torch.randint(1, top + 1, (1,), generator=cpu))) + return out + + +def _random_gen_step(i: int) -> _GenStep: + """Seeded random step i: next_n 1 / 6 alternating, max_seq_len L random among the six + envelope values; the first steps pin the envelope's corners (G 1 and 256, lead 0, the + full 16384-token step, a full batch of 256 behind one context, and 255 contexts ahead of + one request), then every G 1..256 at both next_n, a random number of contexts ahead (none + in one step of five) of random total length up to what the step has left, and the parked + tokens from a real earlier call in every other step pair.""" + cpu = torch.Generator().manual_seed(40_000 + i) + next_n = 6 if i % 2 else 1 + seq = _ENVELOPE_L[int(torch.randint(0, len(_ENVELOPE_L), (1,), generator=cpu))] + if i < len(_ENV_CORNERS): + G, lead = _ENV_CORNERS[i] + else: + G, lead = 1 + ((i - len(_ENV_CORNERS)) // 2) % _ENV_REQS, None # every G at both next_n + room = _ENV_TOKENS - G * next_n + k_max = _ENV_REQS - G + if lead == 0 or k_max == 0 or (lead is None and float(torch.rand(1, generator=cpu)) < 0.2): + ctx = [] + else: + k = ( + 1 + if lead == 1 + else k_max + if lead == -2 + else int(torch.randint(1, k_max + 1, (1,), generator=cpu)) + ) + L = ( + room + if lead in (-1, -2) or float(torch.rand(1, generator=cpu)) < 0.15 + else int(torch.randint(k, room + 1, (1,), generator=cpu)) + ) + ctx = _rand_contexts(min(k, L), L, cpu, seq) + starts = _rand_histories(G, next_n, cpu, seq) + prior = None + if (i // 2) % 2: + # the same requests' previous step: decode at sp - 1, or verify at sp - 1 - accepted + acc = torch.randint(0, next_n, (G,), generator=cpu).tolist() + prior = [max(s - 1 - a, 1) if s > 1 else None for s, a in zip(starts, acc)] + return _GenStep(ctx, starts, next_n, prior, 40_000 + i, seq) + + +def _dump_gen_step(j: int, G: int, n: int, next_n: int = 1, seq: int = 16384) -> _GenStep: + """The dump's (G, next_n, n) at max_seq_len ``seq`` as a step: contexts ahead owning + exactly n - G * ceil(next_n / 2) output rows within the 16384-token budget (random count, + cached parities and order), random histories, parked tokens from a real earlier call in + every other step.""" + cpu = torch.Generator().manual_seed(60_000 + j) + ctx_rows = n - G * math.ceil(next_n / RATIO) + ctx = [] + if ctx_rows: + k = int(torch.randint(1, min(_ENV_REQS - G, ctx_rows) + 1, (1,), generator=cpu)) + cuts = ( + sorted(torch.randperm(ctx_rows - 1, generator=cpu)[: k - 1].add(1).tolist()) + if k > 1 + else [] + ) + rows = [b - a for a, b in zip([0] + cuts, cuts + [ctx_rows])] + odd = [bool(torch.randint(0, 2, (1,), generator=cpu)) for _ in rows] + if sum(2 * r - int(o) for r, o in zip(rows, odd)) + G * next_n > _ENV_TOKENS: + odd = [True] * len(rows) # the shortest contexts owning those rows: 2r - 1 tokens + for r, o in zip(rows, odd): + new = 2 * r - 1 if o else 2 * r + c = 2 * int(torch.randint(0, (seq - new) // 2, (1,), generator=cpu)) + int(o) + ctx.append((c, new)) + starts = _rand_histories(G, next_n, cpu, seq) + prior = [s - 1 if s > 1 else None for s in starts] if j % 2 else None + step = _GenStep(ctx, starts, next_n, prior, 60_000 + j, seq) + assert step.O == ctx_rows and step.O + G * math.ceil(next_n / RATIO) == n, (G, n, step.O) + return step + + +class _EnvGenPool: + """Random distinct kv / score pages of a small pool for every page a request's + positions [lo, hi) touch, in tables ``width`` columns wide; every other table column + points at a trap page.""" + + def __init__( + self, pool: torch.Tensor, spans: Sequence[Tuple[int, int]], seed: int, width: int = 132 + ): + ps = pool.shape[1] + cpu = torch.Generator().manual_seed(seed) + free = (torch.randperm(pool.shape[0] - 1, generator=cpu) + 1).tolist() + self.trap = 0 + self.bt_kv = torch.full((len(spans), width), self.trap, dtype=torch.int32) + self.bt_sc = torch.full_like(self.bt_kv, self.trap) + used = {self.trap} + for b, (lo, hi) in enumerate(spans): + for q in range(lo // ps, (hi - 1) // ps + 1): + self.bt_kv[b, q] = free.pop() + self.bt_sc[b, q] = free.pop() + used |= {int(self.bt_kv[b, q]), int(self.bt_sc[b, q])} + self.bt_kv, self.bt_sc = self.bt_kv.to(DEV), self.bt_sc.to(DEV) + self.mapped = torch.tensor(sorted(used), device=DEV) + + +def _gen_call( + pool_t, bt_kv, bt_sc, kv_score, starts, next_n, lead_rows, lead_out, out, G=None, ape=None +): + """The target's call: generation rows at absolute rows lead_rows + b * next_n of the + step's kv_score, outputs from absolute row lead_out of the step's output buffer.""" + G = len(starts) if G is None else G + n_out = [(s + next_n) // RATIO - s // RATIO for s in starts] + cu_out = [lead_out] + for o in n_out: + cu_out.append(cu_out[-1] + o) + ape = _zero_ape() if ape is None else ape + op( + kv_score, + ape, + pool_t, + pool_t, + bt_kv[:G], + bt_sc[:G], + out, + _i32(s + next_n for s in starts), + _i32(lead_rows + b * next_n for b in range(G + 1)), + _i32(cu_out), + G, + pool_t.shape[1], + HEAD_DIM, + RATIO, + next_n, + ) + return cu_out + + +def _run_gen_step( + step: _GenStep, + pool_t: torch.Tensor, + starts_override=None, + op_lead_shift: int = 0, + pool_cls=None, +): + """Run one envelope step on ``pool_t`` and check it: fp32 rows = reference at the + default tolerances, the target's bf16 output = RNE of the fp32 rows bit for bit, the + contexts' output rows and every row past the packed outputs untouched, the context + rows of kv_score never read (NaN), and the state exactly each position's last value + with nothing else written. Returns (fp32 rows, bf16 rows, reference rows). + ``starts_override`` / ``op_lead_shift`` (controls only) hand the op other starts or + another kv_score lead than the truth's. ``pool_cls`` replaces the table builder + (default ``_EnvGenPool``; the DSpark-tail steps use ``_DSparkTailPool``).""" + G, next_n, ps = step.G, step.next_n, pool_t.shape[1] + gen = _gen(step.seed) + ape = _zero_ape() + prior = step.prior or [None] * G + garbage = 6 if step.prior is None and next_n > 1 else 0 # rejected drafts left behind a rewind + spans = [] + for s, p in zip(step.starts, prior): + lo = min(s - 1, p - 1 if p is not None else s - 1) + hi = max(s + next_n + garbage, p + next_n if p is not None else 0) + spans.append((max(lo, 0), hi)) + env = (pool_cls or _EnvGenPool)(pool_t, spans, step.seed, step.W) + pool_t[env.mapped] = SENTINEL + # Per-request truth arrays over positions [lo, hi) (_Rows): what each position last received. + kvs, scores = [], [] + for lo, hi in spans: + kv, sc = _Rows(lo, hi), _Rows(lo, hi) + kv.t, sc.t = _rand_rows(hi - lo, gen) + kvs.append(kv) + scores.append(sc) + written = [set() for _ in range(G)] + if step.prior is not None: + # The earlier call: each request with a prior start p feeds [p, p + next_n) (at + # next_n 6 some of it past this step's start: drafts this step's rewind overwrites). + idx = [b for b in range(G) if prior[b] is not None] + for b in idx: + p = prior[b] + if p % RATIO: # its own parked token, planted + pos = p - 1 + pool_t[int(env.bt_kv[b, pos // ps]), pos % ps] = kvs[b][pos] + pool_t[int(env.bt_sc[b, pos // ps]), pos % ps] = scores[b][pos] + written[b].add(pos) + if idx: + rows = torch.cat( + [ + torch.cat( + [ + kvs[b][prior[b] : prior[b] + next_n], + scores[b][prior[b] : prior[b] + next_n], + ], + 1, + ) + for b in idx + ] + ).contiguous() + out = torch.full( + (len(idx) * math.ceil(next_n / RATIO) + 1, HEAD_DIM), float("nan"), device=DEV + ) + cu = _gen_call( + pool_t, + env.bt_kv[idx].contiguous(), + env.bt_sc[idx].contiguous(), + rows, + [prior[b] for b in idx], + next_n, + 0, + 0, + out, + ) + torch.cuda.synchronize() + for j, b in enumerate(idx): + p = prior[b] + ref = _pool_windows(kvs[b], scores[b], ape, p // RATIO, (p + next_n) // RATIO) + torch.testing.assert_close( + out[cu[j] : cu[j + 1]], ref, msg=lambda m: f"earlier call b={b}: {m}" + ) + written[b] |= set(range(p, p + next_n)) + # This step's positions get new tokens (at next_n 6 over the earlier call's rejected + # drafts: the rewind); positions past them keep the earlier call's leftovers. + for b, s in enumerate(step.starts): + kvs[b][s : s + next_n], scores[b][s : s + next_n] = _rand_rows(next_n, gen) + for b, s in enumerate(step.starts): + if s % RATIO and (s - 1) not in written[b]: # planted parked token + pos = s - 1 + pool_t[int(env.bt_kv[b, pos // ps]), pos % ps] = kvs[b][pos] + pool_t[int(env.bt_sc[b, pos // ps]), pos % ps] = scores[b][pos] + written[b].add(pos) + for pos in range(s + next_n, s + next_n + garbage): # a rejected draft's leftovers + g_kv, g_sc = _rand_rows(1, gen) + pool_t[int(env.bt_kv[b, pos // ps]), pos % ps] = g_kv[0] + pool_t[int(env.bt_sc[b, pos // ps]), pos % ps] = g_sc[0] + kvs[b][pos], scores[b][pos] = g_kv[0], g_sc[0] + written[b].add(pos) + if garbage: # and garbage over this step's own positions, which it must overwrite + for pos in range(s, s + next_n): + pool_t[int(env.bt_kv[b, pos // ps]), pos % ps] = 1e30 + written[b] |= set(range(s, s + next_n)) + torch.cuda.synchronize() + + # This step's kv_score: L NaN context rows, then the generation rows. + rows = torch.cat( + [ + torch.cat([kv[s : s + next_n], sc[s : s + next_n]], 1) + for kv, sc, s in zip(kvs, scores, step.starts) + ] + ) + kv_score = torch.cat( + [torch.full((step.L, 2 * HEAD_DIM), float("nan"), device=DEV), rows] + ).contiguous() + reserved = G * math.ceil(next_n / RATIO) + op_starts = step.starts if starts_override is None else starts_override + outs = {} + for dt in (torch.float32, torch.bfloat16): + buf = torch.full((step.O + reserved + 1, HEAD_DIM), float("nan"), device=DEV, dtype=dt) + buf[: step.O] = SENTINEL + cu = _gen_call( + pool_t, + env.bt_kv, + env.bt_sc, + kv_score, + op_starts, + next_n, + step.L + op_lead_shift, + step.O, + buf[: step.O + reserved], + ) + outs[dt] = (buf, cu) + torch.cuda.synchronize() + ref = torch.cat( + [ + _pool_windows(kvs[b], scores[b], ape, s // RATIO, (s + next_n) // RATIO) + for b, s in enumerate(step.starts) + ] + ) + N = ref.shape[0] + for dt, (buf, cu) in outs.items(): + assert (buf[: step.O] == SENTINEL).all(), "a context output row was written" + assert torch.isnan(buf[cu[-1] :]).all(), "a row past the packed outputs was written" + assert cu[-1] - step.O == N, (cu[-1], step.O, N) + f32 = outs[torch.float32][0][step.O : step.O + N] + b16 = outs[torch.bfloat16][0][step.O : step.O + N] + torch.testing.assert_close(f32, ref) + assert torch.equal(b16, f32.bfloat16()), "bf16 output != RNE(fp32 output)" + # State: every position a call wrote holds its last value; nothing else changed. + n_slots = 0 + for b in range(G): + pos = torch.tensor(sorted(written[b]), device=DEV) + pk = env.bt_kv[b].long()[pos // ps] + pc = env.bt_sc[b].long()[pos // ps] + assert torch.equal(pool_t[pk, pos % ps], kvs[b][pos]), f"request {b}: kv state" + assert torch.equal(pool_t[pc, pos % ps], scores[b][pos]), f"request {b}: score state" + n_slots += 2 * pos.numel() + assert int((pool_t[env.mapped] != SENTINEL).any(-1).sum()) == n_slots, ( + "a slot outside the named positions changed" + ) + assert (pool_t[env.trap] == SENTINEL).all(), "the trap page was written" + return f32, b16, ref + + +@pytest.mark.parametrize("group", range(_ENV_GROUPS)) +def test_random_partitions_of_the_envelope(group: int): + """Over the engine envelope: seeded random mixed steps -- G in 1..256 + generation requests (every value at both next_n) at next_n 1 and 6 behind 0..256-G + contexts of random total length up to what the 16384-token step leaves (so leads L up to + 16383 and output offsets O up to ~8200), at a random max_seq_len L of the six envelope + values (tables 132 .. 8196 columns wide), random histories up to the envelope's largest + kv length (both parities, page edges and the tail included), parked tokens from a real + earlier call (decode, or a verify whose rejected drafts this step rewinds over) or + planted with rejected-draft garbage behind and over them. Each step: fp32 = reference, + bf16 = RNE(fp32), context and trailing output rows untouched, state exact, nothing else + written (``_run_gen_step``).""" + steps = [_random_gen_step(i) for i in range(group, _ENV_PARTITIONS, _ENV_GROUPS)] + pool_t = torch.empty(_ENV_PAGES, 128, HEAD_DIM, device=DEV) + exact = total = 0 + for step in steps: + _, b16, ref = _run_gen_step(step, pool_t) + exact += int((b16 == ref.bfloat16()).sum()) + total += b16.numel() + frac = exact / max(total, 1) + Gs = sorted({s.G for s in steps}) + print( + f"[envelope random group {group}] {len(steps)} steps, G {Gs[0]}..{Gs[-1]} ({len(Gs)} values), leads " + f"{min(s.L for s in steps)}..{max(s.L for s in steps)}, O " + f"{min(s.O for s in steps)}..{max(s.O for s in steps)}, " + f"contexts up to {max(len(s.ctx) for s in steps)}, widths {sorted({s.W for s in steps})}, histories up to " + f"{max(max(s.starts) for s in steps)}; bf16 vs reference bit-exact {frac:.6f}" + ) + assert frac >= 0.999, frac # the contract's bf16 gate (see test_bf16_output) + + +def test_random_partitions_cover_the_envelope(): + """The seeded steps reach the envelope's corners: every G in 1..256 at both next_n, lead + 0 and the full 16384-token step, up to 255 contexts ahead and full batches of 256, every + max_seq_len (table width) at both next_n, odd and even starts from both a real earlier + call and a planted token, verify rewinds, histories at 1 and at the largest start + (_envelope_kv_max - next_n) of the 1,048,576-token envelope.""" + steps = [_random_gen_step(i) for i in range(_ENV_PARTITIONS)] + for nn in (1, 6): + sub = [s for s in steps if s.next_n == nn] + want = set(range(1, _ENV_REQS + 1)) + assert {s.G for s in sub} == want, (nn, sorted(want - {s.G for s in sub})) + assert any(s.L == 0 for s in sub) and any(s.L + s.G * nn == _ENV_TOKENS for s in sub) + assert {s.W for s in sub} == set(_ENVELOPE_WIDTHS) + assert any(1 in s.starts for s in sub) + assert any( + _envelope_kv_max(_ENVELOPE_L[-1], nn) - nn in s.starts + for s in sub + if s.seq == _ENVELOPE_L[-1] + ) + assert max(len(s.ctx) for s in steps) == _ENV_REQS - 1 and any( + len(s.ctx) + s.G == _ENV_REQS for s in steps + ) + assert len(steps) >= 200 + odd_real = sum( + 1 + for s in steps + if s.prior + for b, x in enumerate(s.starts) + if x % 2 and s.prior[b] is not None + ) + odd_planted = sum(1 for s in steps if not s.prior for x in s.starts if x % 2) + rewinds = sum( + 1 + for s in steps + if s.prior and s.next_n == 6 + for b, x in enumerate(s.starts) + if s.prior[b] is not None and s.prior[b] + 6 > x + ) + assert odd_real and odd_planted and rewinds + print( + f"[coverage] {len(steps)} steps: G 1..{_ENV_REQS} at next_n 1 and 6, leads 0..{max(s.L for s in steps)}, " + f"contexts up to {max(len(s.ctx) for s in steps)}, widths {sorted({s.W for s in steps})}, histories up to " + f"{max(max(s.starts) for s in steps)}; odd starts parked by a real call {odd_real}, planted " + f"{odd_planted}; verify rewinds over rejected drafts {rewinds}" + ) + + +def test_every_dumped_decode_step(): + """The dumps: every (G, next_n, n) the runtime-value dumps recorded for this + op -- the gates' 622 (next_n 1, max_seq_len 16384) and the envelope engines' 62 + (batch 256 at L 32768 with DSpark, L 131072 without, L 1,048,576 with; G up to 256, n up + to 8159) -- built as a step with exactly those values (contexts ahead owning + n - G * ceil(next_n / 2) output rows, histories at that max_seq_len) and run as the + random steps are.""" + pool_t = torch.empty(_ENV_PAGES, 128, HEAD_DIM, device=DEV) + cells = [(G, 1, n, 16384) for G, ns in sorted(_DUMP_DECODE_CALLS.items()) for n in ns] + assert len(cells) == 622 + cells += _ENVELOPE_DUMP_DECODE_CALLS + assert len(_ENVELOPE_DUMP_DECODE_CALLS) == 62 + exact = total = 0 + for j, (G, next_n, n, seq) in enumerate(cells): + _, b16, ref = _run_gen_step(_dump_gen_step(j, G, n, next_n, seq), pool_t) + exact += int((b16 == ref.bfloat16()).sum()) + total += b16.numel() + frac = exact / max(total, 1) + print( + f"[envelope dump] {len(cells)} dumped steps (G up to {max(c[0] for c in cells)}, n up to " + f"{max(c[2] for c in cells)}, max_seq_len up to {max(c[3] for c in cells)}) == reference; " + f"bf16 bit-exact {frac:.6f}" + ) + assert frac >= 0.999, frac + + +def test_kernel_variant_map_over_the_envelope(): + """The profiler's launches over steps spread across the envelope (every + fifth random step plus its corners: G 1..256, next_n 1 and 6, leads 0..16383, tables + 132 .. 8196 columns wide, histories up to 1,048,589, bf16 and fp32 output). Each call + launches exactly one kernel, always the same instantiation + (pagedKvCompressKernel<512, 4, 4, 2, 1, false>: head_dim 512, fp32 kv_score, fp32 state, + ratio 2, one reduction warp, uniform next_n -- the ragged variant runs only when the op is + given `new_tokens_per_seq`, which the entry never passes), with grid (G, 4, 1) and 32 threads -- the host picks the + variant from head_dim, dtypes and ratio only, never from G, next_n, the lead or the + lengths; the output dtype is a runtime argument. Every step profiled here is also + checked against the reference (``_run_gen_step``, outside the profile).""" + import json + import os + import tempfile + + steps = [_random_gen_step(i) for i in range(0, _ENV_PARTITIONS, 5)] + steps += [_random_gen_step(i) for i in range(len(_ENV_CORNERS))] + pool_t = torch.empty(_ENV_PAGES, 128, HEAD_DIM, device=DEV) + launches, expected = [], [] + for step in steps: + _run_gen_step(step, pool_t) + env = _EnvGenPool( + pool_t, [(max(s - 1, 0), s + step.next_n) for s in step.starts], step.seed, step.W + ) + kv_score = torch.randn( + step.L + step.G * step.next_n, 2 * HEAD_DIM, generator=_gen(step.seed), device=DEV + ) + reserved = step.G * math.ceil(step.next_n / RATIO) + for dt in (torch.bfloat16, torch.float32): + out = torch.empty(step.O + reserved, HEAD_DIM, device=DEV, dtype=dt) + launches.append((env, kv_score, step, out)) + expected.append([step.G, 4, 1]) + zero_ape = _zero_ape() # built before the profile starts, so every kernel in it is the op's + torch.cuda.synchronize() + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for env, kv_score, step, out in launches: + _gen_call( + pool_t, + env.bt_kv, + env.bt_sc, + kv_score, + step.starts, + step.next_n, + step.L, + step.O, + out, + ape=zero_ape, + ) + torch.cuda.synchronize() + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "trace.json") + prof.export_chrome_trace(path) + trace = json.load(open(path)) + kernels = [e for e in trace["traceEvents"] if e.get("cat") == "kernel"] + ops = sorted([e for e in kernels if "pagedKvCompress" in e["name"]], key=lambda e: e["ts"]) + names = {e["name"] for e in ops} + grids = [e["args"]["grid"] for e in ops] + blocks = {tuple(e["args"]["block"]) for e in ops} + others = sorted({e["name"] for e in kernels if "pagedKvCompress" not in e["name"]}) + print( + f"[envelope variant map] {len(ops)} op launches over {len(steps)} steps x 2 output dtypes: kernels " + f"{sorted(names)}, blocks {sorted(blocks)}, {len({g[0] for g in grids})} distinct G in grids " + f"{grids[:2]}..{grids[-2:]} (G up to {max(g[0] for g in grids)}); leads {min(s.L for s in steps)}.." + f"{max(s.L for s in steps)}; widths {sorted({s.W for s in steps})}; other kernels {others}" + ) + assert others == [] and len(kernels) == len(launches), ( + "the op launched something besides one kernel per call" + ) + assert len(names) == 1, names + (name,) = names + assert "pagedKvCompressKernel<512,4,4,2,1,false>" in name.replace(" ", ""), name + assert blocks == {(32, 1, 1)}, blocks + assert grids == expected, "grid differs from (G, 4, 1)" + + +def test_envelope_controls(): + """The envelope checks see the call-site defects that matter here: generation rows read + one row early (a lead off by one), starts one position late (kv_lens off by one), and + a verify rewind the op is not told about (the previous step's starts) each fail + ``_run_gen_step`` on a mixed verify step where the correct call passes.""" + pool_t = torch.empty(_ENV_PAGES, 128, HEAD_DIM, device=DEV) + step = _GenStep( + [(0, 1000), (129, 301), (4095, 3)], + [33, 128, 1001, 4094, 16001], + 6, + [30, 125, None, 4090, 15999], + 77_001, + ) + _run_gen_step(step, pool_t) + wrong = { + "lead one row early": dict(op_lead_shift=-1), + "starts one late": dict(starts_override=[s + 1 for s in step.starts]), + "untold rewind": dict( + starts_override=[p if p is not None else s for s, p in zip(step.starts, step.prior)] + ), + } + for name, kw in wrong.items(): + with pytest.raises(AssertionError): + _run_gen_step(step, pool_t, **kw) + print(f"[envelope controls] {sorted(wrong)}: all rejected") + + +# --------------------------------------------------------------------------- +# The DSpark envelope at max_seq_len 16384 +# --------------------------------------------------------------------------- +# With DSpark (max_draft_len 5) and the overlap scheduler on (every gate), the executor widens +# the engine's max_seq_len to 16384 + 5 + 4 + 5 = 16398 (_ENGINE_MAX_SEQ_DSPARK). A request's +# committed tokens stay <= 16384 (the executor keeps the net max_seq_len), so contexts stay below +# position 16384, but a verify step at the end of a full-length request writes generation +# positions up to 16397: logical page 128 of the compressor-state tables, whose width stays the +# manager's max_blocks_per_seq = 132. The steps below put generation requests there. +_DSPARK_MAX_POS = _ENGINE_MAX_SEQ_DSPARK - 1 +_DSPARK_COLS = 132 +_DSPARK_STEPS = 48 + + +class _DSparkTailStep(_GenStep): + """A _GenStep whose generation positions may reach 16397 (contexts still end by 16384).""" + + def __init__(self, ctx, starts, next_n, prior, seed): + self.ctx, self.starts, self.next_n, self.prior, self.seed = ctx, starts, next_n, prior, seed + self.seq, self.W = _ENGINE_MAX_SEQ, _DSPARK_COLS + self.G = len(starts) + self.L = sum(n for _, n in ctx) + self.O = sum((c + n) // RATIO - c // RATIO for c, n in ctx) + assert self.L + self.G * next_n <= _ENV_TOKENS and len(ctx) + self.G <= _ENV_REQS + assert all(c + n <= _ENGINE_MAX_SEQ for c, n in ctx) + assert all(1 <= s and s + next_n - 1 <= _DSPARK_MAX_POS for s in starts) + + +class _DSparkTailPool(_EnvGenPool): + """_EnvGenPool on the manager's 132-column tables: every page a request's positions touch is + mapped to distinct random kv / score pages, except logical page 128 (positions 16384..16511), + which is mapped to the pool's LAST pages (kv P-1-2b, score P-2-2b); every other column + points at the trap page 0.""" + + def __init__( + self, + pool: torch.Tensor, + spans: Sequence[Tuple[int, int]], + seed: int, + width: int = _DSPARK_COLS, + ): + assert width == _DSPARK_COLS + ps = pool.shape[1] + P = pool.shape[0] + cpu = torch.Generator().manual_seed(seed) + far = set(range(P - 2 * len(spans), P)) + free = [p for p in (torch.randperm(P - 1, generator=cpu) + 1).tolist() if p not in far] + self.trap = 0 + self.bt_kv = torch.full((len(spans), _DSPARK_COLS), self.trap, dtype=torch.int32) + self.bt_sc = torch.full_like(self.bt_kv, self.trap) + self.col128 = 0 + for b, (lo, hi) in enumerate(spans): + for q in range(lo // ps, (hi - 1) // ps + 1): + if q == 128: + self.bt_kv[b, q], self.bt_sc[b, q] = P - 1 - 2 * b, P - 2 - 2 * b + self.col128 += 1 + else: + self.bt_kv[b, q] = free.pop() + self.bt_sc[b, q] = free.pop() + self.bt_kv, self.bt_sc = self.bt_kv.to(DEV), self.bt_sc.to(DEV) + self.mapped = torch.tensor( + sorted( + {self.trap} + | set(self.bt_kv.flatten().tolist()) + | set(self.bt_sc.flatten().tolist()) + ), + device=DEV, + ) + _DSparkTailPool.last = self + + +def _dspark_tail_step(i: int) -> _DSparkTailStep: + """Seeded mixed step i at the DSpark tail: next_n 1 / 6 alternating, G 1..256 (32 in the + first two steps), contexts ahead as in the envelope steps, each generation request's start drawn + from the tail [16384 - next_n - 8, 16398 - next_n] (positions up to 16397, both parities) with + probability 3/4 and uniformly otherwise (request 0 always ends at 16397 or 16396), and the + parked tokens from a real earlier call in every other step pair (a verify rewind at next_n 6).""" + cpu = torch.Generator().manual_seed(70_000 + i) + next_n = 6 if i % 2 else 1 + G = 32 if i < 2 else int(torch.randint(1, _ENV_REQS + 1, (1,), generator=cpu)) + room = _ENV_TOKENS - G * next_n + k_max = _ENV_REQS - G + ctx = [] + if k_max and float(torch.rand(1, generator=cpu)) < 0.8: + k = int(torch.randint(1, k_max + 1, (1,), generator=cpu)) + L = int(torch.randint(k, room + 1, (1,), generator=cpu)) + ctx = _rand_contexts(min(k, L), L, cpu) + top = _DSPARK_MAX_POS + 1 - next_n + starts = [] + for _ in range(G): + if float(torch.rand(1, generator=cpu)) < 0.75: + starts.append( + int(torch.randint(_ENGINE_MAX_SEQ - next_n - 8, top + 1, (1,), generator=cpu)) + ) + else: + starts.append(int(torch.randint(1, top + 1, (1,), generator=cpu))) + starts[0] = top - (i // 2) % 2 + prior = None + if (i // 2) % 2: + acc = torch.randint(0, next_n, (G,), generator=cpu).tolist() + prior = [max(s - 1 - a, 1) if s > 1 else None for s, a in zip(starts, acc)] + return _DSparkTailStep(ctx, starts, next_n, prior, 70_000 + i) + + +def test_dspark_tail_mixed_steps(): + """_DSPARK_STEPS seeded mixed steps with generation requests at the DSpark tail (positions up to + 16397 on logical page 128 of 132-column tables, that page at the pool's far end), next_n 1 and + 6, behind random contexts, odd starts parked by a real earlier call (verify rewinds included) + or planted, run and checked as the envelope steps are (``_run_gen_step``): fp32 = reference, bf16 = + RNE(fp32), context and trailing output rows untouched, state exact, nothing else written.""" + steps = [_dspark_tail_step(i) for i in range(_DSPARK_STEPS)] + pool_t = torch.empty(_ENV_PAGES, 128, HEAD_DIM, device=DEV) + exact = total = col128 = tail_rows = 0 + parities, rewinds = set(), 0 + for step in steps: + _, b16, ref = _run_gen_step(step, pool_t, pool_cls=_DSparkTailPool) + col128 += _DSparkTailPool.last.col128 + exact += int((b16 == ref.bfloat16()).sum()) + total += b16.numel() + tail_rows += sum(step.next_n for s in step.starts if s + step.next_n - 1 >= _ENGINE_MAX_SEQ) + parities |= {s % RATIO for s in step.starts if s + step.next_n - 1 >= _ENGINE_MAX_SEQ} + if step.prior and step.next_n == 6: + rewinds += sum( + 1 for s, p in zip(step.starts, step.prior) if p is not None and p + 6 > s + ) + frac = exact / max(total, 1) + top = max(s + st.next_n - 1 for st in steps for s in st.starts) + print( + f"[DSpark tail] {len(steps)} steps, G {min(s.G for s in steps)}..{max(s.G for s in steps)}, " + f"leads {min(s.L for s in steps)}..{max(s.L for s in steps)}, {tail_rows} generation rows at positions " + f">= 16384 (top {top}, start parities {sorted(parities)}), {col128} requests on logical page 128 at the " + f"pool's far end, verify rewinds {rewinds}; bf16 bit-exact {frac:.6f}" + ) + assert top == _DSPARK_MAX_POS and parities == {0, 1} and col128 > 0 and rewinds > 0 + assert frac >= 0.999, frac + + +@pytest.mark.parametrize("next_n", [1, 6]) +def test_dspark_tail_captured_on_the_engine_pool(next_n: int): + """The largest graph batch (G 32) at the DSpark tail, captured on the ENGINE's compressor-state + pool with 132-column tables whose logical page 128 is mapped among the pool's last pages: step + 0 is the eager warm-up (checked), then 3 replays advance every request (request 0 by a full + accepted block, ending at position 16397; the others by 1 + random 0..5 accepted drafts, so + replays rewind over rejected drafts at either parity). Every replay equals the reference and + an eager call bit for bit, and the pool finally holds exactly each position's last value.""" + pool_t = _engine_state_pool() + ape = _zero_ape() + steps, G = 3, 32 + cpu = torch.Generator().manual_seed(71_000 + next_n) + adv = next_n # request 0's advance per step: 1 + (next_n - 1) accepted drafts + c0 = _DSPARK_MAX_POS - (next_n - 1) - steps * adv + cached = [c0] + sorted(torch.randint(c0 - 40, c0 + 1, (G - 1,), generator=cpu).tolist()) + # both start parities from step 0 on, and no request past request 0 (the envelope's end) + cached[1] = c0 if c0 % 2 else c0 - 1 + cached[2] = c0 - 3 if c0 % 2 else c0 - 2 + assert max(cached) == c0 + extra = (steps + 1) * next_n + gen = _gen(71_100 + next_n) + kvs, scores = _truth_rows(cached, extra, gen) + pool = _FarEndPool( + pool_t, + [(max(c - 1, 0), c + extra) for c in cached], + seed=71_200 + next_n, + width=_DSPARK_COLS, + ) + assert pool.bt_kv.shape[1] == _DSPARK_COLS and (pool.bt_kv[:, 128] != pool.trap).any() + for b, c in enumerate(cached): + if c % RATIO: + pool.put(b, c - 1, kvs[b][c - 1], scores[b][c - 1]) + reserved = G * math.ceil(next_n / RATIO) + s_kv_score = torch.zeros(G * next_n, 2 * HEAD_DIM, device=DEV) + s_kv_lens = torch.zeros(G, dtype=torch.int32, device=DEV) + s_cu_seq = torch.arange(G + 1, device=DEV, dtype=torch.int32) * next_n + s_cu_out = torch.zeros(G + 1, dtype=torch.int32, device=DEV) + s_out = torch.zeros(reserved, HEAD_DIM, device=DEV) + e_out = torch.zeros_like(s_out) + + def load(starts): + s_kv_score.copy_( + torch.cat( + [ + torch.cat([kv[s : s + next_n], sc[s : s + next_n]], 1) + for kv, sc, s in zip(kvs, scores, starts) + ] + ) + ) + s_kv_lens.copy_(_i32(s + next_n for s in starts)) + n_out = [(s + next_n) // 2 - s // 2 for s in starts] + s_cu_out.copy_(_i32([0] + torch.tensor(n_out).cumsum(0).tolist())) + return n_out + + def call(out): + op( + s_kv_score, + ape, + pool.kv, + pool.score, + pool.bt_kv, + pool.bt_score, + out, + s_kv_lens, + s_cu_seq, + s_cu_out, + G, + pool.page_size, + HEAD_DIM, + RATIO, + next_n, + ) + + def check(starts, n_out, got): + off = 0 + for b, s in enumerate(starts): + ref = _pool_windows(kvs[b], scores[b], ape, s // 2, (s + next_n) // 2) + torch.testing.assert_close(got[off : off + n_out[b]], ref, msg=lambda m: f"b {b}: {m}") + off += n_out[b] + assert torch.isnan(got[off:]).all(), "rows past the packed outputs written" + + starts = list(cached) + n_out = load(starts) + s_out.fill_(float("nan")) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + call(s_out) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + check(starts, n_out, s_out) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call(s_out) + parities = set() + for _ in range(steps): + accepted = torch.randint(0, next_n, (G,), generator=cpu).tolist() if next_n > 1 else [0] * G + accepted[0] = next_n - 1 + starts = [s + 1 + a for s, a in zip(starts, accepted)] + for b, s in enumerate(starts): + kvs[b][s : s + next_n], scores[b][s : s + next_n] = _rand_rows(next_n, gen) + parities.add(s % RATIO) + n_out = load(starts) + s_out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + check(starts, n_out, s_out) + e_out.fill_(float("nan")) + call(e_out) + torch.cuda.synchronize() + assert torch.equal(s_out.view(torch.int32), e_out.view(torch.int32)), "replay != eager" + assert starts[0] + next_n - 1 == _DSPARK_MAX_POS, starts[0] + last = [s + next_n for s in starts] + pool.check( + [ + (b, kvs[b], scores[b], range(max(c - 1, 0) if c % RATIO else c, last[b])) + for b, c in enumerate(cached) + ] + ) + del graph + print( + f"[DSpark tail capture next_n={next_n}] G 32 at positions up to {_DSPARK_MAX_POS} on 132-column tables " + f"(page 128 at the far end of the [{pool_t.shape[0]}, 128, 512] engine pool): replays == reference " + f"and == eager bitwise; start parities {sorted(parities)}" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_compressor_prefill_reduction.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_compressor_prefill_reduction.py new file mode 100644 index 000000000000..e59486bc4898 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_compressor_prefill_reduction.py @@ -0,0 +1,3761 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the compressor_prefill_reduction catalog entry. + +The entry owns what is certified: `CELLS` is the list, `reference` and `compare` are the gate, and +`test_certified_cells` drives every cell through them, each call of each cell gated against +`op.reference` taken from the pools as they were before the call. The cells are the DeepSeek-V4.1 +ratio-2 compressor of the KV-source layers 2/8/14 -- compress_ratio 2, head_dim 512, fp32 kv_score +rows ``[kv || score]``, fp32 paged state, zero APE (V4.1 has none), fp32 or bf16 output -- over +state page sizes 128 and 256 (the two values DeepseekV4CacheManager accepts for tokens_per_block). +A nonzero APE is covered once so where it enters is measured, not read. + +This file keeps the checks a cell cannot carry: aliased pages behind the window (only the pending +token is read back), the wrong-variant controls, the silent-misuse guards, the undersized +max_outputs hazard, and the engine-size coverage (every row count on the engine's pool, the +target's call on the real managers, offset hazards, the multi-request envelope, every dumped +prefill step, pool-size invariance, zero-row requests). + +Those engine-size checks, and the controls, gate on ``_reference_pool`` / ``_env_ref``: the +checkpoint's own formula (its inference/model.py, Compressor.forward): per channel, the sum over +each consecutive token pair of ``kv * softmax_over_the_pair(score)``, built in fp32 over the WHOLE +sequence at once, independent of the pool layout -- which is what lets it run on the engine's +pools (up to 65 GB) that `op.reference`, which returns whole expected pools, cannot copy. The op +instead runs chunk by chunk over a shared paged pool, so a chunk boundary at an odd position has +to carry the pending token through the pool. + +Pool sizing: the semantics tests use POOL_PAGES, plainly larger than any batch +here needs; the engine's own pools, the multi-request envelope and pool-size +invariance up to the largest pool the gate settings allow are the sections at the end. +Physical pages are a random permutation of the pool, +distinct between the kv and the score tables, and every unused table column +points at a trap page, so slot addressing, a swapped table and an out-of-range +write are all visible. +""" + +import math +import zlib +from typing import Dict, List, Sequence, Tuple + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.compressor_prefill_reduction import ( # noqa: E501 + compressor_prefill_reduction as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "compressor_prefill_reduction requires a CUDA device" + +DEV = torch.device("cuda") +HEAD_DIM = 512 +RATIO = 2 +POOL_PAGES = 1024 +TRAP_PAGE = 0 +KV_SENTINEL = 7777.0 +SCORE_SENTINEL = -7777.0 + +# Chunk plans: sequence -> new-token count per call. Every call is a mixed +# batch whose members sit at different start-position parities, including +# 1-token chunks that only park a pending token (zero outputs for that member). +CHUNK_PLANS: Dict[str, List[List[int]]] = { + "mixed": [ + [1, 2, 97, 1, 200], + [33, 31, 1, 1, 64], + [128, 129, 256, 257, 257], + [1, 1, 1, 1, 1], + [4095, 1, 1, 4094, 1], + ], +} + + +def _gen(seed: int) -> torch.Generator: + g = torch.Generator(device=DEV) + g.manual_seed(seed) + return g + + +def _make_sequence(length: int, gen: torch.Generator) -> Tuple[torch.Tensor, torch.Tensor]: + """fp32 kv and score rows. Scores span near-uniform to near-one-hot pairs; + channel 0 is an exact tie in every pair and channel 1 a gap of 240, where + exp underflows to exactly 0 in fp32.""" + kv = torch.randn(length, HEAD_DIM, generator=gen, device=DEV) * 2.0 + score = torch.randn(length, HEAD_DIM, generator=gen, device=DEV) * 4.0 + score[:, 0] = 1.5 + parity = torch.arange(length, device=DEV) % 2 + score[:, 1] = torch.where(parity == 0, 120.0, -120.0) + return kv, score + + +def _reference_pool(kv: torch.Tensor, score: torch.Tensor, ape: torch.Tensor) -> torch.Tensor: + """model.py's ratio-2 pooling over a whole sequence, fp32: [L, D] -> [L // 2, D].""" + n = kv.shape[0] // RATIO + pos = torch.arange(kv.shape[0], device=DEV) + s = score + ape[pos % RATIO] + kv_g = kv[: n * RATIO].unflatten(0, (n, RATIO)) + s_g = s[: n * RATIO].unflatten(0, (n, RATIO)) + return (kv_g * s_g.softmax(dim=1)).sum(dim=1) + + +class _StatePool: + """Paged fp32 compressor-state pools (kv and score) plus per-sequence block tables.""" + + def __init__(self, lengths: Sequence[int], page_size: int, seed: int): + self.page_size = page_size + self.kv = torch.full((POOL_PAGES, page_size, HEAD_DIM), KV_SENTINEL, device=DEV) + self.score = torch.full((POOL_PAGES, page_size, HEAD_DIM), SCORE_SENTINEL, device=DEV) + n_logical = max(math.ceil(max(lengths) / page_size), 1) + need = len(lengths) * n_logical + assert need < POOL_PAGES, "raise POOL_PAGES" + cpu = torch.Generator().manual_seed(seed) + self.bt_kv = self._table(len(lengths), n_logical, cpu) + self.bt_score = self._table(len(lengths), n_logical, cpu) + + @staticmethod + def _table(batch: int, n_logical: int, cpu: torch.Generator) -> torch.Tensor: + pages = torch.randperm(POOL_PAGES - 1, generator=cpu)[: batch * n_logical] + 1 + table = torch.full((batch, n_logical + 1), TRAP_PAGE, dtype=torch.int32) + table[:, :n_logical] = pages.view(batch, n_logical).to(torch.int32) + return table.to(DEV) + + def slot(self, table: torch.Tensor, b: int, pos: int) -> Tuple[int, int]: + return int(table[b, pos // self.page_size]), pos % self.page_size + + +def _run_call( + pool: _StatePool, + kvs: Sequence[torch.Tensor], + scores: Sequence[torch.Tensor], + starts: Sequence[int], + news: Sequence[int], + ape: torch.Tensor, + out_dtype: torch.dtype, + certify: bool = False, +) -> List[torch.Tensor]: + """One op call over a batch of context chunks; returns each sequence's output rows. + + ``certify`` gates the call on the entry's own ``reference`` (taken from the pools before the + call) and ``compare``, and checks the op returns the three tensors it wrote.""" + rows = [ + torch.cat([kv[s : s + n], sc[s : s + n]], dim=1) + for kv, sc, s, n in zip(kvs, scores, starts, news) + ] + kv_score = torch.cat(rows, dim=0).contiguous() + n_out = [(s + n) // RATIO - s // RATIO for s, n in zip(starts, news)] + cu_seq = [0] + cu_out = [0] + for n, o in zip(news, n_out): + cu_seq.append(cu_seq[-1] + n) + cu_out.append(cu_out[-1] + o) + total = cu_out[-1] + # One guard row past the last output: an out-of-range write lands there. + buf = torch.full((total + 1, HEAD_DIM), float("nan"), device=DEV, dtype=out_dtype) + output = buf[:total] + + def i32(values): + return torch.tensor(values, device=DEV, dtype=torch.int32) + + args = ( + kv_score, + ape, + pool.kv, + pool.score, + pool.bt_kv, + pool.bt_score, + output, + i32([s + n for s, n in zip(starts, news)]), + i32(list(starts)), + i32(cu_seq), + i32(cu_out), + len(news), + pool.page_size, + HEAD_DIM, + RATIO, + max(n_out), + ) + expected = op.reference(*args) if certify else None + result = op(*args) + torch.cuda.synchronize() + if certify: + assert result[0] is output and result[1] is pool.kv and result[2] is pool.score, ( + "the op writes in place and returns (output, paged_kv, paged_score)" + ) + op.compare(result, expected) + assert torch.isnan(buf[total]).all(), "op wrote past the last output row" + return [output[cu_out[b] : cu_out[b + 1]] for b in range(len(news))] + + +def _run_plan( + plan: List[List[int]], + page_size: int, + ape: torch.Tensor, + out_dtype: torch.dtype, + seed: int, + certify: bool = False, +): + gen = _gen(seed) + lengths = [sum(p) for p in plan] + seqs = [_make_sequence(L, gen) for L in lengths] + kvs = [s[0] for s in seqs] + scores = [s[1] for s in seqs] + pool = _StatePool(lengths, page_size, seed) + outs: List[List[torch.Tensor]] = [[] for _ in plan] + starts = [0] * len(plan) + for step in range(len(plan[0])): + news = [p[step] for p in plan] + got = _run_call(pool, kvs, scores, starts, news, ape, out_dtype, certify=certify) + for b, g in enumerate(got): + outs[b].append(g) + starts = [s + n for s, n in zip(starts, news)] + return pool, kvs, scores, [torch.cat(o, dim=0) for o in outs] + + +def _check_state(pool: _StatePool, kvs, scores, ape: torch.Tensor) -> None: + """Every position of every sequence holds its kv and score + ape bit for bit; + every other slot of both pools (the trap page included) is untouched.""" + written_kv = torch.zeros(POOL_PAGES, pool.page_size, dtype=torch.bool, device=DEV) + written_sc = torch.zeros_like(written_kv) + for b, (kv, sc) in enumerate(zip(kvs, scores)): + pos = torch.arange(kv.shape[0], device=DEV) + page_kv = pool.bt_kv[b].long()[pos // pool.page_size] + page_sc = pool.bt_score[b].long()[pos // pool.page_size] + off = pos % pool.page_size + assert torch.equal(pool.kv[page_kv, off], kv), f"seq {b}: paged kv state" + assert torch.equal(pool.score[page_sc, off], sc + ape[pos % RATIO]), ( + f"seq {b}: paged score state" + ) + written_kv[page_kv, off] = True + written_sc[page_sc, off] = True + assert (pool.kv[~written_kv] == KV_SENTINEL).all(), ( + "kv pool written outside the sequences' slots" + ) + assert (pool.score[~written_sc] == SCORE_SENTINEL).all(), ( + "score pool written outside the sequences' slots" + ) + + +def _zero_ape() -> torch.Tensor: + return torch.zeros(RATIO, HEAD_DIM, device=DEV) + + +# --------------------------------------------------------------------------- +# certified surface +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Every call of every cell gated on the entry's reference and compare: the chunk's rows, both + pools bitwise, and the guard row past the last output. A chunk after the first reads the token + the op itself parked. At the end every position of every sequence holds its kv and score + ape + and no other slot of either pool (the trap page included) changed.""" + certifying(op) + spec = cell.spec + seed = zlib.crc32(cell.why.encode()) % 2**31 + if spec["ape"] == "zero": + ape = _zero_ape() + else: + ape = torch.randn(RATIO, HEAD_DIM, generator=_gen(seed + 1), device=DEV) * 2.0 + plan = [list(chunks) for chunks in spec["plan"]] + pool, kvs, scores, _ = _run_plan( + plan, spec["page_size"], ape, spec["out_dtype"], seed=seed, certify=True + ) + _check_state(pool, kvs, scores, ape) + + +def test_only_the_pending_token_is_read_back(): + """Full windows are pooled from the input rows, not re-read from the pool: with + every logical page but the last one aliased to a single scratch page, the + outputs are still exact, and a follow-up call still finds the pending token.""" + ape = _zero_ape() + gen = _gen(59) + length, page_size = 1001, 128 + kv, sc = _make_sequence(length + 3, gen) + pool = _StatePool([length + 3], page_size, seed=59) + last_logical = length // page_size + scratch = int(pool.bt_kv[0, 0]) + pool.bt_kv[0, :last_logical] = scratch + pool.bt_score[0, :last_logical] = int(pool.bt_score[0, 0]) + first = _run_call(pool, [kv], [sc], [0], [length], ape, torch.float32)[0] + second = _run_call(pool, [kv], [sc], [length], [3], ape, torch.float32)[0] + ref = _reference_pool(kv, sc, ape) + torch.testing.assert_close(first, ref[: length // RATIO]) + torch.testing.assert_close(second, ref[length // RATIO :]) + + +# --------------------------------------------------------------------------- +# discrimination: wrong variants must fail the same gate +# --------------------------------------------------------------------------- + + +def test_wrong_variants_fail_the_gate(): + ape = _zero_ape() + plan = CHUNK_PLANS["mixed"] + _, kvs, scores, outs = _run_plan(plan, 128, ape, torch.float32, seed=67) + got = torch.cat(outs, dim=0) + ref = torch.cat([_reference_pool(kv, sc, ape) for kv, sc in zip(kvs, scores)], dim=0) + torch.testing.assert_close(got, ref) + print(f"correct reference: max abs err {(got - ref).abs().max().item():.3e}") + + def pooled(fn): + return torch.cat([fn(kv, sc) for kv, sc in zip(kvs, scores)], dim=0) + + def pairs(kv, sc): + n = kv.shape[0] // RATIO + return kv[: n * RATIO].unflatten(0, (n, RATIO)), sc[: n * RATIO].unflatten(0, (n, RATIO)) + + def softmax_over_channels(kv, sc): + k, s = pairs(kv, sc) + return (k * s.softmax(dim=-1)).sum(dim=1) + + def bf16_operands(kv, sc): + return _reference_pool(kv.bfloat16().float(), sc.bfloat16().float(), ape) + + def uniform_mean(kv, sc): + return pairs(kv, sc)[0].mean(dim=1) + + def pending_token_dropped(): + # A kernel that does not read the pool pools a straddling pair from its + # new token alone (a softmax over one element is 1). + rows = [] + for kv, sc, chunks in zip(kvs, scores, plan): + r = _reference_pool(kv, sc, ape).clone() + start = 0 + for n in chunks[:-1]: + start += n + if start % RATIO == 1: + r[start // RATIO] = kv[start] + rows.append(r) + return torch.cat(rows, dim=0) + + variants = { + "softmax_over_channels": pooled(softmax_over_channels), + "bf16_operands": pooled(bf16_operands), + "uniform_mean": pooled(uniform_mean), + "pending_token_dropped": pending_token_dropped(), + } + for name, wrong in variants.items(): + diff = (got - wrong).abs().max().item() + print(f"control {name}: max abs diff {diff:.3e}") + with pytest.raises(AssertionError): + torch.testing.assert_close(got, wrong) + # The bf16 gate's bit-exact fraction must also see the bf16-operand control. + _, kvs16, scores16, outs16 = _run_plan(plan, 128, ape, torch.bfloat16, seed=67) + got16 = torch.cat(outs16, dim=0) + wrong16 = variants["bf16_operands"].to(torch.bfloat16) + frac = (got16 == wrong16).float().mean().item() + print(f"control bf16_operands vs bf16 output: bit-exact fraction {frac:.6f}") + assert frac < 0.999 + + +# --------------------------------------------------------------------------- +# misuse: what the op accepts silently, and what the wrapper therefore rejects +# --------------------------------------------------------------------------- + + +def _raw_call(kv_score, output, length, page_size=128): + """The bare op, one sequence at start_pos 0, fresh pools.""" + pages = math.ceil(length / page_size) + 1 + pkv = torch.zeros(pages, page_size, HEAD_DIM, device=DEV) + psc = torch.zeros_like(pkv) + bt = torch.arange(pages, device=DEV, dtype=torch.int32).view(1, pages) + n_out = length // RATIO + + def i32(values): + return torch.tensor(values, device=DEV, dtype=torch.int32) + + args = ( + kv_score, + _zero_ape(), + pkv, + psc, + bt, + bt, + output, + i32([length]), + i32([0]), + i32([0, length]), + i32([0, n_out]), + 1, + page_size, + HEAD_DIM, + RATIO, + n_out, + ) + return args, pkv + + +def test_silent_misuse_is_rejected_by_the_wrapper(): + gen = _gen(71) + length = 257 + kv, sc = _make_sequence(length, gen) + kv_score = torch.cat([kv, sc], dim=1).contiguous() + ref = _reference_pool(kv, sc, _zero_ape()) + n_out = length // RATIO + + wide_in = torch.zeros(length, 3 * HEAD_DIM, device=DEV) + wide_in[:, : 2 * HEAD_DIM] = kv_score + wide_out = torch.zeros(n_out, 2 * HEAD_DIM, device=DEV) + cases = { + "fp16_kv_score": (kv_score.half(), torch.zeros(n_out, HEAD_DIM, device=DEV)), + "fp16_output": (kv_score, torch.zeros(n_out, HEAD_DIM, device=DEV, dtype=torch.float16)), + "strided_kv_score": (wide_in[:, : 2 * HEAD_DIM], torch.zeros(n_out, HEAD_DIM, device=DEV)), + "strided_output": (kv_score, wide_out[:, :HEAD_DIM]), + } + for name, (ks, out) in cases.items(): + args, _ = _raw_call(ks, out, length) + op.raw_call(*args) # outside validating: accepted, no error + torch.cuda.synchronize() + diff = (out.float() - ref).abs().max().item() + print(f"raw op, {name}: returned normally, max abs diff vs reference {diff:.3e}") + assert diff > 1.0, f"{name}: expected a silently wrong result from the bare op" + args, _ = _raw_call(ks, torch.zeros_like(out) if out.is_contiguous() else out, length) + match = "float32 or bfloat16" if name.startswith("fp16") else "must be contiguous" + with pytest.raises(AssertionError, match=match): + with validating(op): + op(*args) + + +def test_undersized_max_outputs_silently_skips_rows_and_pending_state(): + """The grid is sized by max_outputs; a sequence with more windows loses the + rows past it AND its trailing pending token (written by its last block).""" + gen = _gen(83) + length = 9 + kv, sc = _make_sequence(length, gen) + kv_score = torch.cat([kv, sc], dim=1).contiguous() + out = torch.full((length // RATIO, HEAD_DIM), float("nan"), device=DEV) + args, pkv = _raw_call(kv_score, out, length) + args = args[:-1] + (2,) + op(*args) + torch.cuda.synchronize() + ref = _reference_pool(kv, sc, _zero_ape()) + torch.testing.assert_close(out[:2], ref[:2]) + assert torch.isnan(out[2:]).all(), "rows past max_outputs were written" + assert torch.equal(pkv[0, 3], kv[3]), "state of a covered window" + assert (pkv[0, 8] == 0).all(), "pending token state was written despite max_outputs" + + +# --------------------------------------------------------------------------- +# The engine's own pool, every row count, the target's call on the +# engine-size manager +# --------------------------------------------------------------------------- +# +# The engine is the DeepSeek-V4.1 core target's accuracy gate: tep4, +# max_seq_len 16384, max_batch_size 32, max_num_tokens 8192, tokens_per_block +# 128, block reuse on, a per-rank KV quota of 72.68959879875183 GiB (the gate's +# log line "KV cache manager v2 device quota set to ..."). The manager built +# below at that quota reproduces the pool shapes the gate's runtime-value dump +# names (asserted in _engine_manager), so its compressor-state pool -- one fp32 +# tensor for the kv and score state of owners 2, 8 and 14 -- is the engine's. + +_ENGINE_QUOTA_BYTES = int(round(72.68959879875183 * (1 << 30))) +_ENGINE_MAX_SEQ = 16384 +# DSpark (max_draft_len 5, overlap on) grows the manager's max_seq_len by 5 + 4 + 5. +_ENGINE_MAX_SEQ_DSPARK = 16398 +# Pool shapes of the gate's runtime-value dump (rank 0). +_ENGINE_SWA_PAGES = 137440 +_ENGINE_COMPRESS_R2_PAGES = {453118, 453119, 453120} +_ENGINE_COMPRESS_R1_PAGES = 151040 +_SENTINEL = 7777.0 +_ENGINE_MANAGERS = {} + +# The engine envelope (per rank): max_num_tokens up to 16384, max_batch_size up to 256 +# (so up to 256 context requests in one call, and max_ctx_compressed_tokens up to 8192), +# max_seq_len L in {16384, 32768, 65536, 131072, 262144, 1048576}, DSpark on (the manager's +# length Lm = L + 14) and off. A context request's kv length reaches Lm (1,048,590 at L 1M with +# DSpark). The compressor-state tables are the manager's max_blocks_per_seq columns wide: 132 / +# 260 / 516 / 1028 / 2052 / 8196 for those L, with and without DSpark (measured on managers of +# the envelope). The largest shared compressor-state pool over the envelope at the 0.5 x +# total-memory quota is the PP2 stage-0 manager's (see _BMAX_PAGES), 247806 pages. +_ENVELOPE_TOKENS = 16384 +_ENVELOPE_REQS = 256 +_ENVELOPE_L = (16384, 32768, 65536, 131072, 262144, 1048576) +_ENVELOPE_WIDTHS = (132, 260, 516, 1028, 2052, 8196) +_ENVELOPE_WIDTH_OF_L = dict(zip(_ENVELOPE_L, _ENVELOPE_WIDTHS)) +_ENVELOPE_LM_1M = 1_048_576 + 14 # the longest context kv length: L 1M with DSpark + + +def _envelope_lm(seq: int, dspark: bool) -> int: + return seq + 14 if dspark else seq + + +def _v41_pretrained_config(): + from types import SimpleNamespace + + text = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=128, + head_dim=HEAD_DIM, + index_head_dim=128, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, + ) + return SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **{k: v for k, v in text.items() if k != "model_type"}, + ) + + +def _engine_manager( + max_draft_len: int, + max_batch_size: int = 32, + quota: int = _ENGINE_QUOTA_BYTES, + storage: str = "bf16", + max_num_tokens: int = 8192, + max_seq_len: int = None, + pp_stage: int = None, +): + """The engine-size V4.1 DeepseekV4CacheManager (one alive at a time; released at module end). + + ``max_batch_size`` 32 is the accuracy gate; 8 is the boot gate's batch size at the + same quota, where the compressor-state pool is larger. ``quota`` defaults to the + accuracy gate's logged per-rank quota; the envelope tests pass the largest quota the + device can give at kv_cache_fraction 0.5. ``storage`` "native" builds it with + TRTLLM_DSV41_KV_STORAGE=native (set for the construction only). ``max_seq_len`` is the + manager's length (default 16384, or 16398 with DSpark). ``pp_stage`` 0 builds the stage-0 + manager of PP2 (``Mapping(world_size=4, rank=0, tp_size=2, pp_size=2)``; the pools do not + depend on tp), built alone in this process, so the manager's + cross-rank allreduce(MIN) of its token capacity runs over this rank only -- the stage-alone + pools, an upper bound on what a PP2 engine's stage allocates. Host tier off + (``host_cache_size=0``): it only adds a host copy for suspend/resume and does not + change the device pools.""" + import os + + if max_seq_len is None: + max_seq_len = _ENGINE_MAX_SEQ_DSPARK if max_draft_len else _ENGINE_MAX_SEQ + key = (max_draft_len, max_batch_size, quota, storage, max_num_tokens, max_seq_len, pp_stage) + if key in _ENGINE_MANAGERS: + return _ENGINE_MANAGERS[key] + _release_engine_managers() + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, + KvCacheConfig, + ) + from tensorrt_llm.mapping import Mapping + + env_before = os.environ.get("TRTLLM_DSV41_KV_STORAGE") + os.environ["TRTLLM_DSV41_KV_STORAGE"] = storage + try: + mgr = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + max_gpu_total_bytes=quota, + host_cache_size=0, + event_buffer_max_size=0, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=40, + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=128, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + if pp_stage is None + else Mapping(world_size=4, rank=2 * pp_stage, tp_size=2, pp_size=2, gpus_per_node=4), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=max_num_tokens, + sparse_attn_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ), + pretrained_config=_v41_pretrained_config(), + spec_config=DSparkDecodingConfig(max_draft_len=max_draft_len, block_size=max_draft_len) + if max_draft_len + else None, + max_cuda_graph_batch_size=max_batch_size, + disable_overlap_scheduler=False, + ) + finally: + if env_before is None: + os.environ.pop("TRTLLM_DSV41_KV_STORAGE") + else: + os.environ["TRTLLM_DSV41_KV_STORAGE"] = env_before + if pp_stage is not None: + assert pp_stage == 0, "only PP2 stage 0 holds the compressor state (owners 2 / 8 / 14)" + assert list(mgr.pp_layers) == list(range(20)), list(mgr.pp_layers) + pool = mgr.get_buffers(2, T.COMPRESSOR_KV) + shapes = { + "swa": tuple(mgr.get_buffers(0, T.SWA).shape), + "compress_r2": sorted({mgr.get_buffers(o, T.COMPRESS).shape[0] for o in (2, 8, 14)}), + "compress_r1": tuple(mgr.get_buffers(20, T.COMPRESS).shape) if pp_stage is None else None, + "compressor_state": tuple(pool.shape), + "swa_dtype": str(mgr.get_buffers(0, T.SWA).dtype), + } + print( + f"[engine manager max_draft_len={max_draft_len} max_batch_size={max_batch_size} quota={quota} " + f"storage={storage} max_num_tokens={max_num_tokens} max_seq_len={max_seq_len} pp_stage={pp_stage} " + f"max_blocks_per_seq={mgr.max_blocks_per_seq}] {shapes}" + ) + assert (shapes["swa_dtype"] == "torch.uint8") == (storage == "native"), shapes + seq = max_seq_len - (14 if max_draft_len else 0) + if seq in _ENVELOPE_WIDTH_OF_L: + assert mgr.max_blocks_per_seq == _ENVELOPE_WIDTH_OF_L[seq], mgr.max_blocks_per_seq + if key == (0, 32, _ENGINE_QUOTA_BYTES, "bf16", 8192, _ENGINE_MAX_SEQ, None): + assert shapes["swa"][0] == _ENGINE_SWA_PAGES, shapes + assert set(shapes["compress_r2"]) == _ENGINE_COMPRESS_R2_PAGES, shapes + assert shapes["compress_r1"][0] == _ENGINE_COMPRESS_R1_PAGES, shapes + for owner in (2, 8, 14): + for role in (T.COMPRESSOR_KV, T.COMPRESSOR_SCORE): + b = mgr.get_buffers(owner, role) + assert ( + b.data_ptr() == pool.data_ptr() + and b.shape == pool.shape + and b.dtype == torch.float32 + ) + assert pool.is_contiguous() and tuple(pool.shape[1:]) == (128, HEAD_DIM) + _ENGINE_MANAGERS[key] = mgr + return mgr + + +def _release_engine_managers(): + for mgr in _ENGINE_MANAGERS.values(): + mgr.shutdown() + _ENGINE_MANAGERS.clear() + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +@pytest.fixture(scope="module", autouse=True) +def _engine_managers_released_at_module_end(): + yield + _release_engine_managers() + + +def _pool_range( + kv: torch.Tensor, score: torch.Tensor, w0: int, w1: int, lo: int = 0 +) -> torch.Tensor: + """model.py's ratio-2 pooling of windows [w0, w1) of per-position arrays that start at + position ``lo``, fp32 (zero ape).""" + p = torch.arange(w0 * RATIO, w1 * RATIO, device=DEV) - lo + k = kv[p].unflatten(0, (w1 - w0, RATIO)) + s = score[p].unflatten(0, (w1 - w0, RATIO)) + return (k * s.softmax(dim=1)).sum(dim=1) + + +def _hazard_pages(num_pages: int) -> List[int]: + """Pages on both sides of every 32-bit offset hazard of an fp32 [pages, 128, 512] + pool (65536 elements, 256 KiB per page): byte offsets 2^31 (page 8192) and 2^32 + (16384), element offsets 2^31 (32768), 2^32 (65536) and 2^33 (131072); those the + pool reaches.""" + return sorted( + c + for c in (8191, 8192, 16383, 16384, 32767, 32768, 65535, 65536, 131071, 131072) + if c < num_pages + ) + + +def _one_sequence_tables( + num_pages: int, ps: int, lo: int, hi: int, seed: int, pages=None, width=None +): + """kv / score tables of one sequence mapping positions [lo, hi) to distinct pages + taken (shuffled) from ``pages`` -- by default the pool's last pages -- and every + other column to a trap page; ``width`` columns (default: one past the last page). + Returns (bt_kv, bt_sc, trap, mapped page ids).""" + logical = list(range(lo // ps, (hi - 1) // ps + 1)) + need = 2 * len(logical) + if pages is None: + pages = list(range(num_pages - need, num_pages)) + pages = list(pages)[:need] + assert len(set(pages)) == need and max(pages) < num_pages + trap = next(t for t in range(num_pages - 1, -1, -1) if t not in set(pages)) + cpu = torch.Generator().manual_seed(seed) + far = [pages[i] for i in torch.randperm(need, generator=cpu).tolist()] + width = logical[-1] + 2 if width is None else width + assert width > logical[-1] + bt_kv = torch.full((1, width), trap, dtype=torch.int32) + bt_sc = torch.full_like(bt_kv, trap) + for q in logical: + bt_kv[0, q], bt_sc[0, q] = far.pop(), far.pop() + mapped = torch.tensor(sorted(set(pages) | {trap}), device=DEV) + return bt_kv.to(DEV), bt_sc.to(DEV), trap, mapped + + +# (start, table width) of the row sweep: from 0 and from an odd 129 on the gate engine's +# 132-column tables, and a chunk ending exactly at the envelope's longest kv length Lm = +# 1,048,590 (L 1,048,576 with DSpark), from an even and an odd start, on its 8196-column tables. +_ROW_SWEEP_N = _ENVELOPE_TOKENS +_ROW_SWEEP = [ + (0, 132), + (129, 132), + (_ENVELOPE_LM_1M - _ROW_SWEEP_N, 8196), + (_ENVELOPE_LM_1M - _ROW_SWEEP_N - 1, 8196), +] + + +@pytest.mark.parametrize("start,width", _ROW_SWEEP, ids=[f"start{s}-w{w}" for s, w in _ROW_SWEEP]) +def test_every_row_count_1_to_16384_on_the_engine_pool(start: int, width: int): + """A context chunk of r new tokens sizes the launch grid (r // 2 windows, + max_outputs = r // 2). On the engine's compressor-state pool, with the sequence's + kv and score pages the pool's LAST pages (shuffled, distinct) and a trap page below + them, every chunk length r = 1..16384 (max_num_tokens 16384) from ``start`` (0; 129, + odd, so the first window pairs the parked token 128 with the chunk's first token; and + 1,032,206 / 1,032,205, whose 16384-token chunk ends at the envelope's longest kv length + 1,048,590 / 1,048,589 on 8196-column tables) returns exactly the first rows of the + r = 16384 call bit for bit, writes nothing past them (NaN guard), and leaves the pool + holding exactly positions [start, start + r) plus the parked token (the calls restore + the same state, so they are idempotent). The 16384-token call matches the reference; + the whole pool outside the named slots is unchanged.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + + pool_t = _engine_manager(0).get_buffers(2, T.COMPRESSOR_KV) + num_pages, ps = pool_t.shape[0], pool_t.shape[1] + N = _ROW_SWEEP_N + lo = start - 1 if start % RATIO else start + gen = _gen(9100 + start % 1000) + # Per-position rows over [lo, start + N), indexed from lo. + kv, sc = _make_sequence(start + N - lo, gen) + # Far-end tables: the pool's last pages cover [lo, start + N) for kv and for score. + bt_kv, bt_sc, trap, mapped = _one_sequence_tables( + num_pages, ps, lo, start + N, seed=9100 + start % 1000, width=width + ) + assert bt_kv.shape[1] == width + assert int(mapped[0]) == trap + region = pool_t[trap:] + region.fill_(_SENTINEL) + outside_sum = pool_t[:trap].view(torch.int32).sum(dtype=torch.int64) + if start % RATIO: + p = start - 1 + pool_t[int(bt_kv[0, p // ps]), p % ps] = kv[p - lo] + pool_t[int(bt_sc[0, p // ps]), p % ps] = sc[p - lo] + parked = 1 if start % RATIO else 0 + + kv_score = torch.cat([kv[start - lo :], sc[start - lo :]], dim=1).contiguous() + ape = _zero_ape() + r_all = torch.arange(1, N + 1, device=DEV, dtype=torch.int32) + kv_lens_all = (start + r_all).view(N, 1).contiguous() + n_all = [(start + r) // RATIO - start // RATIO for r in range(1, N + 1)] + cu_seq_all = torch.stack([torch.zeros_like(r_all), r_all], 1).contiguous() + cu_out_all = torch.stack( + [torch.zeros_like(r_all), torch.tensor(n_all, device=DEV, dtype=torch.int32)], 1 + ) + cu_out_all = cu_out_all.contiguous() + start_pos = _i32_tensor([start]) + buf = torch.full((n_all[-1] + 1, HEAD_DIM), float("nan"), device=DEV) + + def call(r): + n = n_all[r - 1] + op( + kv_score[:r], + ape, + pool_t, + pool_t, + bt_kv, + bt_sc, + buf[:n], + kv_lens_all[r - 1], + start_pos, + cu_seq_all[r - 1], + cu_out_all[r - 1], + 1, + ps, + HEAD_DIM, + RATIO, + n, + ) + return n + + n = call(N) + torch.cuda.synchronize() + full = buf[:n].clone() + torch.testing.assert_close(full, _pool_range(kv, sc, start // RATIO, (start + N) // RATIO, lo)) + buf.fill_(float("nan")) + bad = torch.zeros((), dtype=torch.int64, device=DEV) + for r in range(1, N + 1): + n = call(r) + bad += (buf[:n].view(torch.int32) != full[:n].view(torch.int32)).sum() + bad += ((~torch.isnan(buf)).sum() - n * HEAD_DIM).abs() + buf[:n] = float("nan") + # State so far: kv + score slots of positions [start, start + N) (the first call) and the parked token. + bad += ((region != _SENTINEL).any(-1).sum() - 2 * (N + parked)).abs() + torch.cuda.synchronize() + assert int(bad) == 0, f"{int(bad)} elements / slots differ from the {N}-token call" + pos = torch.arange(lo, start + N, device=DEV) + assert torch.equal(pool_t[bt_kv[0].long()[pos // ps], pos % ps], kv[pos - lo]), "kv state" + assert torch.equal(pool_t[bt_sc[0].long()[pos // ps], pos % ps], sc[pos - lo]), "score state" + assert (region[0] == _SENTINEL).all(), "the trap page was written" + assert int(pool_t[:trap].view(torch.int32).sum(dtype=torch.int64)) == int(outside_sum), ( + "pool written below the region" + ) + print( + f"[row sweep start={start} width={width}] engine pool {tuple(pool_t.shape)}, pages {trap + 1}..{num_pages - 1}" + f", positions up to {start + N - 1}: r 1..{N} bitwise equal to the r = {N} rows" + ) + + +def _i32_tensor(values) -> torch.Tensor: + return torch.tensor(list(values), device=DEV, dtype=torch.int32) + + +# Mixed and context-only steps on the engine-size manager, handed to the op exactly +# as the target's _compress hands them. The dump's prefill calls (T, nc, +# max_ctx_compressed_tokens[2]) are (1018, 1, 508), (1935, 2, 960), (2048, 1, 1024), +# (48, 1, 24), (2, 1, 1), (2, 1, 0) and (1, 1, 0); the other steps add chunks with +# an odd cached count (the parked token read back from the manager's pool), a full +# 8192-token step of four contexts, and a DSpark (max_draft_len 5) step. +# scenario: ([(cached, new)] of the contexts, histories of the generation rows, +# max_draft_len, expected max_ctx_compressed_tokens[2]). +_ENGINE_PREFILL_STEPS = { + "dump_T1018_nc1": ([(0, 1016)], [2047, 16000], 0, 508), + "dump_T1935_nc2": ([(0, 1920), (0, 14)], [5003], 0, 960), + "dump_T2048_nc1": ([(0, 2048)], [], 0, 1024), + "dump_T48_nc1": ([(0, 48)], [], 0, 24), + "dump_T2_nc1_max1": ([(0, 2)], [], 0, 1), + "dump_T2_nc1_max0": ([(0, 1)], [7], 0, 0), + "dump_T1_nc1": ([(0, 1)], [], 0, 0), + "chunks_odd_cached": ([(1023, 1017), (4095, 2001), (1, 1)], [128, 129], 0, 1001), + "T8192_nc4": ([(0, 3000), (511, 2001), (1, 1), (4096, 3190)], [], 0, 1595), + "dspark_T1028_nc1": ([(0, 1016)], [2047, 16000], 5, 508), +} + + +@pytest.mark.parametrize("scenario", list(_ENGINE_PREFILL_STEPS)) +def test_engine_size_manager_target_call(scenario): + """The target's exact call on a step prepared by the engine-size manager: the step's + whole kv_score [T, 1024] (its generation rows NaN: they belong to the decode op), + ``sliding_block_tables[layer_offsets[owner], COMPRESSOR_KV / SCORE][:nc]``, + ``kv_lens_cuda_runtime[:nc]``, ``cached_token_lens_cuda[:nc]`` as start_pos, + ``cu_seq_lens_cuda[:nc + 1]``, ``cu_new_comp_kv_cuda[2][:nc + 1]``, + ``max_ctx_compressed_tokens[2]`` as max_outputs, and the zero-filled output of + ``num_total_compressed_tokens[2]`` rows, bf16 (the target's) and fp32, owners 2/8/14. + The whole pool is NaN except each odd-cached context's parked token, which a + previous chunk leaves. Checks: the context rows match the reference (fp32 default + tolerances; bf16 >= 0.999 bit-exact); the generation rows' reserved output rows stay + zero; the manager gives every new position a slot of its own (asserted: no shared + scratch page in these steps) and each holds its kv / score bit for bit; and no slot + outside the tables' pages for these positions changed.""" + ctx, gen_cached, max_draft_len, expect_max = _ENGINE_PREFILL_STEPS[scenario] + _target_call_on_manager( + _engine_manager(max_draft_len), scenario, ctx, gen_cached, max_draft_len, expect_max + ) + + +def _target_call_on_manager( + manager, + scenario, + ctx, + gen_cached, + max_draft_len, + expect_max, + max_num_tokens=8192, + allow_alias=False, +): + """One step on ``manager`` handed to the op exactly as the target hands it, with the + checks of test_engine_size_manager_target_call. ``allow_alias``: the manager may map + positions of a chunk behind its last two to a shared scratch slot; those slots are then + only counted as written (their content is whichever aliased position the op stored last), + while every slot of the last two positions of each chunk must be the position's own.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + DeepseekV4TrtllmAttentionMetadata, + ) + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + from tensorrt_llm._torch.metadata import KVCacheParams + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from tensorrt_llm.bindings import SamplingConfig + from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig + from tensorrt_llm.mapping import Mapping + + next_n = max_draft_len + 1 + sparse = DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ) + cached = [c for c, _ in ctx] + list(gen_cached) + new = [n for _, n in ctx] + [next_n] * len(gen_cached) + nc, B = len(ctx), len(cached) + requests = [] + try: + for i, (c, n) in enumerate(zip(cached, new)): + req = LlmRequest( + request_id=i, + max_new_tokens=64, + input_tokens=list(range(c + n)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + assert manager.prepare_context(req), f"manager could not hold request {i}" + assert manager.resize_context(req, req.context_chunk_size) + requests.append(req) + md = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor(new, dtype=torch.int), + request_ids=list(range(B)), + max_num_requests=manager.max_batch_size, + num_contexts=nc, + prompt_lens=[c + n for c, n in zip(cached, new)], + max_num_tokens=max_num_tokens, + kv_cache_manager=manager, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + sparse_attention_config=sparse, + ) + md.prepare() + torch.cuda.synchronize() + ctx_tokens = sum(n for _, n in ctx) + num_tokens = sum(new) + ctx_rows = sum((c + n) // RATIO - c // RATIO for c, n in ctx) + total = md.num_total_compressed_tokens[RATIO] + max_out = md.max_ctx_compressed_tokens[RATIO] + assert (md.num_contexts, md.num_seqs, md.num_ctx_tokens) == (nc, B, ctx_tokens) + assert max_out == expect_max, (scenario, max_out) + assert max_out == max((c + n) // RATIO - c // RATIO for c, n in ctx) + assert total == ctx_rows + (B - nc) * math.ceil(next_n / RATIO), (total, ctx_rows) + kv_lens = md.kv_lens_cuda_runtime[:nc] + start_pos = md.cached_token_lens_cuda[:nc] + cu_seq = md.cu_seq_lens_cuda[: nc + 1] + cu_out = md.cu_new_comp_kv_cuda[RATIO][: nc + 1] + assert kv_lens.tolist() == [c + n for c, n in ctx] and start_pos.tolist() == [ + c for c, _ in ctx + ] + assert cu_seq.tolist() == [0] + torch.tensor([n for _, n in ctx]).cumsum(0).tolist() + cu = cu_out.tolist() + assert cu[nc] == ctx_rows + + gen = _gen(9300 + len(scenario)) + truth = [] + kv_score = torch.full((num_tokens, 2 * HEAD_DIM), float("nan"), device=DEV) + row = 0 + for c, n in ctx: + lo = c - 1 if c % RATIO else c + kv, sc = _make_sequence(c + n - lo, gen) # positions [lo, c + n), indexed from lo + kv_score[row : row + n] = torch.cat([kv[c - lo :], sc[c - lo :]], 1) + truth.append((kv, sc, lo)) + row += n + ape = _zero_ape() + page = manager.tokens_per_block + pool = manager.get_buffers(2, T.COMPRESSOR_KV) + pages_seen = set() + aliased = 0 + for owner in (2, 8, 14): + p_kv = manager.get_buffers(owner, T.COMPRESSOR_KV) + p_sc = manager.get_buffers(owner, T.COMPRESSOR_SCORE) + lo_ = manager.layer_offsets[owner] + bt_kv = md.sliding_block_tables[lo_, T.COMPRESSOR_KV.value] + bt_sc = md.sliding_block_tables[lo_, T.COMPRESSOR_SCORE.value] + # Flat slot ids (page * 128 + offset) every new position of every context writes. + writes = [] + for b, (c, n) in enumerate(ctx): + pos = torch.arange(c, c + n, device=DEV) + for tbl in (bt_kv, bt_sc): + pg = tbl[b].long()[pos // page] + assert (pg >= 0).all() and (pg < pool.shape[0]).all(), ( + f"{scenario} owner {owner}: table entry" + ) + writes.append((b, tbl is bt_sc, pos, pg * page + pos % page)) + all_ids = torch.cat([w[3] for w in writes]) + uniq, counts = torch.unique(all_ids, return_counts=True) + shared = int((counts > 1).sum()) + aliased = max(aliased, shared) + # SWA scratch reuse could alias pages behind a window; the gate-size steps get none. + assert allow_alias or shared == 0, ( + f"{scenario} owner {owner}: slots shared between positions" + ) + for b, (c, n) in enumerate(ctx): + for last in range(max(c, c + n - 2), c + n): + for tbl in (bt_kv, bt_sc): + sid = int(tbl[b, last // page]) * page + last % page + assert int(counts[uniq == sid].sum()) == 1, ( + f"{scenario} owner {owner}: last position shared" + ) + pages_seen |= set((uniq // page).tolist()) + for out_dtype in (torch.bfloat16, torch.float32): + pool.fill_(float("nan")) + parked = [] + for b, (c, _) in enumerate(ctx): + if c % RATIO: + p = c - 1 + kv, sc, lo = truth[b] + p_kv[bt_kv[b, p // page].long(), p % page] = kv[p - lo] + p_sc[bt_sc[b, p // page].long(), p % page] = sc[p - lo] + parked += [ + int(bt_kv[b, p // page]) * page + p % page, + int(bt_sc[b, p // page]) * page + p % page, + ] + pooled = torch.zeros(total + 1, HEAD_DIM, device=DEV, dtype=out_dtype) + op( + kv_score, + ape, + p_kv, + p_sc, + bt_kv[:nc], + bt_sc[:nc], + pooled[:total], + kv_lens, + start_pos, + cu_seq, + cu_out, + nc, + page, + HEAD_DIM, + RATIO, + max_out, + ) + torch.cuda.synchronize() + got = torch.cat([pooled[cu[b] : cu[b + 1]].float() for b in range(nc)]) + ref = torch.cat( + [ + _pool_range(kv, sc, c // RATIO, (c + n) // RATIO, lo) + for (kv, sc, lo), (c, n) in zip(truth, ctx) + ] + ) + assert got.shape == ref.shape + if out_dtype == torch.float32: + torch.testing.assert_close( + got, ref, msg=lambda m: f"{scenario} owner {owner}: {m}" + ) + elif got.numel(): + torch.testing.assert_close( + got.bfloat16(), + ref.bfloat16(), + msg=lambda m: f"{scenario} owner {owner}: {m}", + ) + exact = (got.bfloat16() == ref.bfloat16()).float().mean().item() + # Same gate as test_bf16_output. + assert exact >= 0.999, ( + f"{scenario} owner {owner}: bf16 bit-exact fraction {exact}" + ) + assert (pooled[cu[nc] :] == 0).all(), ( + f"{scenario} owner {owner}: rows past the contexts' outputs" + ) + flat = pool.view(-1, HEAD_DIM) + for b, is_score, pos, ids in writes: + kv, sc, lo = truth[b] + single = counts[torch.searchsorted(uniq, ids)] == 1 + want = (sc if is_score else kv)[pos[single] - lo] + assert torch.equal(flat[ids[single]], want), ( + f"{scenario} owner {owner} b {b}: state" + ) + written = int((~torch.isnan(flat)).any(-1).sum()) + expect = int( + torch.unique( + torch.cat([all_ids, torch.tensor(parked, device=DEV, dtype=torch.int64)]) + ).numel() + ) + assert written == expect, ( + f"{scenario} owner {owner}: {written} slots written, {expect} named" + ) + print( + f"[engine target call {scenario}] nc={nc} T={num_tokens} max_outputs={max_out} ctx_rows={ctx_rows} " + f"total={total} pool={tuple(pool.shape)} table_width={bt_kv.shape[-1]} " + f"max_kv_len={max(c + n for c, n in ctx)} distinct pages written={len(pages_seen)} " + f"highest page={max(pages_seen)} slots shared between positions={aliased}" + ) + finally: + for req in requests: + manager.free_resources(req) + + +# Steps on managers of the engine envelope: bf16 storage at 0.5 x total memory +# (_ENVELOPE_QUOTA_BYTES), max_num_tokens 16384, max_batch_size 256, the manager's length Lm (L + 14 +# with DSpark), so the metadata's tables are 260 / 8196 columns wide. The requests are allocated +# at their full length in one resize_context, so their lengths stay short: a 1,048,576-token +# request allocated that way does not fit (resize_context returned False), and at L 1,048,576 +# with max_batch_size 256 the compressor-state pool has 1662 pages. Contexts near position 1M on +# 8196-column tables are the synthetic tests' (the row sweep, the envelope steps). scenario: +# (contexts [(cached, new)], generation histories, max_draft_len, L, expected +# max_ctx_compressed_tokens[2]). The batch-256 engines' dump step (T 16340, nc 39, max 431) at L +# 32768 with DSpark: 39 contexts (one of 862 tokens at the odd cached 1 owning the 431 rows) ahead +# of 217 verify rows, 256 requests; 256 contexts of 64 tokens at odd / even prefixes (a +# 16384-token step of 256 requests); at L 1,048,576 without DSpark the dump's (16384, 1, 8192); +# and with DSpark two contexts (9000 tokens, and 3001 at an odd prefix) ahead of two verify rows. +_ENVELOPE_QUOTA_BYTES = 148_510_474_240 +_ENVELOPE_RNG = torch.Generator().manual_seed(2027) +_ENVELOPE_CTX39 = [(1, 862)] + [ + ((int(c) // 2) * 2 + i % 2, n) + for i, (c, n) in enumerate( + zip( + torch.randint(0, 400, (38,), generator=_ENVELOPE_RNG).tolist(), + [373 + (7 * i) % 13 - 6 + 2 * (i == 0) for i in range(38)], + ) + ) +] +_ENVELOPE_CTX256 = [ + ((int(c) // 2) * 2 + i % 2, 64) + for i, c in enumerate(torch.randint(0, 120, (256,), generator=_ENVELOPE_RNG).tolist()) +] +_ENVELOPE_PREFILL_STEPS = { + "gc_L32k_dump_T16340_nc39": ( + _ENVELOPE_CTX39, + torch.randint(1, 300, (256 - 39,), generator=_ENVELOPE_RNG).tolist(), + 5, + 32768, + 431, + ), + "gc_L32k_nc256_T16384": (_ENVELOPE_CTX256, [], 5, 32768, 32), + "gc_L1m_nods_T16384_nc1": ([(0, 16384)], [], 0, 1_048_576, 8192), + "gc_L1m_T12013_nc2": ([(0, 9000), (1, 3001)], [100, 201], 5, 1_048_576, 4500), +} + + +@pytest.mark.parametrize("scenario", list(_ENVELOPE_PREFILL_STEPS)) +def test_envelope_manager_target_call(scenario): + """The target's exact call (as in test_engine_size_manager_target_call) on steps prepared by + managers of the engine envelope (max_batch_size 256, max_num_tokens 16384, max_seq_len + 32768 + 14 and 1,048,576 (+ 14); the metadata's tables 260 / 8196 columns wide): up to 256 + contexts, T up to 16384, max_ctx_compressed_tokens up to 8192. The manager may map chunk + positions behind the last two to a shared scratch slot (counted and printed); every other + check is the gate-size test's.""" + ctx, gen_cached, max_draft_len, L, expect_max = _ENVELOPE_PREFILL_STEPS[scenario] + lm = _envelope_lm(L, max_draft_len > 0) + next_n = max_draft_len + 1 + assert sum(n for _, n in ctx) + len(gen_cached) * next_n <= _ENVELOPE_TOKENS, scenario + assert len(ctx) + len(gen_cached) <= _ENVELOPE_REQS and all(c + n <= lm for c, n in ctx) + mgr = _engine_manager( + max_draft_len, + _ENVELOPE_REQS, + _ENVELOPE_QUOTA_BYTES, + "bf16", + max_num_tokens=_ENVELOPE_TOKENS, + max_seq_len=lm, + ) + assert mgr.max_blocks_per_seq == _ENVELOPE_WIDTH_OF_L[L] + _target_call_on_manager( + mgr, + scenario, + ctx, + gen_cached, + max_draft_len, + expect_max, + max_num_tokens=_ENVELOPE_TOKENS, + allow_alias=True, + ) + + +# The four engine-size pools: DSpark off / on (max_draft_len 5) at the accuracy gate's +# max_batch_size 32 and at the boot gate's 8, one quota. Their compressor-state pools +# are [22908 | 41088 | 21246 | 39678, 128, 512] fp32; the two at batch 8 pass element +# offset 2^31. Then the engine envelope's largest unsplit pools at 0.5 x total memory: the +# bf16 storage's [120192, 128, 512] (max_num_tokens 16384, batch 1, DSpark) and the native +# [172416, 128, 512] (max_num_tokens 16384, batch 2, DSpark off), past element offset 2^33; and +# the PP2 stage-0 managers at the same tuples: bf16 [188796, 128, 512] and B_max, native +# [247806, 128, 512]. (max_draft_len, max_batch_size, quota, storage, max_num_tokens, pp_stage, +# pages) +_ENGINE_POOLS = [ + (0, 32, _ENGINE_QUOTA_BYTES, "bf16", 8192, None, 22908), + (0, 8, _ENGINE_QUOTA_BYTES, "bf16", 8192, None, 41088), + (5, 32, _ENGINE_QUOTA_BYTES, "bf16", 8192, None, 21246), + (5, 8, _ENGINE_QUOTA_BYTES, "bf16", 8192, None, 39678), + (5, 1, _ENVELOPE_QUOTA_BYTES, "bf16", 16384, None, 120192), + (0, 2, _ENVELOPE_QUOTA_BYTES, "native", 16384, None, 172416), + (5, 1, _ENVELOPE_QUOTA_BYTES, "bf16", 16384, 0, 188796), + (0, 2, _ENVELOPE_QUOTA_BYTES, "native", 16384, 0, 247806), +] + + +@pytest.mark.parametrize( + "max_draft_len,max_batch_size,quota,storage,max_num_tokens,pp_stage,want_pages", + _ENGINE_POOLS, + ids=[ + f"D{d}-bs{b}-{st}-mnt{t}" + (f"-pp2stage{pp}" if pp is not None else "") + f"-p{p}" + for d, b, q, st, t, pp, p in _ENGINE_POOLS + ], +) +def test_offset_hazard_pages_of_the_engine_pools( + max_draft_len: int, + max_batch_size: int, + quota: int, + storage: str, + max_num_tokens: int, + pp_stage, + want_pages: int, +): + """At the far end, run: a 16384-token chunk from start 129 (a parked token + read back) whose kv / score pages are the pages on both sides of every 32-bit offset + hazard the pool reaches (_hazard_pages) and otherwise its last pages. Outputs match + the reference, the pool holds exactly the chunk's state and the parked token, and no + other byte of the pool changed. Control: the same call with every page id at or past + the pool's highest hazard wrapped below it (as a 32-bit offset would wrap) fails that + check.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + + pool_t = _engine_manager( + max_draft_len, + max_batch_size, + quota, + storage, + max_num_tokens=max_num_tokens, + pp_stage=pp_stage, + ).get_buffers(2, T.COMPRESSOR_KV) + num_pages, ps = pool_t.shape[0], pool_t.shape[1] + assert num_pages == want_pages, pool_t.shape + start, N = 129, _ENVELOPE_TOKENS + lo = start - 1 + need = 2 * ((start + N - 1) // ps - lo // ps + 1) + hz = _hazard_pages(num_pages) + pages = hz + [p for p in range(num_pages - need, num_pages) if p not in hz][: need - len(hz)] + gen = _gen(9500 + num_pages) + kv = torch.zeros(start + N, HEAD_DIM, device=DEV) + sc = torch.zeros_like(kv) + kv[lo:], sc[lo:] = _make_sequence(start + N - lo, gen) + kv_score = torch.cat([kv[start:], sc[start:]], dim=1).contiguous() + n = (start + N) // RATIO - start // RATIO + ref = _pool_range(kv, sc, start // RATIO, (start + N) // RATIO) + pos = torch.arange(lo, start + N, device=DEV) + + def run(wrap: int): + bt_kv, bt_sc, trap, mapped = _one_sequence_tables( + num_pages, ps, lo, start + N, seed=9500, pages=pages + ) + pool_t[mapped] = _SENTINEL + outside = _pool_sum(pool_t) - int(pool_t[mapped].view(torch.int32).sum(dtype=torch.int64)) + pool_t[int(bt_kv[0, lo // ps]), lo % ps] = kv[lo] + pool_t[int(bt_sc[0, lo // ps]), lo % ps] = sc[lo] + op_kv = torch.where(bt_kv >= wrap, bt_kv % wrap, bt_kv) if wrap else bt_kv + op_sc = torch.where(bt_sc >= wrap, bt_sc % wrap, bt_sc) if wrap else bt_sc + out = torch.full((n + 1, HEAD_DIM), float("nan"), device=DEV) + op( + kv_score, + _zero_ape(), + pool_t, + pool_t, + op_kv, + op_sc, + out[:n], + _i32_tensor([start + N]), + _i32_tensor([start]), + _i32_tensor([0, N]), + _i32_tensor([0, n]), + 1, + ps, + HEAD_DIM, + RATIO, + n, + ) + torch.cuda.synchronize() + assert torch.isnan(out[n]).all(), "op wrote past the last output row" + + def check(): + torch.testing.assert_close(out[:n], ref) + assert torch.equal(pool_t[bt_kv[0].long()[pos // ps], pos % ps], kv[pos]), "kv state" + assert torch.equal(pool_t[bt_sc[0].long()[pos // ps], pos % ps], sc[pos]), "score state" + assert int((pool_t[mapped] != _SENTINEL).any(-1).sum()) == 2 * pos.numel(), ( + "slot outside the positions" + ) + assert (pool_t[trap] == _SENTINEL).all(), "the trap page was written" + after = _pool_sum(pool_t) - int(pool_t[mapped].view(torch.int32).sum(dtype=torch.int64)) + assert after == outside, "pool written outside the mapped pages" + + return check + + check = run(0) + check() + wrap = max(hz) + print( + f"[hazard pages D={max_draft_len} bs={max_batch_size} {storage} pp_stage={pp_stage}] pool " + f"{tuple(pool_t.shape)} " + f"({pool_t.numel() / 2**31:.2f} x 2^31 elements), hazard pages {hz}: outputs and state exact" + ) + check = run(wrap) + with pytest.raises(AssertionError): + check() + print(f"[hazard pages control] page ids >= {wrap} wrapped mod {wrap}: the check fails") + + +# --------------------------------------------------------------------------- +# The engine configuration's envelope -- multi-request steps, the +# kernel-variant map, and pool-size invariance +# --------------------------------------------------------------------------- +# +# The engine envelope bounds a step at max_num_tokens 16384 tokens and max_batch_size +# 256 requests, each at most the manager's length Lm (max_seq_len L up to 1,048,576, + 14 with +# DSpark). A prefill step hands the op num_contexts = k in 1..256 context chunks packed first +# in kv_score, then the step's generation rows (never read); max_outputs is the target's +# md.max_ctx_compressed_tokens[2] = max_b(kv_len_b // 2 - sp_b // 2), up to 8192; the tables +# are the manager's max_blocks_per_seq (132 .. 8196) columns wide. The row axis of one chunk +# (1..16384 new tokens) is test_every_row_count_1_to_16384_on_the_engine_pool. + +_ENV_TOKENS = _ENVELOPE_TOKENS +_ENV_REQS = _ENVELOPE_REQS +_ENV_PAGES = ( + 4096 # plainly more than a step's pages; pool size itself is test_pool_size_invariance's axis +) +_ENV_SENTINEL = 7777.0 + +# Every (T, num_contexts, max_ctx_compressed_tokens[2]) the core target's runtime-value +# dump recorded for this op: the boot gates' logs and the TEP4 / DEP4 / TP4 accuracy-gate +# previews' logs, parsed line by line (a line is kept only if it is exactly one tuple of +# arity 4). Five garbles were dropped: nc 18 and 4096, two with max above what T allows, and +# (1767, 2, 3) -- once in the TP4 preview's log, where the other ranks print (1767, 2, 465), +# and impossible for two contexts with at most 30 one-token generation rows. T counts the +# step's generation rows too. +_DUMP_PREFILL_CALLS = [ + (1, 1, 0), + (2, 1, 0), + (2, 1, 1), + (14, 1, 7), + (16, 1, 8), + (48, 1, 24), + (524, 1, 246), + (525, 1, 247), + (534, 1, 252), + (560, 1, 265), + (580, 1, 274), + (580, 1, 275), + (581, 1, 275), + (586, 1, 277), + (587, 1, 278), + (593, 1, 281), + (600, 1, 285), + (602, 1, 286), + (604, 1, 286), + (609, 1, 289), + (610, 1, 289), + (612, 1, 290), + (617, 1, 293), + (621, 1, 295), + (624, 1, 296), + (625, 1, 297), + (626, 1, 297), + (628, 1, 298), + (629, 1, 299), + (630, 1, 299), + (632, 1, 301), + (634, 1, 301), + (637, 1, 303), + (640, 1, 304), + (643, 1, 306), + (645, 1, 307), + (647, 1, 308), + (648, 1, 308), + (650, 1, 309), + (654, 1, 311), + (656, 1, 312), + (657, 1, 313), + (658, 1, 313), + (660, 1, 314), + (660, 1, 315), + (661, 1, 315), + (662, 1, 315), + (663, 1, 316), + (665, 1, 317), + (666, 1, 317), + (668, 1, 318), + (671, 1, 320), + (672, 1, 320), + (673, 1, 321), + (674, 1, 321), + (675, 1, 322), + (677, 1, 323), + (680, 1, 324), + (681, 1, 325), + (682, 1, 325), + (683, 1, 326), + (684, 1, 326), + (685, 1, 327), + (687, 1, 328), + (691, 1, 330), + (692, 1, 330), + (693, 1, 331), + (694, 1, 331), + (696, 1, 332), + (697, 1, 333), + (698, 1, 333), + (699, 1, 334), + (701, 1, 335), + (703, 1, 336), + (705, 1, 337), + (706, 1, 337), + (706, 1, 338), + (707, 1, 338), + (710, 1, 339), + (711, 1, 340), + (712, 1, 340), + (713, 1, 341), + (714, 1, 341), + (715, 1, 342), + (716, 1, 342), + (718, 1, 343), + (719, 1, 344), + (720, 1, 344), + (721, 1, 345), + (722, 1, 345), + (723, 1, 346), + (724, 1, 347), + (725, 1, 347), + (727, 1, 348), + (728, 1, 348), + (729, 1, 349), + (730, 1, 349), + (731, 1, 350), + (732, 1, 350), + (733, 1, 351), + (734, 1, 351), + (735, 1, 352), + (737, 1, 353), + (739, 1, 354), + (741, 1, 355), + (742, 1, 355), + (743, 1, 356), + (744, 1, 356), + (745, 1, 357), + (746, 1, 357), + (749, 1, 359), + (750, 1, 359), + (751, 1, 360), + (752, 1, 360), + (753, 1, 361), + (754, 1, 361), + (755, 1, 362), + (756, 1, 362), + (757, 1, 363), + (758, 1, 363), + (759, 1, 364), + (760, 1, 364), + (761, 1, 365), + (763, 1, 366), + (764, 1, 366), + (764, 1, 367), + (765, 1, 367), + (766, 1, 367), + (767, 1, 368), + (768, 1, 368), + (768, 1, 369), + (769, 1, 369), + (770, 1, 369), + (771, 1, 370), + (772, 1, 370), + (773, 1, 371), + (774, 1, 371), + (776, 1, 372), + (776, 1, 373), + (777, 1, 373), + (778, 1, 373), + (779, 1, 374), + (780, 1, 374), + (782, 1, 375), + (783, 1, 376), + (784, 1, 376), + (785, 1, 377), + (786, 1, 377), + (787, 1, 378), + (788, 1, 378), + (789, 1, 379), + (790, 1, 379), + (791, 1, 380), + (792, 1, 380), + (793, 1, 381), + (794, 1, 381), + (795, 1, 382), + (796, 1, 382), + (797, 1, 383), + (798, 1, 383), + (799, 1, 384), + (800, 1, 384), + (801, 1, 385), + (802, 1, 385), + (803, 1, 386), + (804, 1, 386), + (805, 1, 387), + (806, 1, 387), + (807, 1, 388), + (808, 1, 388), + (809, 1, 389), + (810, 1, 389), + (811, 1, 390), + (812, 1, 390), + (813, 1, 391), + (814, 1, 391), + (815, 1, 392), + (816, 1, 392), + (817, 1, 393), + (818, 1, 393), + (819, 1, 394), + (821, 1, 395), + (822, 1, 395), + (823, 1, 396), + (824, 1, 396), + (825, 1, 397), + (826, 1, 397), + (827, 1, 398), + (828, 1, 398), + (829, 1, 399), + (830, 1, 399), + (831, 1, 400), + (832, 1, 400), + (833, 1, 401), + (834, 1, 401), + (835, 1, 402), + (836, 1, 402), + (837, 1, 403), + (838, 1, 403), + (839, 1, 404), + (840, 1, 404), + (841, 1, 405), + (842, 1, 405), + (843, 1, 406), + (844, 1, 406), + (845, 1, 407), + (846, 1, 407), + (847, 1, 408), + (848, 1, 408), + (849, 1, 409), + (850, 1, 409), + (851, 1, 410), + (853, 1, 411), + (854, 1, 411), + (855, 1, 412), + (856, 1, 412), + (856, 1, 413), + (857, 1, 413), + (858, 1, 413), + (859, 1, 414), + (860, 1, 414), + (861, 1, 415), + (862, 1, 415), + (863, 1, 416), + (864, 1, 416), + (865, 1, 417), + (866, 1, 417), + (867, 1, 418), + (868, 1, 418), + (869, 1, 419), + (870, 1, 419), + (871, 1, 420), + (872, 1, 420), + (874, 1, 421), + (875, 1, 422), + (876, 1, 422), + (877, 1, 423), + (878, 1, 423), + (879, 1, 424), + (880, 1, 424), + (881, 1, 425), + (882, 1, 425), + (883, 1, 426), + (884, 1, 426), + (885, 1, 427), + (886, 1, 427), + (887, 1, 428), + (888, 1, 428), + (888, 1, 444), + (889, 1, 429), + (890, 1, 429), + (890, 1, 445), + (891, 1, 430), + (892, 1, 430), + (893, 1, 431), + (894, 1, 431), + (895, 1, 432), + (897, 1, 433), + (898, 1, 433), + (899, 1, 434), + (900, 1, 434), + (901, 1, 435), + (902, 1, 435), + (903, 1, 436), + (904, 1, 436), + (905, 1, 437), + (906, 1, 437), + (907, 1, 438), + (908, 1, 438), + (909, 1, 439), + (910, 1, 439), + (911, 1, 440), + (912, 1, 440), + (913, 1, 441), + (914, 1, 441), + (915, 1, 442), + (916, 1, 442), + (917, 1, 443), + (918, 1, 443), + (919, 1, 444), + (920, 1, 444), + (921, 1, 445), + (922, 1, 445), + (923, 1, 446), + (924, 1, 446), + (925, 1, 447), + (926, 1, 447), + (927, 1, 448), + (928, 1, 448), + (929, 1, 449), + (930, 1, 449), + (931, 1, 450), + (932, 1, 450), + (933, 1, 451), + (934, 1, 451), + (935, 1, 452), + (937, 1, 453), + (938, 1, 453), + (939, 1, 454), + (940, 1, 454), + (941, 1, 455), + (942, 1, 455), + (943, 1, 456), + (944, 1, 456), + (945, 1, 457), + (946, 1, 457), + (946, 1, 458), + (947, 1, 458), + (948, 1, 458), + (949, 1, 459), + (950, 1, 459), + (951, 1, 460), + (952, 1, 460), + (952, 1, 461), + (953, 1, 461), + (954, 1, 461), + (955, 1, 462), + (956, 1, 462), + (957, 1, 463), + (958, 1, 463), + (959, 1, 464), + (960, 1, 464), + (961, 1, 465), + (962, 1, 465), + (963, 1, 466), + (964, 1, 466), + (965, 1, 467), + (966, 1, 467), + (967, 1, 468), + (968, 1, 468), + (969, 1, 469), + (970, 1, 469), + (971, 1, 470), + (972, 1, 470), + (973, 1, 471), + (974, 1, 471), + (975, 1, 472), + (976, 1, 472), + (977, 1, 473), + (978, 1, 473), + (980, 1, 474), + (981, 1, 475), + (982, 1, 475), + (983, 1, 476), + (984, 1, 476), + (985, 1, 477), + (986, 1, 477), + (987, 1, 478), + (988, 1, 478), + (989, 1, 479), + (990, 1, 479), + (991, 1, 480), + (992, 1, 480), + (993, 1, 481), + (994, 1, 481), + (995, 1, 482), + (996, 1, 482), + (997, 1, 483), + (998, 1, 483), + (999, 1, 484), + (1000, 1, 484), + (1001, 1, 485), + (1003, 1, 486), + (1004, 1, 486), + (1004, 1, 487), + (1005, 1, 487), + (1006, 1, 487), + (1007, 1, 488), + (1008, 1, 488), + (1010, 1, 489), + (1011, 1, 490), + (1013, 1, 491), + (1014, 1, 491), + (1015, 1, 492), + (1017, 1, 493), + (1018, 1, 493), + (1018, 1, 508), + (1019, 1, 494), + (1020, 1, 494), + (1021, 1, 495), + (1022, 1, 495), + (1023, 1, 496), + (1024, 1, 496), + (1025, 1, 497), + (1027, 1, 498), + (1028, 1, 498), + (1029, 1, 499), + (1030, 1, 499), + (1031, 1, 500), + (1033, 1, 501), + (1034, 1, 501), + (1035, 1, 502), + (1036, 1, 502), + (1038, 1, 503), + (1039, 1, 504), + (1042, 1, 505), + (1043, 1, 506), + (1044, 1, 506), + (1045, 1, 507), + (1046, 1, 507), + (1047, 1, 508), + (1049, 1, 509), + (1050, 1, 509), + (1051, 1, 510), + (1052, 1, 510), + (1053, 1, 511), + (1054, 1, 511), + (1056, 1, 512), + (1057, 1, 513), + (1059, 1, 514), + (1062, 1, 515), + (1063, 1, 516), + (1065, 1, 517), + (1067, 1, 518), + (1069, 1, 519), + (1070, 1, 519), + (1071, 1, 520), + (1073, 1, 521), + (1074, 1, 521), + (1076, 1, 522), + (1077, 1, 523), + (1079, 1, 524), + (1081, 1, 525), + (1086, 1, 527), + (1088, 1, 528), + (1089, 1, 529), + (1093, 1, 531), + (1095, 1, 532), + (1096, 1, 532), + (1097, 1, 533), + (1098, 1, 533), + (1099, 1, 534), + (1101, 1, 535), + (1103, 1, 536), + (1107, 1, 538), + (1110, 1, 539), + (1111, 1, 540), + (1112, 1, 540), + (1113, 1, 541), + (1115, 1, 542), + (1116, 1, 542), + (1117, 1, 543), + (1118, 1, 543), + (1120, 1, 544), + (1121, 1, 545), + (1122, 1, 545), + (1123, 1, 546), + (1124, 1, 546), + (1125, 1, 547), + (1126, 1, 547), + (1129, 1, 549), + (1130, 1, 549), + (1131, 1, 550), + (1132, 1, 550), + (1135, 1, 552), + (1137, 1, 553), + (1140, 1, 554), + (1142, 1, 555), + (1146, 1, 557), + (1149, 1, 559), + (1150, 1, 559), + (1151, 1, 560), + (1153, 1, 561), + (1156, 1, 562), + (1157, 1, 563), + (1159, 1, 564), + (1160, 1, 564), + (1161, 1, 565), + (1166, 1, 567), + (1167, 1, 568), + (1168, 1, 568), + (1170, 1, 569), + (1171, 1, 570), + (1172, 1, 570), + (1173, 1, 571), + (1174, 1, 571), + (1176, 1, 572), + (1177, 1, 573), + (1180, 1, 574), + (1181, 1, 575), + (1185, 1, 577), + (1187, 1, 593), + (1188, 1, 578), + (1191, 1, 580), + (1201, 1, 585), + (1202, 1, 585), + (1203, 1, 586), + (1208, 1, 588), + (1209, 1, 589), + (1213, 1, 591), + (1216, 1, 592), + (1229, 1, 599), + (1231, 1, 600), + (1232, 1, 600), + (1235, 1, 602), + (1242, 1, 605), + (1248, 1, 608), + (1257, 1, 613), + (1267, 1, 618), + (1268, 1, 618), + (1271, 2, 314), + (1292, 1, 630), + (1304, 1, 636), + (1305, 2, 348), + (1313, 2, 390), + (1329, 1, 649), + (1332, 1, 650), + (1343, 2, 382), + (1363, 2, 335), + (1374, 2, 349), + (1375, 2, 373), + (1385, 2, 349), + (1404, 2, 384), + (1418, 2, 378), + (1421, 2, 365), + (1425, 2, 369), + (1427, 2, 405), + (1433, 2, 370), + (1437, 2, 429), + (1438, 2, 369), + (1440, 2, 427), + (1449, 2, 361), + (1453, 2, 367), + (1454, 2, 365), + (1462, 2, 469), + (1463, 2, 393), + (1469, 2, 364), + (1478, 2, 378), + (1492, 2, 385), + (1492, 2, 387), + (1494, 2, 376), + (1494, 2, 391), + (1496, 2, 412), + (1497, 2, 438), + (1502, 2, 371), + (1503, 2, 388), + (1507, 2, 396), + (1511, 2, 388), + (1512, 2, 394), + (1514, 2, 417), + (1524, 2, 447), + (1526, 2, 427), + (1527, 2, 409), + (1528, 2, 389), + (1532, 1, 750), + (1535, 2, 441), + (1537, 2, 419), + (1538, 2, 423), + (1540, 2, 410), + (1541, 2, 390), + (1542, 2, 392), + (1543, 2, 394), + (1544, 2, 380), + (1544, 2, 422), + (1548, 2, 388), + (1549, 2, 397), + (1550, 2, 393), + (1553, 2, 416), + (1554, 2, 412), + (1555, 2, 415), + (1555, 2, 476), + (1558, 2, 385), + (1563, 2, 390), + (1565, 2, 424), + (1567, 2, 419), + (1568, 2, 387), + (1569, 2, 429), + (1572, 2, 399), + (1573, 2, 427), + (1576, 2, 431), + (1576, 2, 453), + (1577, 2, 392), + (1577, 2, 459), + (1580, 2, 407), + (1581, 2, 404), + (1581, 2, 447), + (1584, 2, 399), + (1584, 2, 425), + (1591, 2, 400), + (1592, 2, 414), + (1592, 2, 420), + (1596, 2, 392), + (1596, 2, 427), + (1597, 2, 394), + (1597, 2, 405), + (1597, 2, 455), + (1598, 2, 407), + (1599, 2, 444), + (1600, 2, 425), + (1601, 2, 415), + (1604, 2, 398), + (1604, 2, 405), + (1604, 2, 417), + (1605, 2, 424), + (1609, 2, 398), + (1610, 2, 427), + (1613, 2, 517), + (1614, 2, 399), + (1614, 2, 413), + (1615, 2, 403), + (1615, 2, 407), + (1621, 2, 402), + (1623, 2, 431), + (1624, 2, 400), + (1624, 2, 437), + (1624, 2, 459), + (1625, 2, 424), + (1625, 2, 435), + (1626, 2, 430), + (1631, 2, 424), + (1631, 2, 454), + (1635, 2, 412), + (1635, 2, 421), + (1636, 2, 418), + (1636, 2, 423), + (1636, 2, 457), + (1638, 2, 436), + (1640, 2, 426), + (1644, 2, 410), + (1645, 2, 429), + (1646, 2, 470), + (1650, 2, 439), + (1650, 2, 440), + (1651, 2, 483), + (1654, 2, 433), + (1655, 2, 434), + (1656, 2, 460), + (1657, 2, 417), + (1657, 2, 465), + (1659, 2, 412), + (1660, 2, 444), + (1666, 2, 441), + (1667, 2, 411), + (1667, 2, 444), + (1667, 2, 456), + (1669, 2, 477), + (1670, 2, 457), + (1671, 2, 434), + (1672, 2, 502), + (1674, 2, 434), + (1676, 2, 436), + (1677, 2, 443), + (1679, 2, 435), + (1684, 2, 462), + (1685, 2, 482), + (1686, 2, 422), + (1691, 2, 471), + (1691, 2, 532), + (1692, 2, 416), + (1692, 2, 471), + (1693, 2, 429), + (1693, 2, 433), + (1693, 2, 435), + (1693, 2, 454), + (1694, 2, 451), + (1694, 2, 481), + (1695, 2, 435), + (1696, 2, 471), + (1697, 2, 455), + (1698, 2, 424), + (1698, 2, 518), + (1699, 2, 448), + (1700, 2, 450), + (1702, 2, 420), + (1702, 2, 427), + (1702, 2, 444), + (1705, 2, 428), + (1705, 2, 435), + (1705, 2, 463), + (1706, 2, 440), + (1708, 2, 452), + (1710, 2, 454), + (1710, 2, 469), + (1712, 2, 465), + (1713, 2, 462), + (1714, 2, 458), + (1715, 2, 455), + (1715, 2, 458), + (1715, 2, 459), + (1716, 2, 462), + (1717, 2, 453), + (1717, 2, 461), + (1718, 2, 425), + (1719, 2, 482), + (1721, 2, 453), + (1723, 2, 498), + (1724, 2, 424), + (1725, 2, 438), + (1725, 2, 554), + (1726, 2, 504), + (1729, 2, 509), + (1730, 2, 432), + (1731, 2, 441), + (1732, 2, 501), + (1733, 2, 459), + (1734, 2, 439), + (1735, 2, 428), + (1735, 2, 438), + (1736, 2, 479), + (1737, 2, 428), + (1739, 2, 428), + (1739, 2, 493), + (1740, 2, 467), + (1740, 2, 564), + (1742, 2, 438), + (1742, 2, 455), + (1743, 2, 453), + (1744, 2, 429), + (1746, 2, 438), + (1746, 2, 456), + (1749, 2, 503), + (1750, 2, 459), + (1750, 2, 460), + (1753, 2, 495), + (1754, 2, 450), + (1755, 2, 462), + (1755, 2, 486), + (1756, 2, 459), + (1758, 2, 435), + (1758, 2, 447), + (1761, 2, 438), + (1761, 2, 451), + (1764, 2, 464), + (1764, 2, 475), + (1765, 2, 499), + (1767, 2, 465), + (1768, 2, 477), + (1770, 2, 501), + (1771, 2, 453), + (1772, 2, 446), + (1772, 2, 475), + (1773, 2, 458), + (1773, 2, 482), + (1774, 2, 502), + (1777, 2, 445), + (1777, 2, 483), + (1777, 2, 495), + (1777, 2, 545), + (1778, 2, 454), + (1779, 2, 471), + (1780, 2, 459), + (1780, 2, 460), + (1780, 2, 463), + (1782, 2, 543), + (1782, 2, 550), + (1783, 2, 496), + (1784, 2, 442), + (1784, 2, 457), + (1785, 2, 444), + (1786, 2, 449), + (1786, 2, 475), + (1787, 2, 493), + (1788, 2, 461), + (1788, 2, 549), + (1789, 2, 511), + (1790, 2, 494), + (1794, 2, 476), + (1796, 2, 467), + (1799, 2, 446), + (1799, 2, 472), + (1800, 2, 480), + (1801, 2, 545), + (1802, 2, 496), + (1802, 2, 524), + (1803, 2, 502), + (1805, 2, 463), + (1805, 2, 479), + (1805, 2, 559), + (1806, 2, 476), + (1810, 2, 464), + (1811, 2, 446), + (1811, 2, 488), + (1811, 2, 509), + (1812, 2, 458), + (1813, 2, 568), + (1815, 2, 462), + (1820, 2, 483), + (1821, 2, 465), + (1821, 2, 478), + (1823, 2, 458), + (1824, 2, 538), + (1827, 2, 599), + (1829, 2, 455), + (1829, 2, 559), + (1831, 2, 456), + (1832, 2, 467), + (1832, 2, 490), + (1832, 2, 573), + (1833, 2, 498), + (1834, 2, 457), + (1835, 2, 475), + (1838, 2, 484), + (1838, 2, 539), + (1840, 2, 533), + (1842, 2, 494), + (1844, 2, 476), + (1844, 2, 491), + (1847, 2, 455), + (1847, 2, 456), + (1847, 2, 470), + (1847, 2, 544), + (1848, 2, 521), + (1849, 2, 488), + (1850, 2, 476), + (1850, 2, 492), + (1850, 2, 501), + (1851, 2, 524), + (1853, 2, 456), + (1854, 2, 481), + (1855, 2, 464), + (1855, 2, 465), + (1856, 2, 533), + (1858, 2, 481), + (1858, 2, 530), + (1862, 2, 571), + (1866, 2, 561), + (1867, 2, 487), + (1868, 2, 510), + (1873, 2, 473), + (1874, 2, 522), + (1880, 2, 480), + (1880, 2, 574), + (1882, 2, 464), + (1885, 2, 487), + (1886, 2, 475), + (1886, 2, 488), + (1886, 2, 561), + (1887, 2, 499), + (1889, 2, 500), + (1891, 2, 477), + (1893, 2, 472), + (1893, 2, 487), + (1894, 2, 537), + (1898, 2, 473), + (1899, 2, 470), + (1899, 2, 487), + (1901, 2, 513), + (1903, 2, 482), + (1903, 2, 554), + (1903, 2, 562), + (1903, 2, 588), + (1905, 2, 561), + (1906, 2, 600), + (1908, 2, 478), + (1912, 2, 585), + (1915, 2, 483), + (1917, 2, 472), + (1921, 2, 568), + (1922, 2, 483), + (1922, 2, 496), + (1925, 2, 501), + (1931, 2, 525), + (1932, 2, 543), + (1935, 2, 960), + (1940, 2, 505), + (1940, 2, 514), + (1941, 2, 484), + (1942, 2, 534), + (1943, 2, 543), + (1944, 2, 512), + (1950, 2, 540), + (1951, 2, 522), + (1951, 2, 535), + (1951, 2, 600), + (1954, 2, 502), + (1955, 2, 577), + (1956, 2, 488), + (1956, 2, 515), + (1956, 2, 549), + (1961, 2, 521), + (1965, 2, 529), + (1967, 2, 491), + (1969, 2, 550), + (1977, 2, 575), + (1982, 2, 504), + (1983, 2, 542), + (1984, 2, 533), + (1986, 2, 499), + (1994, 2, 529), + (1998, 2, 497), + (1998, 2, 545), + (1999, 2, 512), + (1999, 2, 515), + (2002, 2, 538), + (2003, 2, 506), + (2003, 2, 559), + (2005, 2, 527), + (2006, 2, 494), + (2006, 2, 504), + (2007, 2, 559), + (2009, 2, 526), + (2010, 2, 507), + (2010, 2, 518), + (2012, 2, 519), + (2012, 2, 526), + (2014, 2, 547), + (2017, 2, 522), + (2018, 2, 505), + (2023, 2, 512), + (2027, 2, 516), + (2028, 2, 499), + (2034, 2, 524), + (2037, 2, 588), + (2038, 2, 540), + (2039, 2, 565), + (2043, 2, 548), + (2046, 2, 577), + (2047, 2, 519), + (2047, 2, 552), + (2048, 1, 1024), + (2049, 2, 547), + (2049, 2, 562), + (2050, 2, 659), + (2055, 3, 405), + (2056, 2, 559), + (2064, 2, 520), + (2075, 2, 540), + (2075, 2, 546), + (2081, 2, 528), + (2082, 4, 344), + (2084, 2, 549), + (2102, 2, 585), + (2103, 2, 540), + (2107, 2, 533), + (2112, 2, 556), + (2114, 3, 365), + (2122, 2, 599), + (2124, 2, 583), + (2135, 2, 564), + (2137, 2, 529), + (2147, 2, 554), + (2147, 2, 570), + (2151, 2, 541), + (2165, 2, 750), + (2170, 2, 564), + (2176, 2, 618), + (2184, 3, 457), + (2190, 2, 569), + (2193, 2, 545), + (2194, 2, 592), + (2195, 2, 636), + (2209, 2, 568), + (2242, 3, 450), + (2258, 2, 569), + (2261, 3, 417), + (2290, 2, 572), + (2342, 3, 455), + (2362, 3, 418), + (2379, 2, 618), + (2396, 3, 509), + (2398, 3, 445), + (2430, 3, 465), + (2440, 3, 462), + (2443, 2, 649), + (2448, 3, 423), + (2462, 3, 468), + (2486, 3, 448), + (2494, 3, 461), + (2495, 3, 448), + (2509, 3, 468), + (2518, 3, 441), + (2527, 3, 442), + (2535, 3, 511), + (2540, 3, 459), + (2552, 3, 488), + (2570, 3, 468), + (2585, 3, 436), + (2634, 3, 469), + (2655, 3, 532), + (2658, 3, 438), + (2672, 3, 547), + (2673, 3, 500), + (2683, 3, 481), + (2716, 3, 499), + (2722, 3, 472), + (2727, 3, 507), + (2742, 3, 499), + (2744, 3, 471), + (2750, 3, 570), + (2757, 3, 522), + (2778, 3, 588), + (2798, 3, 540), + (2820, 3, 488), + (2842, 3, 585), + (2848, 3, 494), + (2880, 3, 540), + (2893, 3, 522), + (2895, 3, 542), + (2939, 3, 524), + (2997, 3, 630), + (3141, 5, 351), + (3189, 5, 350), + (3234, 4, 557), + (3278, 4, 455), + (3305, 5, 349), + (3317, 4, 433), + (3365, 3, 618), + (3629, 4, 528), + (8087, 12, 416), + (8104, 10, 463), + (8112, 10, 542), + (8118, 9, 548), + (8126, 10, 471), + (8142, 9, 581), + (8142, 10, 475), + (8143, 10, 472), + (8150, 8, 660), + (8172, 10, 659), + (8173, 11, 417), + (8177, 11, 421), + (8178, 8, 570), + (8178, 8, 582), + (8179, 11, 411), + (8192, 1, 4096), +] + +# Every (T, num_contexts, max_ctx_compressed_tokens[2], max_seq_len) the runtime-value dumps of +# the envelope engines recorded for this op: four engines at max_batch_size 256, max_num_tokens +# 16384, L 32768 with DSpark (CUDA graph, eager, TP4, DEP4); two at max_batch_size 8, L 32768; +# one at L 131072 without DSpark; one at L 1,048,576 with DSpark. The DEP4 engine's ranks see +# different steps, so a value one rank recorded is kept. +_ENVELOPE_DUMP_PREFILL_CALLS = [ + (1, 1, 0, 32768), + (1, 1, 0, 131072), + (1, 1, 0, 1048576), + (2, 1, 0, 131072), + (2, 1, 1, 32768), + (2, 1, 1, 131072), + (2, 1, 1, 1048576), + (5, 1, 2, 32768), + (5, 1, 2, 131072), + (5, 1, 2, 1048576), + (7, 1, 0, 32768), + (7, 1, 0, 1048576), + (12, 1, 6, 32768), + (16, 1, 8, 32768), + (29, 2, 8, 131072), + (34, 2, 8, 32768), + (34, 2, 8, 1048576), + (597, 1, 299, 32768), + (610, 1, 299, 32768), + (2252, 1, 1126, 32768), + (4559, 1, 2276, 32768), + (8147, 2, 2938, 32768), + (8192, 1, 4089, 32768), + (8192, 1, 4096, 32768), + (8467, 25, 453, 32768), + (9193, 27, 446, 32768), + (9240, 28, 447, 32768), + (9364, 15, 456, 32768), + (9370, 15, 456, 32768), + (9524, 28, 445, 32768), + (10441, 1, 5220, 32768), + (10444, 1, 5222, 32768), + (10444, 1, 5222, 131072), + (10444, 1, 5222, 1048576), + (10447, 1, 5214, 32768), + (10451, 1, 5218, 32768), + (10451, 1, 5225, 32768), + (16298, 24, 451, 32768), + (16308, 34, 473, 32768), + (16316, 29, 472, 32768), + (16320, 29, 456, 32768), + (16322, 29, 472, 32768), + (16327, 23, 472, 32768), + (16328, 23, 466, 32768), + (16329, 25, 462, 32768), + (16333, 31, 466, 32768), + (16340, 39, 431, 32768), + (16378, 24, 467, 32768), + (16384, 1, 8192, 32768), + (16384, 1, 8192, 131072), + (16384, 1, 8192, 1048576), +] + +_BMAX_DEVICE_TOTAL = 297_020_948_480 # torch.cuda.mem_get_info() total of a GB300 here +# The largest compressor-state pool over the engine envelope at 0.5 x total memory, +# PP2 stages included (measured maxima, max_batch_size sampled): the PP2 stage-0 manager built +# alone (stage 0 holds layers 0..19, so the +# compressor state of owners 2 / 8 / 14 gets a larger share of the quota) with native KV +# storage, max_num_tokens 16384, batch 2, DSpark off, L 16384 -- 247806 pages. Below it: the +# bf16 stage-0 largest 188796 (16384, batch 1, DSpark), the unsplit managers' largest +# (960 tuples per storage mode) 172416 (native, +# 16384, batch 2) and 120192 (bf16, 16384, batch 1, DSpark), the gates' old bound 133248 (native, +# 8192, batch 8), the native gates' largest 94464 and the bf16 gates' 78078. +_BMAX_PAGES = 247806 +_BF16_STAGE0_BMAX_PAGES = 188796 +_UNSPLIT_BMAX_PAGES = 172416 +_GATES_BMAX_PAGES = 133248 +_BF16_ENVELOPE_BMAX_PAGES = 120192 +_BF16_BMAX_PAGES = 78078 +_NATIVE_GATE_BMAX_PAGES = 94464 + + +class _EnvStep: + """One prefill step of the envelope: k context chunks (sp, n) and G generation rows, at + max_seq_len ``seq`` with or without DSpark (the manager's length Lm bounds every context's + kv length; the tables are ``W`` = max_blocks_per_seq columns wide).""" + + def __init__(self, ctx, gen_rows, prior, seed, ape_nonzero=False, seq=16384, dspark=False): + self.ctx = list(ctx) + self.gen_rows = gen_rows + self.prior = prior # True: parked tokens come from an earlier op call; False: planted + self.seed = seed + self.ape_nonzero = ape_nonzero + self.seq, self.dspark = seq, dspark + self.lm = _envelope_lm(seq, dspark) + self.W = _ENVELOPE_WIDTH_OF_L[seq] + assert all(sp + n <= self.lm for sp, n in self.ctx), (seq, dspark) + + @property + def k(self): + return len(self.ctx) + + @property + def tc(self): + return sum(n for _, n in self.ctx) + + @property + def max_outputs(self): + return max((s + n) // RATIO - s // RATIO for s, n in self.ctx) + + +def _compose(rng, total: int, k: int) -> List[int]: + """Random composition of ``total`` into k positive lengths (uniform cuts, or many + 1..3-token chunks and one remainder), in random order.""" + if k == 1: + return [total] + if rng.random() < 0.5 or total - (k - 1) < 1: + cuts = sorted(rng.sample(range(1, total), k - 1)) + return [b - a for a, b in zip([0] + cuts, cuts + [total])] + small = [rng.choice([1, 1, 1, 2, 3]) for _ in range(k - 1)] + if total - sum(small) < 1: + small = [1] * (k - 1) + parts = small + [total - sum(small)] + rng.shuffle(parts) + return parts + + +def _random_start(rng, n: int, lm: int = 16384) -> int: + """Cached prefix of a request with n new tokens (kv_len <= lm): 0, odd or even, and in one + request of ten the chunk ends exactly at lm (the envelope's longest kv length).""" + room = lm - n + u = rng.random() + if room < 1 or u < 0.3: + return 0 + if u > 0.9: + return room + sp = rng.randint(1, room) + if u < 0.6: + sp = sp if sp % 2 else (sp - 1 if sp > 1 else 1) + else: + sp = sp if sp % 2 == 0 else sp - 1 + return sp + + +# (T context tokens, k) edge cases that lead the seeded partitions (step 0 is T 1). +_ENV_EDGES = [ + (1, 1), + (16384, 256), + (256, 256), + (16384, 1), + (257, 256), + (16384, 2), + (16128, 256), + (2, 2), + (16383, 255), + (8192, 32), + (33, 32), + (16384, 128), +] +_ENV_PARTITIONS = len(_ENV_EDGES) + _ENVELOPE_REQS # the edges, then every k 1..256 +_ENV_GROUPS = 8 + + +def _random_step(i: int) -> _EnvStep: + """Seeded step i: the edges first, then k = every value 1..256; context tokens log-uniform + in [k, 16384] (or within 128 of 16384), random compositions and starts, a random max_seq_len + L of the six envelope values with DSpark on or off (contexts up to Lm), generation rows at + next_n 6 (DSpark) or 1 after the contexts, parked tokens from an earlier call in every + other step, a nonzero ape in one of five.""" + import random + + rng = random.Random(23_000 + i) + seq = rng.choice(_ENVELOPE_L) + dspark = rng.random() < 0.5 + if i in ( + 3, + 5, + ): # a 16384-token chunk ending at Lm = 1,048,590 (DSpark) / at 1,048,576 (no DSpark) + seq, dspark = _ENVELOPE_L[-1], i == 3 + lm = _envelope_lm(seq, dspark) + if i == 0: + return _EnvStep([(0, 1)], 0, prior=False, seed=23_000) # T 1, max_outputs 0 + if i < len(_ENV_EDGES): + tc, k = _ENV_EDGES[i] + else: + k = 1 + (i - len(_ENV_EDGES)) % _ENV_REQS + if rng.random() < 0.15: + tc = _ENV_TOKENS - rng.randint(0, 128) + else: + tc = int(round(math.exp(rng.uniform(math.log(k), math.log(_ENV_TOKENS))))) + tc = min(max(tc, k), _ENV_TOKENS) + lens = _compose(rng, tc, k) + ctx = [(_random_start(rng, n, lm), n) for n in lens] + if i in (3, 5): + ctx[0] = (lm - lens[0], lens[0]) + next_n = 6 if dspark else 1 + g = rng.randint(0, _ENV_REQS - k) + g = min(g, (_ENV_TOKENS - tc) // next_n) + return _EnvStep( + ctx, + g * next_n, + prior=(i % 2 == 0), + seed=23_000 + i, + ape_nonzero=(i % 5 == 4), + seq=seq, + dspark=dspark, + ) + + +def _dump_step(j: int, T: int, nc: int, M: int, seq: int = 16384) -> _EnvStep: + """A step with exactly the dumped (T, num_contexts, max_outputs) at max_seq_len ``seq``: + nc contexts, their tokens plus one-token generation rows (at most 256 - nc) summing to + T, one context with exactly M outputs and none with more; random starts.""" + import random + + rng = random.Random(24_000 + j) + g_lo, g_hi = max(0, T - nc * (2 * M + 1)), min(_ENV_REQS - nc, T - nc) + for _ in range(4000 if g_lo <= g_hi else 0): + g = rng.randint(g_lo, g_hi) + tc = T - g + par = [rng.random() < 0.5 for _ in range(nc)] if M > 0 else [False] * nc # True: odd start + caps = [2 * M if odd else 2 * M + 1 for odd in par] + n0 = ( + rng.choice([2 * M - 1, 2 * M]) + if par[0] + else (rng.choice([2 * M, 2 * M + 1]) if M else 1) + ) + rest = tc - n0 + if n0 < 1 or rest < nc - 1 or rest > sum(caps[1:]): + continue + lens = [n0] + [1] * (nc - 1) + left = rest - (nc - 1) + idx = list(range(1, nc)) + while left > 0: + i = rng.choice(idx) + add = min(caps[i] - lens[i], left, rng.randint(1, max(1, left))) + if add <= 0: + idx.remove(i) + continue + lens[i] += add + left -= add + ctx = [] + for n, odd in zip(lens, par): + room = seq - n + sp = rng.randint(0, room // 2) * 2 + if odd: + sp = sp + 1 if sp + 1 <= room else sp - 1 + ctx.append((sp, n)) + order = list(range(nc)) + rng.shuffle(order) + step = _EnvStep([ctx[o] for o in order], g, prior=(j % 2 == 0), seed=24_000 + j, seq=seq) + assert (step.tc + g, step.k, step.max_outputs) == (T, nc, M), (T, nc, M, ctx, g) + return step + raise AssertionError(f"no partition found for the dumped step {(T, nc, M)}") + + +class _EnvPool: + """One fp32 [pages, 128, 512] tensor serving as both the kv and the score pool (as the + manager's compressor-state pool does), with per-request block tables whose pages are + distinct random pages and whose unmapped columns point at a trap page.""" + + def __init__( + self, pool: torch.Tensor, spans, rng_seed: int, pages=None, trap: int = 0, width: int = 132 + ): + self.pool = pool + self.ps = pool.shape[1] + cpu = torch.Generator().manual_seed(rng_seed) + logical = [ + list(range(lo // self.ps, (hi - 1) // self.ps + 1)) if hi > lo else [] + for lo, hi in spans + ] + need = 2 * sum(len(q) for q in logical) + if pages is None: + pages = [p for p in range(pool.shape[0]) if p != trap] + pages = torch.tensor(list(pages)) + assert need <= pages.numel(), (need, pages.numel()) + pick = pages[torch.randperm(pages.numel(), generator=cpu)[:need]].tolist() + self.trap = trap + self.bt_kv = torch.full((_ENV_REQS, width), trap, dtype=torch.int32) + self.bt_sc = torch.full_like(self.bt_kv, trap) + used = {trap} + for b, qs in enumerate(logical): + for q in qs: + pk, pc = pick.pop(), pick.pop() + self.bt_kv[b, q], self.bt_sc[b, q] = pk, pc + used |= {pk, pc} + self.bt_kv = self.bt_kv.to(DEV) + self.bt_sc = self.bt_sc.to(DEV) + self.mapped = torch.tensor(sorted(used), device=DEV) + + +def _env_ref(kv, sc, ape, lo: int, w0: int, w1: int) -> torch.Tensor: + """model.py's ratio-2 pooling of windows [w0, w1) of per-position arrays that start at + (even) position lo; fp32.""" + p = torch.arange(w0 * RATIO, w1 * RATIO, device=DEV) + k = kv[p - lo].unflatten(0, (w1 - w0, RATIO)) + s = (sc[p - lo] + ape[p % RATIO]).unflatten(0, (w1 - w0, RATIO)) + return (k * s.softmax(dim=1)).sum(dim=1) + + +def _env_call( + step, + pool_t, + bt_kv, + bt_sc, + ctx, + rows, + ape, + out_dtype, + reserved: int, + max_outputs=None, + op_starts=None, +): + """The target's call for the contexts ``ctx`` [(sp, n)] whose kv_score rows are ``rows`` + (context rows first, then NaN generation rows): returns (output buffer incl. the + reserved rows and one guard row, cu_out list). ``op_starts(ctx)`` (controls only) + replaces the start_pos the op is handed; every other argument keeps the true ctx.""" + nc = len(ctx) + n_out = [(s + n) // RATIO - s // RATIO for s, n in ctx] + cu_seq, cu_out = [0], [0] + for (_, n), o in zip(ctx, n_out): + cu_seq.append(cu_seq[-1] + n) + cu_out.append(cu_out[-1] + o) + total = cu_out[-1] + reserved + buf = torch.zeros(total + 1, HEAD_DIM, device=DEV, dtype=out_dtype) + buf[total] = float("nan") + mo = max(n_out) if max_outputs is None else max_outputs + op( + rows, + ape, + pool_t, + pool_t, + bt_kv[:nc], + bt_sc[:nc], + buf[:total], + _i32_tensor([s + n for s, n in ctx]), + _i32_tensor([s for s, _ in ctx] if op_starts is None else op_starts(ctx)), + _i32_tensor(cu_seq), + _i32_tensor(cu_out), + nc, + 128, + HEAD_DIM, + RATIO, + mo, + ) + return buf, cu_out + + +def _run_env_step( + step: _EnvStep, + pool_t: torch.Tensor, + env: "_EnvPool" = None, + fill: bool = True, + out_dtypes=(torch.bfloat16, torch.float32), + max_outputs=None, + op_tables=None, + op_starts=None, +): + """Run one envelope step on ``pool_t`` and check it. Returns a dict with the outputs, + the checks' verdicts (raises on the first failure) and the pool snapshot inputs.""" + gen = _gen(step.seed) + ape = ( + torch.randn(RATIO, HEAD_DIM, generator=gen, device=DEV) if step.ape_nonzero else _zero_ape() + ) + # Per request: positions [lo, sp + n) of kv / score; lo is even. With ``prior`` a + # request with sp > 0 gets an earlier chunk [p0, sp) (p0 even) run by an earlier op + # call, which parks the odd token; otherwise the parked token is planted. + import random + + rng = random.Random(step.seed) + spans, data, prior_ctx = [], [], [] + for sp, n in step.ctx: + if step.prior and sp > 0: + p0 = sp - rng.randint(1, min(sp, 128)) + p0 -= p0 % RATIO + else: + p0 = sp - sp % RATIO + spans.append((p0, sp + n)) + prior_ctx.append((p0, sp - p0)) + data.append(_make_sequence(sp + n - p0, gen)) + if env is None: + env = _EnvPool(pool_t, spans, step.seed, width=step.W) + # op_tables (controls only): the tables handed to this step's call, while planting and + # checking keep the true ones. + bt_kv, bt_sc = op_tables(env.bt_kv, env.bt_sc) if op_tables else (env.bt_kv, env.bt_sc) + whole = ( + pool_t.shape[0] <= _ENV_PAGES + ) # small pools are checked whole, large ones on their mapped pages + if fill: + if whole: + pool_t.fill_(_ENV_SENTINEL) + else: + pool_t[env.mapped] = _ENV_SENTINEL + ps = pool_t.shape[1] + + # Earlier chunks (prior mode) or planted parked tokens. + planted = 0 + if step.prior and any(n for _, n in prior_ctx if n): + idx = [b for b, (_, n) in enumerate(prior_ctx) if n] + rows = torch.cat( + [ + torch.cat([data[b][0][: prior_ctx[b][1]], data[b][1][: prior_ctx[b][1]]], 1) + for b in idx + ] + ) + sub_kv = env.bt_kv[idx].contiguous() + sub_sc = env.bt_sc[idx].contiguous() + buf, cu = _env_call( + step, + pool_t, + sub_kv, + sub_sc, + [prior_ctx[b] for b in idx], + rows.contiguous(), + ape, + torch.float32, + 0, + ) + torch.cuda.synchronize() + for j, b in enumerate(idx): + p0, m = prior_ctx[b] + ref = _env_ref(data[b][0], data[b][1], ape, p0, p0 // RATIO, (p0 + m) // RATIO) + torch.testing.assert_close( + buf[cu[j] : cu[j + 1]], ref, msg=lambda s: f"earlier chunk b={b}: {s}" + ) + else: + for b, (sp, n) in enumerate(step.ctx): + if sp % RATIO: + p = sp - 1 + kv, sc = data[b] + lo = spans[b][0] + pool_t[int(env.bt_kv[b, p // ps]), p % ps] = kv[p - lo] + pool_t[int(env.bt_sc[b, p // ps]), p % ps] = sc[p - lo] + ape[p % RATIO] + planted += 1 + + # This step's kv_score: context rows in order, then the generation rows (NaN: never read). + rows = [ + torch.cat([kv[sp - lo :], sc[sp - lo :]], 1) + for (kv, sc), (sp, _), (lo, _) in zip(data, step.ctx, spans) + ] + kv_score = torch.cat( + rows + [torch.full((step.gen_rows, 2 * HEAD_DIM), float("nan"), device=DEV)] + ).contiguous() + reserved = ( + step.gen_rows + ) # output rows past the contexts' (the decode sibling's), must stay zero + outs = {} + for dt in out_dtypes: + buf, cu = _env_call( + step, + pool_t, + bt_kv, + bt_sc, + step.ctx, + kv_score, + ape, + dt, + reserved, + max_outputs, + op_starts, + ) + outs[dt] = (buf, cu) + torch.cuda.synchronize() + + ref = torch.cat( + [ + _env_ref(kv, sc, ape, lo, sp // RATIO, (sp + n) // RATIO) + for (kv, sc), (sp, n), (lo, _) in zip(data, step.ctx, spans) + ] + ) + result = { + "ref": ref, + "outs": outs, + "env": env, + "planted": planted, + "kv_score": kv_score, + "ape": ape, + } + for dt, (buf, cu) in outs.items(): + N = cu[-1] + got = buf[:N] + assert torch.isnan(buf[-1]).all(), "op wrote past the reserved rows" + assert (buf[N:-1] == 0).all(), "op wrote a reserved generation row" + if dt == torch.float32: + torch.testing.assert_close(got, ref) + else: + torch.testing.assert_close(got.float(), ref, atol=1e-5, rtol=2**-8) + if torch.float32 in outs and torch.bfloat16 in outs: + f32, b16 = outs[torch.float32][0], outs[torch.bfloat16][0] + N = outs[torch.float32][1][-1] + # The kernel computes one fp32 quotient and rounds it on store: the bf16 output is + # exactly the RNE of the fp32 output. + assert torch.equal(b16[:N], f32[:N].bfloat16()), "bf16 output != RNE(fp32 output)" + + # State: every position of [lo, sp + n) of every request -- this call's new positions, and + # the earlier chunk or the planted parked token -- holds kv / score + ape bit for bit. + for b, ((kv, sc), (lo, hi)) in enumerate(zip(data, spans)): + pos = torch.arange(lo, hi, device=DEV) + pk = env.bt_kv[b].long()[pos // ps] + pc = env.bt_sc[b].long()[pos // ps] + assert torch.equal(pool_t[pk, pos % ps], kv[pos - lo]), f"request {b}: kv state" + assert torch.equal(pool_t[pc, pos % ps], sc[pos - lo] + ape[pos % RATIO]), ( + f"request {b}: score state" + ) + expect = 2 * sum(hi - lo for lo, hi in spans) + written = int(((pool_t if whole else pool_t[env.mapped]) != _ENV_SENTINEL).any(-1).sum()) + assert written == expect, f"{written} slots written, {expect} named" + assert (pool_t[env.trap] == _ENV_SENTINEL).all(), "the trap page was written" + return result + + +@pytest.mark.parametrize("group", range(_ENV_GROUPS)) +def test_random_partitions_of_the_envelope(group: int): + """Over the engine envelope: seeded random prefill steps -- k in 1..256 + context chunks (every value) of T in 1..16384 context tokens (edge cases T 1 / 2 / 33 / + 256 / 257 / 16128 / 16383 / 16384 at k 1 / 2 / 32 / 128 / 255 / 256 first), random lengths + (uniform cuts, or 1..3-token chunks plus one remainder), at a random max_seq_len L of the + six envelope values with DSpark on or off (tables 132 .. 8196 columns wide), cached + prefixes in [0, Lm - n] (zero, odd, even, and chunks ending exactly at Lm, up to + 1,048,590), random order, random distinct pages for kv and score, generation rows (next_n 6 + with DSpark, else 1; at most 256 requests and 16384 tokens) after the contexts, max_outputs + = max_b(kv_len // 2 - sp // 2) up to 8192. Half the steps get each odd start's parked + token from an earlier multi-request op call (itself checked), half get it planted; one in + five uses a nonzero ape. Each step runs with the target's zero-filled bf16 output and an + fp32 output: fp32 = reference at the default tolerances, bf16 = RNE of the fp32 output bit + for bit, reserved generation rows stay zero, every state slot exact, nothing else + written.""" + steps = [_random_step(i) for i in range(group, _ENV_PARTITIONS, _ENV_GROUPS)] + pool_t = torch.empty(_ENV_PAGES, 128, HEAD_DIM, device=DEV) + exact, total = 0, 0 + for step in steps: + r = _run_env_step(step, pool_t) + b16 = r["outs"][torch.bfloat16][0][: r["ref"].shape[0]] + exact += int((b16 == r["ref"].bfloat16()).sum()) + total += b16.numel() + frac = exact / max(total, 1) + ks = sorted({s.k for s in steps}) + print( + f"[envelope random group {group}] {len(steps)} steps, k {ks[0]}..{ks[-1]} ({len(ks)} values), " + f"T {min(s.tc for s in steps)}..{max(s.tc for s in steps)}, max_outputs " + f"{min(s.max_outputs for s in steps)}..{max(s.max_outputs for s in steps)}, widths " + f"{sorted({s.W for s in steps})}, kv lengths up to {max(sp + n for s in steps for sp, n in s.ctx)}; " + f"bf16 vs reference bit-exact {frac:.6f}" + ) + assert frac >= 0.999, frac # the contract's bf16 gate (see test_bf16_output) + + +def test_random_partitions_cover_the_envelope(): + """The seeded partitions reach the envelope's corners: every k in 1..256, T 1 and 16384, + odd, even and zero starts, 1-token chunks, max_outputs 0 and 8192, every max_seq_len + (table width) with DSpark on and off, chunks ending at 1,048,590 (Lm at L 1M with DSpark) + and at 1,048,576, and at least 200 steps.""" + steps = [_random_step(i) for i in range(_ENV_PARTITIONS)] + ks = {s.k for s in steps} + tcs = {s.tc for s in steps} + assert len(steps) >= 200 + assert ks == set(range(1, _ENV_REQS + 1)), sorted(set(range(1, _ENV_REQS + 1)) - ks) + assert {1, _ENV_TOKENS} <= tcs and all(1 <= t <= _ENV_TOKENS for t in tcs) + starts = [sp for s in steps for sp, _ in s.ctx] + assert ( + any(sp == 0 for sp in starts) + and any(sp % 2 for sp in starts) + and any(sp and sp % 2 == 0 for sp in starts) + ) + assert any(n == 1 for s in steps for _, n in s.ctx) + mos = {s.max_outputs for s in steps} + assert 0 in mos and _ENV_TOKENS // 2 in mos + assert {(s.W, s.dspark) for s in steps} == { + (w, d) for w in _ENVELOPE_WIDTHS for d in (False, True) + } + ends = {sp + n for s in steps for sp, n in s.ctx} + assert _ENVELOPE_LM_1M in ends and _ENVELOPE_L[-1] in ends + assert all( + s.tc + s.gen_rows <= _ENV_TOKENS and s.k + (s.gen_rows > 0) <= _ENV_REQS for s in steps + ) + assert all(sp + n <= s.lm for s in steps for sp, n in s.ctx) + print( + f"[envelope coverage] {len(steps)} partitions: k 1..{_ENV_REQS} all, {len(tcs)} distinct T, max_outputs " + f"{min(mos)}..{max(mos)}, {sum(1 for sp in starts if sp % 2)} odd starts of {len(starts)}, kv lengths up to " + f"{max(ends)}" + ) + + +def test_every_dumped_prefill_step(): + """The dumps: each (T, num_contexts, max_outputs) of the gates' dumps (1014, + max_seq_len 16384) and of the envelope engines' dumps (51: T up to 16384, nc up to 39, + max_outputs up to 8192, at L 32768 / 131072 / 1,048,576) built as a step with exactly + those values at that max_seq_len (random lengths / starts / order, one-token generation + rows filling T) and run as in the random partitions.""" + pool_t = torch.empty(_ENV_PAGES, 128, HEAD_DIM, device=DEV) + cells = [(T, nc, M, 16384) for T, nc, M in _DUMP_PREFILL_CALLS] + _ENVELOPE_DUMP_PREFILL_CALLS + assert len(_ENVELOPE_DUMP_PREFILL_CALLS) == 51 + exact, total = 0, 0 + for j, (T, nc, M, seq) in enumerate(cells): + step = _dump_step(j, T, nc, M, seq) + r = _run_env_step(step, pool_t) + b16 = r["outs"][torch.bfloat16][0][: r["ref"].shape[0]] + exact += int((b16 == r["ref"].bfloat16()).sum()) + total += b16.numel() + frac = exact / max(total, 1) + print( + f"[envelope dump] {len(cells)} dumped steps (nc up to {max(t[1] for t in cells)}, T up to " + f"{max(t[0] for t in cells)}, max_outputs up to {max(t[2] for t in cells)}) == reference; " + f"bf16 bit-exact {frac:.6f}" + ) + assert frac >= 0.999, frac + + +def test_kernel_variant_map_over_the_envelope(): + """The profiler's kernel launches over steps spread across the envelope + (k 1..256, T 1..16384, max_outputs 0 / 1 / up to 8192, tables 132 .. 8196 columns wide, + bf16 and fp32 output): the op launches exactly one kernel per call, always the same + instantiation + (prefillReductionKernel<512, 4, 4, 2, 1>: head_dim 512, fp32 kv_score, fp32 state, + ratio 2, one reduction warp), with grid (num_contexts, max(max_outputs, 1), 1) and + 128 threads; the output dtype is a runtime argument, not a variant. Every step run + here is also checked against the reference (by _run_env_step, outside the profile).""" + import json + import os + import tempfile + + steps = [_random_step(i) for i in range(0, _ENV_PARTITIONS, 6)] + [ + _random_step(i) for i in (1, 2, 3, 8) + ] + steps += [ + _EnvStep([(0, 1)], 0, False, 1), + _EnvStep([(1, 2)], 0, False, 2), + _EnvStep([(0, 8192)], 0, False, 3), + _EnvStep([(0, 16384)], 0, False, 4, seq=32768), + ] + pool_t = torch.empty(_ENV_PAGES, 128, HEAD_DIM, device=DEV) + calls = [] + for step in steps: + r = _run_env_step(step, pool_t) # checked against the reference + calls.append((step, r)) + # Profile the bare op calls again (new rows; the state re-stores itself), every input + # built before the profile starts, so every kernel in it is the op's. + zero_ape = _zero_ape() + launches, expected = [], [] + for step, r in calls: + env = r["env"] + rows = torch.randn( + step.tc + step.gen_rows, 2 * HEAD_DIM, generator=_gen(step.seed), device=DEV + ) + lens = [n for _, n in step.ctx] + outs = [(s + n) // 2 - s // 2 for s, n in step.ctx] + args = ( + _i32_tensor([s + n for s, n in step.ctx]), + _i32_tensor([s for s, _ in step.ctx]), + _i32_tensor([0] + torch.tensor(lens).cumsum(0).tolist()), + _i32_tensor([0] + torch.tensor(outs).cumsum(0).tolist()), + ) + for dt in (torch.bfloat16, torch.float32): + out = torch.empty(sum(outs), HEAD_DIM, device=DEV, dtype=dt) + launches.append( + ( + rows, + env.bt_kv[: step.k], + env.bt_sc[: step.k], + out, + args, + step.k, + step.max_outputs, + ) + ) + expected.append([step.k, max(step.max_outputs, 1), 1]) + torch.cuda.synchronize() + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for rows, bkv, bsc, out, (kvl, spos, cus, cuo), k, mo in launches: + op( + rows, + zero_ape, + pool_t, + pool_t, + bkv, + bsc, + out, + kvl, + spos, + cus, + cuo, + k, + 128, + HEAD_DIM, + RATIO, + mo, + ) + torch.cuda.synchronize() + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "trace.json") + prof.export_chrome_trace(path) + trace = json.load(open(path)) + kernels = [e for e in trace["traceEvents"] if e.get("cat") == "kernel"] + ops = [e for e in kernels if "prefillReduction" in e["name"]] + names = {e["name"] for e in ops} + grids = [e["args"]["grid"] for e in sorted(ops, key=lambda e: e["ts"])] + blocks = {tuple(e["args"]["block"]) for e in ops} + others = sorted({e["name"] for e in kernels if "prefillReduction" not in e["name"]}) + print( + f"[envelope variant map] {len(ops)} op launches over {len(steps)} steps x 2 dtypes: kernels {sorted(names)}, " + f"blocks {sorted(blocks)}; grids (k, max(max_outputs, 1), 1) from {grids[:3]} to {grids[-3:]}; " + f"{len({g[0] for g in grids})} distinct k (max {max(g[0] for g in grids)}), {len({g[1] for g in grids})} " + f"distinct grid.y (max {max(g[1] for g in grids)}); other kernels {others}" + ) + assert others == [] and len(kernels) == len(launches), ( + "the op launched something besides one kernel per call" + ) + assert len(names) == 1, names + (name,) = names + assert "prefillReductionKernel<512,4,4,2,1>" in name.replace(" ", ""), name + assert blocks == {(128, 1, 1)}, blocks + assert grids == expected, "grid differs from (num_contexts, max(max_outputs, 1), 1)" + + +def test_envelope_controls(): + """The envelope checks can see the call-site defects that would matter: an undersized + max_outputs (one below the target's max_ctx_compressed_tokens[2]) and two requests' + kv tables swapped both fail _run_env_step on the steps where they apply.""" + pool_t = torch.empty(_ENV_PAGES, 128, HEAD_DIM, device=DEV) + step = _EnvStep([(0, 1000), (129, 301), (4095, 3)], 5, False, 77) + _run_env_step(step, pool_t) + with pytest.raises(AssertionError): + _run_env_step(step, pool_t, max_outputs=step.max_outputs - 1) + + def swap(kv, sc): + kv = kv.clone() + kv[[0, 1]] = kv[[1, 0]] + return kv, sc + + with pytest.raises(AssertionError): + _run_env_step(step, pool_t, op_tables=swap) + print("[envelope controls] max_outputs - 1 and swapped kv tables: both rejected") + + +def _bmax_pool(): + """The largest compressor-state pool the engine envelope gives at 0.5 x the device's + total memory, PP2 stages included: the live PP2 stage-0 manager (built alone) with the + native KV storage (its smaller SWA / main-KV entries leave the state pool more of the + quota), max_num_tokens 16384, max_batch_size 2, DSpark off, max_seq_len 16384 -- the + largest of the sampled configurations (kv_cache_fraction 0.5 applies to free memory after + the weights, which is below the total, and every role's pages grow monotonically with the + quota).""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + + total = torch.cuda.mem_get_info()[1] + pool_t = _engine_manager( + 0, 2, quota=total // 2, storage="native", max_num_tokens=_ENVELOPE_TOKENS, pp_stage=0 + ).get_buffers(2, T.COMPRESSOR_KV) + if total == _BMAX_DEVICE_TOTAL: + assert tuple(pool_t.shape) == (_BMAX_PAGES, 128, HEAD_DIM), tuple(pool_t.shape) + assert pool_t.is_contiguous() and pool_t.dtype == torch.float32 + return pool_t + + +def _pool_sum(pool_t: torch.Tensor) -> int: + """int64 sum of a pool's int32 bit patterns, 8192 pages at a time: one reduction over a + B_max pool would need a temporary as large as the pool itself.""" + return sum( + int(pool_t[i : i + 8192].view(torch.int32).sum(dtype=torch.int64)) + for i in range(0, pool_t.shape[0], 8192) + ) + + +_INVARIANCE_STEP = _EnvStep( + [ + (0, 1000), + (129, 777), + (4095, 3), + (1, 1), + (16000, 383), + (2, 2000), + (8191, 1), + (0, 2), + (_ENVELOPE_LM_1M - 700, 700), + (_ENVELOPE_LM_1M - 3001, 3000), + ], + 6, + False, + 31_337, + seq=_ENVELOPE_L[-1], + dspark=True, +) + + +def test_pool_size_invariance(): + """One multi-request step (odd and even starts, parked tokens, generation + rows, two contexts ending at 1,048,590 and 1,048,589 on 8196-column tables) on the same + pages -- a random permutation of pages 1..P with trap page 0, so the smallest pool that + holds them has B = P + 1 pages -- run on views pool[:B] of the largest pool of the + engine envelope (the live PP2 stage-0 manager's, _bmax_pool) for B = P + 1, a mid + size, the bf16 gates' largest (78078), the native gates' largest (94464), the bf16 + unsplit largest (120192), the gates' old bound (133248), the native unsplit largest + (172416), the bf16 stage-0 largest (188796) and B_max (247806), and on a separate + exact-size [P + 1, 128, 512] allocation: outputs (bf16 and fp32) and the P + 1 pages of + state are bitwise identical across the ten, equal the reference, and the views' pages + past P + 1 are untouched (checksum).""" + big = _bmax_pool() + B_max = big.shape[0] + step = _INVARIANCE_STEP + spans = [(sp - sp % RATIO, sp + n) for sp, n in step.ctx] + P = 2 * sum((hi - 1) // 128 - lo // 128 + 1 for lo, hi in spans) + small = P + 1 + results = [] + for label, pool_t in [ + ("exact", torch.empty(small, 128, HEAD_DIM, device=DEV)), + ("B_small", big[:small]), + ("B_mid", big[: B_max // 2]), + ("B_bf16_max", big[:_BF16_BMAX_PAGES]), + ("B_native_gate", big[:_NATIVE_GATE_BMAX_PAGES]), + ("B_bf16_envelope_max", big[:_BF16_ENVELOPE_BMAX_PAGES]), + ("B_gates_max", big[:_GATES_BMAX_PAGES]), + ("B_unsplit_max", big[:_UNSPLIT_BMAX_PAGES]), + ("B_bf16_stage0_max", big[:_BF16_STAGE0_BMAX_PAGES]), + ("B_max", big), + ]: + env = _EnvPool(pool_t, spans, 31_337, pages=range(1, small), trap=0, width=step.W) + pool_t[:small] = _ENV_SENTINEL + tail = _pool_sum(pool_t[small:]) if pool_t.shape[0] > small else None + r = _run_env_step(step, pool_t, env=env, fill=False) + if tail is not None: + assert _pool_sum(pool_t[small:]) == tail, f"{label}: tail written" + results.append( + ( + label, + pool_t.shape[0], + {dt: o[0].clone() for dt, o in r["outs"].items()}, + pool_t[:small].clone(), + ) + ) + base = results[0] + for label, B, outs, state in results[1:]: + for dt in outs: + assert torch.equal( + outs[dt].view(torch.int16 if dt == torch.bfloat16 else torch.int32), + base[2][dt].view(torch.int16 if dt == torch.bfloat16 else torch.int32), + ), (label, dt) + assert torch.equal(state.view(torch.int32), base[3].view(torch.int32)), ( + f"{label}: state differs" + ) + print( + f"[envelope pool invariance] pool B in {[r[1] for r in results]} (views of one [{B_max}, 128, 512] fp32 " + f"allocation + an exact-size one): {step.k} contexts on pages 1..{P}, outputs and state bitwise identical" + ) + + +def test_far_end_multi_request_at_b_max(): + """Far end: a multi-request step with odd starts (one context ending at + 1,048,590 on 8196-column tables) whose kv / score pages are the pages on both sides of + every 32-bit offset hazard of the largest pool, B_max = 247806 on the live PP2 stage-0 + manager (byte offsets 2^31, 2^32; + element offsets 2^31, 2^32, 2^33 = pages 8191/8192, 16383/16384, 32767/32768, + 65535/65536, 131071/131072) and otherwise its last pages, the very last page included: + outputs equal the reference, state exact, nothing else in the pool changed (checksum). + Control: the same step with every page id >= 65536 wrapped mod 65536 (a 32-bit + element offset) fails.""" + big = _bmax_pool() + B_max = big.shape[0] + step = _EnvStep( + [ + (129, 4000), + (1, 3000), + (16000, 383), + (4096, 700), + (3, 1), + (0, 107), + (_ENVELOPE_LM_1M - 3001, 3001), + ], + 3, + False, + 4242, + seq=_ENVELOPE_L[-1], + dspark=True, + ) + spans = [(sp - sp % RATIO, sp + n) for sp, n in step.ctx] + need = 2 * sum((hi - 1) // 128 - lo // 128 + 1 for lo, hi in spans) + hz = [ + p + for p in (8191, 8192, 16383, 16384, 32767, 32768, 65535, 65536, 131071, 131072) + if p < B_max + ] + tailp = [p for p in range(B_max - need - 1, B_max) if p not in hz] + pages = hz + tailp[len(tailp) - (need - len(hz)) :] + assert B_max - 1 in pages and len(set(pages)) == need + trap = next(p for p in range(B_max - 1, 0, -1) if p not in set(pages)) + env = _EnvPool(big, spans, 4242, pages=pages, trap=trap, width=step.W) + big[env.mapped] = _ENV_SENTINEL + outside = _pool_sum(big) - int(big[env.mapped].view(torch.int32).sum(dtype=torch.int64)) + _run_env_step(step, big, env=env, fill=False) + after = _pool_sum(big) - int(big[env.mapped].view(torch.int32).sum(dtype=torch.int64)) + assert after == outside, "pool written outside the mapped pages" + print( + f"[far end] pool [{B_max}, 128, 512] fp32 ({big.numel() / 2**32:.2f} x 2^32 elements): {step.k} contexts " + f"on hazard pages {hz} and pages up to {max(pages)}: exact, nothing else written" + ) + + def wrap(kv, sc): + return (torch.where(kv >= 65536, kv % 65536, kv), torch.where(sc >= 65536, sc % 65536, sc)) + + big[env.mapped] = _ENV_SENTINEL + with pytest.raises(AssertionError): + _run_env_step(step, big, env=env, fill=False, op_tables=wrap) + print("[envelope far-end control] page ids >= 65536 wrapped mod 65536: the check fails") + + +# --------------------------------------------------------------------------- +# Encoder replay: context requests with zero compute rows (kv_len == start_pos) +# --------------------------------------------------------------------------- +# Under encoder SWA bounded replay a prefix hit of `reused` tokens (a multiple of 2, the +# compress ratio) starts the request's first chunk R = min(128, reused) positions early, and +# the target's _compress calls this op over the step's COMPUTE rows only: +# kv_score fp32 [rows, 1024] of the compute rows (v41_compute_token_index_cuda), the +# context compute rows first, then the generation rows +# kv_lens metadata.kv_lens_cuda_runtime[:nc] (cached + chunk rows) +# start_pos metadata.v41_compute_start_pos_cuda[:nc] (cached + R = reused) +# cu_seq_lens metadata.v41_compute_cu_seq_lens_cuda[:nc + 1] (the compute rows' offsets) +# cu_kv_comp, max_outputs from the compute start (cu_new_comp_kv_cuda[2], +# max_ctx_compressed_tokens[2]) +# A first chunk that holds exactly its R replay rows -- the runtime admits it: +# resize_context refuses only a chunk shorter than R, and scheduler_v2 rounds a non-final +# chunk down to 128 -- is a context request with kv_len == start_pos: no row, no output, +# at an even start_pos = reused >= 2 (2..16382; a later chunk follows it). Such requests +# sit anywhere in the batch, several per step, beside ordinary chunks and generation rows. +# A step whose every context request is zero-row still calls the op when it has generation +# rows (the target returns early only when no row of the step is a compute row), with +# max_outputs 0. The zero-row chunks' replay rows count against max_num_tokens. + +_REPLAY_WINDOW = 128 +_ZR_PARTITIONS = _ENVELOPE_REQS # every context count 1..256 +_ZR_GROUPS = 8 + + +def _zr_step( + ctx, + gen_rows: int, + next_n: int, + prior: bool, + seed: int, + ape_nonzero: bool = False, + seq: int = 16384, + dspark: bool = False, +) -> _EnvStep: + step = _EnvStep(ctx, gen_rows, prior, seed, ape_nonzero, seq=seq, dspark=dspark) + step.next_n = next_n + return step + + +def _zr_replay_rows(step: _EnvStep) -> int: + """Replay rows of the step's zero-row chunks (R = min(128, reused) each): tokens of the + step that the op does not see.""" + return sum(min(_REPLAY_WINDOW, sp) for sp, n in step.ctx if n == 0) + + +def _zr_sp(rng, lm: int = 16384) -> int: + """start_pos of a zero-row replay chunk = reused, even: a hit of 2..128 tokens (R = + reused, the chunk starts at 0) in about a tenth, a hit within 1024 tokens of the + manager's length lm in a tenth, else 130..lm - 2 (R = 128).""" + u = rng.random() + if u < 0.1: + return rng.randrange(2, _REPLAY_WINDOW + 1, 2) + if u < 0.2: + return rng.randrange(max(_REPLAY_WINDOW + 2, lm - 1024), lm - 1, 2) + return rng.randrange(_REPLAY_WINDOW + 2, lm - 1, 2) + + +def _zero_row_corners() -> List[_EnvStep]: + """The corners, each with parked tokens from an earlier op call and planted.""" + many = [(2 * (65 + 307 * j), 0) for j in range(26)] + most = [(2 * (100 + 257 * j), 0) for j in range(31)] + mixed = [((2 * j + 1) if j % 2 else 100 * j, 260) for j in range(31)] + # 255 short hits (R = sp) beside one chunk filling the 16384-token step (256 contexts) + short = [(2 + 2 * (j % 40), 0) for j in range(255)] + short_chunk = [(1, _ENVELOPE_TOKENS - sum(sp for sp, _ in short))] + # 127 hits of R = 128 up to the 1,048,576-token envelope's far end, beside a 128-token chunk + far = [(2 * (65 + 4127 * j), 0) for j in range(126)] + [(_ENVELOPE_L[-1] - 2, 0)] + shapes = [ + ([(4098, 0), (131, 7), (2, 2)], 5, 1, False, 16384, False), # zero-row first + ([(4098, 6), (136, 0), (2, 2)], 0, 1, False, 16384, False), # zero-row in the middle + ( + [(6, 4), (3, 7), (10, 0)], + 12, + 6, + False, + 16384, + False, + ), # zero-row last, beside two DSpark verify requests + ([(128, 0), (256, 0), (5, 3)], 0, 1, False, 16384, False), # two zero-row requests + ( + [(4098, 0)], + 31, + 1, + False, + 16384, + False, + ), # the only context, beside 31 decode rows (32 requests) + ( + [(16382, 0)], + 30, + 6, + False, + 16384, + False, + ), # the only context at the last start, beside 5 verify requests + ( + many, + 36, + 6, + False, + 16384, + False, + ), # 26 contexts, every one zero-row, beside 6 verify requests + ( + most[:15] + [(1, 4224)] + most[15:], + 0, + 1, + False, + 16384, + False, + ), # 31 zero-row + one chunk: 8192 tokens + ( + mixed[:16] + [(8190, 0)] + mixed[16:], + 0, + 1, + False, + 16384, + False, + ), # one zero-row among 31 chunks + ( + [(2, 0), (0, 300), (16382, 0), (64, 0), (777, 5)], + 6, + 1, + False, + 16384, + False, + ), # sp 2, 16382, 64 (< 128) + ([(4098, 0), (129, 1000)], 0, 1, True, 16384, False), # nonzero ape + ([(7, 0), (4, 4)], 0, 1, False, 16384, False), # odd start_pos: not an engine form + ( + short[:100] + short_chunk + short[100:], + 0, + 1, + False, + 32768, + False, + ), # 256 contexts, 255 zero-row + ( + far[:60] + [(1_000_001, 128)] + far[60:], + 0, + 1, + False, + _ENVELOPE_L[-1], + False, + ), # 127 zero-row to 1,048,574 + ( + [(_ENVELOPE_LM_1M - 2, 0), (_ENVELOPE_LM_1M - 501, 501), (2, 0)], + 12, + 6, + False, + _ENVELOPE_L[-1], + True, + ), # at Lm - 2 = 1,048,588 + ] + out = [] + for c, (ctx, gen_rows, next_n, ape, seq, dspark) in enumerate(shapes): + for prior in (True, False): + out.append( + _zr_step( + ctx, gen_rows, next_n, prior, 67_000 + 2 * c + int(prior), ape, seq, dspark + ) + ) + return out + + +def _zero_row_step(i: int) -> _EnvStep: + """A seeded random step of k = i % 256 + 1 context requests, z of them zero-row (1..k, or + 1..3), at random batch positions, at a random max_seq_len L of the six envelope values with + DSpark on or off (zero-row starts up to Lm - 2; hits shortened to R < 128 where the replay + rows would not fit the 16384-token step); the others a random composition of the tokens + left after the zero-row chunks' replay rows (log-uniform total), at random cached starts; + generation rows (next_n 6 with DSpark, else 1) after the contexts; parked tokens from an + earlier op call in three steps of four; a nonzero ape in one of five.""" + import random + + rng = random.Random(66_000 + i) + seq = rng.choice(_ENVELOPE_L) + dspark = rng.random() < 0.5 + lm = _envelope_lm(seq, dspark) + k = i % _ENV_REQS + 1 + z = rng.randint(1, k) if rng.random() < 0.5 else rng.randint(1, min(3, k)) + zero = [(_zr_sp(rng, lm), 0) for _ in range(z)] + j = 0 + while sum(min(_REPLAY_WINDOW, sp) for sp, _ in zero) > _ENV_TOKENS - (k - z) - 16: + zero[j] = (rng.randrange(2, 16, 2), 0) + j += 1 + budget = _ENV_TOKENS - sum(min(_REPLAY_WINDOW, sp) for sp, _ in zero) + nz = k - z + ctx = [] + if nz: + tc = ( + nz + if budget <= nz + else int(round(math.exp(rng.uniform(math.log(nz), math.log(budget))))) + ) + ctx = [(_random_start(rng, n, lm), n) for n in _compose(rng, min(max(tc, nz), budget), nz)] + tc = sum(n for _, n in ctx) + next_n = 6 if dspark else 1 + g = min(rng.randint(0, _ENV_REQS - k), (budget - tc) // next_n) + if nz == 0 and g == 0: + # zero-row contexts alone call the op only beside generation rows + if k < _ENV_REQS and budget - tc >= next_n: + g = 1 + else: + zero.pop() + ctx = [(_random_start(rng, 1, lm), 1)] + order = ctx + zero + rng.shuffle(order) + return _zr_step(order, g * next_n, next_n, i % 4 != 3, 66_000 + i, i % 5 == 4, seq, dspark) + + +def _zero_row_steps() -> List[_EnvStep]: + return _zero_row_corners() + [_zero_row_step(i) for i in range(_ZR_PARTITIONS)] + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.view(torch.int16 if t.dtype == torch.bfloat16 else torch.int32) + + +def _run_zero_row_pair(step: _EnvStep, pool_t: torch.Tensor) -> dict: + """_run_env_step on the step (reference, state of every request exact, nothing else + written, reserved rows zero), then the step's call again from the resulting pool with + its zero-row requests and without them: the output buffers (bf16 and fp32, reserved and + guard rows included) and the pools after the two calls bitwise equal, and equal to the + first run's output and to the pool before the call. With no non-empty context the + 'without' side is no call at all: the target's zero-filled buffer and an unchanged pool.""" + r = _run_env_step(step, pool_t) + env, kv_score, ape = r["env"], r["kv_score"], r["ape"] + keep = [b for b, (_, n) in enumerate(step.ctx) if n] + snap = pool_t.clone() + same_out = same_pool = True + for dt in (torch.bfloat16, torch.float32): + pool_t.copy_(snap) + with_buf, _ = _env_call( + step, pool_t, env.bt_kv, env.bt_sc, step.ctx, kv_score, ape, dt, step.gen_rows + ) + with_pool = pool_t.clone() + pool_t.copy_(snap) + if keep: + idx = torch.tensor(keep, device=DEV) + wo_buf, _ = _env_call( + step, + pool_t, + env.bt_kv[idx], + env.bt_sc[idx], + [step.ctx[b] for b in keep], + kv_score, + ape, + dt, + step.gen_rows, + ) + else: + wo_buf = torch.zeros_like(with_buf) + wo_buf[-1] = float("nan") + same_out &= torch.equal(_bits(with_buf), _bits(wo_buf)) and torch.equal( + _bits(with_buf), _bits(r["outs"][dt][0]) + ) + same_pool &= torch.equal(_bits(with_pool), _bits(pool_t)) and torch.equal( + _bits(with_pool), _bits(snap) + ) + pool_t.copy_(snap) + return {"same_out": same_out, "same_pool": same_pool, "ref": r["ref"], "outs": r["outs"]} + + +def test_zero_row_requests_at_both_page_sizes(): + """The reviewer's probe plans, frozen, at page sizes 128 and 256 (chunk plans through + _run_plan: a 0 is a zero-row request in that call): a zero-row request first (sp 4098), + in the middle (sp 136), last (sp 10), two in one call, and an odd sp 7 (not an engine + form); fp32 and bf16 outputs match the reference, every position's state exact, no + other slot (trap page included) and no guard row written.""" + plans = { + "zero-row first (sp 4098)": [[4098, 0, 9], [131, 5, 7], [2, 2, 2]], + "zero-row in the middle (sp 136)": [[4098, 6, 9], [131, 5, 0], [2, 2, 2]], + "zero-row last (sp 10)": [[6, 4], [3, 7], [10, 0]], + "two zero-row in one call (sp 128, 256)": [[128, 0, 2], [256, 0, 4], [5, 3, 1]], + "zero-row at odd sp 7 (not an engine form)": [[7, 0, 3], [4, 4, 4]], + } + ape = torch.randn(RATIO, HEAD_DIM, generator=_gen(19), device=DEV) + ran = 0 + for page_size in (128, 256): + for name, plan in plans.items(): + for out_dtype, a in ( + (torch.float32, _zero_ape()), + (torch.bfloat16, _zero_ape()), + (torch.float32, ape), + ): + pool, kvs, scores, outs = _run_plan(plan, page_size, a, out_dtype, seed=31) + for b, (kv, sc, out) in enumerate(zip(kvs, scores, outs)): + ref = _reference_pool(kv, sc, a) + if out_dtype == torch.float32: + torch.testing.assert_close( + out, ref, msg=lambda s: f"{name} page {page_size} seq {b}: {s}" + ) + else: + torch.testing.assert_close(out.float(), ref, atol=1e-5, rtol=2**-8) + _check_state(pool, kvs, scores, a) + ran += 1 + print( + f"[zero-row plans] {len(plans)} plans x pages 128/256 x (fp32, bf16, fp32 + ape): {ran} runs exact" + ) + + +@pytest.mark.parametrize("group", range(_ZR_GROUPS)) +def test_zero_row_requests_over_the_envelope(group: int): + """Zero-row context requests (kv_len == start_pos, the encoder-replay chunk of only its + replay rows) over the engine envelope at the engine's page size 128: the 15 corners + (each with parked tokens from an earlier op call and planted; 255 zero-row requests beside + one chunk, 127 hits of 128 rows up to 1,048,574, one at 1,048,588 beside a chunk ending at + Lm 1,048,590) and 256 seeded random steps (_zero_row_step: every context count 1..256, + 1..k zero-row requests at random batch positions, even start_pos 2..Lm - 2 at a random + max_seq_len, the other chunks random, next_n 1 or 6 generation rows). Each step: + _run_env_step's checks (fp32 = reference, bf16 = RNE(fp32), reserved + rows zero, every request's state exact -- a zero-row request's earlier positions + included -- and nothing else written), then the call with and without its zero-row + requests from the same pool: outputs and pools bitwise equal (_run_zero_row_pair).""" + steps = _zero_row_steps()[group::_ZR_GROUPS] + pool_t = torch.empty(_ENV_PAGES, 128, HEAD_DIM, device=DEV) + diff_out = diff_pool = exact = total = zero = 0 + for step in steps: + r = _run_zero_row_pair(step, pool_t) + diff_out += int(not r["same_out"]) + diff_pool += int(not r["same_pool"]) + zero += sum(1 for _, n in step.ctx if n == 0) + b16 = r["outs"][torch.bfloat16][0][: r["ref"].shape[0]] + exact += int((b16 == r["ref"].bfloat16()).sum()) + total += b16.numel() + frac = exact / max(total, 1) + print( + f"[zero-row group {group}] {len(steps)} steps, {zero} zero-row requests, k " + f"{min(s.k for s in steps)}..{max(s.k for s in steps)}: reference, state and untouched slots exact; " + f"with vs without the zero-row requests: {diff_out} steps' outputs and {diff_pool} steps' pools differ; " + f"bf16 bit-exact {frac:.6f}" + ) + assert diff_out == 0 and diff_pool == 0 + assert frac >= 0.999, frac + + +def test_zero_row_partitions_cover_the_envelope(): + """The zero-row steps reach: every context count 1..256; zero-row requests first, in the + middle and last; several in one step; steps whose every context is zero-row (beside + generation rows); 255 zero-row requests beside one chunk; start_pos 2, below 128 and up + to 1,048,588 (Lm - 2 at L 1M with DSpark); next_n 1 and 6; every max_seq_len (table + width); parked tokens from an earlier call and planted; a step of 16384 tokens with the + replay rows; all inside the envelope.""" + steps = _zero_row_steps() + ks, first, middle, last, multi, only, nn, prior = set(), 0, 0, 0, 0, 0, set(), set() + sps, most = [], 0 + for s in steps: + zs = [b for b, (_, n) in enumerate(s.ctx) if n == 0] + assert zs, "a step without a zero-row request" + ks.add(s.k) + first += int(0 in zs) + last += int(s.k - 1 in zs and s.k > 1) + middle += int(any(0 < b < s.k - 1 for b in zs)) + multi += int(len(zs) >= 2) + only += int(len(zs) == s.k) + most = max(most, len(zs)) + if s.gen_rows: + nn.add(s.next_n) + assert s.gen_rows % s.next_n == 0 + if len(zs) == s.k: + assert s.gen_rows > 0, "zero-row contexts alone make no call" + prior.add(s.prior) + sps += [s.ctx[b][0] for b in zs] + assert s.tc + _zr_replay_rows(s) + s.gen_rows <= _ENV_TOKENS + assert s.k + s.gen_rows // s.next_n <= _ENV_REQS + assert all(sp + n <= s.lm for sp, n in s.ctx) + t_max = max(s.tc + _zr_replay_rows(s) + s.gen_rows for s in steps) + engine_sps = [sp for sp in sps if sp % 2 == 0] + print( + f"[zero-row coverage] {len(steps)} steps, {len(sps)} zero-row requests (first {first}, middle {middle}, " + f"last {last}; >= 2 in {multi} steps, every context in {only}, max {most}); k {min(ks)}..{max(ks)}; " + f"start_pos {min(engine_sps)}..{max(engine_sps)} even, {len(sps) - len(engine_sps)} odd (not an engine " + f"form); next_n {sorted(nn)}; T with the replay rows up to {t_max}" + ) + assert ( + ks == set(range(1, _ENV_REQS + 1)) + and min(first, middle, last, multi, only) > 0 + and most == _ENV_REQS - 1 + ) + assert nn == {1, 6} and prior == {True, False} and t_max == _ENV_TOKENS + assert ( + min(engine_sps) == 2 + and max(engine_sps) == _ENVELOPE_LM_1M - 2 + and any(sp < _REPLAY_WINDOW for sp in engine_sps) + ) + assert {s.W for s in steps} == set(_ENVELOPE_WIDTHS) + assert len(steps) - len(_zero_row_corners()) >= 200 + + +def _wrong_start_in_bounds(step: _EnvStep) -> bool: + """Whether the control call with start_pos = reused - R for the zero-row requests stays + inside kv_score and the output buffer (incl. its guard row): each such request then + pools min(R / 2, grid rows) windows from rows cu_seq_lens[b].. into rows cu_kv_comp[b].. + Steps where it would not are left out of that control rather than run out of bounds.""" + cu, cu_out = [0], [0] + for sp, n in step.ctx: + cu.append(cu[-1] + n) + cu_out.append(cu_out[-1] + (sp + n) // RATIO - sp // RATIO) + rows, out_rows, grid = ( + cu[-1] + step.gen_rows, + cu_out[-1] + step.gen_rows + 1, + max(step.max_outputs, 1), + ) + for b, (sp, n) in enumerate(step.ctx): + k = min(min(_REPLAY_WINDOW, sp) // RATIO, grid) + if n == 0 and (sp % RATIO or cu[b] + RATIO * k > rows or cu_out[b] + k > out_rows): + return False + return True + + +def test_zero_row_controls(): + """The zero-row checks see the call-site defects that would matter: (a) start_pos taken + from cached_token_lens (the replay start, reused - R) instead of the compute start for + the zero-row requests -- the op then pools R rows it was not given -- fails every step + it is run on (the engine-form steps where that wrong call stays in bounds); (b) the + block tables packed without the zero-row requests' rows (moved to the end) fails exactly + the steps where a zero-row request precedes a non-empty one, and passes the others.""" + pool_t = torch.empty(_ENV_PAGES, 128, HEAD_DIM, device=DEV) + steps = _zero_row_steps() + wrong_start = [s for s in steps if _wrong_start_in_bounds(s)] + for step in wrong_start: + with pytest.raises(AssertionError): + _run_env_step( + step, + pool_t, + op_starts=lambda ctx: [ + sp - min(_REPLAY_WINDOW, sp) if n == 0 else sp for sp, n in ctx + ], + ) + assert len(wrong_start) >= 20, len(wrong_start) + corners = _zero_row_corners() + caught = passed = 0 + for step in corners + [_zero_row_step(i) for i in range(0, _ZR_PARTITIONS, 7)]: + zs = [b for b, (_, n) in enumerate(step.ctx) if n == 0] + keep = [b for b in range(step.k) if b not in zs] + order = torch.tensor(keep + zs + list(range(step.k, _ENV_REQS)), device=DEV) + precedes = any(z < max(keep) for z in zs) if keep else False + try: + _run_env_step( + step, + pool_t, + op_tables=lambda kv, sc: (kv[order].contiguous(), sc[order].contiguous()), + ) + ok = True + except AssertionError: + ok = False + assert ok != precedes, (step.ctx, precedes) + caught += int(not ok) + passed += int(ok) + print( + f"[zero-row controls] start_pos = cached tokens: {len(wrong_start)} of {len(wrong_start)} steps rejected " + f"({len(steps) - len(wrong_start)} steps where that call would leave its buffers not run); " + f"tables packed without the zero-row rows: {caught} rejected (a zero-row request before a non-empty one), " + f"{passed} unchanged (none before)" + ) + assert caught > 0 and passed > 0 + + +def test_zero_row_kernel_launches(): + """The zero-row corners launch what every envelope step launches: one + prefillReductionKernel<512, 4, 4, 2, 1> per call, 128 threads, grid (num_contexts, + max(max_outputs, 1), 1) -- a zero-row request only adds a grid row -- and nothing else.""" + import json + import os + import tempfile + + pool_t = torch.empty(_ENV_PAGES, 128, HEAD_DIM, device=DEV) + zero_ape = _zero_ape() + launches, expected = [], [] + for step in _zero_row_corners()[::2]: + r = _run_env_step(step, pool_t) # checked against the reference + env = r["env"] + rows = r["kv_score"] + lens = [n for _, n in step.ctx] + outs = [(s + n) // 2 - s // 2 for s, n in step.ctx] + args = ( + _i32_tensor([s + n for s, n in step.ctx]), + _i32_tensor([s for s, _ in step.ctx]), + _i32_tensor([0] + torch.tensor(lens).cumsum(0).tolist()), + _i32_tensor([0] + torch.tensor(outs).cumsum(0).tolist()), + ) + for dt in (torch.bfloat16, torch.float32): + out = torch.empty(sum(outs) + step.gen_rows, HEAD_DIM, device=DEV, dtype=dt) + launches.append( + ( + rows, + env.bt_kv[: step.k], + env.bt_sc[: step.k], + out, + args, + step.k, + step.max_outputs, + ) + ) + expected.append([step.k, max(step.max_outputs, 1), 1]) + torch.cuda.synchronize() + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for rows, bkv, bsc, out, (kvl, spos, cus, cuo), k, mo in launches: + op( + rows, + zero_ape, + pool_t, + pool_t, + bkv, + bsc, + out, + kvl, + spos, + cus, + cuo, + k, + 128, + HEAD_DIM, + RATIO, + mo, + ) + torch.cuda.synchronize() + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "trace.json") + prof.export_chrome_trace(path) + trace = json.load(open(path)) + kernels = [e for e in trace["traceEvents"] if e.get("cat") == "kernel"] + ops = sorted((e for e in kernels if "prefillReduction" in e["name"]), key=lambda e: e["ts"]) + names = {e["name"] for e in ops} + grids = [e["args"]["grid"] for e in ops] + blocks = {tuple(e["args"]["block"]) for e in ops} + print( + f"[zero-row launches] {len(ops)} launches over {len(launches) // 2} corner steps x 2 dtypes: {sorted(names)}, " + f"blocks {sorted(blocks)}, grids {grids[::2]}" + ) + assert len(kernels) == len(launches) == len(ops) and len(names) == 1 + assert "prefillReductionKernel<512,4,4,2,1>" in next(iter(names)).replace(" ", "") + assert blocks == {(128, 1, 1)} and grids == expected diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_deepgemm_fp8_fp4_mqa_logits.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_deepgemm_fp8_fp4_mqa_logits.py new file mode 100644 index 000000000000..9afb9c62127e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_deepgemm_fp8_fp4_mqa_logits.py @@ -0,0 +1,1949 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the deepgemm_fp8_fp4_mqa_logits catalog entry. + +The entry owns what is certified: `CELLS` is the list, `reference` and `compare` are the gate. +This file turns a cell's spec into tensors and keeps the checks a cell cannot carry: the exact +regime's bit-identity premise, the output stride and allocation, rows past 2^31 bytes, determinism, +empty ranges and extents, CUDA-graph replay, the loud rejections, the wrong-variant controls, the +row and key axes as ranges, the engine envelope's partitions and the row-tiled long-context form. + +Two input regimes, one reference: + +* **exact** -- random packed bytes (E2M1 codes, UE8M0 scale bytes near 127, weights ``n * 2^-10``) + chosen so every head term ``relu(q.k) * w`` and every partial sum of them is exactly + representable in fp32. The op must then be bit-identical to the reference in any accumulation + order; each case asserts that premise before it asserts equality. +* **recipe** -- bf16 values packed with the kernel's own recipe (DeepGEMM ``per_token_cast_to_fp4``, + gran_k 32, packed UE8M0) and fp32 ``randn`` weights; compared at default fp32 tolerance. + +The reference decodes the packed bytes with torch ops (even element in the low nibble, scale byte +``b`` of the int32 word = block ``b``) and computes ``sum_h relu(q_h . k) * w_h`` in float64. The +entry's ``reference`` is that computation, and the cells, ``_check_exact``, the recipe sweeps and +the controls gate through ``op.reference`` / ``op.compare``. ``_reference`` below is its float64 +twin on already-decoded operands: it also returns the visibility mask and the exact-regime premise +bound, takes the wrong-variant flags the controls need, and serves the row-sliced checks (rows +past 2^31 bytes, the engine-envelope partitions); ``_long_context_reference`` decodes keys in chunks for the +2^25-key long-context form. +""" + +import collections +import math +import random +import zlib + +import pytest +import torch +from torch.profiler import ProfilerActivity, profile + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.deepgemm_fp8_fp4_mqa_logits import ( # noqa: E501 + deepgemm_fp8_fp4_mqa_logits as op, +) +from tensorrt_llm.deep_gemm.utils import per_token_cast_to_fp4 + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "deepgemm_fp8_fp4_mqa_logits requires a CUDA device" + +DEV = "cuda" +NEG_INF = float("-inf") +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) +# V4.1 weights fold: weights_proj(x) * index_head_dim^-0.5 * index_n_heads^-0.5, with the GLOBAL +# head count (32) even when a rank holds fewer heads. +GLOBAL_FOLD = 128**-0.5 * 32**-0.5 # == 2^-6 exactly +LOCAL_FOLD_TP4 = 128**-0.5 * 8**-0.5 + + +# ----------------------------------------------------------------------------- reference helpers + + +def _decode(packed, sf, swap_nibbles=False, reverse_sf_bytes=False): + """int8 [..., 64] + int32 [...] -> float64 [..., 128].""" + lut = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + p = packed.view(torch.uint8).long() + lo, hi = lut[p & 0xF], lut[p >> 4] + if swap_nibbles: + lo, hi = hi, lo + v = torch.stack([lo, hi], dim=-1).flatten(-2) + e = sf.contiguous().view(torch.uint8).view(*sf.shape, 4).long() - 127 + if reverse_sf_bytes: + e = e.flip(-1) + scale = ((e + 1023) << 52).view(torch.float64) # 2^e built from the exponent bits: exact + return (v.unflatten(-1, (4, 32)) * scale.unsqueeze(-1)).flatten(-2) + + +def _reference(qd, kd, w, ks, ke, relu_after_sum=False): + """float64 logits [T, S] (-inf outside [ks, ke)), the visibility mask, and max_t,s sum_h |term|.""" + T, S = qd.shape[0], kd.shape[0] + out = torch.empty(T, S, dtype=torch.float64, device=DEV) + abs_sum_max = 0.0 + wd = w.double() + for t0 in range(0, T, 512): + dots = torch.einsum("thd,kd->thk", qd[t0 : t0 + 512], kd) + wt = wd[t0 : t0 + 512].unsqueeze(-1) + if relu_after_sum: + out[t0 : t0 + 512] = (dots * wt).sum(1).clamp_min(0) + else: + terms = dots.clamp_min(0) * wt + out[t0 : t0 + 512] = terms.sum(1) + if S > 0: + abs_sum_max = max(abs_sum_max, float(terms.abs().sum(1).max())) + col = torch.arange(S, device=DEV) + vis = (col[None, :] >= ks[:, None].long()) & (col[None, :] < ke[:, None].long()) + return out.masked_fill(~vis, NEG_INF), vis, abs_sum_max + + +# ----------------------------------------------------------------------------- input builders + + +def _gen(seed): + return torch.Generator(device=DEV).manual_seed(seed) + + +def _rand_packed(rows, g, exp_lo=126, exp_hi=128): + """Random E2M1 byte pairs (every code, both signs) + random UE8M0 bytes in [exp_lo, exp_hi].""" + codes = torch.randint(0, 256, (rows, 64), generator=g, device=DEV, dtype=torch.int32) + sfb = torch.randint(exp_lo, exp_hi + 1, (rows, 4), generator=g, device=DEV, dtype=torch.int32) + return codes.to(torch.uint8).view(torch.int8), sfb.to(torch.uint8).view(torch.int32).view(rows) + + +def _exact_inputs(T, S, H, seed): + """Exact regime: |dot| < 2^15 in units of 2^-4, w = n * 2^-10 with |n| <= 8 (<= 3 significant bits).""" + g = _gen(seed) + q, q_sf = _rand_packed(T * H, g) + kv, kv_sf = _rand_packed(S, g) + w_raw = torch.randint(-8, 9, (T, H), generator=g, device=DEV).float() * 2.0**-4 + return q.view(T, H, 64), q_sf.view(T, H), kv, kv_sf, w_raw + + +def _assert_exact_premise(abs_sum_max): + # Terms are multiples of 2^-14 (dot unit 2^-4 times weight unit 2^-10), each <= 22 significant + # bits. If sum_h |term| < 2^10 every partial sum in any order is an fp32-exact multiple of 2^-14. + assert abs_sum_max < 2.0**10, f"exact-regime premise broken: max sum|term| = {abs_sum_max}" + + +def _causal(T, r, past=0, k_off=0): + pos = torch.arange(past, past + T, device=DEV) + ks = torch.full((T,), k_off, dtype=torch.int32, device=DEV) + ke = (k_off + (pos + 1) // r).to(torch.int32) + return ks, ke + + +def _call(q, q_sf, kv, kv_sf, w, ks, ke, clean=True, max_seqlen_k=0, logits_dtype=torch.float32): + return op(q, q_sf, kv, kv_sf, w, ks, ke, clean, max_seqlen_k, logits_dtype) + + +def _row_stride(S, logits_dtype=torch.float32, block_kv=256): + align = 1024 // torch.empty((), dtype=logits_dtype).element_size() + return -(-(S + block_kv) // align) * align + + +def _check_exact(out, q, q_sf, kv, kv_sf, w, ks, ke): + """The entry's gate, then bit-identity: the exact regime's premise makes the op exact.""" + _, vis, bound = _reference(_decode(q, q_sf), _decode(kv, kv_sf), w, ks, ke) + _assert_exact_premise(bound) + ref = op.reference(q, q_sf, kv, kv_sf, w, ks, ke, True, 0, torch.float32) + assert out.dtype == torch.float32 and out.shape == ref.shape + op.compare(out, ref) + assert torch.equal(out, ref), f"{int((out != ref).sum())} of {out.numel()} elements differ" + return ref, vis + + +# ----------------------------------------------------------------------------- exact regime + + +def _packed(reqs, r): + """(past, new) per request -> per-query ks/ke with each request's keys at its own offset, and S.""" + ks_l, ke_l, S = [], [], 0 + for past, n in reqs: + ks_r, ke_r = _causal(n, r, past=past, k_off=S) + ks_l.append(ks_r) + ke_l.append(ke_r) + S += (past + n) // r + return torch.cat(ks_l), torch.cat(ke_l), S + + +def _logits_bytes(T, S, H): + """The call's one allocation: rows rounded up to the 128 / H query block, row stride align(S+256, 256).""" + return -(-T // (128 // H)) * (128 // H) * _row_stride(S) * 4 + + +def test_exact_output_past_2gib(): + """T 16384 x S 32768: rows from 16257 on start past 2^31 bytes of the logits buffer (2.16 GB). + The first and last 512 rows are checked; the rest is the same kernel at smaller offsets.""" + H, T, r = 32, 16384, 1 + ks, ke, S = _packed([(16384, T)], r) + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, H, seed=99) + w = w_raw * GLOBAL_FOLD + out = _call(q, q_sf, kv, kv_sf, w, ks, ke) + assert out.untyped_storage().nbytes() == _logits_bytes(T, S, H) > 2**31 + kd = _decode(kv, kv_sf) + for rows in (slice(0, 512), slice(T - 512, T)): + ref, _, bound = _reference(_decode(q[rows], q_sf[rows]), kd, w[rows], ks[rows], ke[rows]) + _assert_exact_premise(bound) + assert torch.equal(out[rows], ref.float()) + + +# ----------------------------------------------------------------------------- recipe regime + + +def _recipe_inputs(T, S, H, seed): + g = _gen(seed) + q = torch.randn(T * H, 128, generator=g, device=DEV).bfloat16() + k = torch.randn(S, 128, generator=g, device=DEV).bfloat16() + q_p, q_s = per_token_cast_to_fp4(q, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + k_p, k_s = per_token_cast_to_fp4(k, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + w = torch.randn(T, H, generator=g, device=DEV) * GLOBAL_FOLD + return q_p.view(T, H, 64), q_s.view(T, H), k_p, k_s.view(S), w + + +def _wide_scale_inputs(T, S, H, seed): + """UE8M0 bytes across [90, 160]; per-token / per-key exponents so fp32 stays in range.""" + g = _gen(seed) + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, H, seed=seed) + tq = torch.randint(90, 159, (T, 1, 1), generator=g, device=DEV, dtype=torch.int32) + tk = torch.randint(90, 159, (S, 1), generator=g, device=DEV, dtype=torch.int32) + jq = torch.randint(0, 3, (T, H, 4), generator=g, device=DEV, dtype=torch.int32) + jk = torch.randint(0, 3, (S, 4), generator=g, device=DEV, dtype=torch.int32) + q_sf = (tq + jq).to(torch.uint8).view(torch.int32).view(T, H) + kv_sf = (tk + jk).to(torch.uint8).view(torch.int32).view(S) + return q, q_sf, kv, kv_sf, w_raw * GLOBAL_FOLD + + +def _cell_inputs(spec, seed): + """A cell's tensors and output form -> (q, q_sf, kv, kv_sf, w, ks, ke, clean, max_seqlen_k).""" + H = spec["H"] + clean, msk = True, 0 + if "reqs" in spec: + # per-request (cached, new), each request's keys at its own offset, causal inside + ks, ke, S = _packed(spec["reqs"], spec["r"]) + T = ks.numel() + else: + T, S = spec["T"], spec["S"] + t = torch.arange(T, device=DEV) + if spec["ranges"] == "shifted": # clean_logits=False over shifted per-row ranges + ks = (t % 17).to(torch.int32) + ke = (ks + 1 + t * 2 % (S - 17)).to(torch.int32) + clean = False + else: # "compressed": shifted ranges including empty rows + ks = (t * 3).to(torch.int32) + ke = (ks + (t % 29) + (t % 2) * 300).clamp(max=S).to(torch.int32) + clean, msk = False, int((ke - ks).max()) + if spec["regime"] == "recipe": + q, q_sf, kv, kv_sf, w = _recipe_inputs(T, S, H, seed) + elif spec["regime"] == "wide_scale": + q, q_sf, kv, kv_sf, w = _wide_scale_inputs(T, S, H, seed) + else: + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, H, seed) + w = w_raw * GLOBAL_FOLD + return q, q_sf, kv, kv_sf, w, ks, ke, clean, msk + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + H = spec["H"] + q, q_sf, kv, kv_sf, w, ks, ke, clean, msk = _cell_inputs(spec, zlib.crc32(cell.why.encode())) + T, S = q.shape[0], kv.shape[0] + if spec.get("alloc_check"): + # The first call in a process also allocates a one-time 32 MiB buffer that stays allocated + # (measured); a warm-up call keeps the measurement below independent of test order. + op(q[:1], q_sf[:1], kv, kv_sf, w[:1], ks[:1], ke[:1], clean, msk, torch.float32) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + before = torch.cuda.memory_allocated() + out = op(q, q_sf, kv, kv_sf, w, ks, ke, clean, msk, torch.float32) + torch.cuda.synchronize() + if spec.get("alloc_check"): + # the call allocates exactly the returned logits buffer, no other workspace + assert out.untyped_storage().nbytes() == _logits_bytes(T, S, H) + assert torch.cuda.max_memory_allocated() - before == _logits_bytes(T, S, H) + assert out.dtype == torch.float32 and out.shape == (T, msk or S) + assert out.stride() == ((-(-msk // 256) * 256 if msk else _row_stride(S)), 1) + ref = op.reference(q, q_sf, kv, kv_sf, w, ks, ke, clean, msk, torch.float32) + op.compare(out, ref) + if spec["regime"] == "exact": + # Every head term and partial sum is an fp32-exact multiple of 2^-14: the op is + # bit-identical to the reference on every element it specifies. + _assert_exact_premise(_reference(_decode(q, q_sf), _decode(kv, kv_sf), w, ks, ke)[2]) + specified = ~torch.isnan(ref) + assert torch.equal(out[specified], ref[specified]) + + +def test_deterministic(): + q, q_sf, kv, kv_sf, w = _recipe_inputs(1024, 2048, 32, seed=3) + ks, ke = _causal(1024, 2, past=3072) + first = _call(q, q_sf, kv, kv_sf, w, ks, ke) + for _ in range(10): + assert torch.equal(_call(q, q_sf, kv, kv_sf, w, ks, ke), first) + + +# ----------------------------------------------------------------------------- output forms + + +def test_compressed_form_rejects_clean_logits(): + T, S, H = 200, 1000, 32 + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, H, seed=6) + w = w_raw * GLOBAL_FOLD + t = torch.arange(T, device=DEV) + ks = (t * 3).to(torch.int32) + ke = (ks + (t % 29) + (t % 2) * 300).clamp(max=S).to(torch.int32) + with pytest.raises(RuntimeError, match="not clean_logits"): + _call(q, q_sf, kv, kv_sf, w, ks, ke, clean=True, max_seqlen_k=int((ke - ks).max())) + + +def test_empty_rows_and_empty_extents(): + T, S, H = 16, 16, 32 + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, H, seed=7) + w = w_raw * GLOBAL_FOLD + ks = torch.full((T,), 4, dtype=torch.int32, device=DEV) + ke = torch.where(torch.arange(T, device=DEV) % 2 == 0, 4, 9).to(torch.int32) + out = _call(q, q_sf, kv, kv_sf, w, ks, ke) + _check_exact(out, q, q_sf, kv, kv_sf, w, ks, ke) + assert bool(torch.isneginf(out[0::2]).all()) + # T == 0 or S == 0 cannot build the kernel's TMA descriptors; the error is not sticky. + with pytest.raises(RuntimeError, match="CUDA_ERROR_INVALID_VALUE"): + _call(q[:0], q_sf[:0], kv, kv_sf, w[:0], ks[:0], ke[:0]) + with pytest.raises(RuntimeError, match="CUDA_ERROR_INVALID_VALUE"): + _call(q, q_sf, kv[:0], kv_sf[:0], w, torch.zeros_like(ks), torch.zeros_like(ke)) + out = _call(q, q_sf, kv, kv_sf, w, ks, ke) + torch.cuda.synchronize() + _check_exact(out, q, q_sf, kv, kv_sf, w, ks, ke) + + +def test_cuda_graph_replay(): + T, S, H = 256, 512, 32 + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, H, seed=8) + w = w_raw * GLOBAL_FOLD + ks, ke = _causal(T, 2, past=S * 2 - T) + _call(q, q_sf, kv, kv_sf, w, ks, ke) # JIT-compile outside the capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = _call(q, q_sf, kv, kv_sf, w, ks, ke) + for seed in (9, 10): + q2, q_sf2, kv2, kv_sf2, w_raw2 = _exact_inputs(T, S, H, seed=seed) + ke2 = (ke - seed).clamp(min=0) + for dst, src in ( + (q, q2), + (q_sf, q_sf2), + (kv, kv2), + (kv_sf, kv_sf2), + (w, w_raw2 * GLOBAL_FOLD), + (ke, ke2), + ): + dst.copy_(src) + graph.replay() + torch.cuda.synchronize() + _check_exact(out, q2, q_sf2, kv2, kv_sf2, w_raw2 * GLOBAL_FOLD, ks, ke2) + + +# ----------------------------------------------------------------------------- loud rejections + + +def test_rejections(): + T, S, H = 8, 16, 32 + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, H, seed=12) + w = w_raw * GLOBAL_FOLD + ks, ke = _causal(T, 1, past=S - T) + bad = { + "q_fp.scalar_type()": dict(q=q.view(torch.uint8)), + "head_dim == _head_dim": dict(kv=kv.view(torch.uint8)), + "t.dim() == N": dict(kv_sf=kv_sf.view(S, 1)), + "seq_len == _seq_len and num_heads == _num_heads": dict(q_sf=q_sf.view(T * H, 1)), + "q_fp.is_contiguous()": dict(q=torch.cat([q, q], 2)[:, :, :64]), + "kv_fp.is_contiguous()": dict(kv=torch.cat([kv, kv], 1)[:, :64]), + "cu_seq_len_k_start.scalar_type()": dict(ks=ks.long()), + "weights.scalar_type() != torch::kBFloat16 or logits_dtype == torch::kBFloat16": dict( + w=w.bfloat16() + ), + } + args = dict(q=q, q_sf=q_sf, kv=kv, kv_sf=kv_sf, w=w, ks=ks, ke=ke) + for msg, override in bad.items(): + a = {**args, **override} + with pytest.raises(RuntimeError, match=msg.replace("(", r"\(").replace(")", r"\)")): + _call(a["q"], a["q_sf"], a["kv"], a["kv_sf"], a["w"], a["ks"], a["ke"]) + # Head counts: any multiple of 4 up to 128 is taken (24 and 128 are measured against the + # reference in test_head_counts_off_the_cells_match_the_reference); 6 is refused. + q2, q_sf2, _, _, w2 = _exact_inputs(T, S, 6, seed=13) + with pytest.raises(RuntimeError, match=r"num_heads % 4 == 0"): + _call(q2, q_sf2, kv, kv_sf, w2, ks, ke) + + +@pytest.mark.parametrize("heads", [24, 128]) +def test_head_counts_off_the_cells_match_the_reference(heads): + """The kernel accepts any head count that is a multiple of 4 up to 128 (128 / H query rows per + block). Two such counts outside the cells, one that does not divide 128 and the largest, are + held to the entry's gate on a causal ratio-1 step. Not a certification of these counts: a + measurement that the acceptance is not a silent wrong answer.""" + T, S = 64, 256 + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, heads, seed=15 + heads) + w = w_raw * GLOBAL_FOLD + ks, ke = _causal(T, 1, past=S - T) + out = _call(q, q_sf, kv, kv_sf, w, ks, ke) + op.compare(out, op.reference(q, q_sf, kv, kv_sf, w, ks, ke, True, 0, torch.float32)) + + +# ----------------------------------------------------------------------------- wrong-variant controls + + +def test_wrong_variants_fail_the_gate(): + """The entry's gate catches each of these mistakes (counts printed for the receipt notes).""" + T, S, H, r = 512, 256, 32, 2 + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, H, seed=14) + w = w_raw * GLOBAL_FOLD + ks, ke = _causal(T, r) + out = _call(q, q_sf, kv, kv_sf, w, ks, ke) + qd, kd = _decode(q, q_sf), _decode(kv, kv_sf) + _, vis, _ = _reference(qd, kd, w, ks, ke) + good = op.reference(q, q_sf, kv, kv_sf, w, ks, ke, True, 0, torch.float32) + op.compare(out, good) + assert torch.equal(out, good) + controls = { + "relu after the head sum": _reference(qd, kd, w, ks, ke, relu_after_sum=True)[0], + "local-head weight fold (8 heads)": _reference(qd, kd, w_raw * LOCAL_FOLD_TP4, ks, ke)[0], + # Swapping the nibble order of BOTH operands permutes both vectors alike and leaves every dot + # product unchanged (measured: 0 differences), so the control swaps k only. + "k nibbles swapped": _reference(qd, _decode(kv, kv_sf, swap_nibbles=True), w, ks, ke)[0], + "scale bytes reversed": _reference( + _decode(q, q_sf, reverse_sf_bytes=True), + _decode(kv, kv_sf, reverse_sf_bytes=True), + w, + ks, + ke, + )[0], + "causal end + 1": _reference(qd, kd, w, ks, ke + 1)[0], + } + for name, wrong in controls.items(): + wrong32 = wrong.float() + n_diff = int((out != wrong32).sum()) + n_tol = int( + ( + ~torch.isclose(out, wrong32, rtol=1.3e-6, atol=1e-5, equal_nan=False) + & ~(torch.isneginf(out) & torch.isneginf(wrong32)) + ).sum() + ) + print( + f"control [{name}]: {n_diff} of {out.numel()} elements differ bitwise, " + f"{n_tol} outside default fp32 tolerance ({int(vis.sum())} visible)" + ) + assert n_tol > 0 + with pytest.raises(AssertionError): + op.compare(out, wrong32) + + +def test_wrong_variants_fail_the_gate_long_context(): + """At S 32576 (packed, r = 1) the gate sees the mistakes a long key range can hide: a kernel that + stops at the previously certified S 16384, one whose key row index wraps at 2^14, and one that + ignores the per-request start. Checked on the last 1024 query rows (all in the third request, + whose keys lie at rows 23480 .. 32575).""" + H, r = 32, 1 + ks, ke, S = _packed([(16384, 4096), (0, 3000), (8000, 1096)], r) + T = ks.numel() + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, H, seed=15) + w = w_raw * GLOBAL_FOLD + out = _call(q, q_sf, kv, kv_sf, w, ks, ke) + rows = slice(T - 1024, T) + qd, kd = _decode(q[rows], q_sf[rows]), _decode(kv, kv_sf) + wr, ksr, ker = w[rows], ks[rows], ke[rows] + _, vis, bound = _reference(qd, kd, wr, ksr, ker) + _assert_exact_premise(bound) + got = out[rows] + good = op.reference(q[rows], q_sf[rows], kv, kv_sf, wr, ksr, ker, True, 0, torch.float32) + op.compare(got, good) + assert torch.equal(got, good) + wrapped = kd[torch.arange(S, device=DEV) % 2**14] + controls = { + "keys past 16384 dropped": _reference(qd, kd, wr, ksr, ker.clamp(max=16384))[0], + "key row index wrapped at 2^14": _reference(qd, wrapped, wr, ksr, ker)[0], + "per-request start ignored (ks = 0)": _reference(qd, kd, wr, torch.zeros_like(ksr), ker)[0], + } + for name, wrong in controls.items(): + wrong32 = wrong.float() + n_tol = int( + ( + ~torch.isclose(got, wrong32, rtol=1.3e-6, atol=1e-5) + & ~(torch.isneginf(got) & torch.isneginf(wrong32)) + ).sum() + ) + print( + f"control [{name}]: {n_tol} of {got.numel()} elements outside default fp32 tolerance " + f"({int(vis.sum())} visible)" + ) + assert n_tol > 0 + with pytest.raises(AssertionError): + op.compare(got, wrong32) + + +# ----------------------------------------------------------------------------- row and key axes +# +# The target passes any context chunk T in 1..8192 (max_num_tokens) over any compressed key count S in +# 1..32768. These tests certify both axes as ranges, not as sampled points: a maximal run against the +# float64 reference, then every smaller extent bitwise equal to the matching slice of that run, plus +# one compiled kernel pair serving every extent (profiler kernel names). +# +# Harness sighting. Each sweep alternates two input sets with the same ranges (A for even extents, B for +# odd), walks extents downwards so each result fits in the block the previous one freed, and fills every +# returned buffer with NaN before freeing it. An element the op left unwritten therefore reads NaN or the +# other set's value, never a stale copy of the right one; each sweep asserts that most results landed on +# a NaN-filled block, and runs a one-row (or one-key) shift control that must fail. + +_HEADS = 32 +_S_MAX = 32768 +# The one kernel a full-width (clean_logits=True) call launches at every extent: the logits kernel, +# which also writes -inf to the columns outside each row's range (no separate clean pass). +_LOGITS_KERNEL = "void deep_gemm::sm100_mqa_logits<32u, 128u, true, false, true, false, 4u, 256u," +# Every S = 0, 1, 2, 3 or -1 mod the kernel's 256-key tile up to 32768. That includes S = 0, 1, -1 mod +# 8192 columns and every residue mod 4 (the key-scale TMA extent is rounded to 4). +_S_EDGES = sorted( + {k * 256 + d for k in range(129) for d in (-1, 0, 1, 2, 3) if 1 <= k * 256 + d <= _S_MAX} +) + + +def _chunk(T, S, r): + """One request's context chunk whose last row sees key S - 1: ks = 0, ke = (p + 1) // r for the + absolute positions p = S*r - T .. S*r - 1, clamped to [0, S] (a row before the first compressed key + sees none). This is the target's single-request call form.""" + pos = torch.arange(S * r - T, S * r, device=DEV) + ke = torch.div(pos + 1, r, rounding_mode="floor").clamp(0, S).to(torch.int32) + return torch.zeros(T, dtype=torch.int32, device=DEV), ke + + +def _packed3(T, S): + """Three requests packed as the target packs a context step: each request's keys at its own offset + (the second request starts at an odd key), causal inside each request, rows before a request's first + key see none.""" + rows = (T // 2, T // 4 + 1, T - T // 2 - T // 4 - 1) + s0 = min(S, (S * 5 // 8) | 1) + s1 = (S - s0) // 3 + ks_l, ke_l, off = [], [], 0 + for n, s in zip(rows, (s0, s1, S - s0 - s1)): + t = torch.arange(n, device=DEV) + ks_l.append(torch.full((n,), off, dtype=torch.int32, device=DEV)) + ke_l.append((off + (s - n + t + 1).clamp(0, s)).to(torch.int32)) + off += s + return torch.cat(ks_l), torch.cat(ke_l) + + +def _axis_ranges(form, T, S): + return _packed3(T, S) if form == "packed" else _chunk(T, S, int(form[1])) + + +def _axis_full(regime, T, S, ks, ke, seed): + """Inputs of one set plus the maximal run, checked against the float64 reference.""" + if regime == "exact": + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, _HEADS, seed) + w = w_raw * GLOBAL_FOLD + else: + q, q_sf, kv, kv_sf, w = _recipe_inputs(T, S, _HEADS, seed) + full = _call(q, q_sf, kv, kv_sf, w, ks, ke) + if regime == "exact": + _check_exact(full, q, q_sf, kv, kv_sf, w, ks, ke) + else: + op.compare(full, op.reference(q, q_sf, kv, kv_sf, w, ks, ke, True, 0, torch.float32)) + return q, q_sf, kv, kv_sf, w, full + + +def _poison(out): + """NaN over the whole storage of a returned logits buffer (the padded rows and columns included).""" + out.as_strided((out.untyped_storage().nbytes() // 4,), (1,), 0).fill_(float("nan")) + + +@pytest.mark.parametrize( + "regime,form", + [("exact", "r1"), ("exact", "r2"), ("exact", "packed"), ("recipe", "packed")], +) +@pytest.mark.parametrize("T_max,S", [(8192, 16384), (4096, 32768)]) +def test_row_axis_every_T(T_max, S, regime, form): + """Every T in 1..T_max: the first T rows of q, q_sf, weights, ks and ke over the same keys give + exactly the first T rows of the T_max call, which itself matches float64 (exact regime: bit-identical; + recipe regime: default fp32 tolerance, and the prefix rows are still bitwise equal).""" + ks, ke = _axis_ranges(form, T_max, S) + sets = [ + _axis_full(regime, T_max, S, ks, ke, seed=T_max + S + 7 * k + len(form)) for k in (0, 1) + ] + torch.cuda.synchronize() + bad = torch.zeros((), dtype=torch.int64, device=DEV) + shift = {} + poisoned, reused = set(), 0 + for T in range(T_max, 0, -1): + q, q_sf, kv, kv_sf, w, full = sets[T % 2] + out = _call(q[:T], q_sf[:T], kv, kv_sf, w[:T], ks[:T], ke[:T]) + assert out.shape == (T, S) and out.stride() == (_row_stride(S), 1) + reused += out.data_ptr() in poisoned + bad += (out != full[:T]).sum() + if T in (1, 5, 128, 129, T_max // 2 + 1, T_max - 1): + shift[T] = (out != full[1 : T + 1]).sum() # control: the gate must see a one-row shift + _poison(out) + poisoned.add(out.data_ptr()) + del out + shift = {T: int(n) for T, n in shift.items()} + print( + f"row axis {regime}/{form} T 1..{T_max} x S {S}: {int(bad)} mismatching elements; " + f"{reused}/{T_max} results on a NaN-filled block; one-row-shift control mismatches {shift}" + ) + assert int(bad) == 0 + assert reused >= T_max // 2 + assert all(n > 0 for n in shift.values()) + + +@pytest.mark.parametrize( + "regime,form,T,every_S", + [ + ("exact", "r1", 8192, False), + ("exact", "r2", 8192, False), + ("exact", "packed", 8192, False), + ("recipe", "packed", 8192, False), + ("exact", "r1", 67, True), + ("exact", "r2", 67, True), + ("exact", "packed", 67, True), + ], +) +def test_key_axis(regime, form, T, every_S): + """At fixed T, S in _S_EDGES (T 8192) or every S in 1..32768 (T 67, which is not a multiple of the + 4-row query block): the first S key rows with the ranges of this form at S give, on every visible + element, exactly the T x 32768 call in which every row sees every key, and -inf everywhere else.""" + ks0 = torch.zeros(T, dtype=torch.int32, device=DEV) + ke0 = torch.full((T,), _S_MAX, dtype=torch.int32, device=DEV) + sets = [_axis_full(regime, T, _S_MAX, ks0, ke0, seed=T + 3 * k + len(form)) for k in (0, 1)] + torch.cuda.synchronize() + s_values = range(_S_MAX, 0, -1) if every_S else _S_EDGES[::-1] + bad = torch.zeros((), dtype=torch.int64, device=DEV) + n_vis = torch.zeros((), dtype=torch.int64, device=DEV) + shift = {} + poisoned, reused = set(), 0 + for S in s_values: + q, q_sf, kv, kv_sf, w, full = sets[S % 2] + ks, ke = _axis_ranges(form, T, S) + out = _call(q, q_sf, kv[:S], kv_sf[:S], w, ks, ke) + assert out.shape == (T, S) and out.stride() == (_row_stride(S), 1) + reused += out.data_ptr() in poisoned + col = torch.arange(S, device=DEV) + vis = (col[None, :] >= ks[:, None]) & (col[None, :] < ke[:, None]) + bad += ((out != full[:, :S]) & vis).sum() + (~torch.isneginf(out) & ~vis).sum() + n_vis += vis.sum() + if S in (255, 256, 257, 8191, 8193, 32767): + shift[S] = ( + (out != full[:, 1 : S + 1]) & vis + ).sum() # control: a one-key shift must fail + _poison(out) + poisoned.add(out.data_ptr()) + del out + shift = {S: int(n) for S, n in shift.items()} + print( + f"key axis {regime}/{form} T {T}, {len(s_values)} values of S: {int(bad)} mismatching elements " + f"({int(n_vis)} visible checked); {reused}/{len(s_values)} results on a NaN-filled block; " + f"one-key-shift control mismatches {shift}" + ) + assert int(bad) == 0 + assert reused >= len(s_values) // 2 + assert all(n > 0 for n in shift.values()) + + +# Every (T, S, ratio) of the core target's runtime-value dump (boot and accuracy gates), and the 24K +# passkey prompt's three chunks (8192, 8192, 7674 tokens; the last one is in the dump) at both ratios. +_DUMPED = [ + (1, 1, 1), + (2, 1, 2), + (2, 2, 1), + (14, 8, 2), + (14, 16, 1), + (16, 8, 2), + (16, 16, 1), + (48, 24, 2), + (48, 48, 1), + (888, 1469, 2), + (888, 2938, 1), + (890, 1469, 2), + (890, 2938, 1), + (1016, 1469, 2), + (1016, 2938, 1), + (1934, 969, 2), + (1934, 1938, 1), + (2048, 1024, 2), + (2048, 1025, 2), + (2048, 2048, 1), + (2048, 2050, 1), + (8192, 8192, 1), + (8192, 4096, 2), + (8192, 16384, 1), + (8192, 8192, 2), + (7674, 24058, 1), + (7674, 12029, 2), +] + + +@pytest.mark.parametrize("T,S,r", _DUMPED) +def test_dumped_values(T, S, r): + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, _HEADS, seed=7 * T + S + r) + w = w_raw * GLOBAL_FOLD + ks, ke = _chunk(T, S, r) + out = _call(q, q_sf, kv, kv_sf, w, ks, ke) + _check_exact(out, q, q_sf, kv, kv_sf, w, ks, ke) + assert out.stride() == (_row_stride(S), 1) + + +def test_one_compiled_variant(): + """T and S are runtime arguments only: across T 1..8192 and S 1..32768, every form, one kernel runs + per call -- the logits kernel, which also cleans the out-of-range columns -- with the same template + arguments at every extent.""" + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(8192, _S_MAX, _HEADS, seed=41) + w = w_raw * GLOBAL_FOLD + ts = (1, 2, 3, 4, 5, 7, 8, 9, 67, 127, 128, 129, 1283, 4095, 4096, 4097, 7674, 8191, 8192) + ss = (1, 2, 3, 4, 5, 255, 256, 257, 969, 8191, 8192, 8193, 16384, 24058, 32767, 32768) + cases = [(T, S, _axis_ranges(f, T, S)) for T in ts for S in ss for f in ("r1", "r2", "packed")] + _call( + q[:4], q_sf[:4], kv, kv_sf, w[:4], *_chunk(4, _S_MAX, 1) + ) # JIT, if any, outside the profile + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for T, S, (ks, ke) in cases: + _call(q[:T], q_sf[:T], kv[:S], kv_sf[:S], w[:T], ks, ke) + torch.cuda.synchronize() + names = collections.Counter( + e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA + ) + for name, n in names.items(): + print(f"kernel x{n}: {name}") + assert set(names.values()) == {len(cases)}, names + assert len(names) == 1, names + (name,) = names + assert name.startswith(_LOGITS_KERNEL), name + + +# ----------------------------------------------------------------------------- multi-request steps: the engine envelope +# +# One engine step passes T <= 8192 context rows from k <= 32 context requests, each with cached + new <= 16384 +# tokens. The metadata (``_prepare_v41_ctx_index_k``) gathers every request's compressed keys, cached ones +# included, request after request, so S = sum_b kv_len_b // r reaches 524288 (r 1) / 262144 (r 2), and +# gives request b's token at position p the range [cu_b, cu_b + (p + 1) // r) with cu_b = sum_{j= 1 and c + n <= _ENV_LEN for c, n in reqs) + res = _run_partition(reqs, r, seed=5000 + 7 * i + r) + if res is None: + stats["skipped (S 0)"] += 1 + continue + T, S, n_vis = res + stats["partitions"] += 1 + stats["rows"] += T + stats["visible checked"] += n_vis + stats["max S"] = max(stats["max S"], S) + stats["max k"] = max(stats["max k"], len(reqs)) + print(f"random partitions block {block} r {r}: {dict(stats)}") + assert stats["partitions"] + stats["skipped (S 0)"] == per + + +def test_multi_request_wrong_ranges_are_seen(): + """The partition check must fail when one request's range is shifted by one key, or when two requests' + key blocks are swapped (the ranges built as if request 2's keys preceded request 1's).""" + reqs = [(3000, 700), (0, 257), (9000, 1043)] + r = 1 + ks, ke, S = _packed(reqs, r) + T = ks.numel() + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, _HEADS, seed=77) + w = w_raw * GLOBAL_FOLD + out = _call(q, q_sf, kv, kv_sf, w, ks, ke) + bounds = _bounds(reqs) + _check_partition(out, q, q_sf, kv, kv_sf, w, ks, ke, bounds) + shift = torch.zeros_like(ks) + shift[:700] = 1 + # Swapped: the same rows, but the ranges built as if request 2's keys preceded request 1's. + spans = [(c + n) // r for c, n in reqs] + offs = {0: 0, 2: spans[0], 1: spans[0] + spans[2]} + sw = [_causal(n, r, past=c, k_off=offs[b]) for b, (c, n) in enumerate(reqs)] + swapped_ks, swapped_ke = torch.cat([a for a, _ in sw]), torch.cat([e for _, e in sw]) + wrong = { + "request 0 range shifted by one key": (ks + shift, ke + shift), + "requests 1 and 2 key blocks swapped": (swapped_ks, swapped_ke), + } + for name, (wks, wke) in wrong.items(): + with pytest.raises(AssertionError): + _check_partition(out, q, q_sf, kv, kv_sf, w, wks, wke, bounds) + print(f"control [{name}]: the partition check fails") + + +def _maximal(T, k, r): + """k requests sharing T new rows as evenly as possible, each cached up to 16384 total.""" + new = [T // k + (i < T % k) for i in range(k)] + return [(_ENV_LEN - n, n) for n in new] + + +@pytest.mark.parametrize("r", [1, 2]) +def test_multi_request_maximal_cell(r): + """T 8192 in 32 requests of 16384 tokens each: S 524288 at r 1 (a 17.2 GB logits buffer whose last rows + start past element 2^32) and 262144 at r 2. Every request block checked in full.""" + T, S, n_vis = _run_partition(_maximal(_ENV_T, _ENV_K, r), r, seed=31 + r) + assert (T, S) == (_ENV_T, _ENV_K * _ENV_LEN // r) + assert _row_stride(S) * (T - 1) > 2**32 or r == 2 + print(f"maximal cell r {r}: T {T} x S {S}, {n_vis} visible elements checked") + + +def test_multi_request_kernel_variant_map(): + """Kernel names over the envelope: T in 1..8192, k in 1..32 requests, cached prefix 0 or up to 16384, + ratios 1 and 2, S up to 524288. Every call launches the same one kernel as the row and key axes' + variant map (``test_one_compiled_variant``).""" + ts, ks_ = (1, 5, 129, 4097, 8192), (1, 2, 17, 32) + cells = [] + for T in ts: + for k in ks_: + if k > T: + continue + new = [T // k + (i < T % k) for i in range(k)] + for cached in ("none", "full"): + reqs = [((_ENV_LEN - n) if cached == "full" else 0, n) for n in new] + for r in (1, 2): + ks, ke, S = _packed(reqs, r) + if S: + cells.append((T, S, ks, ke)) + S_max = max(S for _, S, _, _ in cells) + assert S_max == _ENV_K * _ENV_LEN + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(_ENV_T, S_max, _HEADS, seed=43) + w = w_raw * GLOBAL_FOLD + _call( + q[:4], q_sf[:4], kv, kv_sf, w[:4], *_chunk(4, S_max, 1) + ) # JIT, if any, outside the profile + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for T, S, ks, ke in cells: + out = _call(q[:T], q_sf[:T], kv[:S], kv_sf[:S], w[:T], ks, ke) + del out + torch.cuda.synchronize() + names = collections.Counter( + e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA + ) + for name, n in names.items(): + print(f"kernel x{n}: {name}") + print(f"variant map: {len(cells)} calls, S {min(S for _, S, _, _ in cells)}..{S_max}") + assert set(names.values()) == {len(cells)}, names + assert len(names) == 1, names + (name,) = names + assert name.startswith(_LOGITS_KERNEL), name + + +# Every distinct (tc, S, r) of the mqa_logits records in the core target's runtime-value dumps (boot and +# accuracy gates; 1928 cells, none outside the envelope), as "tc.S.r" when one request carries all tc rows, +# else "tc.S.r.nc.max_q_len" joined from a trtllm_sparse_mla_dsv4.varlen record with the same tc. +# ``_dumped_partition`` turns each into a consistent partition. +_DUMPED_CELLS = """ +1.1.1 2.1.2 2.2.1 14.8.2 14.16.1 16.8.2 16.16.1 48.24.2 48.48.1 493.342.2 493.685.1 494.336.2 494.672.1 +503.379.2 503.758.1 529.356.2 529.712.1 549.274.2 549.349.2 549.549.1 549.698.1 550.436.2 550.872.1 555.379.2 +555.759.1 556.368.2 556.737.1 562.281.2 562.562.1 569.422.2 569.844.1 571.407.2 571.814.1 573.383.2 573.767.1 +578.289.2 578.578.1 579.357.2 579.715.1 581.290.2 581.581.1 586.293.2 586.382.2 586.586.1 586.764.1 590.295.2 +590.417.2 590.590.1 590.835.1 593.362.2 593.725.1 594.377.2 594.754.1 595.430.2 595.861.1 597.298.2 597.597.1 +598.406.2 598.813.1 599.299.2 599.599.1 601.368.2 601.375.2 601.736.1 601.750.1 603.369.2 603.739.1 606.303.2 +606.606.1 609.304.2 609.609.1 612.306.2 612.612.1 614.374.2 614.749.1 616.308.2 616.616.1 617.376.2 617.753.1 +619.309.2 619.619.1 623.311.2 623.623.1 625.312.2 625.625.1 626.313.2 626.626.1 627.313.2 627.627.1 629.314.2 +629.402.2 629.629.1 629.804.1 630.315.2 630.630.1 631.315.2 631.631.1 632.316.2 632.632.1 634.317.2 634.430.2 +634.634.1 634.861.1 635.317.2 635.635.1 637.318.2 637.637.1 640.320.2 640.640.1 641.386.2 641.773.1 642.387.2 +642.774.1 643.321.2 643.643.1 644.322.2 644.644.1 646.323.2 646.418.2 646.646.1 646.837.1 649.324.2 649.649.1 +650.325.2 650.650.1 651.325.2 651.651.1 652.326.2 652.393.2 652.652.1 652.787.1 653.326.2 653.653.1 654.327.2 +654.654.1 656.328.2 656.656.1 660.330.2 660.660.1 661.330.2 661.396.2 661.661.1 661.793.1 662.331.2 662.420.2 +662.662.1 662.841.1 663.331.2 663.663.1 665.332.2 665.457.2 665.665.1 665.915.1 666.333.2 666.666.1 667.333.2 +667.412.2 667.667.1 667.825.1 668.334.2 668.668.1 670.335.2 670.427.2 670.670.1 670.854.1 672.336.2 672.672.1 +674.337.2 674.674.1 675.337.2 675.412.2 675.675.1 675.824.1 676.338.2 676.676.1 679.339.2 679.679.1 680.340.2 +680.680.1 681.340.2 681.681.1 682.435.2 682.445.2 682.871.1 682.890.1 683.341.2 683.683.1 684.342.2 684.684.1 +685.342.2 685.685.1 687.343.2 687.687.1 688.344.2 688.445.2 688.688.1 688.891.1 689.344.2 689.689.1 690.345.2 +690.429.2 690.690.1 690.858.1 691.345.2 691.691.1 692.346.2 692.692.1 693.414.2 693.828.1 694.347.2 694.694.1 +696.348.2 696.696.1 697.348.2 697.697.1 698.349.2 698.698.1 699.349.2 699.699.1 700.350.2 700.700.1 701.350.2 +701.701.1 702.351.2 702.702.1 703.351.2 703.703.1 704.352.2 704.704.1 706.353.2 706.706.1 708.452.2 708.905.1 +710.355.2 710.710.1 711.355.2 711.711.1 712.356.2 712.423.2 712.712.1 712.847.1 713.356.2 713.713.1 714.357.2 +714.714.1 715.357.2 715.715.1 718.359.2 718.718.1 719.359.2 719.719.1 720.360.2 720.478.2 720.720.1 720.956.1 +721.360.2 721.721.1 722.361.2 722.722.1 723.361.2 723.723.1 724.362.2 724.487.2 724.724.1 724.975.1 725.362.2 +725.725.1 726.363.2 726.726.1 727.363.2 727.727.1 728.364.2 728.728.1 729.364.2 729.729.1 730.365.2 730.730.1 +732.366.2 732.732.1 733.366.2 733.493.2 733.733.1 733.986.1 734.367.2 734.458.2 734.734.1 734.916.1 735.367.2 +735.735.1 736.368.2 736.736.1 737.368.2 737.466.2 737.737.1 737.932.1 738.369.2 738.738.1 739.369.2 739.739.1 +740.370.2 740.740.1 741.370.2 741.741.1 742.371.2 742.742.1 743.371.2 743.743.1 745.372.2 745.440.2 745.745.1 +745.880.1 746.373.2 746.746.1 747.373.2 747.747.1 748.374.2 748.748.1 749.374.2 749.749.1 751.375.2 751.751.1 +752.376.2 752.752.1 753.376.2 753.753.1 754.377.2 754.754.1 755.377.2 755.755.1 756.378.2 756.493.2 756.756.1 +756.987.1 757.378.2 757.757.1 758.379.2 758.758.1 759.379.2 759.759.1 760.380.2 760.760.1 761.380.2 761.448.2 +761.761.1 761.897.1 762.381.2 762.762.1 763.381.2 763.763.1 764.382.2 764.764.1 765.382.2 765.765.1 766.383.2 +766.766.1 767.383.2 767.767.1 768.384.2 768.768.1 769.384.2 769.769.1 770.385.2 770.770.1 771.385.2 771.771.1 +772.386.2 772.476.2 772.772.1 772.952.1 773.386.2 773.773.1 774.387.2 774.774.1 775.387.2 775.775.1 776.388.2 +776.776.1 777.388.2 777.777.1 778.389.2 778.778.1 779.389.2 779.779.1 780.390.2 780.493.2 780.780.1 780.986.1 +781.390.2 781.781.1 782.391.2 782.782.1 783.391.2 783.783.1 784.392.2 784.459.2 784.784.1 784.918.1 785.392.2 +785.785.1 786.393.2 786.786.1 787.393.2 787.787.1 788.394.2 788.788.1 790.395.2 790.790.1 791.395.2 791.791.1 +792.396.2 792.792.1 793.396.2 793.793.1 794.397.2 794.794.1 795.397.2 795.795.1 796.398.2 796.796.1 797.398.2 +797.797.1 798.399.2 798.605.2 798.798.1 798.1211.1 799.399.2 799.799.1 800.400.2 800.800.1 801.400.2 801.801.1 +802.401.2 802.471.2 802.802.1 802.943.1 803.401.2 803.803.1 804.402.2 804.804.1 805.402.2 805.805.1 806.403.2 +806.806.1 807.403.2 807.807.1 808.404.2 808.808.1 809.404.2 809.809.1 810.405.2 810.810.1 811.405.2 811.811.1 +812.406.2 812.491.2 812.812.1 812.982.1 813.406.2 813.499.2 813.813.1 813.999.1 814.407.2 814.814.1 815.407.2 +815.815.1 816.408.2 816.816.1 817.408.2 817.817.1 818.409.2 818.818.1 819.409.2 819.819.1 820.410.2 820.820.1 +822.411.2 822.822.1 823.411.2 823.823.1 824.412.2 824.824.1 825.412.2 825.493.2 825.825.1 825.986.1 826.413.2 +826.826.1 827.413.2 827.827.1 828.414.2 828.828.1 829.414.2 829.829.1 830.415.2 830.830.1 831.415.2 831.831.1 +832.416.2 832.832.1 833.416.2 833.833.1 834.417.2 834.834.1 835.417.2 835.835.1 836.418.2 836.836.1 837.418.2 +837.519.2 837.837.1 837.1039.1 838.419.2 838.838.1 839.419.2 839.839.1 840.420.2 840.840.1 841.420.2 841.841.1 +843.421.2 843.843.1 844.422.2 844.844.1 845.422.2 845.845.1 846.423.2 846.846.1 847.423.2 847.847.1 848.424.2 +848.848.1 849.424.2 849.849.1 850.425.2 850.850.1 851.425.2 851.851.1 852.426.2 852.852.1 853.426.2 853.567.2 +853.853.1 853.1135.1 854.427.2 854.511.2 854.580.2 854.854.1 854.1023.1 854.1160.1 855.427.2 855.855.1 856.428.2 +856.856.1 857.428.2 857.857.1 858.429.2 858.530.2 858.858.1 858.1060.1 859.429.2 859.859.1 860.430.2 860.860.1 +861.430.2 861.861.1 862.431.2 862.522.2 862.862.1 862.1045.1 863.431.2 863.863.1 864.432.2 864.864.1 866.433.2 +866.866.1 867.433.2 867.867.1 868.434.2 868.868.1 869.434.2 869.869.1 870.435.2 870.870.1 871.435.2 871.871.1 +872.436.2 872.872.1 873.436.2 873.873.1 874.437.2 874.874.1 875.437.2 875.875.1 876.438.2 876.876.1 877.438.2 +877.877.1 878.439.2 878.878.1 879.439.2 879.879.1 880.440.2 880.880.1 881.440.2 881.881.1 882.441.2 882.882.1 +883.441.2 883.883.1 884.442.2 884.884.1 885.442.2 885.885.1 886.443.2 886.886.1 887.443.2 887.887.1 888.444.2 +888.888.1 888.1469.2 888.2938.1 889.444.2 889.889.1 890.445.2 890.890.1 890.1469.2 890.2938.1 891.445.2 +891.891.1 892.446.2 892.892.1 893.446.2 893.893.1 894.447.2 894.894.1 895.447.2 895.895.1 896.448.2 896.896.1 +897.448.2 897.897.1 898.449.2 898.898.1 899.449.2 899.899.1 900.450.2 900.900.1 901.450.2 901.901.1 902.451.2 +902.902.1 903.451.2 903.903.1 904.452.2 904.904.1 906.453.2 906.906.1 907.453.2 907.907.1 908.454.2 908.908.1 +909.454.2 909.909.1 910.455.2 910.910.1 911.455.2 911.911.1 912.456.2 912.912.1 913.456.2 913.913.1 914.457.2 +914.914.1 915.457.2 915.556.2 915.915.1 915.1112.1 916.458.2 916.916.1 917.458.2 917.917.1 918.459.2 918.918.1 +919.459.2 919.919.1 920.460.2 920.920.1 921.460.2 921.547.2 921.921.1 921.1094.1 922.461.2 922.922.1 923.461.2 +923.923.1 924.462.2 924.924.1 925.462.2 925.925.1 926.463.2 926.926.1 927.463.2 927.927.1 928.464.2 928.928.1 +929.464.2 929.929.1 930.465.2 930.930.1 931.465.2 931.931.1 932.466.2 932.932.1 933.466.2 933.933.1 934.467.2 +934.934.1 935.467.2 935.935.1 936.468.2 936.936.1 937.468.2 937.937.1 938.469.2 938.938.1 939.469.2 939.939.1 +940.470.2 940.940.1 941.470.2 941.941.1 942.471.2 942.942.1 943.471.2 943.943.1 944.472.2 944.944.1 945.472.2 +945.945.1 946.473.2 946.946.1 947.473.2 947.947.1 949.474.2 949.949.1 950.475.2 950.950.1 951.475.2 951.951.1 +952.476.2 952.952.1 953.476.2 953.953.1 954.477.2 954.954.1 955.477.2 955.955.1 956.478.2 956.956.1 957.478.2 +957.957.1 958.479.2 958.958.1 959.479.2 959.959.1 960.480.2 960.960.1 961.480.2 961.961.1 962.481.2 962.962.1 +963.481.2 963.963.1 964.482.2 964.964.1 965.482.2 965.965.1 966.483.2 966.966.1 967.483.2 967.967.1 968.484.2 +968.968.1 969.484.2 969.969.1 970.485.2 970.970.1 972.486.2 972.972.1 973.486.2 973.571.2 973.973.1 973.1142.1 +974.487.2 974.974.1 975.487.2 975.975.1 976.488.2 976.976.1 977.488.2 977.977.1 979.489.2 979.979.1 980.490.2 +980.980.1 982.491.2 982.982.1 983.491.2 983.983.1 984.492.2 984.984.1 986.493.2 986.986.1 987.493.2 987.987.1 +988.494.2 988.988.1 989.494.2 989.989.1 990.495.2 990.990.1 991.495.2 991.991.1 992.496.2 992.575.2 992.992.1 +992.1150.1 993.496.2 993.993.1 994.497.2 994.994.1 996.498.2 996.996.1 997.498.2 997.997.1 998.499.2 998.998.1 +999.499.2 999.999.1 1000.500.2 1000.1000.1 1002.501.2 1002.1002.1 1003.501.2 1003.1003.1 1004.502.2 1004.1004.1 +1005.502.2 1005.1005.1 1007.503.2 1007.1007.1 1008.504.2 1008.1008.1 1011.505.2 1011.1011.1 1012.506.2 +1012.580.2 1012.1012.1 1012.1161.1 1013.506.2 1013.1013.1 1014.507.2 1014.1014.1 1015.507.2 1015.1015.1 +1016.508.2 1016.1016.1 1016.1469.2 1016.2938.1 1018.509.2 1018.1018.1 1019.509.2 1019.1019.1 1020.510.2 +1020.1020.1 1021.510.2 1021.1021.1 1022.511.2 1022.1022.1 1023.511.2 1023.1023.1 1025.512.2 1025.1025.1 +1026.513.2 1026.1026.1 1028.514.2 1028.1028.1 1031.515.2 1031.1031.1 1032.516.2 1032.1032.1 1034.517.2 +1034.1034.1 1036.518.2 1036.1036.1 1038.519.2 1038.1038.1 1039.519.2 1039.1039.1 1040.520.2 1040.1040.1 +1042.521.2 1042.1042.1 1043.521.2 1043.1043.1 1045.522.2 1045.1045.1 1046.523.2 1046.1046.1 1048.524.2 +1048.1048.1 1050.525.2 1050.1050.1 1055.527.2 1055.1055.1 1057.528.2 1057.1057.1 1058.529.2 1058.1058.1 +1062.531.2 1062.1062.1 1064.532.2 1064.1064.1 1065.532.2 1065.1065.1 1066.533.2 1066.1066.1 1067.533.2 +1067.1067.1 1068.534.2 1068.1068.1 1070.535.2 1070.1070.1 1072.536.2 1072.1072.1 1076.538.2 1076.1076.1 +1079.539.2 1079.1079.1 1080.540.2 1080.1080.1 1081.540.2 1081.1081.1 1082.541.2 1082.1082.1 1084.542.2 +1084.1084.1 1085.542.2 1085.1085.1 1086.543.2 1086.1086.1 1087.543.2 1087.1087.1 1089.544.2 1089.1089.1 +1090.545.2 1090.1090.1 1091.545.2 1091.1091.1 1092.546.2 1092.1092.1 1093.546.2 1093.1093.1 1094.547.2 +1094.1094.1 1095.547.2 1095.1095.1 1098.549.2 1098.1098.1 1099.549.2 1099.1099.1 1100.550.2 1100.1100.1 +1101.550.2 1101.1101.1 1104.552.2 1104.1104.1 1106.553.2 1106.1106.1 1109.554.2 1109.1109.1 1111.555.2 +1111.1111.1 1115.557.2 1115.1115.1 1118.559.2 1118.1118.1 1119.559.2 1119.1119.1 1120.560.2 1120.1120.1 +1122.561.2 1122.1122.1 1125.562.2 1125.1125.1 1126.563.2 1126.1126.1 1128.564.2 1128.1128.1 1129.564.2 +1129.1129.1 1130.565.2 1130.1130.1 1135.567.2 1135.1135.1 1136.568.2 1136.1136.1 1137.568.2 1137.1137.1 +1139.569.2 1139.1139.1 1140.570.2 1140.1140.1 1141.570.2 1141.1141.1 1142.571.2 1142.1142.1 1143.571.2 +1143.1143.1 1145.572.2 1145.1145.1 1146.573.2 1146.1146.1 1149.574.2 1149.1149.1 1150.575.2 1150.1150.1 +1154.577.2 1154.1154.1 1157.578.2 1157.1157.1 1160.580.2 1160.1160.1 1170.585.2 1170.1170.1 1171.585.2 +1171.1171.1 1172.586.2 1172.1172.1 1177.588.2 1177.1177.1 1178.589.2 1178.1178.1 1182.591.2 1182.1182.1 +1185.592.2 1185.1185.1 1187.593.2 1187.1187.1 1198.599.2 1198.1198.1 1200.600.2 1200.1200.1 1201.600.2 +1201.1201.1 1204.602.2 1204.1204.1 1211.605.2 1211.1211.1 1217.608.2 1217.1217.1 1226.613.2 1226.1226.1 +1236.618.2 1236.1236.1 1237.618.2 1237.1237.1 1241.620.2.2.629 1241.1241.1.2.629 1261.630.2 1261.1261.1 +1273.636.2 1273.1273.1 1275.637.2.2.697 1275.1275.1.2.697 1283.769.2.2.780 1283.1538.1.2.780 1298.649.2 +1298.1298.1 1301.650.2 1301.1301.1 1313.656.2.2.764 1313.1313.1.2.764 1333.666.2.2.670 1333.1333.1.2.670 +1344.672.2.2.698 1344.1344.1.2.698 1345.672.2.2.746 1345.1345.1.2.746 1355.830.2.2.698 1355.1661.1.2.698 +1374.687.2.2.768 1374.1374.1.2.768 1388.693.2.2.757 1388.1388.1.2.757 1391.695.2.2.730 1391.1391.1.2.730 +1395.850.2.2.738 1395.1701.1.2.738 1397.698.2.2.811 1397.1397.1.2.811 1403.701.2.2.740 1403.1403.1.2.740 +1407.778.2.2.858 1407.1556.1.2.858 1408.704.2.2.738 1408.1408.1.2.738 1410.806.2 1410.1614.1.2.855 +1419.709.2.2.723 1423.711.2.2.734 1423.1423.1.2.734 1424.712.2.2.730 1424.1424.1.2.730 1432.805.2.2.938 +1432.1610.1.2.938 1433.716.2.2.787 1433.1433.1.2.787 1439.719.2.2.729 1439.1439.1.2.729 1448.724.2.2.756 +1448.1448.1.2.756 1462.731.2.2.774 1462.1462.1.2.774 1464.731.2.2.753 1464.836.2.2.782 1464.1464.1.2.782 +1464.1672.1.2.782 1466.798.2 1466.1598.1.2.825 1467.733.2.2.877 1467.1467.1.2.877 1472.736.2.2.742 +1473.736.2.2.777 1473.1473.1.2.777 1477.738.2.2.792 1477.1477.1.2.792 1481.740.2.2.777 1481.1481.1.2.777 +1482.741.2.2.788 1482.1482.1.2.788 1484.741.2.2.835 1484.1484.1.2.835 1494.746.2.2.895 1494.1494.1.2.895 +1496.814.2.2.854 1496.1628.1.2.854 1497.748.2.2.818 1497.1497.1.2.818 1498.748.2.2.779 1498.1498.1.2.779 +1501.750.2 1501.1501.1 1505.752.2.2.882 1505.1505.1.2.882 1507.753.2.2.839 1507.1507.1.2.839 1508.753.2.2.847 +1508.1508.1.2.847 1510.755.2.2.820 1510.1510.1.2.820 1511.755.2.2.780 1511.1511.1.2.780 1512.756.2.2.784 +1512.1512.1.2.784 1513.756.2.2.788 1513.1513.1.2.788 1514.843.2.2.844 1514.849.2.2.844 1514.1687.1.2.759 +1514.1698.1.2.759 1518.759.2.2.776 1518.1518.1.2.776 1519.824.2.2.793 1519.1648.1.2.793 1520.760.2.2.786 +1520.1520.1.2.786 1523.761.2.2.832 1523.1523.1.2.832 1524.761.2.2.825 1524.1524.1.2.825 1525.762.2.2.831 +1525.859.2.2.831 1525.1525.1.2.831 1525.1719.1.2.831 1528.763.2.2.771 1528.1528.1.2.771 1533.766.2.2.780 +1533.1533.1.2.780 1535.767.2.2.848 1535.1535.1.2.848 1537.768.2.2.839 1537.1537.1.2.839 1538.768.2.2.775 +1538.1538.1.2.775 1539.769.2.2.859 1539.1539.1.2.859 1542.770.2.2.799 1542.1542.1.2.799 1543.771.2.2.854 +1543.1543.1.2.854 1546.772.2.2.863 1546.773.2.2.906 1546.1546.1.2.863 1547.773.2.2.784 1547.881.2.2.784 +1547.1547.1.2.784 1547.1762.1.2.784 1550.774.2.2.815 1550.1550.1.2.815 1551.775.2.2.809 1551.1551.1.2.809 +1554.776.2.2.799 1554.777.2.2.850 1554.1554.1.2.799 1561.780.2.2.801 1561.1561.1.2.801 1562.780.2.2.841 +1562.872.2.2.828 1562.1562.1.2.828 1562.1744.1.2.828 1566.783.2.2.784 1566.1566.1.2.784 1567.783.2.2.911 +1567.1567.1.2.911 1568.783.2.2.815 1568.1568.1.2.815 1569.784.2.2.888 1569.1569.1.2.888 1570.784.2.2.851 +1570.1570.1.2.851 1571.785.2.2.831 1571.1571.1.2.831 1574.786.2.2.811 1574.787.2.2.834 1574.1574.1.2.834 +1575.787.2.2.848 1575.1575.1.2.848 1579.789.2.2.796 1579.1579.1.2.796 1580.790.2.2.854 1580.1580.1.2.854 +1583.866.2.2.1034 1583.1732.1.2.1034 1584.791.2.2.799 1584.792.2.2.826 1584.1584.1.2.799 1585.792.2.2.807 +1585.1585.1.2.807 1591.795.2.2.804 1591.1591.1.2.804 1593.796.2.2.862 1593.1593.1.2.862 1594.796.2.2.801 +1594.797.2.2.918 1594.1594.1.2.918 1595.797.2.2.848 1595.1595.1.2.848 1596.798.2.2.860 1596.1596.1.2.860 +1601.800.2.2.909 1601.1601.1.2.909 1605.802.2.2.824 1605.1605.1.2.824 1606.802.2.2.847 1606.803.2.2.914 +1606.1606.1.2.847 1608.803.2.2.873 1608.1608.1.2.873 1610.804.2.2.853 1610.1610.1.2.853 1614.807.2.2.820 +1614.1614.1.2.820 1615.807.2.2.858 1615.1615.1.2.858 1616.882.2 1616.1765.1.2.941 1620.810.2.2.878 +1620.1620.1.2.878 1621.810.2.2.967 1621.1621.1.2.967 1624.812.2.2.866 1624.1624.1.2.866 1625.812.2.2.869 +1625.1625.1.2.869 1626.813.2.2.920 1626.1626.1.2.920 1627.813.2.2.834 1627.1627.1.2.834 1629.814.2.2.825 +1629.1629.1.2.825 1630.882.2 1630.1765.1.2.889 1636.817.2.2.883 1636.1636.1.2.883 1637.818.2.2.822 +1637.1637.1.2.822 1639.1639.1.2.954 1640.820.2.2.914 1640.1640.1.2.914 1641.820.2.2.868 1641.1641.1.2.868 +1642.820.2.2.1005 1642.1642.1.2.1005 1644.821.2.2.869 1644.1644.1.2.869 1646.822.2.2.873 1646.1646.1.2.873 +1647.823.2.2.887 1647.1647.1.2.887 1649.824.2.2.871 1649.1649.1.2.871 1654.827.2.2.924 1654.1654.1.2.924 +1655.827.2.2.964 1655.1655.1.2.964 1656.827.2.2.845 1656.1656.1.2.845 1661.830.2.2.1064 1661.1661.1.2.1064 +1662.830.2.2.833 1662.1662.1.2.833 1663.831.2.2.866 1663.896.2.2.866 1663.1663.1.2.866 1663.1792.1.2.866 +1664.832.2.2.962 1664.899.2.2.962 1664.1664.1.2.962 1664.1800.1.2.962 1665.832.2.2.871 1665.1665.1.2.871 +1666.832.2.2.943 1666.1666.1.2.943 1667.833.2.2.910 1667.1667.1.2.910 1668.833.2.2.849 1668.834.2.2.1036 +1668.1668.1.2.1036 1669.834.2.2.896 1669.1669.1.2.896 1670.835.2.2.900 1670.1670.1.2.900 1672.835.2.2.841 +1672.836.2.2.854 1672.920.2.2.854 1672.1672.1.2.841 1672.1841.1.2.841 1675.837.2.2.856 1675.1675.1.2.856 +1676.838.2.2.880 1676.1676.1.2.880 1678.839.2.2.904 1678.1678.1.2.904 1680.839.2.2.939 1680.840.2.2.908 +1680.1680.1.2.939 1682.841.2.2.930 1682.1682.1.2.930 1683.841.2.2.924 1683.1683.1.2.924 1684.842.2.2.916 +1684.1684.1.2.916 1685.842.2.2.916 1685.1685.1.2.916 1686.842.2.2.925 1686.1686.1.2.925 1687.843.2.2.907 +1687.1687.1.2.907 1688.944.2 1688.1890.1.2.851 1689.844.2.2.965 1689.1689.1.2.965 1691.912.2.2.907 +1691.1825.1.2.907 1693.846.2.2.996 1693.1693.1.2.996 1694.847.2.2.848 1694.1694.1.2.848 1695.847.2.2.876 +1695.936.2.2.876 1695.1695.1.2.876 1695.1873.1.2.876 1696.848.2.2.1008 1696.1696.1.2.1008 1699.849.2.2.1019 +1699.1699.1.2.1019 1700.850.2.2.864 1700.1700.1.2.864 1701.850.2.2.883 1701.1701.1.2.883 1702.850.2.2.1003 +1702.1702.1.2.1003 1703.851.2.2.919 1703.1703.1.2.919 1704.932.2 1704.1865.1.2.879 1705.852.2.2.876 +1705.1705.1.2.876 1706.852.2.2.959 1706.1706.1.2.959 1707.853.2.2.856 1707.1707.1.2.856 1709.854.2.2.987 +1709.1709.1.2.987 1710.854.2.2.1129 1710.855.2.2.934 1710.1710.1.2.934 1712.855.2.2.877 1712.856.2.2.910 +1712.1712.1.2.877 1713.856.2.2.906 1713.1713.1.2.906 1714.857.2.2.858 1714.1714.1.2.858 1716.857.2.2.877 +1716.858.2.2.912 1716.1716.1.2.877 1719.859.2.2.1007 1719.1719.1.2.1007 1720.859.2.2.917 1720.967.2 +1720.1720.1.2.917 1720.1935.1.2.917 1723.861.2.2.990 1723.1723.1.2.990 1724.862.2.2.900 1724.1724.1.2.900 +1725.862.2.2.925 1725.1725.1.2.925 1726.863.2.2.918 1726.1726.1.2.918 1728.863.2.2.895 1728.1728.1.2.895 +1731.865.2.2.877 1731.1731.1.2.877 1734.866.2.2.929 1734.867.2.2.950 1734.1734.1.2.929 1735.867.2.2.998 +1735.1735.1.2.998 1737.868.2.2.930 1737.1737.1.2.930 1738.869.2.2.954 1738.1738.1.2.954 1740.967.2 +1740.1935.1.2.1003 1741.870.2.2.907 1741.1741.1.2.907 1742.870.2.2.893 1742.871.2.2.950 1742.1742.1.2.950 +1743.871.2.2.917 1743.1743.1.2.917 1744.872.2.2.1004 1744.1744.1.2.1004 1747.873.2.2.966 1747.1747.1.2.966 +1748.873.2.2.909 1748.1748.1.2.909 1749.874.2.2.942 1749.1749.1.2.942 1750.874.2.2.921 1750.875.2.2.918 +1750.1750.1.2.918 1752.875.2.2.1087 1752.876.2.2.1100 1752.1752.1.2.1100 1753.876.2.2.993 1753.1753.1.2.993 +1754.876.2.2.915 1754.877.2.2.884 1754.1754.1.2.884 1755.877.2.2.888 1755.1755.1.2.888 1756.877.2.2.951 +1756.878.2.2.898 1756.1756.1.2.898 1757.878.2.2.986 1757.1757.1.2.986 1758.878.2.2.923 1758.879.2.2.1098 +1758.1758.1.2.923 1759.879.2.2.1023 1759.1759.1.2.1023 1760.880.2.2.988 1760.1760.1.2.988 1764.882.2.2.952 +1764.1764.1.2.952 1766.882.2.2.935 1766.1766.1.2.935 1769.884.2.2.893 1769.1769.1.2.893 1770.885.2.2.960 +1770.1770.1.2.960 1771.885.2.2.1091 1771.1771.1.2.1091 1772.989.2.2.992 1772.1011.2.2.992 1772.1978.1.2.992 +1772.2023.1.2.992 1773.886.2.2.1005 1773.1773.1.2.1005 1775.887.2.2.926 1775.1775.1.2.926 1776.887.2.2.953 +1776.1776.1.2.953 1780.889.2.2.929 1780.1780.1.2.929 1781.890.2.2.893 1781.1781.1.2.893 1782.891.2.2.916 +1782.1782.1.2.916 1783.891.2.2.1137 1783.1783.1.2.1137 1785.892.2.2.924 1785.1785.1.2.924 1790.895.2.2.966 +1790.1790.1.2.966 1791.895.2.2.930 1791.1791.1.2.930 1793.995.2.2.915 1793.1990.1.2.915 1794.897.2.2.1076 +1794.1794.1.2.1076 1797.898.2.2.1198 1797.1797.1.2.1198 1799.899.2.2.910 1799.1799.1.2.910 1801.900.2.2.912 +1801.1801.1.2.912 1802.901.2.2.934 1802.1802.1.2.934 1803.901.2.2.996 1804.902.2.2.914 1804.1804.1.2.914 +1805.902.2.2.951 1805.1805.1.2.951 1808.903.2.2.1079 1808.1808.1.2.1079 1810.904.2.2.1067 1810.1810.1.2.1067 +1812.906.2.2.988 1812.1812.1.2.988 1814.906.2.2.953 1814.1814.1.2.953 1817.908.2.2.910 1817.1817.1.2.910 +1818.908.2.2.1043 1818.1818.1.2.1043 1819.909.2.2.977 1819.1819.1.2.977 1820.909.2.2.1003 1820.910.2.2.952 +1820.1820.1.2.1003 1821.910.2.2.1048 1821.1821.1.2.1048 1823.911.2.2.912 1823.1823.1.2.912 1824.911.2.2.963 +1824.1824.1.2.963 1825.912.2.2.929 1825.1825.1.2.929 1826.999.2 1826.1999.1.2.1067 1828.914.2.2.962 +1828.1828.1.2.962 1832.915.2.2.1143 1832.1832.1.2.1143 1836.918.2.2.1122 1836.1836.1.2.1122 1837.918.2.2.974 +1837.1837.1.2.974 1838.918.2.2.1021 1838.1838.1.2.1021 1843.921.2.2.947 1843.1843.1.2.947 1844.921.2.2.1045 +1844.1844.1.2.1045 1850.924.2.2.961 1850.1850.1.2.961 1852.926.2.2.928 1852.1852.1.2.928 1855.927.2.2.975 +1855.1855.1.2.975 1856.927.2.2.977 1856.928.2.2.1122 1856.1856.1.2.977 1857.928.2.2.998 1857.1857.1.2.998 +1859.929.2.2.1000 1859.1859.1.2.1000 1861.930.2.2.955 1861.1861.1.2.955 1863.931.2.2.945 1863.1863.1.2.945 +1864.932.2.2.1074 1864.1864.1.2.1074 1868.1020.2 1868.2041.1.2.947 1869.934.2.2.974 1869.1869.1.2.974 +1871.935.2.2.1026 1871.1871.1.2.1026 1873.936.2.2.1176 1873.1873.1.2.1176 1875.937.2.2.1122 1875.1875.1.2.1122 +1876.937.2.2.1201 1876.1876.1.2.1201 1878.938.2.2.957 1878.1878.1.2.957 1882.941.2.2.1170 1882.1882.1.2.1170 +1885.942.2.2.966 1885.1885.1.2.966 1887.943.2.2.945 1887.1887.1.2.945 1891.945.2.2.1136 1891.1891.1.2.1136 +1892.946.2.2.966 1892.1025.2.2.966 1892.1892.1.2.966 1892.2050.1.2.966 1895.947.2.2.1003 1895.1895.1.2.1003 +1901.950.2.2.1050 1901.1901.1.2.1050 1902.950.2.2.1087 1902.1902.1.2.1087 1910.954.2.2.1011 1910.955.2.2.1028 +1910.1910.1.2.1011 1911.955.2.2.969 1911.1911.1.2.969 1912.956.2.2.1068 1912.1912.1.2.1068 1913.956.2.2.1086 +1913.1913.1.2.1086 1914.956.2.2.1025 1914.1914.1.2.1025 1920.959.2.2.1081 1920.1920.1.2.1081 1921.960.2.2.1201 +1921.1921.1.2.1201 1924.962.2.2.1004 1924.1924.1.2.1004 1925.962.2.2.1154 1925.1925.1.2.1154 1926.962.2.2.977 +1926.1926.1.2.977 1931.965.2.2.1043 1931.1931.1.2.1043 1934.969.2.2.1920 1934.1938.1.2.1920 1935.967.2.2.1058 +1935.1935.1.2.1058 1937.968.2.2.982 1937.1937.1.2.982 1939.969.2.2.1100 1939.1939.1.2.1100 1947.973.2.2.1150 +1947.1947.1.2.1150 1952.976.2.2.1008 1952.1952.1.2.1008 1953.976.2.2.1085 1953.1953.1.2.1085 1954.977.2.2.1066 +1954.1954.1.2.1066 1956.977.2.2.999 1956.1956.1.2.999 1964.982.2.2.1058 1964.1964.1.2.1058 1968.983.2.2.1091 +1968.984.2.2.994 1968.1968.1.2.994 1969.984.2.2.1024 1969.1969.1.2.1024 1972.986.2.2.1076 1972.1972.1.2.1076 +1973.986.2.2.1012 1973.1060.2.2.1012 1973.1973.1.2.1012 1973.2122.1.2.1012 1975.987.2.2.1055 1975.1975.1.2.1055 +1976.988.2.2.988 1976.1976.1.2.988 1977.988.2.2.1119 1977.1977.1.2.1119 1979.989.2.2.1053 1979.1979.1.2.1053 +1980.989.2.2.1015 1980.990.2.2.1036 1980.1980.1.2.1015 1982.990.2.2.1053 1982.991.2.2.1038 1982.1982.1.2.1038 +1984.992.2.2.1094 1984.1984.1.2.1094 1987.993.2.2.1045 1987.1987.1.2.1045 1988.993.2.2.1011 1988.1988.1.2.1011 +1993.996.2.2.1024 1993.1993.1.2.1024 1997.998.2.2.1032 1997.1997.1.2.1032 1998.998.2.2.999 1998.1998.1.2.999 +2004.1001.2.2.1049 2004.2004.1.2.1049 2007.1003.2.2.1176 2007.2007.1.2.1176 2008.1004.2.2.1080 +2008.2008.1.2.1080 2009.1004.2.2.1130 2009.2009.1.2.1130 2013.1006.2.2.1097 2013.2013.1.2.1097 +2016.1008.2.2.1154 2016.2016.1.2.1154 2017.1008.2.2.1038 2017.2017.1.2.1038 2019.1009.2.2.1094 +2019.2019.1.2.1094 2020.1009.2.2.1319 2020.2020.1.2.1319 2026.1100.2.2.1118 2026.2026.1.2.1118 +2026.2201.1.2.1118 2034.1017.2.2.1040 2034.2034.1.2.1040 2045.1022.2.2.1080 2045.2045.1.2.1080 2048.1024.2 +2048.1025.2 2048.2048.1 2048.2050.1 2051.1025.2.2.1057 2051.2051.1.2.1057 2054.1026.2.2.1099 2054.1350.2.4.688 +2054.2054.1.2.1099 2054.2702.1.2.1099 2072.1036.2.2.1170 2072.2072.1.2.1170 2073.1036.2.2.1080 +2073.2073.1.2.1080 2077.1038.2.2.1066 2077.2077.1.2.1066 2082.1041.2.2.1112 2082.2082.1.2.1112 2085.1042.2.3.730 +2085.2085.1.3.730 2092.1046.2.2.1198 2092.2092.1.2.1198 2094.1122.2 2094.2246.1.2.1167 2105.1052.2.2.1128 +2107.1053.2.2.1058 2107.2107.1.2.1058 2117.1058.2.2.1109 2117.2117.1.2.1109 2121.1060.2.2.1082 +2121.2121.1.2.1082 2135.1067.2.2.1501 2135.2135.1.2.1501 2140.1070.2.2.1128 2140.2140.1.2.1128 +2146.1073.2.2.1236 2146.2146.1.2.1236 2155.1077.2.3.914 2155.2155.1.3.914 2160.1079.2.2.1139 2160.2160.1.2.1139 +2163.1081.2.2.1091 2163.2163.1.2.1091 2164.1081.2.2.1185 2164.2164.1.2.1185 2165.1082.2.2.1273 +2165.2165.1.2.1273 2179.1089.2.2.1136 2179.2179.1.2.1136 2193.1096.2.2.1177 2193.2193.1.2.1177 2213.1106.2.3.900 +2213.2213.1.3.900 2228.1113.2.2.1139 2228.2228.1.2.1139 2232.1237.2 2232.2477.1.3.835 2260.1129.2.2.1145 +2260.2260.1.2.1145 2313.1251.2.3.910 2313.2504.1.3.910 2333.1166.2.3.836 2333.2333.1.3.836 2349.1174.2.2.1237 +2349.2349.1.2.1237 2367.1249.2.3.1018 2367.2499.1.3.1018 2369.1184.2.3.891 2369.2369.1.3.891 2401.1307.2.3.930 +2401.2616.1.3.930 2411.1204.2.3.925 2411.2411.1.3.925 2413.1206.2.2.1298 2413.2413.1.2.1298 2419.1209.2.3.846 +2419.2419.1.3.846 2433.1216.2.3.936 2433.2433.1.3.936 2457.1228.2.3.896 2457.2457.1.3.896 2465.1232.2.3.923 +2465.2465.1.3.923 2466.1233.2.3.896 2466.2466.1.3.896 2480.1239.2.3.936 2480.2480.1.3.936 2489.1244.2.3.883 +2489.2489.1.3.883 2498.1249.2.3.884 2498.2498.1.3.884 2506.1320.2 2506.2641.1.3.1023 2511.2511.1.3.919 +2523.1261.2.3.977 2523.2523.1.3.977 2541.1270.2.3.936 2541.2541.1.3.936 2556.1277.2.3.872 2556.2556.1.3.872 +2605.1302.2.3.938 2605.2605.1.3.938 2626.1312.2.3.1065 2626.2626.1.3.1065 2629.1313.2.3.877 2629.2629.1.3.877 +2643.1320.2.3.1095 2643.2643.1.3.1095 2644.1321.2.3.1000 2644.2644.1.3.1000 2654.1326.2.3.963 2654.2654.1.3.963 +2687.1343.2.3.999 2687.2687.1.3.999 2693.1346.2.3.945 2693.2693.1.3.945 2698.1348.2.3.1015 2698.2698.1.3.1015 +2713.1356.2.3.999 2713.2713.1.3.999 2715.1357.2.3.942 2715.2715.1.3.942 2721.1360.2.3.1140 2721.2721.1.3.1140 +2728.1431.2 2728.2863.1.3.1045 2749.1464.2.3.1177 2749.2930.1.3.1177 2769.1384.2.3.1080 2769.2769.1.3.1080 +2791.1394.2.3.977 2791.2791.1.3.977 2813.1406.2.3.1170 2813.2813.1.3.1170 2819.1408.2.3.989 2819.2819.1.3.989 +2851.1425.2.3.1081 2851.2851.1.3.1081 2864.1431.2.3.1045 2864.2864.1.3.1045 2866.1432.2.3.1085 +2866.2866.1.3.1085 2910.1454.2.3.1049 2910.2910.1.3.1049 2968.1483.2.3.1261 2968.2968.1.3.1261 3114.1689.2.5.702 +3114.3380.1.5.702 3162.1713.2 3162.3428.1.5.701 3206.1927.2.4.1114 3206.3854.1.4.1114 3250.1623.2.4.911 +3250.3250.1.4.911 3278.1708.2 3278.3417.1.5.697 3289.1643.2.4.867 3289.3289.1.4.867 3336.1667.2.3.1237 +3336.3336.1.3.1237 3601.1799.2.4.1057 3601.3601.1.4.1057 8071.4430.2.12.833 8071.8866.1.12.833 +8095.4119.2.10.927 8095.8243.1.10.927 8102.4059.2 8102.8123.1.10.1085 8117.4065.2.9.1097 8117.8135.1.9.1097 +8119.4452.2.10.942 8119.8906.1.10.942 8135.4523.2.10.950 8135.9049.1.10.950 8136.4523.2.10.944 +8136.9052.1.10.944 8141.4077.2.9.1163 8141.8159.1.9.1163 8150.4073.2.8.1320 8150.8150.1.8.1320 +8153.4212.2.10.1318 8153.8430.1.10.1318 8156.4087.2 8156.8179.1.11.835 8160.4087.2.11.842 8160.8182.1.11.842 +8161.4219.2.11.823 8161.8445.1.11.823 8178.4086.2.8.1141 8178.4087.2.8.1141 8178.8178.1.8.1141 8192.4096.2 +8192.8192.1 +""" + + +def _dumped_partition(tc, S, r, nc, mq): + """(cached, new) per request consistent with one dumped cell: nc requests, the first with mq new + tokens, the others splitting tc - mq (each in 1..mq); at ratio 2 as many of them odd as S needs; + then cached prefixes add the missing S keys, request after request, up to 16384 tokens each. + None if no such partition exists.""" + if nc == 1: + if mq != tc: + return None + new = [tc] + else: + m, rest = nc - 1, tc - mq + if not m <= rest <= m * mq: + return None + new = [mq] + [rest // m + (i < rest % m) for i in range(m)] + if r == 2: + want = max(tc - 2 * S, tc % 2) - mq % 2 # odd lengths wanted among the others + if want < 0 or want > m: + return None + for _ in range(2 * m): + odd = [i for i in range(1, nc) if new[i] % 2] + even = [i for i in range(1, nc) if new[i] % 2 == 0] + if len(odd) == want: + break + src = even if len(odd) < want else odd + up = [i for i in src if new[i] < mq] + down = [i for i in src if new[i] > 1 and i != (up[0] if up else -1)] + if not up or not down: + return None + new[up[0]] += 1 + new[down[0]] -= 1 + if sum(n % 2 for n in new[1:]) != want: + return None + extra = S - sum(n // r for n in new) + if extra < 0: + return None + reqs = [] + for n in new: + e = min(extra, (16384 - n) // r) + extra -= e + reqs.append((r * e, n)) + if extra: + return None + return reqs + + +def _dumped_cells(): + return [tuple(int(x) for x in tok.split(".")) for tok in _DUMPED_CELLS.split()] + + +@pytest.mark.parametrize("r", [1, 2]) +def test_multi_request_dumped_partitions(r): + """One consistent partition per dumped cell at this ratio, vs float64.""" + cells = [c for c in _dumped_cells() if c[2] == r] + stats = collections.Counter() + for c in cells: + tc, S, _, *rest = c + nc, mq = rest if rest else (1, tc) + reqs = _dumped_partition(tc, S, r, nc, mq) + assert reqs is not None and len(reqs) == nc and sum(n for _, n in reqs) == tc + assert sum((cached + n) // r for cached, n in reqs) == S and max(n for _, n in reqs) == mq + T, S_run, n_vis = _run_partition(reqs, r, seed=tc + 3 * S + r) + assert (T, S_run) == (tc, S) + stats["cells"] += 1 + stats["multi-request"] += nc > 1 + stats["rows"] += T + stats["visible checked"] += n_vis + print(f"dumped partitions r {r}: {dict(stats)}") + assert stats["cells"] == len(cells) > 0 + + +# ----------------------------------------------------------------------------- long context: row tiles, long gathers +# +# The V4.1 target's prefill index pass (``_index_prefill`` in the target's modeling.py) calls this op once per ROW +# TILE: rows [a, b) of the step's tc <= 16384 (max_num_tokens) context rows, passed as the row-offset views q[a:b], +# q_sf[a:b], weights[a:b], ks[a:b], ke[a:b], against ALL s_len keys the step gathered (full width, clean_logits=True, +# fp32 logits). ``_index_row_tiles`` sizes the tiles so rows x align(s_len + 256, 256) <= 2^28 fp32 elements: one +# tile of all rows when they fit, otherwise tiles of 2^28 // width rows, rounded down to a multiple of 64 when that +# is >= 64. The decoder-replay pass makes the same call on its tail rows (``_tail_cand_masks``: the last 128 +# positions of each bounded request, ratio-1 keys). s_len is the sum over the step's context requests of +# kv_len // r, and one request reaches Lm = 1,048,590 tokens (max_seq_len 1,048,576 plus DSpark's 14). s_len is +# certified here up to 2^25: the op that gathers the keys (indexer_k_cache_gather) misaddresses rows from 2^25 on. +# +# Two references. ``_long_context_reference`` is the float64 reference at any extent (keys decoded 2^16 rows at a +# time). ``_LongContextPalette`` draws the query rows and key rows of a call from palettes, so one float64 [Q, P] +# table serves every call: a call's expected logits are that table gathered by its draws, -inf outside each row's +# range. A kernel that reads the wrong key (query) row lands on another palette entry with probability 1 - 1/P +# (1 - 1/Q) per element; ``test_long_context_wrong_answers_are_seen`` shows the check failing for each such mistake. + + +_LONG_CONTEXT_TC = 16384 # context rows per step (max_num_tokens) +_LONG_CONTEXT_K = 256 # context requests per step (max_batch_size) +_LONG_CONTEXT_LM = 1048590 # the longest request: max_seq_len 1,048,576 + 14 with DSpark +_LONG_CONTEXT_BUDGET = 1 << 28 # fp32 logits elements per call (the target's _INDEX_TILE_ELEMS) +_LONG_CONTEXT_S_MAX = 1 << 25 # gathered keys per step, certified up to here + + +def _long_context_rows_max(s_len): + """The tallest tile at s_len: all of a step's rows go in one call while rows <= 2^28 // width.""" + return min(_LONG_CONTEXT_TC, max(1, _LONG_CONTEXT_BUDGET // _row_stride(s_len))) + + +def _long_context_tiles(rows, s_len): + """The target's row tiles (``_index_row_tiles``) of a step with ``rows`` context rows over s_len keys.""" + per = max(1, _LONG_CONTEXT_BUDGET // _row_stride(s_len)) + if per >= rows: + return [(0, rows)] + if per >= 64: + per -= per % 64 + return [(a, min(rows, a + per)) for a in range(0, rows, per)] + + +def _long_context_mask_(out, ks, ke): + """-inf outside each row's [ks, ke), in place, 2^26 elements at a time.""" + T, S = out.shape + col = torch.arange(S, device=DEV) + step = max(1, (1 << 26) // S) + for t0 in range(0, T, step): + vis = (col[None, :] >= ks[t0 : t0 + step, None].long()) & ( + col[None, :] < ke[t0 : t0 + step, None].long() + ) + out[t0 : t0 + step].masked_fill_(~vis, NEG_INF) + return out + + +def _long_context_reference(q, q_sf, w, kv, kv_sf, ks, ke): + """The float64 reference at any extent, rounded once to fp32: [T, S], -inf outside [ks, ke); and + max_t,s sum_h |term| (the exact regime's premise). Keys decoded 2^16 at a time, query rows taken 64 at a time.""" + T, S = q.shape[0], kv.shape[0] + qd, wd = _decode(q, q_sf), w.double() + out = torch.empty(T, S, dtype=torch.float32, device=DEV) + bound = torch.zeros((), dtype=torch.float64, device=DEV) + for s0 in range(0, S, 1 << 16): + kd = _decode(kv[s0 : s0 + (1 << 16)], kv_sf[s0 : s0 + (1 << 16)]) + for t0 in range(0, T, 64): + terms = ( + torch.einsum("thd,kd->thk", qd[t0 : t0 + 64], kd) + .clamp_min_(0) + .mul_(wd[t0 : t0 + 64, :, None]) + ) + out[t0 : t0 + 64, s0 : s0 + kd.shape[0]] = terms.sum(1) + bound = torch.maximum(bound, terms.abs_().sum(1).max()) + del terms + return _long_context_mask_(out, ks, ke), float(bound) + + +class _LongContextPalette: + """Q query rows (with their head weights) and P key rows in the exact regime, and the float64 logits of every + (query row, key row) pair, rounded to fp32: ``table[i, j]``. A call's inputs are rows drawn from the palettes; + its expected logits are ``table`` gathered by the draws, -inf outside each row's range.""" + + def __init__(self, seed, n_q=2048, n_k=8192): + self.q, self.q_sf, self.kv, self.kv_sf, w_raw = _exact_inputs(n_q, n_k, _HEADS, seed) + self.w = w_raw * GLOBAL_FOLD + end = torch.full((n_q,), n_k, dtype=torch.int32, device=DEV) + table, _, bound = _reference( + _decode(self.q, self.q_sf), + _decode(self.kv, self.kv_sf), + self.w, + torch.zeros_like(end), + end, + ) + _assert_exact_premise(bound) + self.table = table.float() + self.n_q, self.n_k = n_q, n_k + + def draw(self, T, S, g): + """Row draws, and the inputs they make: (q, q_sf, kv, kv_sf, weights) for T query rows and S keys.""" + qi = torch.randint(0, self.n_q, (T,), generator=g, device=DEV) + ki = torch.randint(0, self.n_k, (S,), generator=g, device=DEV) + return qi, ki, (self.q[qi], self.q_sf[qi], self.kv[ki], self.kv_sf[ki], self.w[qi]) + + def expected(self, qi, ki, ks, ke): + return _long_context_mask_(self.table.index_select(0, qi).index_select(1, ki), ks, ke) + + +def _long_context_check_tiles(pal, qi, ki, inputs, ks, ke, tiles): + """Each tile called as the target calls it (row-offset views of the step's tensors, all S keys) and compared, + every element, with the palette reference. Returns the visible elements checked.""" + q, q_sf, kv, kv_sf, w = inputs + S = kv.shape[0] + n_vis = 0 + for a, b in tiles: + assert (b - a) * _row_stride(S) <= _LONG_CONTEXT_BUDGET or b - a == 1 + out = _call(q[a:b], q_sf[a:b], kv, kv_sf, w[a:b], ks[a:b], ke[a:b]) + assert out.shape == (b - a, S) and out.stride() == (_row_stride(S), 1) + assert out.untyped_storage().nbytes() == _logits_bytes(b - a, S, _HEADS) + exp = pal.expected(qi[a:b], ki, ks[a:b], ke[a:b]) + bad = int((out != exp).sum()) + assert bad == 0, f"tile [{a}, {b}) over S {S}: {bad} of {out.numel()} elements differ" + n_vis += int((ke[a:b] - ks[a:b]).clamp_min(0).sum()) + _poison(out) + del out, exp + return n_vis + + +def _long_context_sample_tiles(tiles, rng, n=16): + """All tiles when there are at most n; else the first, the last and n - 2 others drawn at random.""" + if len(tiles) <= n: + return tiles + mid = sorted(rng.sample(range(1, len(tiles) - 1), n - 2)) + return [tiles[0]] + [tiles[i] for i in mid] + [tiles[-1]] + + +# Every distinct (rows, s_len, r) of the deepgemm_fp8_fp4_mqa_logits and .tail_cand records in the long-context +# runs' runtime-value dumps (tep4, DSpark, max_num_tokens 8192, max_seq_len 131072 / 262144 / 1048576), as +# "rows.s_len.r". The four tail_cand cells (128 rows at s_len 130773, 261860, 523951, 1048219, ratio 1) are among +# them. +_LONG_CONTEXT_DUMPED = """ +1.1.1 2.1.2 2.2.1 64.32768.1 64.32768.2 128.8192.1 128.65536.1 128.65536.2 128.98304.1 128.98304.2 128.130773.1 +128.192512.2 128.196608.1 128.196608.2 128.221184.1 128.221184.2 128.225280.2 128.229376.1 128.229376.2 +128.261860.1 128.282624.2 128.286720.1 128.286720.2 128.290816.2 128.294912.1 128.294912.2 128.299008.2 +128.421888.2 128.425984.1 128.425984.2 128.430080.2 128.434176.1 128.434176.2 128.438272.2 128.442368.1 +128.442368.2 128.446464.2 128.450560.1 128.450560.2 128.454656.2 128.458752.1 128.458752.2 128.462848.2 +128.523951.1 128.524288.1 128.532480.1 128.540672.1 128.548864.1 128.557056.1 128.565248.1 128.573440.1 +128.581632.1 128.589824.1 128.598016.1 128.606208.1 128.614400.1 128.622592.1 128.630784.1 128.638976.1 +128.647168.1 128.655360.1 128.663552.1 128.671744.1 128.679936.1 128.688128.1 128.696320.1 128.1048219.1 +155.1048219.1 175.261975.2 175.523951.1 192.163840.1 192.163840.2 192.704512.1 192.712704.1 192.720896.1 +192.729088.1 192.737280.1 192.745472.1 192.753664.1 192.761856.1 192.770048.1 192.778240.1 192.786432.1 +192.794624.1 192.802816.1 192.811008.1 192.819200.1 192.827392.1 192.835584.1 219.524109.2 256.131072.1 +256.131072.2 256.843776.1 256.851968.1 256.860160.1 256.868352.1 256.876544.1 256.884736.1 256.892928.1 +256.901120.1 256.909312.1 256.917504.1 256.925696.1 256.933888.1 256.942080.1 256.950272.1 256.958464.1 +256.966656.1 256.974848.1 256.983040.1 256.991232.1 256.999424.1 256.1007616.1 256.1015808.1 256.1024000.1 +256.1032192.1 256.1040384.1 256.1048219.1 320.704512.1 320.712704.1 320.720896.1 320.729088.1 320.737280.1 +320.745472.1 320.753664.1 320.761856.1 320.770048.1 320.778240.1 320.786432.1 320.794624.1 320.802816.1 +320.811008.1 320.819200.1 320.827392.1 320.835584.1 384.606208.1 384.614400.1 384.622592.1 384.630784.1 +384.638976.1 384.647168.1 384.655360.1 384.663552.1 384.671744.1 384.679936.1 384.688128.1 384.696320.1 +448.352256.1 448.352256.2 448.356352.2 448.360448.1 448.360448.2 448.364544.2 448.368640.1 448.368640.2 +448.372736.2 448.376832.1 448.376832.2 448.380928.2 448.524109.2 448.524288.1 448.532480.1 448.540672.1 +448.548864.1 448.557056.1 448.565248.1 448.573440.1 448.581632.1 448.589824.1 448.598016.1 512.69632.2 +512.102400.2 512.135168.2 512.139264.1 512.139264.2 512.167936.2 512.172032.1 512.172032.2 512.200704.2 +512.204800.1 512.204800.2 512.208896.2 512.262144.1 512.262144.2 512.266240.2 512.270336.1 512.270336.2 +512.274432.2 512.278528.1 512.278528.2 512.323584.2 512.327680.1 512.327680.2 512.331776.2 512.335872.1 +512.335872.2 512.339968.2 512.344064.1 512.344064.2 512.348160.2 512.385024.1 512.385024.2 512.389120.2 +512.393216.1 512.393216.2 512.397312.2 512.401408.1 512.401408.2 512.405504.2 512.409600.1 512.409600.2 +512.413696.2 512.417792.1 512.417792.2 512.466944.1 512.466944.2 512.471040.2 512.475136.1 512.475136.2 +512.479232.2 512.483328.1 512.483328.2 512.487424.2 512.491520.1 512.491520.2 512.495616.2 512.499712.1 +512.499712.2 512.503808.2 512.507904.1 512.507904.2 512.512000.2 512.516096.1 512.516096.2 512.520192.2 +512.523951.1 576.233472.2 576.237568.1 576.237568.2 576.241664.2 576.245760.1 576.245760.2 576.421888.2 +576.425984.1 576.425984.2 576.430080.2 576.434176.1 576.434176.2 576.438272.2 576.442368.1 576.442368.2 +576.446464.2 576.450560.1 576.450560.2 576.454656.2 576.458752.1 576.458752.2 576.462848.2 640.385024.1 +640.385024.2 640.389120.2 640.393216.1 640.393216.2 640.397312.2 640.401408.1 640.401408.2 640.405504.2 +640.409600.1 640.409600.2 640.413696.2 640.417792.1 640.417792.2 704.106496.1 704.106496.2 704.303104.1 +704.303104.2 704.307200.2 704.311296.1 704.311296.2 704.315392.2 704.319488.1 704.319488.2 704.352256.1 +704.352256.2 704.356352.2 704.360448.1 704.360448.2 704.364544.2 704.368640.1 704.368640.2 704.372736.2 +704.376832.1 704.376832.2 704.380928.2 740.261860.1 768.143360.2 768.323584.2 768.327680.1 768.327680.2 +768.331776.2 768.335872.1 768.335872.2 768.339968.2 768.344064.1 768.344064.2 768.348160.2 832.176128.2 +832.180224.1 832.180224.2 832.303104.1 832.303104.2 832.307200.2 832.311296.1 832.311296.2 832.315392.2 +832.319488.1 832.319488.2 896.212992.1 896.212992.2 896.217088.2 896.282624.2 896.286720.1 896.286720.2 +896.290816.2 896.294912.1 896.294912.2 896.299008.2 960.261975.2 960.262144.1 960.262144.2 960.266240.2 +960.270336.1 960.270336.2 960.274432.2 960.278528.1 960.278528.2 1024.36864.2 1024.73728.1 1024.73728.2 +1024.147456.1 1024.147456.2 1024.249856.2 1024.253952.1 1024.253952.2 1024.258048.2 1024.261860.1 1088.110592.2 +1088.233472.2 1088.237568.1 1088.237568.2 1088.241664.2 1088.245760.1 1088.245760.2 1152.184320.2 1152.188416.1 +1152.188416.2 1152.221184.1 1152.221184.2 1152.225280.2 1152.229376.1 1152.229376.2 1216.212992.1 1216.212992.2 +1216.217088.2 1280.114688.1 1280.114688.2 1280.151552.2 1280.200704.2 1280.204800.1 1280.204800.2 1280.208896.2 +1344.192512.2 1344.196608.1 1344.196608.2 1408.77824.2 1408.184320.2 1408.188416.1 1408.188416.2 1472.118784.2 +1472.176128.2 1472.180224.1 1472.180224.2 1536.155648.1 1536.155648.2 1536.159744.2 1536.167936.2 1536.172032.1 +1536.172032.2 1600.163840.1 1600.163840.2 1664.81920.1 1664.81920.2 1664.122880.1 1664.122880.2 1664.155648.1 +1664.155648.2 1664.159744.2 1728.40960.1 1728.40960.2 1728.151552.2 1749.130773.1 1792.147456.1 1792.147456.2 +1856.143360.2 1920.135168.2 1920.139264.1 1920.139264.2 1956.130930.2 1984.130930.2 1984.131072.1 1984.131072.2 +2048.86016.2 2048.126976.2 2048.130773.1 2176.122880.1 2176.122880.2 2240.118784.2 2304.45056.2 2304.90112.1 +2304.90112.2 2304.114688.1 2304.114688.2 2368.110592.2 2496.106496.1 2496.106496.2 2560.94208.2 2560.102400.2 +2688.98304.1 2688.98304.2 2816.49152.1 2816.49152.2 2816.94208.2 2944.90112.1 2944.90112.2 3072.86016.2 +3200.53248.2 3264.81920.1 3264.81920.2 3392.77824.2 3584.57344.1 3584.57344.2 3584.73728.1 3584.73728.2 +3840.69632.2 3861.65386.2 3904.61440.2 4032.65386.2 4032.65536.1 4032.65536.2 4288.61440.2 4608.57344.1 +4608.57344.2 4992.53248.2 5376.49152.1 5376.49152.2 5888.45056.2 6464.40960.1 6464.40960.2 7168.36864.2 +8128.32768.1 8128.32768.2 8192.4096.2 8192.8192.1 8192.8192.2 8192.12288.2 8192.16384.1 8192.16384.2 +8192.20480.2 8192.24576.1 8192.24576.2 8192.28672.2 +""" + + +def _long_context_dumped_cells(): + return [tuple(int(x) for x in tok.split(".")) for tok in _LONG_CONTEXT_DUMPED.split()] + + +# The s_len values of the row-axis sweeps: the tile-height changes at the 16K envelope (16128 is the longest s_len +# whose 16384 rows fit one call), the long-context runs' lengths (131072 / 262144 / 1048576, and the passkey +# prompts' 261860 / 523951 / 1048219), Lm itself (tallest tile 255, tiles of 192), both sides of the 64-row rounding +# (2^28 // width = 128 | 127 at 2096896 | 2097152, 64 | 63 at 4194048 | 4194304), and on to 2^25 (tiles of 7). +_LONG_CONTEXT_S_ROWS = ( + 16128, + 16384, + 32768, + 65536, + 131072, + 261860, + 262144, + 523951, + 524288, + 1048219, + _LONG_CONTEXT_LM, + 2096896, + 2097152, + 4194048, + 4194304, + 1 << 23, + 1 << 24, + _LONG_CONTEXT_S_MAX, +) +_LONG_CONTEXT_S_RECIPE = (16128, 131072, _LONG_CONTEXT_LM, 4194304, _LONG_CONTEXT_S_MAX) + + +def test_long_context_tile_heights(): + """The tiling at the lengths the long-context runs reached; every tile covers the rows once within the budget; + every dumped (rows, s_len) is a tile the tiling makes.""" + assert _long_context_tiles(_LONG_CONTEXT_TC, 16128) == [(0, _LONG_CONTEXT_TC)] + assert _long_context_rows_max(16384) == 16131 and _long_context_tiles(_LONG_CONTEXT_TC, 16384)[ + :2 + ] == [ + (0, 16128), + (16128, _LONG_CONTEXT_TC), + ] + assert _long_context_tiles(_LONG_CONTEXT_TC, 523951)[0] == (0, 512) and _long_context_tiles( + _LONG_CONTEXT_TC, 1048219 + )[0] == (0, 256) + assert _long_context_rows_max(_LONG_CONTEXT_LM) == 255 and _long_context_tiles( + _LONG_CONTEXT_TC, _LONG_CONTEXT_LM + )[0] == (0, 192) + assert [_long_context_rows_max(S) for S in (2096896, 2097152, 4194048, 4194304)] == [ + 128, + 127, + 64, + 63, + ] + assert _long_context_tiles(_LONG_CONTEXT_TC, 2097152)[0] == (0, 64) and _long_context_tiles( + _LONG_CONTEXT_TC, 4194304 + )[0] == (0, 63) + assert ( + _long_context_rows_max(_LONG_CONTEXT_S_MAX) == 7 + and len(_long_context_tiles(_LONG_CONTEXT_TC, _LONG_CONTEXT_S_MAX)) == 2341 + ) + for S in _LONG_CONTEXT_S_ROWS + (1, 16383, 3_000_000, 12_345_677): + for rows in (1, 7, 63, 64, 65, 255, 256, 4097, _LONG_CONTEXT_TC - 1, _LONG_CONTEXT_TC): + tiles = _long_context_tiles(rows, S) + assert tiles[0][0] == 0 and tiles[-1][1] == rows + assert all(t[1] == u[0] for t, u in zip(tiles, tiles[1:])) + assert all( + 0 < b - a <= _long_context_rows_max(S) + and (b - a) * _row_stride(S) <= _LONG_CONTEXT_BUDGET + for a, b in tiles + ) + for rows, S, _ in _long_context_dumped_cells(): + assert rows <= _long_context_rows_max(S), (rows, S) + + +def _long_context_full(regime, T, S, ks, ke, seed): + """Inputs of one set and the T x S call, checked against the float64 reference.""" + if regime == "exact": + q, q_sf, kv, kv_sf, w_raw = _exact_inputs(T, S, _HEADS, seed) + w = w_raw * GLOBAL_FOLD + else: + q, q_sf, kv, kv_sf, w = _recipe_inputs(T, S, _HEADS, seed) + full = _call(q, q_sf, kv, kv_sf, w, ks, ke) + ref, bound = _long_context_reference(q, q_sf, w, kv, kv_sf, ks, ke) + if regime == "exact": + _assert_exact_premise(bound) + assert torch.equal(full, ref), ( + f"{int((full != ref).sum())} of {full.numel()} elements differ" + ) + else: + assert torch.equal(torch.isneginf(full), torch.isneginf(ref)) + torch.testing.assert_close(full, ref) + del ref + return q, q_sf, kv, kv_sf, w, full + + +_LONG_CONTEXT_ROW_CASES = ( + [("exact", "r1", S) for S in _LONG_CONTEXT_S_ROWS] + + [("exact", "packed", S) for S in _LONG_CONTEXT_S_ROWS] + + [("exact", "r2", 16128)] + + [("recipe", "packed", S) for S in _LONG_CONTEXT_S_RECIPE] +) + + +@pytest.mark.parametrize("regime,form,S", _LONG_CONTEXT_ROW_CASES) +def test_long_context_row_axis_every_tile_height(regime, form, S): + """Every tile height the target makes at s_len S, 1 .. min(16384, 2^28 // width): the first T rows (views of the + tallest tile's inputs) over the same S keys give exactly the first T rows of the tallest tile, which itself + matches float64. ``r1``: one request whose last row sees key S - 1; ``packed``: three requests, keys at their own + offsets. Harness sighting as in ``test_row_axis_every_T``.""" + T_max = _long_context_rows_max(S) + ks, ke = _axis_ranges(form, T_max, S) + sets = [ + _long_context_full(regime, T_max, S, ks, ke, seed=S % 100_003 + 7 * k + len(form)) + for k in (0, 1) + ] + torch.cuda.synchronize() + bad = torch.zeros((), dtype=torch.int64, device=DEV) + shift = {} + poisoned, reused = set(), 0 + for T in range(T_max, 0, -1): + q, q_sf, kv, kv_sf, w, full = sets[T % 2] + out = _call(q[:T], q_sf[:T], kv, kv_sf, w[:T], ks[:T], ke[:T]) + assert out.shape == (T, S) and out.stride() == (_row_stride(S), 1) + assert out.untyped_storage().nbytes() == _logits_bytes(T, S, _HEADS) + reused += out.data_ptr() in poisoned + bad += (out != full[:T]).sum() + if ( + T < T_max + and T in (1, 2, 5, T_max // 2 + 1, T_max - 1) + and bool((ke[: T + 1] > ks[: T + 1]).any()) + ): + shift[T] = (out != full[1 : T + 1]).sum() # control: the gate must see a one-row shift + _poison(out) + poisoned.add(out.data_ptr()) + del out + shift = {T: int(n) for T, n in shift.items()} + print( + f"tile heights {regime}/{form} S {S}: T 1..{T_max}, {int(bad)} mismatching elements; " + f"{reused}/{T_max} results on a NaN-filled block; one-row-shift control mismatches {shift}" + ) + assert int(bad) == 0 + assert reused >= T_max // 2 + assert shift and all(n > 0 for n in shift.values()) + + +def _long_context_key_sample(S_max, seed): + """S in (16384, S_max] (``test_key_axis`` covers 1..32768): S = 0, 1, 2, 3 or -1 mod 256 within 2048 of every + power of two from 2^15, the row-axis lengths, the dumped s_len, S_max, and 200 seeded log-uniform draws.""" + s = {S_max} + for k in range(15, S_max.bit_length()): + for c in range((2**k - 2048) // 256, (2**k + 2048) // 256 + 1): + s.update(c * 256 + d for d in (-1, 0, 1, 2, 3)) + s.update(_LONG_CONTEXT_S_ROWS) + s.update(S for _, S, _ in _long_context_dumped_cells()) + rng = random.Random(seed) + s.update(int(2 ** rng.uniform(14, math.log2(S_max))) for _ in range(200)) + return sorted(x for x in s if 16384 < x <= S_max) + + +@pytest.mark.parametrize("form", ["r1", "packed"]) +@pytest.mark.parametrize("T,S_max", [(7, _LONG_CONTEXT_S_MAX), (255, _LONG_CONTEXT_LM)]) +def test_long_context_key_axis(T, S_max, form): + """At the tallest tile of 2^25 (7 rows) and of Lm (255 rows): the first S keys, with this form's ranges at S, + give on every visible element exactly the T x S_max call in which every row sees every key (itself bit-identical + to float64), and -inf everywhere else.""" + ks0 = torch.zeros(T, dtype=torch.int32, device=DEV) + ke0 = torch.full((T,), S_max, dtype=torch.int32, device=DEV) + sets = [ + _long_context_full("exact", T, S_max, ks0, ke0, seed=T + 3 * k + len(form)) for k in (0, 1) + ] + torch.cuda.synchronize() + s_values = _long_context_key_sample(S_max, seed=T)[::-1] + probe = { + s_values[1], + s_values[2], + s_values[len(s_values) // 2], + s_values[-1], + 1 << 24, + (1 << 24) + 1, + } + bad = torch.zeros((), dtype=torch.int64, device=DEV) + n_vis = torch.zeros((), dtype=torch.int64, device=DEV) + shift = {} + poisoned, reused = set(), 0 + for S in s_values: + q, q_sf, kv, kv_sf, w, full = sets[S % 2] + ks, ke = _axis_ranges(form, T, S) + out = _call(q, q_sf, kv[:S], kv_sf[:S], w, ks, ke) + assert out.shape == (T, S) and out.stride() == (_row_stride(S), 1) + reused += out.data_ptr() in poisoned + col = torch.arange(S, device=DEV) + vis = (col[None, :] >= ks[:, None]) & (col[None, :] < ke[:, None]) + bad += ((out != full[:, :S]) & vis).sum() + (~torch.isneginf(out) & ~vis).sum() + n_vis += vis.sum() + if S in probe: + shift[S] = ( + (out != full[:, 1 : S + 1]) & vis + ).sum() # control: a one-key shift must fail + _poison(out) + poisoned.add(out.data_ptr()) + del out, vis + shift = {S: int(n) for S, n in shift.items()} + print( + f"key axis {form} T {T}, {len(s_values)} values of S in {s_values[-1]}..{s_values[0]}: {int(bad)} " + f"mismatching elements ({int(n_vis)} visible checked); {reused}/{len(s_values)} results on a NaN-filled " + f"block; one-key-shift control mismatches {shift}" + ) + assert int(bad) == 0 + assert reused >= len(s_values) // 2 + assert shift and all(n > 0 for n in shift.values()) + + +@pytest.mark.parametrize("r", [1, 2]) +def test_long_context_row_logits_do_not_depend_on_the_call(r): + """A row's logits do not depend on the other query rows or key rows of the call, bitwise, in the recipe regime + (where fp32 rounding is not exact). Rows: the last 16384 positions of one Lm-token request at ratio r. The same + row gives the same logits (1) in one untiled call of all 16384 rows (64 GiB of logits at r 1), (2) in its tile + of the target's tiling, (3) called alone over just its own keys, and (4) with the request's keys shifted inside + a larger gather of other keys: ending a 2^25-key gather (7-row tiles), and at an odd offset (1,234,567) of a + 3,000,000-key one (64-row tiles); the last 256 rows as that step's tail rows, some also alone there.""" + S, R = _LONG_CONTEXT_LM // r, _LONG_CONTEXT_TC + q, q_sf, kv, kv_sf, w = _recipe_inputs(R, S, _HEADS, seed=60 + r) + ks, ke = _chunk(R, S, r) + whole = _call(q, q_sf, kv, kv_sf, w, ks, ke) + assert whole.untyped_storage().nbytes() == _logits_bytes(R, S, _HEADS) + last = slice(R - 64, R) + ref, _ = _long_context_reference(q[last], q_sf[last], w[last], kv, kv_sf, ks[last], ke[last]) + assert torch.equal(torch.isneginf(whole[last]), torch.isneginf(ref)) + torch.testing.assert_close(whole[last], ref) + del ref + counts = collections.Counter() + tiles = _long_context_tiles(R, S) + for a, b in tiles: + out = _call(q[a:b], q_sf[a:b], kv, kv_sf, w[a:b], ks[a:b], ke[a:b]) + counts["tiles"] += int((out != whole[a:b]).sum()) + del out + rng = random.Random(70 + r) + alone = {0, 1, 2, 3, 4, 5, R - 5, R - 4, R - 3, R - 2, R - 1} + alone |= {t for a, _ in tiles[1:] for t in (a - 1, a)} + alone |= set(rng.sample(range(R), 48)) + for t in sorted(alone): + n = int(ke[t]) # ks is 0: the row's own keys are the first n + out = _call( + q[t : t + 1], + q_sf[t : t + 1], + kv[:n], + kv_sf[:n], + w[t : t + 1], + ks[t : t + 1], + ke[t : t + 1], + ) + counts["alone"] += int((out[0] != whole[t, :n]).sum()) + del out + tail = R - 256 + for s_big, off in ((_LONG_CONTEXT_S_MAX, _LONG_CONTEXT_S_MAX - S), (3_000_000, 1_234_567)): + pre, pre_sf = _rand_packed(off, _gen(80 + r)) + post, post_sf = _rand_packed(s_big - off - S, _gen(90 + r)) + big, big_sf = torch.cat([pre, kv, post]), torch.cat([pre_sf, kv_sf, post_sf]) + del pre, pre_sf, post, post_sf + ks2, ke2 = ks + off, ke + off + for a, b in _long_context_tiles(256, s_big): + a, b = a + tail, b + tail + out = _call(q[a:b], q_sf[a:b], big, big_sf, w[a:b], ks2[a:b], ke2[a:b]) + counts[f"shifted {s_big}"] += int((out[:, off : off + S] != whole[a:b]).sum()) + counts[f"shifted {s_big}"] += int((~torch.isneginf(out[:, :off])).sum()) + counts[f"shifted {s_big}"] += int((~torch.isneginf(out[:, off + S :])).sum()) + del out + for t in (tail, tail + 1, R - 2, R - 1): + out = _call( + q[t : t + 1], + q_sf[t : t + 1], + big, + big_sf, + w[t : t + 1], + ks2[t : t + 1], + ke2[t : t + 1], + ) + counts[f"shifted {s_big}, alone"] += int((out[0, off : off + S] != whole[t]).sum()) + del out + del big, big_sf + print( + f"ratio {r}: S {S}, {len(tiles)} tiles of {tiles[0][1]}, {len(alone)} rows alone; mismatching " + f"elements vs the untiled call: {dict(counts)}" + ) + assert set(counts.values()) == {0}, counts + + +_LONG_CONTEXT_N_PARTS, _LONG_CONTEXT_PART_BLOCKS = 200, 4 + + +def _long_context_random_partition(rng): + """A context step of the long-context form. T in 1..16384 (half uniform, half log-uniform) split at random + cuts into k in 1..min(256, T) requests (half uniform, half log-uniform), in random order; ratio 1 or 2. A + request is a fresh prompt a third of the time, else cached + new is log-uniform in new..Lm, or, in a quarter of + the steps (long-context steps), uniform in Lm/2..Lm. When the step's s_len would pass 2^25, every cached prefix + is scaled down to fit.""" + T = ( + rng.randint(1, _LONG_CONTEXT_TC) + if rng.random() < 0.5 + else min(_LONG_CONTEXT_TC, int(2 ** rng.uniform(0, 14))) + ) + kmax = min(_LONG_CONTEXT_K, T) + k = ( + rng.randint(1, kmax) + if rng.random() < 0.5 + else min(kmax, int(2 ** rng.uniform(0, math.log2(kmax + 1)))) + ) + cuts = sorted(rng.sample(range(1, T), k - 1)) + new = [b - a for a, b in zip([0] + cuts, cuts + [T])] + rng.shuffle(new) + r = rng.choice((1, 2)) + long_ctx = rng.random() < 1 / 4 + reqs = [] + for n in new: + if rng.random() < 1 / 3: + total = n + elif long_ctx: + total = rng.randint(max(n, _LONG_CONTEXT_LM // 2), _LONG_CONTEXT_LM) + else: + total = max(n, int(2 ** rng.uniform(math.log2(n), math.log2(_LONG_CONTEXT_LM)))) + reqs.append((total - n, n)) + S = sum((c + n) // r for c, n in reqs) + if S > _LONG_CONTEXT_S_MAX: + f = (_LONG_CONTEXT_S_MAX * r - sum(n for _, n in reqs) - r * k) / sum(c for c, _ in reqs) + reqs = [(int(c * f), n) for c, n in reqs] + return reqs, r + + +@pytest.mark.parametrize("block", range(_LONG_CONTEXT_PART_BLOCKS)) +def test_long_context_random_partitions(block): + """200 seeded random context steps (50 per block) of the long-context form, each tiled as the target tiles it; + up to 16 tiles per step (all when there are at most 16; else the first, the last and 14 at random) called as the + target calls them and checked, every element, against the palette reference.""" + pal = _LongContextPalette(seed=900 + block) + per = _LONG_CONTEXT_N_PARTS // _LONG_CONTEXT_PART_BLOCKS + stats = collections.Counter() + for i in range(block * per, (block + 1) * per): + rng = random.Random(7000 + i) + reqs, r = _long_context_random_partition(rng) + assert ( + 1 <= len(reqs) <= _LONG_CONTEXT_K and 1 <= sum(n for _, n in reqs) <= _LONG_CONTEXT_TC + ) + assert all(n >= 1 and c >= 0 and c + n <= _LONG_CONTEXT_LM for c, n in reqs) + ks, ke, S = _packed(reqs, r) + assert S <= _LONG_CONTEXT_S_MAX + if S == 0: + stats["skipped (S 0)"] += 1 + continue + T = ks.numel() + qi, ki, inputs = pal.draw(T, S, _gen(8000 + i)) + tiles = _long_context_tiles(T, S) + chk = _long_context_sample_tiles(tiles, rng) + stats["visible checked"] += _long_context_check_tiles(pal, qi, ki, inputs, ks, ke, chk) + stats["steps"] += 1 + stats[f"ratio {r}"] += 1 + stats["rows"] += T + stats["tiles"] += len(tiles) + stats["tiles checked"] += len(chk) + stats["rows checked"] += sum(b - a for a, b in chk) + stats["steps with S > 1M"] += S > 1 << 20 + stats["steps with S > 2^24"] += S > 1 << 24 + stats["max S"] = max(stats["max S"], S) + stats["max k"] = max(stats["max k"], len(reqs)) + stats["max tiles"] = max(stats["max tiles"], len(tiles)) + del inputs + print(f"long-context partitions block {block}: {dict(stats)}") + assert stats["steps"] + stats["skipped (S 0)"] == per + + +def _long_context_tail_step(rng): + """A decoder-replay step's context requests in batch order, (kv_len, has tail rows): k in 1..128 bounded + requests finishing their prompt here (kv_len = the prompt, log-uniform in 129..Lm, or in a quarter of the steps + uniform in Lm/2..Lm), with m in 0..32 more context requests that have no tail rows here (kv_len log-uniform in + 1..Lm); scaled down to s_len <= 2^25.""" + k = rng.randint(1, 128) if rng.random() < 0.5 else min(128, int(2 ** rng.uniform(0, 7.01))) + m = rng.randint(0, 32) if rng.random() < 0.5 else 0 + if rng.random() < 1 / 4: + reqs = [(rng.randint(_LONG_CONTEXT_LM // 2, _LONG_CONTEXT_LM), True) for _ in range(k)] + else: + reqs = [ + (max(129, int(2 ** rng.uniform(math.log2(129), math.log2(_LONG_CONTEXT_LM)))), True) + for _ in range(k) + ] + reqs += [(int(2 ** rng.uniform(0, math.log2(_LONG_CONTEXT_LM))), False) for _ in range(m)] + rng.shuffle(reqs) + total = sum(n for n, _ in reqs) + if total > _LONG_CONTEXT_S_MAX: + f = (_LONG_CONTEXT_S_MAX - 129 * len(reqs)) / total + reqs = [(max(129 if tail else 1, int(n * f)), tail) for n, tail in reqs] + return reqs + + +def _long_context_tail_rows(reqs): + """The tail rows' ranges: each bounded request's last 128 positions p see [cu_b, cu_b + p + 1) of the gather.""" + ks, ke, off = [], [], 0 + for n, tail in reqs: + if tail: + p = torch.arange(n - 128, n, device=DEV) + ks.append(torch.full((128,), off, dtype=torch.int32, device=DEV)) + ke.append((off + p + 1).to(torch.int32)) + off += n + return torch.cat(ks), torch.cat(ke), off + + +@pytest.mark.parametrize("block", range(2)) +def test_long_context_tail_cand_steps(block): + """The decoder-replay tail form (``deepgemm_fp8_fp4_mqa_logits.tail_cand``): 100 seeded steps (50 per block), 128 + rows per bounded request against every ratio-1 key of the step, tiled as the target tiles them; up to 16 tiles + per step checked, every element, against the palette reference.""" + pal = _LongContextPalette(seed=930 + block) + stats = collections.Counter() + for i in range(block * 50, (block + 1) * 50): + rng = random.Random(9500 + i) + reqs = _long_context_tail_step(rng) + ks, ke, S = _long_context_tail_rows(reqs) + assert S <= _LONG_CONTEXT_S_MAX and all(n <= _LONG_CONTEXT_LM for n, _ in reqs) + T = ks.numel() + qi, ki, inputs = pal.draw(T, S, _gen(9600 + i)) + tiles = _long_context_tiles(T, S) + chk = _long_context_sample_tiles(tiles, rng) + stats["visible checked"] += _long_context_check_tiles(pal, qi, ki, inputs, ks, ke, chk) + stats["steps"] += 1 + stats["tail requests"] += T // 128 + stats["tiles checked"] += len(chk) + stats["rows checked"] += sum(b - a for a, b in chk) + stats["max rows"] = max(stats["max rows"], T) + stats["steps with S > 2^24"] += S > 1 << 24 + stats["max S"] = max(stats["max S"], S) + del inputs + print(f"tail-cand steps block {block}: {dict(stats)}") + assert stats["steps"] == 50 + + +def test_long_context_dumped_tiles(): + """Every (rows, s_len, r) of the long-context runs' dumps as one tile: the last ``rows`` positions of a request + whose last row sees key s_len - 1, against the palette reference.""" + pal = _LongContextPalette(seed=950) + cells = _long_context_dumped_cells() + n_vis = 0 + for i, (rows, S, r) in enumerate(cells): + ks, ke = _chunk(rows, S, r) + qi, ki, inputs = pal.draw(rows, S, _gen(9900 + i)) + n_vis += _long_context_check_tiles(pal, qi, ki, inputs, ks, ke, [(0, rows)]) + del inputs + print( + f"dumped tiles: {len(cells)} cells, s_len up to {max(S for _, S, _ in cells)}, {n_vis} visible checked" + ) + assert len(cells) == 471 + + +def test_long_context_kernel_variant_map(): + """The tiles of the long-context form launch the same one kernel, with the same template arguments, as the + row and key axes' variant map: heights 1, 4, 7, the tallest and one less, and the last tile of a 16384-row step, + at every row-axis s_len up to 2^25, packed ranges.""" + q, q_sf, kv, kv_sf, w_raw = _exact_inputs( + _LONG_CONTEXT_TC, _LONG_CONTEXT_S_MAX, _HEADS, seed=44 + ) + w = w_raw * GLOBAL_FOLD + cases = [] + for S in _LONG_CONTEXT_S_ROWS: + a, b = _long_context_tiles(_LONG_CONTEXT_TC, S)[-1] + for T in sorted( + {1, 4, 7, _long_context_rows_max(S), max(1, _long_context_rows_max(S) - 1), b - a} + ): + cases.append((T, S, *_packed3(T, S))) + _call( + q[:4], q_sf[:4], kv, kv_sf, w[:4], *_chunk(4, _LONG_CONTEXT_S_MAX, 1) + ) # JIT, if any, outside the profile + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for T, S, ks, ke in cases: + out = _call(q[:T], q_sf[:T], kv[:S], kv_sf[:S], w[:T], ks, ke) + del out + torch.cuda.synchronize() + names = collections.Counter( + e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA + ) + for name, n in names.items(): + print(f"kernel x{n}: {name}") + print( + f"long-context variant map: {len(cases)} calls, S {min(c[1] for c in cases)}..{max(c[1] for c in cases)}" + ) + assert set(names.values()) == {len(cases)}, names + assert len(names) == 1, names + (name,) = names + assert name.startswith(_LOGITS_KERNEL), name + + +def test_long_context_wrong_answers_are_seen(): + """The palette check fails for each mistake a long gather could hide. The step: the last 63 positions of an + Lm-token request whose keys end a 2^25-key gather; its last 7-row tile. Wrong expectations: ranges one key late + (only the range edges differ), keys read one row late, key index wrapped at 2^24, the request's gather offset + ignored, the tile's row offset ignored.""" + pal = _LongContextPalette(seed=960) + S, off = _LONG_CONTEXT_S_MAX, _LONG_CONTEXT_S_MAX - _LONG_CONTEXT_LM + p = torch.arange(_LONG_CONTEXT_LM - 63, _LONG_CONTEXT_LM, device=DEV) + ks = torch.full((63,), off, dtype=torch.int32, device=DEV) + ke = (off + p + 1).to(torch.int32) + qi, ki, inputs = pal.draw(63, S, _gen(961)) + a, b = _long_context_tiles(63, S)[-1] + assert (a, b) == (56, 63) + q, q_sf, kv, kv_sf, w = inputs + out = _call(q[a:b], q_sf[a:b], kv, kv_sf, w[a:b], ks[a:b], ke[a:b]) + assert torch.equal(out, pal.expected(qi[a:b], ki, ks[a:b], ke[a:b])) + wrapped = ki[torch.arange(S, device=DEV) % (1 << 24)] + wrong = { + "ranges one key late": pal.expected(qi[a:b], ki, ks[a:b] + 1, ke[a:b] + 1), + "keys read one row late": pal.expected(qi[a:b], ki.roll(1), ks[a:b], ke[a:b]), + "key index wrapped at 2^24": pal.expected(qi[a:b], wrapped, ks[a:b], ke[a:b]), + "gather offset ignored": pal.expected(qi[a:b], ki, ks[a:b] - off, ke[a:b] - off), + "tile row offset ignored": pal.expected(qi[: b - a], ki, ks[a:b], ke[a:b]), + } + for name, exp in wrong.items(): + n = int((out != exp).sum()) + print( + f"control [{name}]: {n} of {out.numel()} elements differ ({int((ke[a:b] - ks[a:b]).sum())} visible)" + ) + assert n > 0 + del exp diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_deepgemm_fp8_fp4_paged_mqa_logits.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_deepgemm_fp8_fp4_paged_mqa_logits.py new file mode 100644 index 000000000000..9fca52fb53f6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_deepgemm_fp8_fp4_paged_mqa_logits.py @@ -0,0 +1,2547 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the deepgemm_fp8_fp4_paged_mqa_logits and deepgemm_paged_mqa_logits_metadata entries. + +The entries own what is certified: `CELLS` is the list, `reference` and `compare` are the gate. +The schedule has no plain-torch reference; it is certified at the logits cells, as their input, +and its values are gated by the logits they drive. This file turns a cell's spec into tensors and +keeps the checks a cell cannot carry: the exact head-reduction order, input immutability, CUDA-graph +capture, the runtime widths, pools and mixed-step views, every batch and width as a prefix, the +loud rejections, the wrong-layout controls, the pool-size bounds and the long-context envelope. +``_Step.reference`` is the float64 twin of the entry's reference that reads the keys from their +logical storage (so it also serves pools held as views of a larger allocation) and takes the +wrong variants the controls need; ``_check`` applies the entry's gate wherever the step holds its +pool. + +The DeepSeek-V4.1 decode / DSpark-verify indexer logits over the native MXFP4 index-K pool: +``score[b*next_n + j, s] = sum_h relu(q[b, j, h] . k[s]) * w[b*next_n + j, h]`` for the compressed +positions ``s < (p_b + j + 1) // ratio`` of request ``b`` whose first query token sits at absolute +position ``p_b``. Every input is packed here from native torch ops (no DeepGEMM helper, no other +catalog entry), and the reference decodes those same bytes and computes the score in fp64. +""" + +import re +import time +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.deepgemm_fp8_fp4_paged_mqa_logits import ( # noqa: E501 + deepgemm_fp8_fp4_paged_mqa_logits as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.deepgemm_fp8_fp4_paged_mqa_logits import ( # noqa: E501 + deepgemm_paged_mqa_logits_metadata as metadata_op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "deepgemm_fp8_fp4_paged_mqa_logits requires a CUDA device" + +DEV = "cuda" +HEADS = 32 # index_n_heads, replicated on every rank +DIM = 128 # index_head_dim +KEY_BYTES = 68 # 64 E2M1 code bytes + 4 UE8M0 scale bytes per key +# The pool's page count is the cache manager's number, never the batch's: each case allocates +# POOL_FACTOR x the pages its batch needs (+ POOL_EXTRA) and scatters the batch's pages over it +# with a random permutation, so "the page the block table names" and "the page this request +# happens to own" are different numbers. Chosen here, not taken from a runtime dump. +POOL_FACTOR = 3 +POOL_EXTRA = 8 +# Logits rows are padded to a 1024-byte multiple of a 256-column split (fp32: 256 columns). +SPLIT_KV = 256 + +_E2M1 = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float64, + device=DEV, +) +_MIDPOINTS = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], device=DEV) + + +# ─── packing (native torch; the reference recipe of the checkpoint's fp4_act_quant) ─────────── + + +def _quant_mxfp4(x: torch.Tensor): + """x [..., 128] -> (int8 [..., 64], int32 [...]). + + Per 32-block: scale = 2^ceil(log2(max(amax, 6*2^-126) / 6)) stored as a biased UE8M0 byte + (byte b of the int32 = block b); codes = E2M1(x / scale) rounded to nearest-even, element 2i + in the low nibble of byte i. + """ + xf = x.float().unflatten(-1, (4, 32)) + amax = xf.abs().amax(-1).clamp_min(6 * 2.0**-126) + bits = (amax * (1.0 / 6.0)).contiguous().view(torch.int32) + e = ((bits >> 23) & 0xFF) - 127 + ((bits & 0x7FFFFF) != 0).to(torch.int32) + scale = ((e + 127) << 23).view(torch.float32) + y = (xf / scale.unsqueeze(-1)).clamp(-6.0, 6.0) + a = y.abs() + code = (a.unsqueeze(-1) > _MIDPOINTS).sum(-1) + code = code + ((a == 0.75) | (a == 1.75) | (a == 3.5)).to(code.dtype) # ties to even code + code = (code + 8 * ((y < 0) & (code > 0)).to(code.dtype)).flatten(-2) + packed = (code[..., 0::2] | (code[..., 1::2] << 4)).to(torch.uint8).view(torch.int8) + sf = (e + 127).to(torch.uint8).contiguous().view(torch.int32).squeeze(-1) + return packed.contiguous(), sf.contiguous() + + +def _decode(packed: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """int8 [..., 64] + int32 [...] -> fp64 [..., 128], exact.""" + p = packed.view(torch.uint8).long() + v = torch.stack([_E2M1[p & 0xF], _E2M1[p >> 4]], dim=-1).flatten(-2) + sb = sf.contiguous().view(torch.uint8).view(*sf.shape, 4).long() + # 2^(byte - 127) built from its fp64 bit pattern: exact (torch.pow(2.0, e) on CUDA is not, + # it was measured 1 ulp low for some e) + scale = ((sb + (1023 - 127)) << 52).view(torch.float64) + return (v.unflatten(-1, (4, 32)) * scale.unsqueeze(-1)).flatten(-2) + + +def _write_pool(codes: torch.Tensor, sf: torch.Tensor, layout: str = "paged") -> torch.Tensor: + """codes int8 [P, bk, 64] + sf int32 [P, bk] -> uint8 [P, bk, 1, 68]. + + "paged" is the kernel's layout: a page holds its bk keys' 64 code bytes back to back, then + their bk 4-byte scale words. "interleaved" (a control) stores each key's 68 bytes together. + """ + P, bk, _ = codes.shape + if layout == "paged": + flat = torch.cat( + [ + codes.view(torch.uint8).reshape(P, bk * 64), + sf.view(torch.uint8).reshape(P, bk * 4), + ], + dim=1, + ) + else: + flat = torch.cat([codes.view(torch.uint8), sf.view(torch.uint8).view(P, bk, 4)], dim=2) + return flat.contiguous().view(P, bk, 1, KEY_BYTES) + + +# ─── one decode step ────────────────────────────────────────────────────────────────────────── + + +class _Step: + """A decode / DSpark-verify step of one indexer layer. + + ``positions[b]`` is the absolute position of request b's first query token, or None for a + CUDA-graph padding request (no pages, context length 0). + """ + + def __init__( + self, + positions, + next_n, + ratio, + block_kv, + heads=HEADS, + seed=0, + extra_ctx=37, + max_ctx=None, + pool_pages=None, + ): + """``max_ctx`` fixes the static width (default: largest length + ``extra_ctx``). + ``pool_pages`` sizes the pool to a runtime page count: the batch's pages are then drawn at + random from the whole pool (always including its last page), and every other page holds + 0xFF bytes, whose E8M0 scale is NaN, so reading a page the table does not name turns a + visible element into NaN.""" + g = torch.Generator(device=DEV) + g.manual_seed(seed) + self.B, self.n, self.H, self.bk, self.ratio = len(positions), next_n, heads, block_kv, ratio + lens = [ + [0] * next_n if p is None else [(p + j + 1) // ratio for j in range(next_n)] + for p in positions + ] + self.lens = torch.tensor(lens, dtype=torch.int32, device=DEV) + # the static width a captured decode step would bake in: wider than this batch needs + self.max_ctx = int(self.lens.max()) + extra_ctx if max_ctx is None else max_ctx + assert self.max_ctx >= int(self.lens.max()) + pages = [(max(row) + block_kv - 1) // block_kv for row in lens] + width = (self.max_ctx + block_kv - 1) // block_kv + 2 + if pool_pages is None: + self.used = None + self.num_pages = POOL_FACTOR * sum(pages) + POOL_EXTRA + self.num_pages += self.num_pages % 2 # even, so a control can re-view it as 2x pages + self.k_codes = torch.empty( + self.num_pages, block_kv, DIM // 2, dtype=torch.int8, device=DEV + ) + self.k_sf = torch.empty(self.num_pages, block_kv, dtype=torch.int32, device=DEV) + self.refill_keys(g) + perm = torch.randperm(self.num_pages, generator=g, device=DEV).to(torch.int32) + else: + self.num_pages = pool_pages + perm = torch.randperm(pool_pages, generator=g, device=DEV)[: sum(pages)].to(torch.int32) + if sum(pages) and not bool((perm == pool_pages - 1).any()): + perm[-1] = pool_pages - 1 + self.used = perm.long() + self.k_codes = torch.full( + (pool_pages, block_kv, DIM // 2), -1, dtype=torch.int8, device=DEV + ) + self.k_sf = torch.full((pool_pages, block_kv), -1, dtype=torch.int32, device=DEV) + self.refill_keys(g) + self.block_table = torch.full((self.B, width), -1, dtype=torch.int32, device=DEV) + off = 0 + for b, npg in enumerate(pages): + self.block_table[b, :npg] = perm[off : off + npg] + off += npg + # Keys past a request's length inside its last page are never visible; poison their + # scale bytes with 0xFF (the E8M0 NaN encoding), as stale or unwritten slots may be. + sf_bytes = self.k_sf.view(torch.uint8).view(self.num_pages, block_kv, 4) + for b, row in enumerate(lens): + L = max(row) + if L % block_kv: + sf_bytes[self.block_table[b, L // block_kv].long(), L % block_kv :] = 0xFF + self.pool = _write_pool(self.k_codes, self.k_sf) + q = torch.randn(self.B, next_n, heads, DIM, generator=g, device=DEV) * 2.0 + self.q, self.q_sf = _quant_mxfp4(q) + # weights_proj(x) * 128^-0.5 * 32^-0.5 is a bf16 value; the kernel takes it as fp32 + w = torch.randn(self.B * next_n, heads, generator=g, device=DEV) * 2.0**-4 + self.weights = w.to(torch.bfloat16).float() + + def refill_keys(self, g): + if self.used is not None: + k = torch.randn(self.used.numel(), self.bk, DIM, generator=g, device=DEV) + self.k_codes[self.used], self.k_sf[self.used] = _quant_mxfp4(k) + return + for c in range(0, self.num_pages, 4096): + k = torch.randn(min(4096, self.num_pages - c), self.bk, DIM, generator=g, device=DEV) + self.k_codes[c : c + k.shape[0]], self.k_sf[c : c + k.shape[0]] = _quant_mxfp4(k) + + def run(self, **over): + a = dict( + q=self.q, + q_sf=self.q_sf, + kv_cache=self.pool, + weights=self.weights, + context_lens=self.lens, + block_table=self.block_table, + max_context_len=self.max_ctx, + ) + block_kv = over.pop("meta_block_kv", self.bk) + a.update(over) + meta = metadata_op(a["context_lens"], block_kv) + return op(schedule_meta=meta, **a) + + def _keys(self, b): + L = int(self.lens[b].max()) + npg = (L + self.bk - 1) // self.bk + idx = self.block_table[b, :npg].long() + return _decode(self.k_codes[idx], self.k_sf[idx]).reshape(-1, DIM)[:L] + + def _dots(self, b): + """exact per-head q.k in fp64: [n, H, L]""" + return torch.einsum("jhd,sd->jhs", _decode(self.q[b], self.q_sf[b]), self._keys(b)) + + def reference(self, relu=True, weights=None): + """fp64 [B*n, max_ctx]; -inf at every column row (b, j) must not see.""" + w = (self.weights if weights is None else weights).double() + out = torch.full( + (self.B * self.n, self.max_ctx), float("-inf"), dtype=torch.float64, device=DEV + ) + for b in range(self.B): + if int(self.lens[b].max()) == 0: + continue + d = self._dots(b) + d = d.clamp_min(0) if relu else d + score = (d * w[b * self.n : (b + 1) * self.n].unsqueeze(-1)).sum(1) + for j in range(self.n): + L = int(self.lens[b, j]) + out[b * self.n + j, :L] = score[j, :L] + return out + + def emulate(self): + """fp32 emulation of the kernel's head reduction: per head a = q.k (exact in fp32); + four FMA chains c_l = fma(a_h + |a_h|, w_h, c_l) over heads h = l, l+4, ...; then + ((c_0 + c_2) + (c_1 + c_3)) / 2. fma is formed exactly in fp64 and rounded once.""" + out = torch.full((self.B * self.n, self.max_ctx), float("-inf"), device=DEV) + for b in range(self.B): + if int(self.lens[b].max()) == 0: + continue + a = self._dots(b).float() + two_relu = (a + a.abs()).double() + w = self.weights[b * self.n : (b + 1) * self.n].double() + chains = [] + for lane in range(4): + acc = torch.zeros(self.n, a.shape[-1], device=DEV) + for h in range(lane, self.H, 4): + acc = (two_relu[:, h] * w[:, h : h + 1] + acc.double()).float() + chains.append(acc) + red = ((chains[0] + chains[2]) + (chains[1] + chains[3])) / 2 + for j in range(self.n): + L = int(self.lens[b, j]) + out[b * self.n + j, :L] = red[j, :L] + return out + + +def _aligned_width(max_ctx: int, dtype: torch.dtype) -> int: + split = -(-max_ctx // SPLIT_KV) * SPLIT_KV + per_1k = 1024 // torch.tensor([], dtype=dtype).element_size() + return -(-split // per_1k) * per_1k + + +def _check(step: _Step, out: torch.Tensor): + """Shape/stride contract, the entry's gate where the step holds its pool, default fp32 + assert_close vs the fp64 twin, and the exact reduction order.""" + rows = step.B * step.n + assert out.shape == (rows, step.max_ctx) and out.dtype == torch.float32 + assert out.stride() == (_aligned_width(step.max_ctx, torch.float32), 1) + if step.pool is not None: + op.compare( + out, + op.reference( + step.q, + step.q_sf, + step.pool, + step.weights, + step.lens, + step.block_table, + None, + step.max_ctx, + ), + ) + ref = step.reference() + vis = torch.isfinite(ref) + assert int(vis.sum()) == int(step.lens.sum()) + got = out[vis] + torch.testing.assert_close(got, ref.float()[vis]) + # Stronger than the tolerance: the kernel's fp32 result is exactly the documented reduction + # (measured bit-identical on every visible element on sm_103). + assert torch.equal(got, step.emulate()[vis]) + return ref, vis + + +def _outside_fraction(got: torch.Tensor, ref: torch.Tensor) -> float: + """Fraction of elements outside the default fp32 assert_close gate.""" + ref32 = ref.float() + tol = 1e-5 + 1.3e-6 * ref32.abs() + return ((got - ref32).abs() > tol).float().mean().item() + + +# ─── cases ──────────────────────────────────────────────────────────────────────────────────── + +# (ratio, page): the INDEXER_COMPRESS page holds tokens_per_block / ratio entries; +# tokens_per_block 128 gives (2, 64) and (1, 128); tokens_per_block 256 gives (2, 128). +PAGES = [(2, 64), (1, 128), (2, 128)] + + +def _positions(regime, B, ratio, block_kv, seed): + g = torch.Generator() + g.manual_seed(seed) + if regime == "boundary": + # first-row compressed lengths on page and 256-column split edges + lens = [1, 2, block_kv - 1, block_kv, block_kv + 1, 255, 256, 257, 511, 513, 700, 1000] + pos = [lens[i % len(lens)] * ratio - 1 + (i // len(lens)) for i in range(B)] + else: + lo, hi = { + "gate": (1024, 16384), # the accuracy gates' contexts + "long": (20480, 40960), # the long-context run (>= 20K tokens) + "64k": (65536 * ratio - 8192, 65536 * ratio - 8), # ~64K compressed positions + }[regime] + pos = torch.randint(lo, hi, (B,), generator=g).tolist() + if B >= 8: + pos[B // 2] = None # a CUDA-graph padding request in the middle of the batch + return pos + + +def _cell_step(spec, seed): + """A cell's decode step: first-row positions from a regime, or the cell's own.""" + ratio, block_kv = spec["ratio"], spec["block_kv"] + if "regime" in spec: + positions = _positions(spec["regime"], spec["batch"], ratio, block_kv, seed) + else: + positions = spec["positions"] + return _Step(positions, spec["next_n"], ratio, block_kv, heads=spec["heads"], seed=seed) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + # The schedule is certified at the logits cells, as their input: both entries are armed, and + # the schedule's values are gated by the logits they drive. + assert [c.spec for c in metadata_op.CELLS] == [c.spec for c in op.CELLS] + certifying(op, metadata_op) + spec = cell.spec + dtype = spec.get("logits_dtype", torch.float32) + step = _cell_step(spec, zlib.crc32(cell.why.encode())) + before = [t.clone() for t in (step.q, step.q_sf, step.weights, step.lens, step.block_table)] + pool_before = step.pool.clone() if step.pool.numel() < (1 << 26) else None + meta = metadata_op(step.lens, step.bk) + metadata_op.compare(meta, metadata_op.reference(step.lens, step.bk)) + args = (step.q, step.q_sf, step.pool, step.weights, step.lens, step.block_table, meta) + out = op(*args, step.max_ctx, dtype) + op.compare(out, op.reference(*args, step.max_ctx, dtype)) + assert out.shape == (step.B * step.n, step.max_ctx) and out.dtype == dtype + assert out.stride() == (_aligned_width(step.max_ctx, dtype), 1) + # Stronger than the gate: the kernel's fp32 result is exactly the documented head reduction + # (bf16 logits: that value rounded once), on every visible element. + vis = torch.arange(step.max_ctx, device=DEV)[None, :] < step.lens.view(-1, 1) + assert torch.equal(out[vis], step.emulate()[vis].to(dtype)) + # inputs are not modified + for t, t0 in zip((step.q, step.q_sf, step.weights, step.lens, step.block_table), before): + assert torch.equal(t, t0) + if pool_before is not None: + assert torch.equal(step.pool, pool_before) + + +def test_cuda_graph_replay(): + """Metadata + logits captured together. Each replay reads that step's lengths, block-table + rows, pool contents, queries and weights from the captured buffers; the captured static + width (max_context_len) stays.""" + capacity = [5000, 12000, 700, 3000, None, 200, 64, 9000] + step = _Step(capacity, 6, 2, 64, seed=11, extra_ctx=64) + eager = step.run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + with torch.cuda.graph(graph): + out = step.run() + torch.cuda.current_stream().wait_stream(stream) + graph.replay() + torch.cuda.synchronize() + vis = torch.isfinite(step.reference()) + assert torch.equal(out[vis], eager[vis]) + g = torch.Generator(device=DEV) + g.manual_seed(12) + rows = step.block_table.clone() + for perm in ([1, 0, 2, 3, 4, 5, 6, 7], [7, 1, 4, 3, 0, 6, 5, 2]): + # request b now owns table row perm[b]; its new position fits that row's pages + positions = [] + for b, r in enumerate(perm): + cap = capacity[r] + positions.append(None if cap is None or b == 4 else [cap - 3, 1, cap // 2][b % 3]) + new = _Step(positions, 6, 2, 64, seed=int(perm[0]) + 100) + step.lens.copy_(new.lens) + step.q.copy_(new.q) + step.q_sf.copy_(new.q_sf) + step.weights.copy_(new.weights) + step.block_table.copy_(rows[perm]) + step.refill_keys(g) + step.pool.copy_(_write_pool(step.k_codes, step.k_sf)) + graph.replay() + torch.cuda.synchronize() + _check(step, out) + + +# ─── the DeepSeek-V4.1 core target's runtime values ────────────────────────────────────────── +# From the core target's runtime-value dump (boot/eval/GSM8K runs, tep4/dep4/tp4): +# - the engine's index-K pools: (453120, 64, 1, 68) at ratio 2 and (151040, 128, 1, 68) at ratio 1; +# - max_context_len is the metadata's static max_compressed_kv_lens = max_positions // ratio, with +# max_positions = max_seq_len + max_draft_tokens: 8192 / 16384 at max_seq_len 16384 and +# 16384 / 32768 at max_seq_len 32768 (ratio 2 / ratio 1), far above most rows' lengths; +# - decode batches (generation requests per step, next_n = 1) 1-10, 16-19, 24 and 27-32; the +# ones below are those the boundary/gate/long/64k sweep does not enumerate. +RUNTIME_POOL_PAGES = {64: 453120, 128: 151040} +STATIC = [(2, 64, 8192), (1, 128, 16384), (2, 64, 16384), (1, 128, 32768)] +DUMP_BATCHES = [2, 3, 4, 5, 6, 7, 9, 10, 16, 17, 18, 19, 24, 27, 28, 29, 30, 31] + + +def _static_positions(B, next_n, ratio, max_ctx, seed): + """First-row positions under a static max_context_len: request 0 sits at the cap (its last row + sees exactly max_ctx entries), request 1 at position 1, the rest drawn from GSM8K-to-passkey + contexts, so most rows are far shorter than the width.""" + g = torch.Generator() + g.manual_seed(seed) + cap = max_ctx * ratio - next_n + pos = torch.randint(2, min(cap, 24576), (B,), generator=g).tolist() + pos[0] = cap + if B > 1: + pos[1] = 1 + return pos + + +@pytest.mark.parametrize("batch", DUMP_BATCHES) +@pytest.mark.parametrize( + "ratio,block_kv,max_ctx", STATIC, ids=[f"r{r}_page{p}_w{w}" for r, p, w in STATIC] +) +def test_static_width_decode(ratio, block_kv, max_ctx, batch): + """Decode at the target's static width over a runtime-sized pool: every visible element is + exact (_check), and no row reads a page its table row does not name (those pages are NaN).""" + seed = zlib.crc32(f"static/{ratio}/{block_kv}/{max_ctx}/{batch}".encode()) + pos = _static_positions(batch, 1, ratio, max_ctx, seed) + step = _Step( + pos, 1, ratio, block_kv, seed=seed, max_ctx=max_ctx, pool_pages=RUNTIME_POOL_PAGES[block_kv] + ) + out = step.run() + _check(step, out) + assert int(step.lens.max()) == max_ctx + assert out.is_contiguous() # the static widths are multiples of 256 columns + + +@pytest.mark.parametrize("batch", [3, 10, 19, 31]) +@pytest.mark.parametrize("ratio,block_kv,max_seq_len", [(2, 64, 16384), (1, 128, 32768)]) +def test_static_width_dspark_verify(ratio, block_kv, max_seq_len, batch): + """next_n = 6 under the DSpark static width (max_seq_len + 5 draft tokens) // ratio, which is + not a multiple of 256: the rows are padded (stride 8448 / 33024).""" + next_n, max_ctx = 6, (max_seq_len + 5) // ratio + seed = zlib.crc32(f"dspark/{ratio}/{block_kv}/{max_ctx}/{batch}".encode()) + pos = _static_positions(batch, next_n, ratio, max_ctx, seed) + step = _Step( + pos, + next_n, + ratio, + block_kv, + seed=seed, + max_ctx=max_ctx, + pool_pages=RUNTIME_POOL_PAGES[block_kv], + ) + _check(step, step.run()) + assert int(step.lens.max()) == max_ctx + + +@pytest.mark.parametrize("ratio,block_kv,pool_pages", [(2, 64, 600000), (1, 128, 300000)]) +def test_pool_past_2gib(ratio, block_kv, pool_pages): + """A pool of 2.61 GB, about 1.3x the runtime's 1.97 GB (453120 x 64 pages): the pages the batch + owns are drawn from the whole pool, so about a sixth of them start past byte 2^31.""" + max_ctx = 16384 * 2 // ratio + seed = zlib.crc32(f"pool2g/{ratio}/{block_kv}".encode()) + pos = _static_positions(8, 1, ratio, max_ctx, seed) + step = _Step(pos, 1, ratio, block_kv, seed=seed, max_ctx=max_ctx, pool_pages=pool_pages) + page_bytes = block_kv * KEY_BYTES + assert step.pool.numel() > 2**31 and int((step.used * page_bytes >= 2**31).sum()) > 0 + _check(step, step.run()) + + +def _embed(x, offset, fill): + """x placed at row ``offset`` of a larger buffer; returns that row-offset view.""" + buf = torch.full((offset + x.shape[0] + 3, *x.shape[1:]), fill, dtype=x.dtype, device=DEV) + buf[offset : offset + x.shape[0]] = x + return buf[offset : offset + x.shape[0]] + + +def _mixed_views(step, tc, nc, pad): + """The views a mixed (context + generation) step hands the op: q / q_sf / weights / + context_lens start at token row tc of the step's buffers, the table at request row nc of + the metadata's [max_num_sequences, max_blocks_per_seq] table (here ``pad`` columns wider).""" + rows = step.B * step.n + width = step.block_table.shape[1] + table = torch.zeros(nc + step.B + 2, width + pad, dtype=torch.int32, device=DEV) + table[nc : nc + step.B, :width] = step.block_table + return dict( + q=_embed(step.q.view(rows, step.H, DIM // 2), tc, 0).view(step.B, step.n, step.H, DIM // 2), + q_sf=_embed(step.q_sf.view(rows, step.H), tc, 0).view(step.B, step.n, step.H), + weights=_embed(step.weights, tc, 0.0), + context_lens=_embed(step.lens.view(-1), tc, 0).view(step.B, step.n), + block_table=table[nc : nc + step.B], + ) + + +@pytest.mark.parametrize("tc,nc,pad", [(1, 1, 1), (573, 3, 0), (7674, 2, 1)]) +def test_mixed_step_row_offset_views(tc, nc, pad): + """Row-offset views, including context_lens and table bases only 4-byte aligned (tc 1 / 573; + nc 1 at a 131-column table) and a table row stride wider than the rows need.""" + pos = [100, 3000, 7000, 16383, 555, 1, 2, 9000, 4000, 12000] + step = _Step(pos, 1, 2, 64, seed=tc + nc, max_ctx=8192, pool_pages=RUNTIME_POOL_PAGES[64]) + views = _mixed_views(step, tc, nc, pad) + assert views["block_table"].stride(0) == step.block_table.shape[1] + pad + _check(step, step.run(**views)) + + +# The union of the gates' capture sets, per rank: every decode batch 1..8 (boot engine, padding +# off) and 1/2/4/8/16/24/32 (accuracy engine, padding on), at next_n 1. +CORE_GRAPH_BATCHES = [1, 2, 3, 4, 5, 6, 7, 8, 16, 24, 32] + + +def _poison_tails(step): + """0xFF (E8M0 NaN) scale bytes on every slot past each request's last-row length in its last + page: such slots are never visible, whatever the pool holds there.""" + sf_bytes = step.k_sf.view(torch.uint8).view(step.num_pages, step.bk, 4) + for b in range(step.B): + L = int(step.lens[b].max()) + if L % step.bk: + sf_bytes[step.block_table[b, L // step.bk].long(), L % step.bk :] = 0xFF + + +def _capture(step): + """Metadata + logits captured together on a side stream, as the target's step captures them.""" + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + with torch.cuda.graph(graph): + out = step.run() + torch.cuda.current_stream().wait_stream(stream) + return graph, out + + +def _replay_and_check(step, graph, out, tag, prev): + """One replay into an output poisoned with NaN: exact vs the fp64 reference and bit-exact vs + the documented reduction (_check), bitwise equal to an eager call on the same inputs, and + rejected by the previous step's reference (a replay that re-read the captured step is seen).""" + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + ref, vis = _check(step, out) + again = step.run() + assert torch.equal(out[vis], again[vis]) + prev_ref, prev_vis = prev + both = vis & prev_vis + stale = _outside_fraction(out[both], prev_ref[both]) + print( + f"[{tag}] visible {int(vis.sum())}, previous step's reference: {stale:.4f} of " + f"{int(both.sum())} outside the gate" + ) + assert stale > 0.5 + return ref, vis + + +@pytest.mark.parametrize("batch", CORE_GRAPH_BATCHES) +@pytest.mark.parametrize( + "ratio,block_kv,max_ctx", STATIC, ids=[f"r{r}_page{p}_w{w}" for r, p, w in STATIC] +) +def test_cuda_graph_static_width(ratio, block_kv, max_ctx, batch): + """The target captures metadata + logits per decode graph batch with the static width baked + in. Every request's table row covers the full width. The first replay equals the eager + warm-up call bitwise; each of 3 more replays reads new lengths (incl. 0, 1 and the cap), + queries, weights and keys (tails poisoned NaN) and is checked by _replay_and_check.""" + seed = zlib.crc32(f"graph/{ratio}/{block_kv}/{max_ctx}/{batch}".encode()) + step = _Step( + [max_ctx * ratio - 1] * batch, + 1, + ratio, + block_kv, + seed=seed, + max_ctx=max_ctx, + pool_pages=RUNTIME_POOL_PAGES[block_kv], + ) + eager = step.run() + torch.cuda.synchronize() + graph, out = _capture(step) + graph.replay() + torch.cuda.synchronize() + prev = _check(step, out) + assert torch.equal(out[prev[1]], eager[prev[1]]) + g = torch.Generator(device=DEV) + g.manual_seed(seed) + for rep in range(3): + lens = torch.randint(0, max_ctx + 1, (batch, 1), generator=g, device=DEV, dtype=torch.int32) + lens[0] = max_ctx + if batch > 1: + lens[-1] = rep % 2 # a padding-like row: 0 or 1 entries + step.lens.copy_(lens) + new = _Step([1] * batch, 1, ratio, block_kv, seed=seed + rep + 1) + step.q.copy_(new.q) + step.q_sf.copy_(new.q_sf) + step.weights.copy_(new.weights) + step.refill_keys(g) + _poison_tails(step) + step.pool.copy_(_write_pool(step.k_codes, step.k_sf)) + prev = _replay_and_check( + step, graph, out, f"graph r{ratio} w{max_ctx} B{batch} rep{rep}", prev + ) + + +# ─── DSpark verify under CUDA graph at the target's graph batches ───────────────────────────── +# With DSpark on, the target's generation-only steps are captured with next_n = 6 rows per +# request at the union of the gates' capture sets: every G in 1..8 (the DSpark boot +# engine: max_batch_size 8, padding off; its runtime-value dump records G 3, 5, 6 and 7 at the +# graph-copy widths 8201 / 16403) and G 1, 2, 4, 8, 16, 24, 32 (the accuracy engine, padding on). +# max_context_len is the metadata's max_compressed_kv_lens[r] = (max_seq_len + max_draft_tokens) +# // r: max_draft_tokens is 0 in the eager metadata and can be 5 in its CUDA-graph copy (which +# re-runs __post_init__ after the spec-dec update), so both widths are covered, at max_seq_len +# 16384 (gates) and 32768 (the long-context run). +DSPARK_GRAPH_BATCHES = [1, 2, 3, 4, 5, 6, 7, 8, 16, 24, 32] +DSPARK_STATIC = [ + (2, 64, 8192), + (2, 64, 8194), + (1, 128, 16384), + (1, 128, 16389), + (2, 64, 16384), + (2, 64, 16386), + (1, 128, 32768), + (1, 128, 32773), + # The DSpark engine's own widths: the runtime raises the engine's max_seq_len + # to 16398, so the eager / graph-copy widths are 16398 // r and 16403 // r. + (2, 64, 8199), + (2, 64, 8201), + (1, 128, 16398), + (1, 128, 16403), +] + + +def _dspark_lens(B, ratio, max_ctx, rep, g): + """Per-row lengths (p + j + 1) // ratio of a verify step: request 0 at the cap (its last row + sees exactly max_ctx entries), request 1 at position rep % 2 (first rows see 0 or 1 entries), + on even reps a graph padding request (all 0) last, the rest random below the cap.""" + cap = max_ctx * ratio - 6 + pos = torch.randint(0, cap + 1, (B,), generator=g).tolist() + pos[0] = cap + if B > 1: + pos[1] = rep % 2 + if B > 2 and rep % 2 == 0: + pos[-1] = None + lens = [[0] * 6 if p is None else [(p + j + 1) // ratio for j in range(6)] for p in pos] + return torch.tensor(lens, dtype=torch.int32, device=DEV) + + +@pytest.mark.parametrize("batch", DSPARK_GRAPH_BATCHES) +@pytest.mark.parametrize( + "ratio,block_kv,max_ctx", DSPARK_STATIC, ids=[f"r{r}_page{p}_w{w}" for r, p, w in DSPARK_STATIC] +) +def test_cuda_graph_dspark_verify(ratio, block_kv, max_ctx, batch): + """next_n = 6 captured at every DSpark graph batch with each static width: metadata + logits + in one graph, as the target captures them. The first replay equals the eager warm-up call + bitwise; 3 more replays read new lengths (the cap, 0/1, a padding request), queries, weights + and keys (tails poisoned NaN), each checked by _replay_and_check.""" + seed = zlib.crc32(f"graph-dspark/{ratio}/{block_kv}/{max_ctx}/{batch}".encode()) + step = _Step( + [max_ctx * ratio - 6] * batch, + 6, + ratio, + block_kv, + seed=seed, + max_ctx=max_ctx, + pool_pages=RUNTIME_POOL_PAGES[block_kv], + ) + assert int(step.lens.max()) == max_ctx + eager = step.run() + torch.cuda.synchronize() + graph, out = _capture(step) + graph.replay() + torch.cuda.synchronize() + prev = _check(step, out) + assert torch.equal(out[prev[1]], eager[prev[1]]) + g = torch.Generator(device=DEV) + g.manual_seed(seed) + gh = torch.Generator() + gh.manual_seed(seed) + for rep in range(3): + step.lens.copy_(_dspark_lens(batch, ratio, max_ctx, rep, gh)) + new = _Step([1] * batch, 6, ratio, block_kv, seed=seed + rep + 1) + step.q.copy_(new.q) + step.q_sf.copy_(new.q_sf) + step.weights.copy_(new.weights) + step.refill_keys(g) + _poison_tails(step) + step.pool.copy_(_write_pool(step.k_codes, step.k_sf)) + prev = _replay_and_check( + step, graph, out, f"graph dspark r{ratio} w{max_ctx} G{batch} rep{rep}", prev + ) + + +# ─── every generation batch 1..32 at the static widths ─────────────────────────────────────── +# Rows are independent: each (b, j) row reads only its own request's q, weights, length and +# table row, so a batch's result is the first B requests of a 32-request call. One 32-request +# step per (width, next_n) is checked against the reference, and every B in 1..31 over the first +# B requests (prefix views of the same buffers, schedule rebuilt from those lengths) must be +# bitwise equal to the first B * next_n rows of that call on every visible element. +PREFIX_CASES = [(r, p, w, 1) for r, p, w in STATIC] + [(r, p, w, 6) for r, p, w in DSPARK_STATIC] + + +@pytest.mark.parametrize( + "ratio,block_kv,max_ctx,next_n", + PREFIX_CASES, + ids=[f"r{r}_page{p}_w{w}_nextn{n}" for r, p, w, n in PREFIX_CASES], +) +def test_every_batch_is_a_prefix(ratio, block_kv, max_ctx, next_n): + seed = zlib.crc32(f"prefix/{ratio}/{block_kv}/{max_ctx}/{next_n}".encode()) + pos = _static_positions(32, next_n, ratio, max_ctx, seed) + pos[5] = None # a graph padding request: length 0, all -1 table row + step = _Step( + pos, + next_n, + ratio, + block_kv, + seed=seed, + max_ctx=max_ctx, + pool_pages=RUNTIME_POOL_PAGES[block_kv], + ) + full = step.run() + _, vis = _check(step, full) + for B in range(1, 32): + rows = B * next_n + out = step.run( + q=step.q[:B], + q_sf=step.q_sf[:B], + weights=step.weights[:rows], + context_lens=step.lens[:B], + block_table=step.block_table[:B], + ) + assert out.shape == (rows, max_ctx) + v = vis[:rows] + assert torch.equal(out[v], full[:rows][v]), f"batch {B} differs from the prefix" + + +# ─── deliberately wrong controls: the same gate must reject each ───────────────────────────── + + +def _control_step(): + return _Step([130, 999, 4100, 20001, 777, 3, 12000, 64], 6, 2, 64, seed=21) + + +def test_control_wrong_block_table(): + """Each request reads its neighbour's pages (block-table rows rolled by one).""" + step = _control_step() + ref = step.reference() + vis = torch.isfinite(ref) + rolled = step.block_table.roll(1, dims=0).contiguous() + got = step.run(block_table=rolled)[vis] + frac = _outside_fraction(got, ref[vis]) + print(f"[control wrong_block_table] outside the gate: {frac:.4f}") + assert frac > 0.5 + with pytest.raises(AssertionError): + torch.testing.assert_close(got, ref.float()[vis]) + + +def test_control_wrong_page_size(): + """Pages written 64 entries wide, read as 128-entry pages (the table halved accordingly).""" + step = _control_step() + ref = step.reference() + vis = torch.isfinite(ref) + pool128 = step.pool.view(-1, 128, 1, KEY_BYTES) + bt128 = (step.block_table // 2).contiguous() + got = step.run(kv_cache=pool128, block_table=bt128, meta_block_kv=128)[vis] + frac = _outside_fraction(got, ref[vis]) + print(f"[control wrong_page_size] outside the gate: {frac:.4f}") + assert frac > 0.5 + with pytest.raises(AssertionError): + torch.testing.assert_close(got, ref.float()[vis]) + + +def test_control_key_interleaved_layout(): + """The pool written key-interleaved ([64 codes | 4 scales] per key) instead of paged.""" + step = _control_step() + ref = step.reference() + vis = torch.isfinite(ref) + pool = _write_pool(step.k_codes, step.k_sf, layout="interleaved") + got = step.run(kv_cache=pool)[vis] + frac = _outside_fraction(got, ref[vis]) + print(f"[control key_interleaved_layout] outside the gate: {frac:.4f}") + assert frac > 0.5 + with pytest.raises(AssertionError): + torch.testing.assert_close(got, ref.float()[vis]) + + +def test_control_reference_without_relu(): + """The gate sees the relu: the op's output against a relu-free score.""" + step = _control_step() + ref = step.reference(relu=False) + vis = torch.isfinite(ref) + got = step.run()[vis] + frac = _outside_fraction(got, ref[vis]) + print(f"[control reference_without_relu] outside the gate: {frac:.4f}") + assert frac > 0.5 + with pytest.raises(AssertionError): + torch.testing.assert_close(got, ref.float()[vis]) + + +def test_control_heads_reversed_weights(): + """Head weights applied in reverse head order.""" + step = _control_step() + ref = step.reference() + vis = torch.isfinite(ref) + got = step.run(weights=step.weights.flip(-1).contiguous())[vis] + frac = _outside_fraction(got, ref[vis]) + print(f"[control heads_reversed_weights] outside the gate: {frac:.4f}") + assert frac > 0.5 + with pytest.raises(AssertionError): + torch.testing.assert_close(got, ref.float()[vis]) + + +# ─── loud rejections (host-side asserts; the CUDA context survives) ─────────────────────────── + + +def _raises(fn, text): + with pytest.raises(RuntimeError, match=re.escape(text)): + fn() + + +def test_rejections(): + step = _Step([130, 999], 1, 2, 64, seed=31) + meta = metadata_op(step.lens, 64) + base = dict( + q=step.q, + q_sf=step.q_sf, + kv_cache=step.pool, + weights=step.weights, + context_lens=step.lens, + block_table=step.block_table, + schedule_meta=meta, + max_context_len=step.max_ctx, + ) + + def call(**over): + return op(**{**base, **over}) + + page_check = "block_kv == 32 or block_kv == 64 or block_kv == 128" + # tokens_per_block 256 at ratio 1 would give 256-entry index-K pages + _raises(lambda: metadata_op(step.lens, 256), page_check) + pool256 = torch.zeros(4, 256, 1, KEY_BYTES, dtype=torch.uint8, device=DEV) + _raises(lambda: call(kv_cache=pool256), f"(arch_major == 10 and ({page_check}))") + # Any multiple of 4 up to 128 heads is taken (24 is measured against the reference in + # test_head_count_off_the_cells_matches_the_reference); 6 is refused. + _raises( + lambda: call( + q=step.q[:, :, :6].contiguous(), + q_sf=step.q_sf[:, :, :6].contiguous(), + weights=step.weights[:, :6].contiguous(), + ), + "num_heads > 0 and num_heads <= 128 and num_heads % 4 == 0", + ) + _raises(lambda: call(context_lens=step.lens.view(-1)), "context_lens.dim() == 2") + _raises( + lambda: metadata_op(step.lens.view(-1), 64), + "context_lens.dim() == 2", + ) + other_sms = metadata_op(step.lens, 64, num_sms=8) + _raises( + lambda: call(schedule_meta=other_sms), + "_schedule_meta_size == num_sms / num_ctas_per_cluster + 1", + ) + _raises( + lambda: call(weights=step.weights.bfloat16()), + "weights.scalar_type() != torch::kBFloat16 or logits_dtype == torch::kBFloat16", + ) + _raises( + lambda: call(logits_dtype=torch.float16), + "logits_dtype == torch::kFloat32 or logits_dtype == torch::kBFloat16", + ) + wt = torch.empty(HEADS, step.B, device=DEV).t() + wt.copy_(step.weights) + _raises(lambda: call(weights=wt), "weights.stride(1) == 1") + _raises( + lambda: call(block_table=step.block_table.long()), + "block_table.scalar_type() == torch::kInt", + ) + _raises( + lambda: call(context_lens=step.lens.long()), "context_lens.scalar_type() == torch::kInt" + ) + _raises( + lambda: call(q=step.q.view(torch.uint8)), + "q_fp.scalar_type() == (is_fp4 ? kPackedFP4 : torch::kFloat8_e4m3fn)", + ) + _raises( + lambda: call(kv_cache=step.pool.view(torch.int8)), + "fused_kv_cache.scalar_type() == torch::kByte", + ) + # the context survived every rejection + _check(step, step.run()) + + +def test_head_count_off_the_cells_matches_the_reference(): + """24 heads, which the cells do not cover and which does not divide 128: the kernel takes it, + and its logits are held to the entry's gate and the fp64 twin. A measurement that the + acceptance is not a silent wrong answer, not a certification of 24 heads.""" + step = _Step([130, 999], 1, 2, 64, heads=24, seed=32) + out = step.run() + rows = step.B * step.n + assert out.shape == (rows, step.max_ctx) and out.dtype == torch.float32 + op.compare( + out, + op.reference( + step.q, + step.q_sf, + step.pool, + step.weights, + step.lens, + step.block_table, + None, + step.max_ctx, + ), + ) + ref = step.reference() + vis = torch.isfinite(ref) + torch.testing.assert_close(out[vis], ref.float()[vis]) + + +# ─── controls for the runtime-value cases above ─────────────────────────────────────────────── + + +def test_control_table_names_unowned_page(): + """The runtime-sized pool makes a stray page read visible: one table entry per request moved + to a page the batch does not own (0xFF bytes, NaN scale) fails the gate with NaN.""" + step = _Step( + [100, 3000, 7000, 16383, 555, 9000, 4000, 12000], + 1, + 2, + 64, + seed=41, + max_ctx=8192, + pool_pages=RUNTIME_POOL_PAGES[64], + ) + ref = step.reference() + vis = torch.isfinite(ref) + owned = torch.zeros(step.num_pages, dtype=torch.bool, device=DEV) + owned[step.used] = True + stray = int(torch.nonzero(~owned)[0]) + table = step.block_table.clone() + table[:, 0] = stray + got = step.run(block_table=table)[vis] + ref32 = ref.float()[vis] + n_out = int( + (~((got - ref32).abs() <= 1e-5 + 1.3e-6 * ref32.abs())).sum() + ) # NaN counts as outside + n_nan = int(torch.isnan(got).sum()) + print( + f"[control table_names_unowned_page] outside the gate: {n_out} of {got.numel()} visible " + f"(the entries of each request's first page), NaN: {n_nan}" + ) + assert n_nan > 0 + with pytest.raises(AssertionError): + torch.testing.assert_close(got, ref.float()[vis]) + + +def test_control_mixed_step_offset_lost(): + """A mixed step whose query-side views lost their token offset (each generation row reads + the row tc - 1 before it, i.e. the last context token for row 0 and its neighbour for the + rest) fails the gate.""" + tc = 573 + step = _Step( + [100, 3000, 7000, 16383, 555, 1, 2, 9000, 4000, 12000], + 1, + 2, + 64, + seed=42, + max_ctx=8192, + pool_pages=RUNTIME_POOL_PAGES[64], + ) + views = _mixed_views(step, tc, 1, 1) + rows = step.B * step.n + for k in ("q", "q_sf", "weights"): + base = views[k].reshape(rows, *views[k].shape[2:]) if k != "weights" else views[k] + shifted = base.as_strided(base.shape, base.stride(), base.storage_offset() - base.stride(0)) + views[k] = shifted.reshape(views[k].shape) + ref = step.reference() + vis = torch.isfinite(ref) + got = step.run(**views)[vis] + frac = _outside_fraction(got, ref[vis]) + print(f"[control mixed_step_offset_lost] outside the gate: {frac:.4f}") + assert frac > 0.5 + with pytest.raises(AssertionError): + torch.testing.assert_close(got, ref.float()[vis]) + + +# ─── pool size: any B <= B_BOUND[ratio] blocks ───────────────────────────────────────────────── +# The cache manager's INDEXER_COMPRESS storage at B blocks (tokens_per_block 128): ratio 1 is one +# uint8 [B, 128, 1, 68] pool; ratio 2 is ONE allocation of 3B 64-entry pages shared by the three +# ratio-2 owners, owner k (k = 0, 1, 2) seeing the view alloc[k:] of 3B - k pages, its logical +# block b at view page 3b (its table holds 3 x block). B depends on free memory. For a manager that +# holds every layer, 298,013,687,808 B of device memory over the 113,920 persistent bytes per block +# of native KV storage (890 B/token x 128) gives B_MAX = 2,615,990; the bf16-storage engines +# (349,440 B/block) stay below 852,832. A pipeline stage's manager holds fewer bytes per block, so +# that derivation does not bound it; its measured maxima follow. +B_MAX = 2_615_990 +B_CORE_BF16 = 852_832 +# The engine envelope's managers (max_num_tokens up to 16384, max_batch_size up to 256, max_seq_len up +# to 1,048,576), measured at 0.5 x the GB300's memory as KV quota over a sampled (max_num_tokens, +# max_batch_size, DSpark, max_seq_len) grid: the unsplit manager's index-K pools reach 1,295,549 +# blocks (ratio 1; ratio 2: 3,886,647 pages). A PP2 stage manager holds half the layers: +# stage 1 (layers 20-39, the ratio-1 index source) reaches 3,245,959 ratio-1 blocks built alone and +# 2,159,957 in a PP2 engine (whose ranks agree on the smaller stage's token capacity); stage 0 (the +# ratio-2 sources) reaches 6,480,780 ratio-2 pages = 3 x 2,160,260 blocks, below 3 x B_MAX. +B_ENV_UNSPLIT = 1_295_549 +B_PP2_R1_ENGINE = 2_159_957 +B_PP2_R2_STAGE = 2_160_260 +B_PP2_R1_STAGE = 3_245_959 +B_BOUND = {1: B_PP2_R1_STAGE, 2: B_MAX} # the certified ceiling per ratio, in blocks +INV_BLOCKS = 160 # the invariance batches use blocks 0..159 only: B_min = 160 +POOL_SIZES = { + 1: [ + INV_BLOCKS, + 481, + 1493, + 151040, + B_CORE_BF16, + B_ENV_UNSPLIT, + B_PP2_R1_ENGINE, + B_MAX, + B_PP2_R1_STAGE, + ], + 2: [INV_BLOCKS, 481, 1493, 151040, B_CORE_BF16, B_ENV_UNSPLIT, B_PP2_R2_STAGE, B_MAX], +} +GATE_WIDTHS = [(2, 64, 8192), (1, 128, 16384)] + + +def _alloc_index_k(B, ratio): + """The manager's storage at B blocks, every byte 0xFF (NaN scale): an exact-size allocation.""" + pages, bk = (B, 128) if ratio == 1 else (3 * B, 64) + return torch.full((pages, bk, 1, KEY_BYTES), 0xFF, dtype=torch.uint8, device=DEV) + + +def _owners(ratio): + return [0] if ratio == 1 else [0, 1, 2] + + +class _PoolStep(_Step): + """A step whose keys live in a caller-owned pool: ``block_table`` holds view page ids and + ``slots`` [B, width] the row of ``codes`` / ``sf`` each named page holds (-1 past a request's + pages). The reference and the emulation decode the keys from those rows.""" + + def __init__(self, positions, next_n, ratio, block_kv, pages_of, seed, max_ctx): + g = torch.Generator(device=DEV) + g.manual_seed(seed) + self.B, self.n, self.H, self.bk, self.ratio = len(positions), next_n, HEADS, block_kv, ratio + lens = [ + [0] * next_n if p is None else [(p + j + 1) // ratio for j in range(next_n)] + for p in positions + ] + self.lens = torch.tensor(lens, dtype=torch.int32, device=DEV) + self.max_ctx = max_ctx + self.pool = None # every call passes its owner view + width = (max_ctx + block_kv - 1) // block_kv + 2 + self.block_table = torch.full((self.B, width), -1, dtype=torch.int32, device=DEV) + for b, row in enumerate(lens): + npg = (max(row) + block_kv - 1) // block_kv + if npg: + self.block_table[b, :npg] = torch.tensor(pages_of(b, npg), dtype=torch.int32) + named = self.block_table[self.block_table >= 0] + self.pages = torch.unique(named).long() # sorted view page ids the batch names + self.slots = torch.full_like(self.block_table, -1) + self.slots[self.block_table >= 0] = torch.searchsorted(self.pages, named.long()).to( + torch.int32 + ) + k = torch.randn(self.pages.numel(), block_kv, DIM, generator=g, device=DEV) + self.k_codes, self.k_sf = _quant_mxfp4(k) + self.page_bytes = _write_pool(self.k_codes, self.k_sf) + q = torch.randn(self.B, next_n, HEADS, DIM, generator=g, device=DEV) * 2.0 + self.q, self.q_sf = _quant_mxfp4(q) + w = torch.randn(self.B * next_n, HEADS, generator=g, device=DEV) * 2.0**-4 + self.weights = w.to(torch.bfloat16).float() + + def _keys(self, b): + L = int(self.lens[b].max()) + npg = (L + self.bk - 1) // self.bk + idx = self.slots[b, :npg].long() + return _decode(self.k_codes[idx], self.k_sf[idx]).reshape(-1, DIM)[:L] + + def physical(self, k): + """Allocation pages of the named view pages for the owner whose view starts k pages in.""" + return self.pages + k + + def run_on(self, view, **over): + return self.run(kv_cache=view, **over) + + +def _written_span(step, out): + """Every element the kernel writes: request b's rows up to ceil256 of its last row's length + (at most the width). Returned as int32 bit patterns so NaN compares bitwise.""" + bits = out.view(torch.int32) + parts = [] + for b in range(step.B): + span = min(-(-int(step.lens[b].max()) // SPLIT_KV) * SPLIT_KV, step.max_ctx) + parts.append(bits[b * step.n : (b + 1) * step.n, :span].reshape(-1)) + return torch.cat(parts) + + +def _invariance_steps(ratio, block_kv, max_ctx): + """next_n 1 and 6 at batches 1, 7 and 32 (request 0 at the width, request 1 at position 1, a + padding request at 32), every request's pages a random choice of blocks 0..INV_BLOCKS-1 (pages + alias across requests, as prefix reuse does), block INV_BLOCKS-1 always named.""" + mult = 1 if ratio == 1 else 3 + steps = {} + for next_n in (1, 6): + for batch in (1, 7, 32): + seed = zlib.crc32(f"pool-inv/{ratio}/{max_ctx}/{next_n}/{batch}".encode()) + pos = _static_positions(batch, next_n, ratio, max_ctx, seed) + if batch == 32: + pos[5] = None + gh = torch.Generator() + gh.manual_seed(seed) + + def pages_of(b, npg, gh=gh): + blocks = torch.randperm(INV_BLOCKS, generator=gh)[:npg] + if b == 0 and not bool((blocks == INV_BLOCKS - 1).any()): + blocks[-1] = ( + INV_BLOCKS - 1 + ) # the smallest pool is exactly the one that holds it + return (blocks * mult).tolist() + + steps[(next_n, batch)] = _PoolStep( + pos, next_n, ratio, block_kv, pages_of, seed, max_ctx + ) + return steps + + +@pytest.mark.parametrize( + "ratio,block_kv,max_ctx", GATE_WIDTHS, ids=[f"r{r}_page{p}_w{w}" for r, p, w in GATE_WIDTHS] +) +def test_pool_size_invariance(ratio, block_kv, max_ctx): + """The same queries, weights, lengths, block tables and page bytes at every pool size B in + POOL_SIZES[ratio] (an exact-size allocation each, one at a time; every page the batch does not + own is 0xFF): the result is bitwise identical across B on every element the kernel writes (each + request's rows up to ceil256 of its last row's length), and passes the reference at B_min. + Ratio 2 runs each owner's view (0, 1 and 2 pages into the shared allocation).""" + steps = _invariance_steps(ratio, block_kv, max_ctx) + first = {} + sizes = POOL_SIZES[ratio] + assert max(sizes) == B_BOUND[ratio] + for B in sizes: + alloc = _alloc_index_k(B, ratio) + for k in _owners(ratio): + view = alloc[k:] + for key, step in steps.items(): + phys = step.physical(k) + alloc[phys] = step.page_bytes + out = step.run_on(view) + alloc[phys] = 0xFF + span = _written_span(step, out) + if B == sizes[0]: + assert int(step.pages.max()) // (3 if ratio == 2 else 1) == INV_BLOCKS - 1 + _check(step, out) + first[(k, key)] = span + else: + assert torch.equal(span, first[(k, key)]), ( + f"B {B}, owner view +{k}, next_n/batch {key}: differs from B {sizes[0]}" + ) + print( + f"[pool invariance r{ratio} w{max_ctx}] B {B}: {alloc.numel() / 1e9:.2f} GB, " + f"{len(steps) * len(_owners(ratio))} calls bitwise equal to B {sizes[0]}" + ) + del alloc, view + torch.cuda.empty_cache() + + +def _far_pages(ratio, B, k): + """Owner k's view pages at the far end of the B-block storage: its last page, and the pages + around byte offsets 2^31, 2^32, 2^33 and 2^34 of the allocation (the page holding each offset + straddles it; ratio 2 takes the owner's pages among the five around it).""" + bk = 128 if ratio == 1 else 64 + pb = bk * KEY_BYTES + alloc_pages = B if ratio == 1 else 3 * B + mult = 1 if ratio == 1 else 3 + phys = [alloc_pages - 3 + k if ratio == 2 else alloc_pages - 1] + straddling = [] + for n in (31, 32, 33, 34): + p = 2**n // pb + assert 2**n % pb and p < alloc_pages + straddling.append(p) + phys += [q for q in range(p - 2, p + 3) if q % mult == k] + view = [q - k for q in phys] + assert all(v % mult == 0 for v in view) + return view, [p - k for p in straddling if p % mult == k] + + +# (ratio, page, width, blocks): both ratios at B_MAX at the gate widths, and ratio 1 at the PP2 +# stage-1 ceiling at the gate width and at the largest DSpark graph width (8193 pages per request). +FAR_CASES = [ + (2, 64, 8192, B_MAX), + (1, 128, 16384, B_MAX), + (1, 128, 16384, B_PP2_R1_STAGE), + (1, 128, 1048595, B_PP2_R1_STAGE), +] + + +@pytest.mark.parametrize( + "ratio,block_kv,max_ctx,B", + FAR_CASES, + ids=[f"r{r}_page{p}_w{w}_B{b}" for r, p, w, b in FAR_CASES], +) +def test_pool_far_end_addressing(ratio, block_kv, max_ctx, B): + """At B blocks (ratio 1 at B_MAX: 22.77 GB, at B_PP2_R1_STAGE: 28.25 GB; ratio 2 at B_MAX: + 7,847,970 pages, 34.15 GB), each owner view: a batch of 8 whose request 0 (at the width) names + the owner's last page and its pages around byte offsets 2^31..2^34 first, and whose other pages + are drawn from the whole view, passes the reference at next_n 1 and 6. Two wrong tables must + fail the same check: every entry moved to its neighbouring page, and page ids whose byte offset + wraps at 32 bits.""" + alloc = _alloc_index_k(B, ratio) + pb = block_kv * KEY_BYTES + mult = 1 if ratio == 1 else 3 + for k in _owners(ratio): + view = alloc[k:] + far, straddle = _far_pages(ratio, B, k) + for next_n in (1, 6): + key = f"pool-far/{ratio}/{k}/{next_n}" + ("" if B == B_MAX else f"/{B}/{max_ctx}") + seed = zlib.crc32(key.encode()) + pos = _static_positions(8, next_n, ratio, max_ctx, seed) + gh = torch.Generator() + gh.manual_seed(seed) + + def pages_of(b, npg, gh=gh, far=far): + rnd = (torch.randint(0, B, (npg,), generator=gh) * mult).tolist() + return (far + rnd)[:npg] if b == 0 else rnd + + step = _PoolStep(pos, next_n, ratio, block_kv, pages_of, seed, max_ctx) + assert set(far) <= set(step.block_table[0].tolist()) + named = step.pages * pb + assert far[0] == mult * (B - 1) # the owner's last page (block B - 1) + print( + f"[pool far r{ratio} B{B} w{max_ctx} view+{k} next_n {next_n}] {step.pages.numel()} pages, " + f"last view page {far[0]}, straddling 2^31..2^34: {straddle}, past 2^32: " + f"{int((named + k * pb >= 2**32).sum())}, past 2^34: {int((named + k * pb >= 2**34).sum())}" + ) + phys = step.physical(k) + alloc[phys] = step.page_bytes + ref, vis = _check(step, step.run_on(view)) + ref32 = ref.float()[vis] + # wrong control 1: each entry names its neighbouring page (the last page its predecessor) + bt = step.block_table + last = view.shape[0] - 1 + shifted = torch.where(bt < 0, bt, torch.where(bt >= last, bt - 1, bt + 1)) + # wrong control 2: a 32-bit byte offset (page * page_bytes mod 2^32) + wrapped = torch.where(bt < 0, bt, ((bt.long() * pb) % 2**32 // pb).to(torch.int32)) + assert int((wrapped != bt).sum()) > 0 + for tag, table in (("neighbour page", shifted), ("32-bit wrap", wrapped)): + got = step.run_on(view, block_table=table.contiguous())[vis] + n_out = int((~((got - ref32).abs() <= 1e-5 + 1.3e-6 * ref32.abs())).sum()) + print( + f"[control pool far r{ratio} B{B} w{max_ctx} view+{k} next_n {next_n} {tag}] outside the " + f"gate: {n_out} of {got.numel()} visible" + ) + assert n_out > 0 + with pytest.raises(AssertionError): + torch.testing.assert_close(got, ref32) + alloc[phys] = 0xFF + del alloc, view + torch.cuda.empty_cache() + + +# The PP2 stage-1 pool under CUDA graph: DSpark off at L 1,048,576 (next_n 1, width 1,048,576) is the +# configuration that attains B_PP2_R1_STAGE; DSpark on captures next_n 6 at 1,048,595. +FAR_GRAPH = [(1, 1048576), (6, 1048595)] + + +@pytest.mark.parametrize("next_n,max_ctx", FAR_GRAPH, ids=[f"n{n}_w{w}" for n, w in FAR_GRAPH]) +def test_pool_far_end_capture(next_n, max_ctx): + """At B_PP2_R1_STAGE ratio-1 blocks (28.25 GB): metadata + logits captured over a batch of 8 + whose request 0 (at the width) names the pool's last page and the pages around byte offsets + 2^31..2^34 first, the rest of its 8193 pages and the other requests' pages drawn from the whole + pool. The first replay equals the eager warm-up bitwise; 3 more replays each run a new step + written into the captured buffers (new lengths, queries, weights, table rows and page bytes, the + previous step's pages back to 0xFF, the output poisoned NaN): exact against the fp64 reference + and bitwise the documented reduction (_check), bitwise an eager call, and far from the previous + step's reference.""" + ratio, block_kv, B = 1, 128, B_PP2_R1_STAGE + alloc = _alloc_index_k(B, ratio) + far, _ = _far_pages(ratio, B, 0) + + def make(rep): + seed = zlib.crc32(f"pool-far-graph/{next_n}/{max_ctx}/{rep}".encode()) + pos = _static_positions(8, next_n, ratio, max_ctx, seed) + gh = torch.Generator() + gh.manual_seed(seed) + + def pages_of(b, npg): + rnd = torch.randint(0, B, (npg,), generator=gh).tolist() + return (far + rnd)[:npg] if b == 0 else rnd + + st = _PoolStep(pos, next_n, ratio, block_kv, pages_of, seed, max_ctx) + assert set(far) <= set(st.block_table[0].tolist()) + return st + + step = make(0) + alloc[step.physical(0)] = step.page_bytes + eager = step.run_on(alloc) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + with torch.cuda.graph(graph): + out = step.run_on(alloc) + torch.cuda.current_stream().wait_stream(stream) + graph.replay() + torch.cuda.synchronize() + prev = _check(step, out) + assert torch.equal(out[prev[1]], eager[prev[1]]) + del eager + cur = step + for rep in range(1, 4): + new = make(rep) + alloc[cur.physical(0)] = 0xFF + alloc[new.physical(0)] = new.page_bytes + for dst, src in ( + (step.lens, new.lens), + (step.q, new.q), + (step.q_sf, new.q_sf), + (step.weights, new.weights), + (step.block_table, new.block_table), + ): + dst.copy_(src) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + ref, vis = _check(new, out) + again = new.run_on(alloc) + assert torch.equal(out[vis], again[vis]) + del again + both = vis & prev[1] + stale = _outside_fraction(out[both], prev[0][both]) + print( + f"[pool far graph B{B} w{max_ctx} n{next_n} rep{rep}] {new.pages.numel()} pages, visible " + f"{int(vis.sum())}, previous step's reference: {stale:.4f} of {int(both.sum())} outside the gate" + ) + assert stale > 0.5 + prev, cur = (ref, vis), new + del alloc + torch.cuda.empty_cache() + + +# ─── mixed steps at both gate widths, the target's view form ───────────────────────────────── +# In a step with tc context tokens (nc context requests) ahead of G generation requests of next_n +# rows, the target hands the op: q / q_sf / weights = split(buf, [tc, gen])[1] of the step's +# token-major [T, ...] buffers (the rows before tc hold the context tokens' own values); +# context_lens = compressed_visible_lens_cuda[r][tc:T] of a max_num_tokens (8192) buffer, viewed +# [G, next_n]; block_table = indexer_compress_block_tables[r][nc:nc+G] of the metadata's +# [max_num_sequences, max_blocks_per_seq] table, 132 columns at max_seq_len 16384. It does this at +# ratio 1 (128-entry pages, static width 16384; layers 20-36) and ratio 2 (64-entry pages, 8192; +# layers 2-18). A CUDA graph is captured only for steps with no context request (the scheduler's +# can_run_cuda_graph is num_context_requests == 0), so the captured form is tc = nc = 0: prefix +# views of the same buffers. +VIEW_WIDTHS = [(1, 128, 16384), (2, 64, 8192)] +VIEW_MAX_TOKENS = 8192 +VIEW_TABLE_COLS = 132 +VIEW_TABLE_ROWS = 40 + + +def _view_cells(next_n): + """(tc, nc, G): odd and even tc and nc from 1 up to the largest steps (T 8192 with 32 requests; + T 8191 with 31 at next_n 6).""" + return [ + (1, 1, 10), + (2, 2, 10), + (573, 3, 10), + (574, 4, 10), + (7673, 5, 10), + (7674, 2, 10), + (8191 - 2 * next_n, 29, 3 if next_n == 1 else 2), + (8192 - 2 * next_n, 30, 2), + ] + + +VIEW_MIXED = [ + (r, p, w, n, tc, nc, G) + for r, p, w in VIEW_WIDTHS + for n in (1, 6) + for tc, nc, G in _view_cells(n) +] + + +def _view_full_table(step, g, rows, cols, first): + """A [rows, cols] metadata table with step.block_table at rows first..first+B: every entry is a + page of the step's pool (the -1s past each request's pages and the rows of other requests get + random pool pages, which the kernel must not read).""" + table = torch.randint( + 0, step.num_pages, (rows, cols), generator=g, device=DEV, dtype=torch.int64 + ) + table = table.to(torch.int32) + own = table[first : first + step.B, : step.block_table.shape[1]] + table[first : first + step.B, : step.block_table.shape[1]] = torch.where( + step.block_table >= 0, step.block_table, own + ) + return table + + +def _view_token_buffers(step, tc, g): + """The step's token-major buffers with tc context rows ahead of the generation rows.""" + rows = step.B * step.n + ctx_q, ctx_sf = _quant_mxfp4(torch.randn(tc, step.H, DIM, generator=g, device=DEV) * 2.0) + q = torch.cat([ctx_q, step.q.view(rows, step.H, DIM // 2)]) + q_sf = torch.cat([ctx_sf, step.q_sf.view(rows, step.H)]) + w = torch.cat([(torch.randn(tc, step.H, generator=g, device=DEV) * 2.0**-4), step.weights]) + lens = torch.randint( + 0, step.max_ctx + 1, (VIEW_MAX_TOKENS,), generator=g, device=DEV, dtype=torch.int32 + ) + lens[tc : tc + rows] = step.lens.view(-1) + return q, q_sf, w, lens + + +def _view_outside(got, ref): + """Fraction of elements outside the default fp32 gate, NaN counted as outside.""" + ref32 = ref.float() + return (~((got - ref32).abs() <= 1e-5 + 1.3e-6 * ref32.abs())).float().mean().item() + + +def _target_views(step, bufs, table, tc, nc): + q, q_sf, w, lens = bufs + rows = step.B * step.n + return dict( + q=q.split([tc, rows], 0)[1].reshape(step.B, step.n, step.H, DIM // 2), + q_sf=q_sf.split([tc, rows], 0)[1].reshape(step.B, step.n, step.H), + weights=w.split([tc, rows], 0)[1], + context_lens=lens[tc : tc + rows].reshape(step.B, step.n), + block_table=table[nc : nc + step.B], + ) + + +@pytest.mark.parametrize( + "ratio,block_kv,max_ctx,next_n,tc,nc,G", + VIEW_MIXED, + ids=[f"r{r}_w{w}_n{n}_tc{tc}_nc{nc}_G{G}" for r, p, w, n, tc, nc, G in VIEW_MIXED], +) +def test_mixed_step_views(ratio, block_kv, max_ctx, next_n, tc, nc, G): + """The target's mixed-step call at the ratio's static width over the runtime-sized pool: views + at token row tc and table row nc, table 132 columns wide (row bases 16-byte aligned) and 133 + (4-byte aligned). Generation request 0 is at the width, request 1 at position 1. Every + visible element passes the fp64 reference and equals the documented reduction bitwise, the + inputs are unchanged, and the result is bitwise the contiguous call's. Controls: the query-side + views one row short of tc, and the table view one row short of nc (row 0 reads the last context + request's pages), each fail the gate.""" + assert tc + G * next_n <= VIEW_MAX_TOKENS and nc + G <= 32 + seed = zlib.crc32(f"c10-mixed/{ratio}/{max_ctx}/{next_n}/{tc}/{nc}/{G}".encode()) + pos = _static_positions(G, next_n, ratio, max_ctx, seed) + step = _Step( + pos, + next_n, + ratio, + block_kv, + seed=seed, + max_ctx=max_ctx, + pool_pages=RUNTIME_POOL_PAGES[block_kv], + ) + assert int(step.lens.max()) == max_ctx + g = torch.Generator(device=DEV) + g.manual_seed(seed) + contiguous = step.run() + ref, vis = _check(step, contiguous) + bufs = _view_token_buffers(step, tc, g) + for cols in (VIEW_TABLE_COLS, VIEW_TABLE_COLS + 1): + table = _view_full_table(step, g, VIEW_TABLE_ROWS, cols, nc) + views = _target_views(step, bufs, table, tc, nc) + assert ( + views["block_table"].stride(0) == cols + and views["q"].storage_offset() == tc * step.H * DIM // 2 + ) + assert views["context_lens"].storage_offset() == tc + before = [t.clone() for t in (*bufs, table)] + out = step.run(**views) + _check(step, out) + assert torch.equal(out[vis], contiguous[vis]) + assert all(torch.equal(a, b) for a, b in zip(before, (*bufs, table))) + # Control 1: q / q_sf / weights views one row short (row 0 reads the last context token). + short = _target_views(step, bufs, table, tc, nc) + for k, buf in (("q", bufs[0]), ("q_sf", bufs[1]), ("weights", bufs[2])): + v = buf[tc - 1 : tc - 1 + step.B * step.n] + short[k] = v.reshape(short[k].shape) + frac_q = _view_outside(step.run(**short)[vis], ref[vis]) + # Control 2: the table view one row short of nc (each request reads the previous one's pages). + shifted = _target_views(step, bufs, table, tc, nc) + shifted["block_table"] = table[nc - 1 : nc - 1 + step.B] + frac_t = _view_outside(step.run(**shifted)[vis], ref[vis]) + print( + f"[view mixed r{ratio} w{max_ctx} n{next_n} tc{tc} nc{nc} G{G}] visible {int(vis.sum())}; " + f"controls outside the gate: q offset lost {frac_q:.4f}, table offset lost {frac_t:.4f}" + ) + assert frac_q > 0.5 and frac_t > 0.5 + + +VIEW_GRAPH = [(r, p, w, 1, b) for r, p, w in VIEW_WIDTHS for b in CORE_GRAPH_BATCHES] + [ + (r, p, w, 6, b) for r, p, w in VIEW_WIDTHS for b in DSPARK_GRAPH_BATCHES +] + + +@pytest.mark.parametrize( + "ratio,block_kv,max_ctx,next_n,batch", + VIEW_GRAPH, + ids=[f"r{r}_w{w}_n{n}_G{b}" for r, p, w, n, b in VIEW_GRAPH], +) +def test_cuda_graph_view_form(ratio, block_kv, max_ctx, next_n, batch): + """The captured form of the same call (tc = nc = 0) at every graph batch of CORE_GRAPH_BATCHES + (the gates' capture sets): context_lens a prefix view of the 8192-entry lengths buffer, the table rows 0..G-1 of a + [40, 132] table. Metadata + logits captured after an eager warm-up; the first replay equals + it bitwise; 3 more replays read new lengths (the cap, 0/1, a padding request), queries, + weights, keys and re-assigned table rows (rolled by one request), each checked by + _replay_and_check.""" + seed = zlib.crc32(f"c10-graph/{ratio}/{max_ctx}/{next_n}/{batch}".encode()) + step = _Step( + [max_ctx * ratio - next_n] * batch, + next_n, + ratio, + block_kv, + seed=seed, + max_ctx=max_ctx, + pool_pages=RUNTIME_POOL_PAGES[block_kv], + ) + assert int(step.lens.max()) == max_ctx + g = torch.Generator(device=DEV) + g.manual_seed(seed) + lens_buf = torch.randint( + 0, max_ctx + 1, (VIEW_MAX_TOKENS,), generator=g, device=DEV, dtype=torch.int32 + ) + rows = batch * next_n + lens_buf[:rows] = step.lens.view(-1) + table = _view_full_table(step, g, VIEW_TABLE_ROWS, VIEW_TABLE_COLS, 0) + step.lens = lens_buf[:rows].view(batch, next_n) + step.block_table = table[:batch] + eager = step.run() + torch.cuda.synchronize() + graph, out = _capture(step) + graph.replay() + torch.cuda.synchronize() + prev = _check(step, out) + assert torch.equal(out[prev[1]], eager[prev[1]]) + gh = torch.Generator() + gh.manual_seed(seed) + for rep in range(3): + if next_n == 1: + lens = torch.randint( + 0, max_ctx + 1, (batch, 1), generator=g, device=DEV, dtype=torch.int32 + ) + lens[0] = max_ctx + if batch > 1: + lens[-1] = rep % 2 + else: + lens = _dspark_lens(batch, ratio, max_ctx, rep, gh) + step.lens.copy_(lens) + step.block_table.copy_(step.block_table.roll(1, dims=0)) + new = _Step([1] * batch, next_n, ratio, block_kv, seed=seed + rep + 1) + step.q.copy_(new.q) + step.q_sf.copy_(new.q_sf) + step.weights.copy_(new.weights) + step.refill_keys(g) + _poison_tails(step) + step.pool.copy_(_write_pool(step.k_codes, step.k_sf)) + prev = _replay_and_check( + step, graph, out, f"view graph r{ratio} w{max_ctx} n{next_n} G{batch} rep{rep}", prev + ) + + +# ─── the DSpark engine's widths and the static-width axis ──────────────────────────────────── +# With DSpark on, the runtime raises the engine's max_seq_len M from the user's 16384 to +# 16384 + 5 (overlap scheduler, tokens_per_gen_step - 1) + 4 (get_num_extra_kv_tokens, +# max_draft_len - 1) + 5 (tokens_per_gen_step - 1) = 16398 (py_executor_creator.py). The static +# width max_compressed_kv_lens[r] = (M + max_draft_tokens) // r is then 16398 // r eagerly and +# 16403 // r in the metadata's CUDA-graph copy: 8199 / 8201 at ratio 2, 16398 / 16403 at ratio 1. +# Positions reach M - 1 = 16397, so lengths reach 16398 (r1) / 8199 (r2) and a request's pages +# reach table column 128 of the 132-column tables. +DSPARK_ENGINE_WIDTHS = [(2, 64, 8199), (2, 64, 8201), (1, 128, 16398), (1, 128, 16403)] +DSPARK_EAGER_WIDTHS = [(1, 128, 16398), (2, 64, 8199)] +DSPARK_GRAPH_COPY_WIDTHS = [(1, 128, 16403), (2, 64, 8201)] + + +@pytest.mark.parametrize("batch", [3, 10, 19, 31]) +@pytest.mark.parametrize("next_n", [1, 6], ids=["decode", "dspark_verify"]) +@pytest.mark.parametrize( + "ratio,block_kv,max_ctx", + DSPARK_ENGINE_WIDTHS, + ids=[f"r{r}_page{p}_w{w}" for r, p, w in DSPARK_ENGINE_WIDTHS], +) +def test_dspark_engine_width(ratio, block_kv, max_ctx, next_n, batch): + """Each DSpark engine width over the runtime-sized pool: request 0 exactly at the width (its + pages reach table column 128), request 1 at position 1, the rest random. Every visible element + passes the fp64 reference and equals the documented reduction bitwise (_check); the row stride + is ceil256(width).""" + seed = zlib.crc32(f"c12-engine/{ratio}/{block_kv}/{max_ctx}/{next_n}/{batch}".encode()) + pos = _static_positions(batch, next_n, ratio, max_ctx, seed) + step = _Step( + pos, + next_n, + ratio, + block_kv, + seed=seed, + max_ctx=max_ctx, + pool_pages=RUNTIME_POOL_PAGES[block_kv], + ) + out = step.run() + _check(step, out) + assert int(step.lens.max()) == max_ctx + assert (int(step.lens[0].max()) - 1) // block_kv == 128 # the last page is table column 128 + assert out.stride(0) == -(-max_ctx // 256) * 256 + + +DSPARK_DECODE_GRAPH = [(r, p, w, b) for r, p, w in DSPARK_ENGINE_WIDTHS for b in CORE_GRAPH_BATCHES] + + +@pytest.mark.parametrize( + "ratio,block_kv,max_ctx,batch", + DSPARK_DECODE_GRAPH, + ids=[f"r{r}_page{p}_w{w}_B{b}" for r, p, w, b in DSPARK_DECODE_GRAPH], +) +def test_dspark_cuda_graph_decode_width(ratio, block_kv, max_ctx, batch): + """next_n 1 captured at every decode batch of CORE_GRAPH_BATCHES at each DSpark engine width (the body of + test_cuda_graph_static_width: first replay = eager warm-up bitwise, 3 replays on new data).""" + test_cuda_graph_static_width(ratio, block_kv, max_ctx, batch) + + +DSPARK_MIXED = [ + (r, p, w, n, tc, nc, G) + for r, p, w in DSPARK_EAGER_WIDTHS + for n in (1, 6) + for tc, nc, G in _view_cells(n) +] + + +@pytest.mark.parametrize( + "ratio,block_kv,max_ctx,next_n,tc,nc,G", + DSPARK_MIXED, + ids=[f"r{r}_w{w}_n{n}_tc{tc}_nc{nc}_G{G}" for r, p, w, n, tc, nc, G in DSPARK_MIXED], +) +def test_dspark_mixed_step_views(ratio, block_kv, max_ctx, next_n, tc, nc, G): + """Mixed steps are eager, so they carry the eager DSpark widths 16398 (r1) / 8199 (r2): the + body of test_mixed_step_views (views at token row tc and table row nc of 132- and + 133-column tables, both controls) at those widths.""" + test_mixed_step_views(ratio, block_kv, max_ctx, next_n, tc, nc, G) + + +DSPARK_VIEW_GRAPH = [ + (r, p, w, 1, b) for r, p, w in DSPARK_GRAPH_COPY_WIDTHS for b in CORE_GRAPH_BATCHES +] + [(r, p, w, 6, b) for r, p, w in DSPARK_GRAPH_COPY_WIDTHS for b in DSPARK_GRAPH_BATCHES] + + +@pytest.mark.parametrize( + "ratio,block_kv,max_ctx,next_n,batch", + DSPARK_VIEW_GRAPH, + ids=[f"r{r}_w{w}_n{n}_G{b}" for r, p, w, n, b in DSPARK_VIEW_GRAPH], +) +def test_dspark_cuda_graph_view_form(ratio, block_kv, max_ctx, next_n, batch): + """The captured view form (lengths a prefix of the 8192-entry buffer, table rows of a [40, 132] + table) at the graph-copy widths 16403 (r1) / 8201 (r2): the body of + test_cuda_graph_view_form.""" + test_cuda_graph_view_form(ratio, block_kv, max_ctx, next_n, batch) + + +# The static-width axis. The width W reaches DeepGEMM only as the output allocation +# [rows, ceil256(W)] (csrc/apis/attention.hpp: align(align(W, 256), 1024 B / 4 B)), its slice to W +# and the kernel's runtime logits_stride argument; the JIT template of sm100_paged_mqa_logits +# (csrc/jit_kernels/impls/sm100_mqa_logits.hpp, generate_impl) does not contain it, and the work +# schedule is built from the lengths alone. So a call at width W with lengths <= W must give, on +# every visible element, the bits a call at any larger width gives. The test runs EVERY W from 1 +# to the 132-column table's reach (r1: 16640 = ceil256(16403), 130 pages; r2: 8448 = 132 pages of +# 64, = ceil256(8201)) against one checked call at that maximum, and again on the wider tables of +# a max_seq_len 32768 engine up to r1 33024 / r2 16640 (258 / 260 pages; that engine's widths and their +# DSpark forms, (32768 + 14 + {0, 5}) // r, all lie below). +WIDTH_AXIS = [(1, 128, 16640, 132), (2, 64, 8448, 132), (1, 128, 33024, 260), (2, 64, 16640, 260)] + + +@pytest.mark.parametrize("next_n", [1, 6], ids=["decode", "dspark_verify"]) +@pytest.mark.parametrize( + "ratio,block_kv,w_max,cols", + WIDTH_AXIS, + ids=[f"r{r}_page{p}_wmax{w}_cols{c}" for r, p, w, c in WIDTH_AXIS], +) +def test_every_static_width_is_a_prefix(ratio, block_kv, w_max, cols, next_n): + """32 requests (request 0 at w_max, request 1 at position 1, request 5 a padding request, the + rest random) on the runtime-sized pool (unowned pages 0xFF) and a 132- (or 260-) column table whose + unused entries name random pool pages. The w_max call is checked (_check). Then for every W in + 1..w_max, lengths min(len, W) (request 0 always exactly at W), schedule rebuilt: shape + (rows, W), row stride ceil256(W), and every visible element bitwise equal to the w_max call's. + Control: the same output read at row stride W (a consumer ignoring the padding) differs.""" + seed = zlib.crc32(f"c12-axis/{ratio}/{block_kv}/{w_max}/{next_n}".encode()) + pos = _static_positions(32, next_n, ratio, w_max, seed) + pos[5] = None + step = _Step( + pos, + next_n, + ratio, + block_kv, + seed=seed, + max_ctx=w_max, + pool_pages=RUNTIME_POOL_PAGES[block_kv], + ) + assert int(step.lens.max()) == w_max + assert bool((step.block_table[:, cols:] == -1).all()) + step.block_table = step.block_table[:, :cols].contiguous() + g = torch.Generator(device=DEV) + g.manual_seed(seed) + step.block_table = _view_full_table(step, g, step.B, cols, 0) + assert step.block_table.shape == (32, cols) + full = step.run() + _check(step, full) + full_bits = full.view(torch.int32) + rows = step.B * step.n + lens_max = step.lens.clone() + col = torch.arange(w_max, device=DEV, dtype=torch.int32) + bad = torch.zeros((), dtype=torch.int64, device=DEV) + seen = torch.zeros((), dtype=torch.int64, device=DEV) + ctl_rows = {} + slowest = 0.0 + t0 = time.time() + for W in range(1, w_max + 1): + lens = lens_max.clamp(max=W) + out, call_s = _timed(step.run, context_lens=lens, max_context_len=W) + slowest = max(slowest, call_s) + assert out.shape == (rows, W) and out.stride() == (_aligned_width(W, torch.float32), 1) + vis = col[:W] < lens.view(-1, 1) + bad += ((out.view(torch.int32) != full_bits[:, :W]) & vis).sum() + seen += vis.sum() + if W in (8199, 16398, 32787) or (W == w_max - 1 and W % 256): + # control: the buffer read at row stride W, as if it were contiguous + wrong = torch.as_strided(out, (rows, W), (W, 1)) + ctl_rows[W] = (((wrong.view(torch.int32) != full_bits[:, :W]) & vis).sum(1) > 0)[ + step.n : + ] + torch.cuda.synchronize() + dt = time.time() - t0 + L = lens_max.view(-1).long() + expect = int((L * (L + 1) // 2 + L * (w_max - L)).sum()) + print( + f"[static-width axis r{ratio} wmax{w_max} cols{cols} n{next_n}] {w_max} widths in {dt:.1f}s, visible elements " + f"compared {int(seen)}, mismatches {int(bad)}; slowest call {slowest * 1e3:.1f} ms" + ) + assert int(seen) == expect and int(bad) == 0 + assert slowest < JIT_BOUND_S + for W, r in ctl_rows.items(): + live = int((lens_max.view(-1)[step.n :] > 0).sum()) + print( + f"[static-width axis control W {W}] rows read at stride W that differ: {int(r.sum())} of {live} live rows" + ) + assert int(r.sum()) >= live // 2 + + +# ─── the engine envelope: G <= 256, max_seq_len up to 1,048,576 ────────────────────────────── +# Graph batches: every G in 1..32 plus TRT-LLM's padded list at cuda_graph_config +# max_batch_size 256 (llm_args.py, _generate_cuda_graph_batch_sizes: 1, 2, 4, 8, 16..128 step 8, +# 192, 256); eager steps hold any G in 1..256 generation requests. Per-request length: the user's +# max_seq_len L in ENV_L. With DSpark on, the engine's max_seq_len is Lm = L + 14 and the +# metadata's CUDA-graph copy adds max_draft_tokens 5, so the static widths +# max_compressed_kv_lens[r] are Lm // r (eager) and (Lm + 5) // r (graph); with DSpark off both are +# L // r. The index-K block tables have ceil(Lm / 128) + 3 columns (132 at Lm 16398, 8196 at +# 1048590) and the lengths buffer has max_num_tokens = 16384 entries. +ENV_GSET = list(range(1, 33)) + [40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256] +ENV_L = [16384, 32768, 65536, 131072, 262144, 1048576] +ENV_MAX_NUM_TOKENS = 16384 +ENV_TABLE_ROWS = 264 # max_num_sequences >= max_batch_size 256 +ENV_PAGE = {1: 128, 2: 64} # index-K entries per page at tokens_per_block 128 +# The reference works on chunks of at most this many (row, column) elements (1 GiB of fp64 dots +# per 4M); the bitwise-only comparisons use the larger budget. +ENV_REF_BUDGET = 1 << 22 +ENV_CMP_BUDGET = 1 << 26 + + +def _env_widths(): + """(kind, L, ratio, W, table columns) of every static width the envelope's engines pass.""" + out = [] + for L in ENV_L: + lm = L + 14 + for r in (1, 2): + out.append(("on_eager", L, r, lm // r, -(-lm // 128) + 3)) + out.append(("on_graph", L, r, (lm + 5) // r, -(-lm // 128) + 3)) + out.append(("off", L, r, L // r, -(-L // 128) + 3)) + return out + + +ENV_WIDTHS = _env_widths() + +_E2M1_X2 = torch.tensor( + [0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12], dtype=torch.float32, device=DEV +) + + +def _codes_x2(packed: torch.Tensor) -> torch.Tensor: + """packed E2M1 [..., 64] -> fp32 [..., 128]: twice each element's value, an integer in [-12, 12].""" + p = packed.view(torch.uint8).int() + return torch.stack([_E2M1_X2[p & 0xF], _E2M1_X2[p >> 4]], dim=-1).flatten(-2) + + +def _exp2(e: torch.Tensor) -> torch.Tensor: + """2^e in fp64 from its bit pattern (exact).""" + return ((e.long() + 1023) << 52).view(torch.float64) + + +def _exact_dots(q, q_sf, codes, sf_bytes): + """Exact per-head q.k in fp64: q int8 [Ba, n, H, 64] + q_sf int32 [Ba, n, H]; keys as their pool + bytes, codes uint8 [Ba, C, 64] + scale bytes uint8 [Ba, C, 4] -> [Ba, n, H, C]. + + Per 32-element block the products of the doubled E2M1 codes sum to an integer of magnitude at + most 32 * 144 = 4608, which an fp32 GEMM forms exactly in any order; each block's sum times its + two power-of-two scales (2^(bq - 127) * 2^(bk - 127) / 4) is exact in fp64, and so is the sum of + the four blocks. test_env_exact_dots_equal_the_decode_reference pins this against _decode.""" + Ba, n, H, _ = q.shape + C = codes.shape[1] + qx = _codes_x2(q).view(Ba, n * H, 4, 32) + qe = q_sf.contiguous().view(torch.uint8).view(Ba, n * H, 4).int() + kx = _codes_x2(codes).view(Ba, C, 4, 32) + ke = sf_bytes.int() + out = None + for blk in range(4): + s = torch.bmm(qx[:, :, blk], kx[:, :, blk].transpose(1, 2)) + term = s.double() * _exp2(qe[:, :, blk, None] + ke[:, None, :, blk] - 256) + out = term if out is None else out + term + return out.view(Ba, n, H, C) + + +def _emulate_heads(dots, w): + """The kernel's documented fp32 head reduction over exact dots [Ba, n, H, C] and fp64 weights + [Ba, n, H] (the same recipe as _Step.emulate).""" + Ba, n, H, C = dots.shape + a = dots.float() + two = (a + a.abs()).double() + chains = [] + for lane in range(4): + acc = torch.zeros(Ba, n, C, device=DEV) + for h in range(lane, H, 4): + acc = (two[:, :, h] * w[:, :, h, None] + acc.double()).float() + chains.append(acc) + return ((chains[0] + chains[2]) + (chains[1] + chains[3])) / 2 + + +def _env_chunks(lens, budget): + """(requests, c0, c1) tiles covering every visible element of a [G, next_n] step exactly once: + the requests whose last row reaches column c0, over a column range sized to the budget.""" + G, n = lens.shape + req_len = lens.max(1).values.cpu() + maxlen = int(req_len.max()) if G else 0 + c0 = 0 + while c0 < maxlen: + bs = torch.nonzero(req_len > c0).flatten() + C = max(SPLIT_KV, budget // (bs.numel() * n) // SPLIT_KV * SPLIT_KV) + c1 = min(c0 + C, maxlen) + yield (None if bs.numel() == G else bs.to(DEV)), c0, c1 + c0 = c1 + + +def _tile(x3, bs, c0, c1): + t = x3[:, :, c0:c1] + return t if bs is None else t[bs] + + +def _outside(got, ref): + """Elements outside the default fp32 assert_close gate; NaN counts as outside.""" + return int((~((got - ref).abs() <= 1e-5 + 1.3e-6 * ref.abs())).sum()) + + +def _env_compare(out, q, q_sf, weights, lens, table, pool, W, same=(), prev=None): + """Every visible element (row (b, j), column c < lens[b, j]) of ``out`` against the exact + per-head dots of the keys ``table`` names in ``pool``: outside the default fp32 gate of the fp64 + score, and not bitwise equal to the documented fp32 reduction. ``same``: outputs that must equal + ``out`` bitwise on every visible element. ``prev``: (output, lens) of the previous step, compared + through the same gate on the elements both steps see. Works on request x column tiles, so a + [1536, 1048595] output never needs a dense fp64 copy. Returns the counts.""" + G, n = lens.shape + H = q.shape[2] + P, bk = pool.shape[0], pool.shape[1] + flat = pool.view(P, bk * KEY_BYTES) + assert out.shape == (G * n, W) and out.dtype == torch.float32 + assert out.stride() == (_aligned_width(W, torch.float32), 1) + out3 = out.view(G, n, W) + w3 = weights.view(G, n, H).double() + lens64 = lens.long() + cnt = dict(seen=0, outside=0, not_emu=0, not_same=0, stale_shared=0, stale_outside=0) + for bs, c0, c1 in _env_chunks(lens, ENV_REF_BUDGET): + idx = slice(None) if bs is None else bs + p0 = c0 // bk + pages = table[idx, p0 : -(-c1 // bk)].long().clamp(0, P - 1) + Ba = pages.shape[0] + raw = flat[pages] + o = c0 - p0 * bk + codes = raw[..., : bk * 64].reshape(Ba, -1, 64)[:, o : o + c1 - c0] + sfb = raw[..., bk * 64 :].reshape(Ba, -1, 4)[:, o : o + c1 - c0] + dots = _exact_dots(q[idx], q_sf[idx], codes, sfb) + wb = w3[idx] + ref = (dots.clamp_min(0) * wb[..., None]).sum(2).float() + emu = _emulate_heads(dots, wb) + del dots + vis = torch.arange(c0, c1, device=DEV) < lens64[idx][..., None] + got = _tile(out3, bs, c0, c1)[vis] + cnt["seen"] += got.numel() + cnt["outside"] += _outside(got, ref[vis]) + cnt["not_emu"] += int((got.view(torch.int32) != emu[vis].view(torch.int32)).sum()) + for s in same: + sv = _tile(s.view(G, n, W), bs, c0, c1)[vis] + cnt["not_same"] += int((sv.view(torch.int32) != got.view(torch.int32)).sum()) + if prev is not None: + pout, plens = prev + both = ( + torch.arange(c0, c1, device=DEV) + < torch.minimum(lens64, plens.long())[idx][..., None] + ) + cnt["stale_shared"] += int(both.sum()) + cnt["stale_outside"] += _outside( + _tile(out3, bs, c0, c1)[both], _tile(pout.view(G, n, W), bs, c0, c1)[both] + ) + return cnt + + +def _env_assert_exact(cnt, lens, tag): + assert cnt["seen"] == int(lens.long().sum()), f"{tag}: visible elements {cnt}" + assert cnt["outside"] == 0 and cnt["not_emu"] == 0 and cnt["not_same"] == 0, f"{tag}: {cnt}" + + +def _env_visible_mismatch(a, b, lens): + """(mismatches, compared): ``a`` against ``b`` bitwise on every element visible under ``lens`` + (the two may have different widths; every visible column lies below both).""" + G, n = lens.shape + a3, b3 = a.view(G, n, a.shape[1]), b.view(G, n, b.shape[1]) + lens64 = lens.long() + bad = torch.zeros((), dtype=torch.int64, device=DEV) + seen = torch.zeros((), dtype=torch.int64, device=DEV) + for bs, c0, c1 in _env_chunks(lens, ENV_CMP_BUDGET): + idx = slice(None) if bs is None else bs + vis = torch.arange(c0, c1, device=DEV) < lens64[idx][..., None] + diff = _tile(a3, bs, c0, c1).view(torch.int32) != _tile(b3, bs, c0, c1).view(torch.int32) + bad += (diff & vis).sum() + seen += vis.sum() + return int(bad), int(seen) + + +# The one-compiled-variant probe. DeepGEMM builds a kernel the first time a process launches a new +# template instantiation, and a build of either kernel took 1.18 s (metadata) and 2.1 s (logits) +# here; a launch of a built kernel takes microseconds. So a call at a new batch or width that +# returns within JIT_BOUND_S (synchronized wall time, output allocation included) built nothing. The +# on-disk cache cannot serve as the probe: it is shared with every process on the node, and its +# listing was seen missing an existing entry while other processes renamed new ones in. +JIT_BOUND_S = 0.5 +_WARM_FAMILIES = set() + + +def _reserve(rows, W): + """Leave a free block of exactly a [rows, W] fp32 output's size in the caching allocator, so the + timed call that follows reuses it instead of timing a device allocation.""" + torch.empty(rows * _aligned_width(W, torch.float32), device=DEV) + + +def _timed(fn, **kw): + torch.cuda.synchronize() + t0 = time.perf_counter() + out = fn(**kw) + torch.cuda.synchronize() + return out, time.perf_counter() - t0 + + +def _warm_family(next_n, block_kv): + """Build (or load) the (next_n, page) variants of both kernels once per process, at a one-request + 256-column call, so a later timed call measures only what its own shape adds.""" + if (next_n, block_kv) in _WARM_FAMILIES: + return + lens = torch.ones(1, next_n, dtype=torch.int32, device=DEV) + q = torch.zeros(1, next_n, HEADS, DIM // 2, dtype=torch.int8, device=DEV) + q_sf = torch.full((1, next_n, HEADS), 0x7F7F7F7F, dtype=torch.int32, device=DEV) + pool = torch.zeros(2, block_kv, 1, KEY_BYTES, dtype=torch.uint8, device=DEV) + w = torch.zeros(next_n, HEADS, device=DEV) + table = torch.zeros(1, 4, dtype=torch.int32, device=DEV) + meta = metadata_op(lens, block_kv) + op(q, q_sf, pool, w, lens, table, meta, SPLIT_KV) + torch.cuda.synchronize() + _WARM_FAMILIES.add((next_n, block_kv)) + + +class _EnvStep: + """A generation step in the target's own buffer layout at an envelope shape. + + q / q_sf / weights are the rows after ``tc`` context tokens of token-major buffers, + context_lens is ``lens_buf[tc : tc + G * next_n]`` of a max_num_tokens lengths buffer, and the + block table is rows ``nc .. nc + G`` of an [ENV_TABLE_ROWS, cols] table (tc = nc = 0 is the + captured form). ``set_step`` gives every request fresh pages drawn at random from the whole + ``pool_pages`` pool (its last page always among them) with fresh keys, and fresh queries and + weights; every page no request owns holds 0xFF bytes (NaN scale), the slots past each request's + length in its last page are NaN-poisoned, and every table entry past a request's pages names a + random pool page. The buffers are static, so a captured graph reads each step's values.""" + + def __init__(self, G, next_n, ratio, W, cols, pool_pages, seed, tc=0, nc=0): + self.G, self.n, self.ratio, self.bk, self.W, self.cols = ( + G, + next_n, + ratio, + ENV_PAGE[ratio], + W, + cols, + ) + self.g = torch.Generator(device=DEV) + self.g.manual_seed(seed) + self.gh = torch.Generator() + self.gh.manual_seed(seed) + self.pool = torch.full( + (pool_pages, self.bk, 1, KEY_BYTES), 0xFF, dtype=torch.uint8, device=DEV + ) + self.flat = self.pool.view(pool_pages, self.bk * KEY_BYTES) + self.table_buf = torch.randint( + 0, pool_pages, (ENV_TABLE_ROWS, cols), generator=self.g, device=DEV, dtype=torch.int32 + ) + self.lens_buf = torch.randint( + 0, W + 1, (ENV_MAX_NUM_TOKENS,), generator=self.g, device=DEV, dtype=torch.int32 + ) + rows = G * next_n + assert tc + rows <= ENV_MAX_NUM_TOKENS and nc + G <= ENV_TABLE_ROWS + cq, csf = _quant_mxfp4( + torch.randn(tc + rows, HEADS, DIM, generator=self.g, device=DEV) * 2.0 + ) + self.q_buf, self.q_sf_buf = cq, csf + self.w_buf = ( + (torch.randn(tc + rows, HEADS, generator=self.g, device=DEV) * 2.0**-4) + .to(torch.bfloat16) + .float() + ) + self.tc, self.nc = tc, nc + self.q = self.q_buf[tc:].view(G, next_n, HEADS, DIM // 2) + self.q_sf = self.q_sf_buf[tc:].view(G, next_n, HEADS) + self.weights = self.w_buf[tc:] + self.lens = self.lens_buf[tc : tc + rows].view(G, next_n) + self.block_table = self.table_buf[nc : nc + G] + self.owned = None + + def positions(self, rep, far=True): + """First-query positions: request 0 at the cap (its last row sees exactly W entries), + request 1 at position rep % 2, on even reps a padding request last, one request in the upper + half of the range, the rest at gate-like contexts below min(cap, 24576).""" + G, cap = self.G, self.W * self.ratio - self.n + pos = torch.randint(0, min(cap, 24576) + 1, (G,), generator=self.gh).tolist() + pos[0] = cap + if G > 1: + pos[1] = rep % 2 + if G > 2 and rep % 2 == 0: + pos[-1] = None + if far and G > 3: + pos[G // 2] = int(torch.randint(cap // 2, cap + 1, (1,), generator=self.gh)) + return pos + + def lens_of(self, positions): + n, r = self.n, self.ratio + return [ + [0] * n if p is None else [max(p + j + 1, 0) // r for j in range(n)] for p in positions + ] + + def set_step(self, positions): + n, bk, P = self.n, self.bk, self.pool.shape[0] + if self.owned is not None: + self.flat[self.owned] = 0xFF + lens = self.lens_of(positions) + assert max(max(row) for row in lens) <= self.W + self.lens.copy_(torch.tensor(lens, dtype=torch.int32)) + npg = [-(-max(row) // bk) for row in lens] + total = sum(npg) + assert total <= P + perm = torch.randperm(P, generator=self.g, device=DEV)[:total] + if total and not bool((perm == P - 1).any()): + perm[-1] = P - 1 + self.block_table.copy_( + torch.randint( + 0, P, (self.G, self.cols), generator=self.g, device=DEV, dtype=torch.int32 + ) + ) + off = 0 + for b, k in enumerate(npg): + self.block_table[b, :k] = perm[off : off + k].to(torch.int32) + off += k + for c in range(0, total, 4096): + pg = perm[c : c + 4096] + k, ks = _quant_mxfp4(torch.randn(pg.numel(), bk, DIM, generator=self.g, device=DEV)) + self.flat[pg] = _write_pool(k, ks).view(pg.numel(), bk * KEY_BYTES) + for b, row in enumerate(lens): + L = max(row) + if L % bk: + self.flat[self.block_table[b, L // bk].long(), bk * 64 + (L % bk) * 4 :] = 0xFF + self.owned = perm + q, qs = _quant_mxfp4(torch.randn(self.G, n, HEADS, DIM, generator=self.g, device=DEV) * 2.0) + self.q.copy_(q) + self.q_sf.copy_(qs) + w = torch.randn(self.G * n, HEADS, generator=self.g, device=DEV) * 2.0**-4 + self.weights.copy_(w.to(torch.bfloat16).float()) + + def run(self, **over): + a = dict( + q=self.q, + q_sf=self.q_sf, + kv_cache=self.pool, + weights=self.weights, + context_lens=self.lens, + block_table=self.block_table, + max_context_len=self.W, + ) + a.update(over) + meta = metadata_op(a["context_lens"], self.bk) + return op(schedule_meta=meta, **a) + + def compare(self, out, **kw): + return _env_compare( + out, + self.q, + self.q_sf, + self.weights, + self.lens, + self.block_table, + self.pool, + self.W, + **kw, + ) + + +def test_env_exact_dots_equal_the_decode_reference(): + """_exact_dots (integer-code GEMMs per 32-block, fp64 scaling) is bitwise the fp64 einsum of the + decoded values the rest of this file uses, with keys over a 2^-8..2^8 magnitude range.""" + g = torch.Generator(device=DEV) + g.manual_seed(7) + q, qs = _quant_mxfp4(torch.randn(3, 6, HEADS, DIM, generator=g, device=DEV) * 2.0) + k, ks = _quant_mxfp4( + torch.randn(3, 5000, DIM, generator=g, device=DEV) + * torch.logspace(-8, 8, 5000, base=2.0, device=DEV)[:, None] + ) + ref = torch.einsum("bjhd,bsd->bjhs", _decode(q, qs), _decode(k, ks)) + got = _exact_dots( + q, qs, k.view(torch.uint8), ks.contiguous().view(torch.uint8).view(3, 5000, 4) + ) + assert torch.equal(got, ref) + + +def test_env_checker_agrees_with_check_and_sees_wrong_tables(): + """_env_compare on a _Step case that _check passes: every visible element counted, none outside, + all bitwise equal to the emulation. With the block-table rows rolled by one request (a wrong + output) it reports almost every visible element outside the gate.""" + step = _Step( + [130, 999, 4100, 20001, 777, 3, 12000, 64], + 6, + 2, + 64, + seed=21, + max_ctx=16393, + pool_pages=RUNTIME_POOL_PAGES[64], + ) + out = step.run() + _check(step, out) + args = (step.q, step.q_sf, step.weights, step.lens, step.block_table, step.pool, step.max_ctx) + cnt = _env_compare(out, *args) + _env_assert_exact(cnt, step.lens, "env checker") + wrong = step.run(block_table=step.block_table.roll(1, dims=0).contiguous()) + bad = _env_compare(wrong, *args) + print( + f"[env checker control] rolled table rows: {bad['outside']} of {bad['seen']} outside the gate" + ) + assert bad["outside"] > bad["seen"] // 2 + + +# ─── capture over the envelope: every graph batch at every width ───────────────────────────── +# Captured widths are the graph copies: DSpark on, (Lm + 5) // r at next_n 6 (verify) and next_n 1; +# DSpark off, L // r at next_n 1. Each at every G in ENV_GSET, in the captured view form (tc = nc = 0: +# lengths a prefix of the 16384-entry buffer, table rows 0 .. G of the [264, cols] table), over the +# runtime-sized pool. +ENV_GRAPH = [ + (L, r, W, c, n) + for kind, L, r, W, c in ENV_WIDTHS + for n in ((6, 1) if kind == "on_graph" else (1,) if kind == "off" else ()) +] +ENV_GRAPH_CASES = [(L, r, W, c, n, G) for L, r, W, c, n in ENV_GRAPH for G in ENV_GSET] + + +@pytest.mark.parametrize( + "L,ratio,W,cols,next_n,G", + ENV_GRAPH_CASES, + ids=[f"L{L}_r{r}_w{W}_c{c}_n{n}_G{G}" for L, r, W, c, n, G in ENV_GRAPH_CASES], +) +def test_env_cuda_graph(L, ratio, W, cols, next_n, G): + """Metadata + logits captured after an eager warm-up; the first replay equals it bitwise; then 3 + replays, each on a new step written into the captured buffers (lengths with request 0 at the + width, 0/1-long rows, a padding request on even steps; fresh pages, keys, table rows, queries + and weights; the output poisoned NaN): exact against the fp64 score and bitwise the documented + reduction on every visible element, bitwise equal to an eager call on the same inputs, and far + from the previous step's result (a replay that re-read the captured step is seen). The + eager warm-up at the new (G, W) returns within JIT_BOUND_S (no kernel build).""" + seed = zlib.crc32(f"env-graph/{L}/{ratio}/{W}/{next_n}/{G}".encode()) + step = _EnvStep(G, next_n, ratio, W, cols, RUNTIME_POOL_PAGES[ENV_PAGE[ratio]], seed) + step.set_step(step.positions(0)) + _warm_family(next_n, step.bk) + eager, dt = _timed(step.run) + assert dt < JIT_BOUND_S, f"eager call at a new (G, W) took {dt:.3f}s" + graph, out = _capture(step) + graph.replay() + torch.cuda.synchronize() + _env_assert_exact(step.compare(out, same=(eager,)), step.lens, "first replay") + del eager + prev = (out.clone(), step.lens.clone()) + for rep in range(1, 4): + step.set_step(step.positions(rep)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + again = step.run() + cnt = step.compare(out, same=(again,), prev=prev) + _env_assert_exact(cnt, step.lens, f"replay {rep}") + stale = cnt["stale_outside"] / max(cnt["stale_shared"], 1) + print( + f"[env graph L{L} r{ratio} w{W} n{next_n} G{G} rep{rep}] visible {cnt['seen']}, previous " + f"step's result: {stale:.4f} of {cnt['stale_shared']} outside the gate; eager warm-up {dt * 1e3:.1f} ms" + ) + assert stale > 0.5 + prev = (out.clone(), step.lens.clone()) + del again + + +# ─── eager: every batch 1..256 at every width ──────────────────────────────────────────────── +# Rows are independent, so a batch is a prefix of a larger one: one 256-request step per (width, +# next_n) is checked exactly, and every B in 1..255 (prefix views of the same buffers, schedule +# rebuilt) must equal its first B * next_n rows bitwise on every visible element. next_n 1 at all +# 36 widths; next_n 6 at the 24 DSpark ones (eager and graph copy). +ENV_EAGER = [ + (L, r, W, c, n) + for kind, L, r, W, c in ENV_WIDTHS + for n in ((1, 6) if kind.startswith("on") else (1,)) +] + + +@pytest.mark.parametrize( + "L,ratio,W,cols,next_n", + ENV_EAGER, + ids=[f"L{L}_r{r}_w{W}_c{c}_n{n}" for L, r, W, c, n in ENV_EAGER], +) +def test_env_every_batch_is_a_prefix(L, ratio, W, cols, next_n): + seed = zlib.crc32(f"env-prefix/{L}/{ratio}/{W}/{next_n}".encode()) + step = _EnvStep(256, next_n, ratio, W, cols, RUNTIME_POOL_PAGES[ENV_PAGE[ratio]], seed) + pos = step.positions(1) + pos[5] = None # a padding request inside the batch + step.set_step(pos) + full = step.run() + _env_assert_exact(step.compare(full), step.lens, "256-request step") + assert int(step.lens.max()) == W + bad = seen = 0 + slowest, slowest_b = 0.0, 0 + for B in range(1, 256): + rows = B * next_n + _reserve(rows, W) + out, dt = _timed( + step.run, + q=step.q[:B], + q_sf=step.q_sf[:B], + weights=step.weights[:rows], + context_lens=step.lens[:B], + block_table=step.block_table[:B], + ) + if dt > slowest: + slowest, slowest_b = dt, B + assert out.shape == (rows, W) and out.stride() == (_aligned_width(W, torch.float32), 1) + m, s = _env_visible_mismatch(out, full[:rows], step.lens[:B]) + bad, seen = bad + m, seen + s + print( + f"[env prefix L{L} r{ratio} w{W} c{cols} n{next_n}] B 1..255: {seen} visible elements compared, " + f"{bad} mismatches; slowest call {slowest * 1e3:.1f} ms (B {slowest_b})" + ) + assert bad == 0 + assert slowest < JIT_BOUND_S + + +# ─── eager mixed steps at the envelope (tokens x batch x width) ────────────────────────────── +# The target's mixed-step views at T <= 16384 and nc + G <= 256, at the eager widths of L 65536 and +# 1048576 (8196-column tables at Lm 1048590): (tc, nc, G) from 1 context token and 255 generation +# requests up to T 16384 and 255 context requests. +def _env_mixed_cells(n): + return [ + (1, 1, 255), + (ENV_MAX_NUM_TOKENS - 255 * n, 1, 255), + (8192, 128, 128), + (ENV_MAX_NUM_TOKENS - n, 255, 1), + ] + + +ENV_MIXED = [ + (L, r, W, c, n, tc, nc, G) + for kind, L, r, W, c in ENV_WIDTHS + if L in (65536, 1048576) and kind in ("on_eager", "off") + for n in ((1, 6) if kind == "on_eager" else (1,)) + for tc, nc, G in _env_mixed_cells(n) +] + + +@pytest.mark.parametrize( + "L,ratio,W,cols,next_n,tc,nc,G", + ENV_MIXED, + ids=[f"L{L}_r{r}_w{W}_c{c}_n{n}_tc{tc}_nc{nc}_G{G}" for L, r, W, c, n, tc, nc, G in ENV_MIXED], +) +def test_env_mixed_step_views(L, ratio, W, cols, next_n, tc, nc, G): + """The views at token row tc and table row nc (tables cols and cols + 1 wide): exact on every + visible element, bitwise the contiguous call, inputs unchanged. Controls: the query-side views + one row short of tc and the table view one row short of nc each fail the gate.""" + seed = zlib.crc32(f"env-mixed/{L}/{ratio}/{W}/{next_n}/{tc}/{nc}/{G}".encode()) + step = _EnvStep( + G, next_n, ratio, W, cols, RUNTIME_POOL_PAGES[ENV_PAGE[ratio]], seed, tc=tc, nc=nc + ) + step.set_step(step.positions(1)) + assert int(step.lens.max()) == W + rows = G * next_n + contiguous = step.run( + q=step.q.clone(), + q_sf=step.q_sf.clone(), + weights=step.weights.clone(), + context_lens=step.lens.clone(), + block_table=step.block_table.clone(), + ) + _env_assert_exact(step.compare(contiguous), step.lens, "contiguous") + wide = torch.randint( + 0, + step.pool.shape[0], + (ENV_TABLE_ROWS, cols + 1), + generator=step.g, + device=DEV, + dtype=torch.int32, + ) + wide[:, :cols] = step.table_buf + for table in (step.table_buf, wide): + view = table[nc : nc + G] + assert view.stride(0) == table.shape[1] and step.q.storage_offset() == tc * HEADS * DIM // 2 + before = [t.clone() for t in (step.q_buf, step.q_sf_buf, step.w_buf, step.lens_buf, table)] + out = step.run(block_table=view) + _env_assert_exact( + step.compare(out, same=(contiguous,)), step.lens, f"views, {table.shape[1]} columns" + ) + assert all( + torch.equal(a, b) + for a, b in zip(before, (step.q_buf, step.q_sf_buf, step.w_buf, step.lens_buf, table)) + ) + short = step.run( + q=step.q_buf[tc - 1 : tc - 1 + rows].view(G, next_n, HEADS, DIM // 2), + q_sf=step.q_sf_buf[tc - 1 : tc - 1 + rows].view(G, next_n, HEADS), + weights=step.w_buf[tc - 1 : tc - 1 + rows], + ) + shifted = step.run(block_table=step.table_buf[nc - 1 : nc - 1 + G]) + c_q, c_t = step.compare(short), step.compare(shifted) + frac_q, frac_t = c_q["outside"] / c_q["seen"], c_t["outside"] / c_t["seen"] + print( + f"[env mixed L{L} r{ratio} w{W} n{next_n} tc{tc} nc{nc} G{G}] visible {c_q['seen']}; controls outside " + f"the gate: q offset lost {frac_q:.4f}, table offset lost {frac_t:.4f}" + ) + assert frac_q > 0.5 and frac_t > 0.5 + + +# ─── the memory edge: every request of G 256 at the largest widths ─────────────────────────── +# All 256 requests' rows at the far end: at next_n 6 and W 1,048,595 the output is [1536, 1048832] +# fp32 (6.44 GB, element offsets up to 1.61e9 and byte offsets past 2^32), over a pool of +# 256 x 8193 pages (ratio 1: 2,097,416 blocks, 18.3 GB; below the ratio-1 bound B_BOUND[1]). +ENV_EDGE = [ + (1, 1048595, 8196, 6), + (2, 524297, 8196, 6), + (1, 1048576, 8195, 1), + (2, 524288, 8195, 1), +] + + +@pytest.mark.parametrize( + "ratio,W,cols,next_n", ENV_EDGE, ids=[f"r{r}_w{W}_c{c}_n{n}" for r, W, c, n in ENV_EDGE] +) +def test_env_edge_every_request_at_the_width(ratio, W, cols, next_n): + """Eager step checked exactly; captured after it (first replay bitwise the eager call); then 2 + replays on new steps (every request within 2048 positions of the cap, fresh pages, keys, table + rows, queries, weights), each exact, bitwise an eager call, and far from the previous result.""" + G, bk = 256, ENV_PAGE[ratio] + cap = W * ratio - next_n + pool_pages = G * -(-W // bk) + 8 + seed = zlib.crc32(f"env-edge/{ratio}/{W}/{next_n}".encode()) + step = _EnvStep(G, next_n, ratio, W, cols, pool_pages, seed) + step.set_step([cap] * G) + eager = step.run() + cnt = step.compare(eager) + _env_assert_exact(cnt, step.lens, "eager") + assert int(step.lens[:, -1].min()) == W + print( + f"[env edge r{ratio} w{W} n{next_n} eager] {cnt['seen']} visible elements exact; output " + f"{tuple(eager.shape)} row stride {eager.stride(0)}; pool {pool_pages} pages" + ) + graph, out = _capture(step) + graph.replay() + torch.cuda.synchronize() + m, s = _env_visible_mismatch(out, eager, step.lens) + assert m == 0 and s == int(step.lens.long().sum()) + del eager + prev = (out.clone(), step.lens.clone()) + for rep in range(1, 3): + step.set_step(torch.randint(cap - 2048, cap + 1, (G,), generator=step.gh).tolist()) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + again = step.run() + cnt = step.compare(out, same=(again,), prev=prev) + _env_assert_exact(cnt, step.lens, f"edge replay {rep}") + stale = cnt["stale_outside"] / max(cnt["stale_shared"], 1) + print( + f"[env edge r{ratio} w{W} n{next_n} rep{rep}] visible {cnt['seen']}, previous step's result: " + f"{stale:.4f} of {cnt['stale_shared']} outside the gate; peak " + f"{torch.cuda.max_memory_allocated() / 2**30:.1f} GiB" + ) + assert stale > 0.5 + prev = (out.clone(), step.lens.clone()) + del again + + +# ─── the length axis up to the largest widths ──────────────────────────────────────────────── +# A row's value at column c depends on its query, weights and the key at c, not on its length, so a +# row of length t is the first t columns of the same row at full length. One 256-request step with +# every row at the width is checked exactly; then the lengths alone change: at next_n 1 EVERY length +# 1..W is some request's length in exactly one call; at next_n 6 every last-row length 1..256 and +# every one within 1 of a multiple of 64 (the 64- and 128-entry pages and the 256-column splits), +# rows (p + j + 1) // r. Each call's visible elements must equal the full step's bitwise. +ENV_LENGTH_AXIS = [(1, 1048595, 8196), (2, 524297, 8196)] + + +@pytest.mark.parametrize("next_n", [1, 6], ids=["decode", "dspark_verify"]) +@pytest.mark.parametrize( + "ratio,W,cols", ENV_LENGTH_AXIS, ids=[f"r{r}_w{W}_c{c}" for r, W, c in ENV_LENGTH_AXIS] +) +def test_env_every_length(ratio, W, cols, next_n): + G, bk = 256, ENV_PAGE[ratio] + cap = W * ratio - next_n + seed = zlib.crc32(f"env-length/{ratio}/{W}/{next_n}".encode()) + step = _EnvStep(G, next_n, ratio, W, cols, G * -(-W // bk) + 8, seed) + step.set_step([cap] * G) + full = step.run() + _env_assert_exact(step.compare(full), step.lens, "full-length step") + if next_n == 1: + targets = list(range(1, W + 1)) + else: + targets = sorted( + set(range(1, 257)) + | {m + d for m in range(64, W + 2, 64) for d in (-1, 0, 1) if 1 <= m + d <= W} + ) + order = torch.randperm(len(targets), generator=step.gh).tolist() + targets = [targets[i] for i in order] + targets += targets[: -len(targets) % G] + bad = seen = 0 + slowest = 0.0 + t0 = time.time() + for c in range(0, len(targets), G): + lens = torch.tensor( + step.lens_of([t * ratio - next_n for t in targets[c : c + G]]), + dtype=torch.int32, + device=DEV, + ) + assert [int(x) for x in lens[:, -1].tolist()] == targets[c : c + G] + out, dt = _timed(step.run, context_lens=lens) + slowest = max(slowest, dt) + m, s = _env_visible_mismatch(out, full, lens) + bad, seen = bad + m, seen + s + torch.cuda.synchronize() + print( + f"[env length axis r{ratio} w{W} n{next_n}] {len(set(targets))} last-row lengths in " + f"{len(targets) // G} calls, {time.time() - t0:.1f}s: {seen} visible elements compared, {bad} mismatches; " + f"slowest call {slowest * 1e3:.1f} ms" + ) + assert bad == 0 + assert slowest < JIT_BOUND_S + + +# ─── every output row stride up to the largest widths ──────────────────────────────────────── +# The width W reaches DeepGEMM only as the output allocation [rows, ceil256(W)], its slice to W and +# the kernel's runtime logits_stride = ceil256(W); the JIT template does not contain it and the +# schedule is built from the lengths alone. So any W runs the launch of the other widths with the +# same ceil256(W). This runs every stride class from 256 to ceil256 of the ratio's largest width, +# each at its lowest (W = 256k + 1), second-highest and highest width, plus every W in 1..256 and +# every envelope width of the ratio: 32 requests (request 0 at the maximum, request 1 at position 1, +# request 5 a padding request), lengths min(len, W), schedule rebuilt, every visible element +# bitwise equal to the checked call at the maximum. Control: an output read at row stride W (a +# consumer ignoring the padding) differs on the live rows. +ENV_STRIDE_AXIS = [(1, 1048595, 8196), (2, 524297, 8196)] + + +@pytest.mark.parametrize("next_n", [1, 6], ids=["decode", "dspark_verify"]) +@pytest.mark.parametrize( + "ratio,W,cols", ENV_STRIDE_AXIS, ids=[f"r{r}_w{W}_c{c}" for r, W, c in ENV_STRIDE_AXIS] +) +def test_env_every_row_stride(ratio, W, cols, next_n): + seed = zlib.crc32(f"env-stride/{ratio}/{W}/{next_n}".encode()) + step = _EnvStep(32, next_n, ratio, W, cols, RUNTIME_POOL_PAGES[ENV_PAGE[ratio]], seed) + pos = step.positions(1) + pos[5] = None + step.set_step(pos) + full = step.run() + _env_assert_exact(step.compare(full), step.lens, "maximal width") + lens_max = step.lens.clone() + top = _aligned_width(W, torch.float32) + engine = {w for _, _, r, w, _ in ENV_WIDTHS if r == ratio and w <= W} + widths = sorted( + set(range(1, 257)) + | engine + | {w for m in range(256, top + 1, 256) for w in (m - 255, m - 1, m)} + ) + controls = sorted(w for w in engine if w % SPLIT_KV) + rows = step.G * next_n + bad = seen = expect = 0 + ctl = {} + slowest = 0.0 + t0 = time.time() + for Wp in widths: + lens = lens_max.clamp(max=Wp) + out, dt = _timed(step.run, context_lens=lens, max_context_len=Wp) + slowest = max(slowest, dt) + assert out.shape == (rows, Wp) and out.stride() == (_aligned_width(Wp, torch.float32), 1) + m, s = _env_visible_mismatch(out, full, lens) + bad, seen, expect = bad + m, seen + s, expect + int(lens.long().sum()) + if Wp in controls: + wrong = torch.as_strided(out, (rows, Wp), (Wp, 1)) + live = lens.view(-1)[next_n:] > 0 + vis = torch.arange(Wp, device=DEV) < lens.view(-1, 1) + differs = ((wrong.view(torch.int32) != full[:, :Wp].view(torch.int32)) & vis).any(1) + ctl[Wp] = (int((differs[next_n:] & live).sum()), int(live.sum())) + torch.cuda.synchronize() + print( + f"[env stride axis r{ratio} w{W} n{next_n}] {len(widths)} widths (strides 256..{top}) in " + f"{time.time() - t0:.1f}s: {seen} visible elements compared, {bad} mismatches; slowest call " + f"{slowest * 1e3:.1f} ms" + ) + assert bad == 0 and seen == expect + assert slowest < JIT_BOUND_S + for Wp, (m, live) in ctl.items(): + print( + f"[env stride axis control W {Wp}] rows read at stride W that differ: {m} of {live} live rows" + ) + assert m >= live // 2 diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_dsv4_local_to_global_indices.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_dsv4_local_to_global_indices.py new file mode 100644 index 000000000000..a2a00537d485 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_dsv4_local_to_global_indices.py @@ -0,0 +1,3005 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the dsv4_local_to_global_indices catalog entry. + +Certified column (DeepSeek-V4.1 values): a 128-column SWA segment over 128-token +pages, and a 512-column compressed top-k segment over 128 // ratio entries per page +for ratio 1 (128) and ratio 2 (64); ratio-0 layers call it without a compressed +segment. The original cases use pools of POOL_PAGES = 4096 pages -- chosen, plainly +larger than any batch here needs, before a runtime-value dump existed -- and block +tables of TABLE_ROWS = 64 rows with the batch at scattered rows, so a table-row or +page addressing defect cannot hide behind row == batch index. The cases after those +use the core engine's pool sizes and then every pool, width and batch an engine can +be configured with. + +The entry owns what is certified: `CELLS` is the list, `reference` and `compare` are +the gate. This file turns a cell's spec into tensors and keeps the checks a cell +cannot carry: strided views, past-the-table indices, a -1 page, T = 0, PDL, capture +and replay, the in-tree V4 launcher, the wrong-variant controls, the op's own +rejections, and the engine-size sweeps. +""" + +import math +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.dsv4_local_to_global_indices import ( # noqa: E501 + dsv4_local_to_global_indices as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "dsv4_local_to_global_indices requires a CUDA device" + +DEV = "cuda" +WINDOW = 128 +TOPK = 512 +SWA_TPB = 128 +POOL_PAGES = 4096 +TABLE_ROWS = 64 + + +def _reference(req_id, table, local, tpb, offset): + """One segment's pool rows: `op.reference` on that segment alone, as int32 like the op's + output (every row these cases produce is far below 2**31).""" + return op.reference(req_id, table, local, tpb, offset).to(torch.int32) + + +def _table(width, g): + """[TABLE_ROWS, width] int32; each row holds distinct pages of a POOL_PAGES pool.""" + return torch.stack( + [torch.randperm(POOL_PAGES, generator=g, device=DEV)[:width] for _ in range(TABLE_ROWS)] + ).to(torch.int32) + + +def _window(positions): + """The causal window of each query: positions max(0, p-127) .. p in order, then -1.""" + start = (positions[:, None] - WINDOW + 1).clamp(min=0) + cols = start + torch.arange(WINDOW, device=DEV, dtype=torch.int32) + return torch.where(cols > positions[:, None], -1, cols).to(torch.int32) + + +def _topk_local(positions, ratio, g): + """V4.1 index top-k shape: min(512, (p+1)//ratio) distinct compressed positions in + [0, (p+1)//ratio), in score order (not sorted by position), then -1.""" + visible = (positions.long() + 1) // ratio + nmax = int(visible.max()) if positions.numel() else 0 + out = torch.full((positions.shape[0], TOPK), -1, dtype=torch.int32, device=DEV) + if nmax == 0: + return out + scores = torch.rand(positions.shape[0], nmax, generator=g, device=DEV) + col = torch.arange(nmax, device=DEV) + scores = scores.masked_fill(col[None, :] >= visible[:, None], float("-inf")) + k = min(TOPK, nmax) + idx = scores.topk(k, dim=1).indices + rank = torch.arange(k, device=DEV) + out[:, :k] = torch.where(rank[None, :] < visible[:, None], idx, -1).to(torch.int32) + return out + + +def _batch(ratio, seqs, max_seq_len, seed, replay_tail=None): + """Tokens of requests given as (cached, new); returns the op's inputs plus positions.""" + g = torch.Generator(device=DEV).manual_seed(seed) + rows = torch.randperm(TABLE_ROWS, generator=g, device=DEV)[: len(seqs)].to(torch.int32) + req_id = torch.cat([rows[b].repeat(new) for b, (_, new) in enumerate(seqs)]).contiguous() + positions = torch.cat( + [torch.arange(c, c + new, device=DEV, dtype=torch.int32) for c, new in seqs] + ) + swa_table = _table(math.ceil(max_seq_len / SWA_TPB), g) + swa_local = _window(positions) + if replay_tail is not None: + # Bounded replay: the window is truncated to the replayed segment. + swa_local = torch.where(swa_local < replay_tail, -1, swa_local).to(torch.int32) + if ratio == 0: + return req_id, swa_table, swa_local, None, None, 0, positions + tpb_c = SWA_TPB // ratio + comp_table = _table(math.ceil(max_seq_len // ratio / tpb_c), g) + comp_local = _topk_local(positions, ratio, g) + return req_id, swa_table, swa_local, comp_table, comp_local, tpb_c, positions + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + ratio = spec["ratio"] + assert spec["pool_pages"] == POOL_PAGES and spec["table_rows"] == TABLE_ROWS + assert spec["swa_tokens_per_block"] == SWA_TPB + req_id, swa_table, swa_local, comp_table, comp_local, tpb_c, positions = _batch( + ratio, + spec["seqs"], + spec["max_seq_len"], + seed=zlib.crc32(cell.why.encode()), + replay_tail=spec["replay_tail"], + ) + assert tpb_c == spec["compressed_tokens_per_block"] + args = ( + req_id, + swa_table, + swa_local, + spec["swa_tokens_per_block"], + spec["swa_row_offset"], + comp_table, + comp_local, + spec["compressed_tokens_per_block"], + spec["compressed_row_offset"], + ) + out = op(*args) + t = positions.shape[0] + assert out.shape == (t, WINDOW + (TOPK if ratio else 0)) and out.is_contiguous() + ref = op.reference(*args) + # The inputs are the V4.1 window: min(128, p + 1) valid rows per query (fewer under + # replay), and min(512, (p + 1) // ratio) valid compressed entries. + if spec["replay_tail"] is None: + want_swa = (positions.long() + 1).clamp(max=WINDOW) + assert int((ref[:, :WINDOW] >= 0).sum()) == int(want_swa.sum()) + if ratio: + want_c = ((positions.long() + 1) // ratio).clamp(max=TOPK) + assert int((ref[:, WINDOW:] >= 0).sum()) == int(want_c.sum()) + op.compare(out, ref) + + +@pytest.mark.parametrize("ratio", [1, 2]) +def test_strided_views_are_read_through_their_strides(ratio): + """Block tables as column slices of wider tables, index rows as non-contiguous views.""" + req_id, swa_table, swa_local, comp_table, comp_local, tpb_c, _ = _batch( + ratio, [(0, 700), (5000, 3), (130, 1)], 16384, seed=7 + ) + t = req_id.shape[0] + wide_swa = torch.full((TABLE_ROWS, swa_table.shape[1] + 37), -7, dtype=torch.int32, device=DEV) + wide_swa[:, : swa_table.shape[1]] = swa_table + wide_c = torch.full((TABLE_ROWS, comp_table.shape[1] + 5), -7, dtype=torch.int32, device=DEV) + wide_c[:, 5:] = comp_table + swa_t = torch.empty((WINDOW, t), dtype=torch.int32, device=DEV) + swa_t.copy_(swa_local.t()) + comp_wide = torch.full((t, WINDOW + TOPK), -9, dtype=torch.int32, device=DEV) + comp_wide[:, WINDOW:] = comp_local + swa_view, comp_view = swa_t.t(), comp_wide[:, WINDOW:] + bt_swa_view, bt_c_view = wide_swa[:, : swa_table.shape[1]], wide_c[:, 5:] + assert not swa_view.is_contiguous() and not comp_view.is_contiguous() + assert not bt_swa_view.is_contiguous() and not bt_c_view.is_contiguous() + out = op(req_id, bt_swa_view, swa_view, SWA_TPB, 0, bt_c_view, comp_view, tpb_c, 0) + assert torch.equal(out[:, :WINDOW], _reference(req_id, swa_table, swa_local, SWA_TPB, 0)) + assert torch.equal(out[:, WINDOW:], _reference(req_id, comp_table, comp_local, tpb_c, 0)) + + +def test_local_index_past_the_table_is_minus_one(): + g = torch.Generator(device=DEV).manual_seed(11) + req_id = torch.tensor([3, 3, 60], dtype=torch.int32, device=DEV) + swa_table = _table(4, g) # covers local indices 0 .. 511 + comp_table = _table(2, g) # covers 0 .. 127 at 64 per page + swa_local = torch.tensor([[511, 512, 1000, -1] * 32] * 3, dtype=torch.int32, device=DEV) + comp_local = torch.tensor([[127, 128, 5000, -1] * 128] * 3, dtype=torch.int32, device=DEV) + out = op(req_id, swa_table, swa_local, 128, 0, comp_table, comp_local, 64, 0) + assert bool((out[:, 0:WINDOW:4] >= 0).all()) and bool((out[:, WINDOW::4] >= 0).all()) + for c in (1, 2, 3): + assert bool((out[:, c:WINDOW:4] == -1).all()) + assert bool((out[:, WINDOW + c :: 4] == -1).all()) + + +def test_unallocated_page_is_not_masked(): + """A -1 page in a table is used as page -1: the op does not validate page ids.""" + req_id = torch.tensor([0], dtype=torch.int32, device=DEV) + table = torch.full((1, 4), -1, dtype=torch.int32, device=DEV) + local = torch.arange(WINDOW, dtype=torch.int32, device=DEV)[None, :] + 5 + out = op(req_id, table, local, SWA_TPB, 1000) + expect = 1000 - SWA_TPB + (local % SWA_TPB) + assert torch.equal(out, expect.to(torch.int32)) + + +@pytest.mark.parametrize("ratio", [0, 2]) +def test_zero_tokens(ratio): + req_id, swa_table, swa_local, comp_table, comp_local, tpb_c, _ = _batch( + ratio, [(0, 4)], 1024, seed=3 + ) + out = op( + req_id[:0], + swa_table, + swa_local[:0], + SWA_TPB, + 0, + comp_table, + None if comp_local is None else comp_local[:0], + tpb_c, + 0, + ) + assert out.shape == (0, WINDOW + (TOPK if ratio else 0)) and out.dtype == torch.int32 + + +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_pdl_launch_gives_the_same_table(ratio): + req_id, swa_table, swa_local, comp_table, comp_local, tpb_c, _ = _batch( + ratio, [(0, 1000), (9000, 6)], 16384, seed=5 + ) + args = (req_id, swa_table, swa_local, SWA_TPB, 0, comp_table, comp_local, tpb_c, 0) + assert torch.equal( + op(*args, enable_pdl=True), + op(*args, enable_pdl=False), + ) + + +@pytest.mark.parametrize("enable_pdl", [False, True]) +@pytest.mark.parametrize("ratio", [1, 2]) +def test_cuda_graph_capture_and_replay(ratio, enable_pdl): + """Capture a decode-shaped call; replay it over new requests, positions and tables.""" + seqs = [(c, 6) for c in range(100, 16000, 530)] + first = _batch(ratio, seqs, 16384, seed=21) + static = [x.clone() if isinstance(x, torch.Tensor) else x for x in first[:5]] + tpb_c = first[5] + + def call(): + return op( + static[0], + static[1], + static[2], + SWA_TPB, + 0, + static[3], + static[4], + tpb_c, + 0, + enable_pdl=enable_pdl, + ) + + call() # JIT outside the capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = call() + for seed in (22, 23): + new = _batch(ratio, [(c + seed, 6) for c, _ in seqs], 16384, seed=seed) + for dst, src in zip(static, new[:5]): + dst.copy_(src) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out[:, :WINDOW], _reference(new[0], new[1], new[2], SWA_TPB, 0)) + assert torch.equal(out[:, WINDOW:], _reference(new[0], new[3], new[4], tpb_c, 0)) + + +def test_matches_the_in_tree_v4_launcher_where_it_applies(): + """Same kernel: ratio 0 and ratio 2 agree bit for bit with the V4 Python launcher. + + (That launcher gates its compressed segment on compress_ratio > 1, so it cannot + express a 1:1 compressed pool; the contract records the probe.) + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.kernels import ( + deepseek_v4_local_to_global_indices, + ) + + for ratio in (0, 2): + req_id, swa_table, swa_local, comp_table, comp_local, tpb_c, _ = _batch( + ratio, [(0, 900), (7000, 6), (15000, 1)], 16384, seed=31 + ) + ours = op(req_id, swa_table, swa_local, SWA_TPB, 256, comp_table, comp_local, tpb_c, 64) + theirs = deepseek_v4_local_to_global_indices( + req_id=req_id, + block_table_swa=swa_table, + swa_local_indices=swa_local, + swa_pool_base_ptr=0, + swa_buffer_ptr=256, + tokens_per_block=SWA_TPB, + token_stride=1, + block_table_compressed=comp_table, + compressed_local_indices=comp_local, + compress_pool_base_ptr=0, + compressed_buffer_ptr=64, + compress_ratio=max(ratio, 1), + num_compressed_indices=TOPK if ratio else 0, + ) + assert torch.equal(ours, theirs) + + +def test_deliberately_wrong_controls_are_rejected(): + """Each wrong variant is run through the op and fails the same exact-equality gate.""" + report = {} + for ratio in (1, 2): + req_id, swa_table, swa_local, comp_table, comp_local, tpb_c, positions = _batch( + ratio, [(0, 1000), (0, 129), (9000, 6), (16000, 1)], 16384, seed=41 + ) + swa_off, comp_off = 3 * POOL_PAGES * SWA_TPB, 2 * POOL_PAGES * tpb_c + ref = torch.cat( + [ + _reference(req_id, swa_table, swa_local, SWA_TPB, swa_off), + _reference(req_id, comp_table, comp_local, tpb_c, comp_off), + ], + dim=1, + ) + good = op( + req_id, swa_table, swa_local, SWA_TPB, swa_off, comp_table, comp_local, tpb_c, comp_off + ) + assert torch.equal(good, ref) + + # Off-by-one window: positions p-128 .. p-1 (the query's own row missing). + shifted = _window(positions - 1) + variants = { + "window_off_by_one": ( + req_id, + swa_table, + shifted, + SWA_TPB, + swa_off, + comp_table, + comp_local, + tpb_c, + comp_off, + ), + # The other ratio's compressed page size: 64 on a ratio-1 layer, 128 on a ratio-2 layer. + "wrong_ratio_page_size": ( + req_id, + swa_table, + swa_local, + SWA_TPB, + swa_off, + comp_table, + comp_local, + 64 if ratio == 1 else 128, + comp_off, + ), + "neighbouring_table_row": ( + (req_id + 1) % TABLE_ROWS, + swa_table, + swa_local, + SWA_TPB, + swa_off, + comp_table, + comp_local, + tpb_c, + comp_off, + ), + "pool_offsets_swapped": ( + req_id, + swa_table, + swa_local, + SWA_TPB, + comp_off, + comp_table, + comp_local, + tpb_c, + swa_off, + ), + } + for name, args in variants.items(): + wrong = op(*args) + bad = int((wrong != ref).sum()) + report[(ratio, name)] = (bad, ref.numel()) + assert not torch.equal(wrong, ref), f"control {name} at ratio {ratio} was not rejected" + for (ratio, name), (bad, total) in report.items(): + print(f"control ratio={ratio} {name}: {bad} of {total} elements differ") + + +def test_rejections_are_loud(): + req_id, swa_table, swa_local, comp_table, comp_local, tpb_c, _ = _batch( + 2, [(0, 300)], 16384, seed=51 + ) + with pytest.raises(ValueError, match="must be int32 tensors"): + op(req_id.long(), swa_table, swa_local, SWA_TPB, 0) + with pytest.raises(ValueError, match="must be int32 tensors"): + op(req_id, swa_table, swa_local, SWA_TPB, 0, comp_table, comp_local.long(), tpb_c, 0) + with pytest.raises(ValueError, match="W_compressed must be a power of two, got 500"): + op(req_id, swa_table, swa_local, SWA_TPB, 0, comp_table, comp_local[:, :500], tpb_c, 0) + with pytest.raises(ValueError, match="W_swa must be a power of two, got 127"): + op(req_id, swa_table, swa_local[:, :127], SWA_TPB, 0) + with pytest.raises(ValueError, match="given together or not at all"): + op(req_id, swa_table, swa_local, SWA_TPB, 0, comp_table, None, tpb_c, 0) + with pytest.raises(ValueError, match="req_id must be a contiguous 1-D tensor"): + wide = torch.stack([req_id, req_id], dim=1) + op(wide[:, 0], swa_table, swa_local, SWA_TPB, 0) + with pytest.raises(ValueError, match="compressed_tokens_per_block must be positive"): + op(req_id, swa_table, swa_local, SWA_TPB, 0, comp_table, comp_local, 0, 0) + with pytest.raises(ValueError, match=r"swa_row_offset must be in \[0, 2\*\*31\)"): + op(req_id, swa_table, swa_local, SWA_TPB, -1) + + +# ─── The DeepSeek-V4.1 core call at engine-size tables ───────────────────────────────────────── +# The core target calls this op once per layer per step: ratio 0 as +# (req_idx_per_token[:T], SWA table, swa_local_indices_cuda[:T], 128, 0) and ratios 1 / 2 with the +# ratio's compress block table, the layer's [T, 512] index top-k and 128 // ratio entries per page, +# offsets 0. T is every row count 1..8192 (max_num_tokens) in eager steps, and in captured decode +# steps every batch 1..8 (boot, padding off), {1, 2, 4, 8, 16, 24, 32} (accuracy, padding on) and +# 6 x those batches (DSpark target verify). T enters the kernel only as the grid size. +# Tables are sized from the engine's pools: SWA 188160 pages x 128 rows, COMPRESS 453120 x 64 +# (ratio 2) and 151040 x 128 (ratio 1); block tables are 64 rows x 129 pages (chosen: the gate's +# max_seq_len 16384 needs 128 pages of 128 tokens per request, plus one), and every page id comes +# from the top 64 * 129 pages of its pool, so every produced row lies at the pool's far end. +import tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.kernels as _l2g_kernels # noqa: E402 + +ENGINE_SWA_PAGES = 188160 +ENGINE_COMP_PAGES = {1: 151040, 2: 453120} +ENGINE_ROWS = 64 +ENGINE_BLOCKS = 129 + + +def _engine_table(pages, g): + """[ENGINE_ROWS, ENGINE_BLOCKS] page ids, a random arrangement of the pool's top + ENGINE_ROWS * ENGINE_BLOCKS pages (the last page included).""" + band = ENGINE_ROWS * ENGINE_BLOCKS + ids = pages - band + torch.randperm(band, generator=g, device=DEV) + return ids.view(ENGINE_ROWS, ENGINE_BLOCKS).to(torch.int32) + + +def _topk_local_with_newest(positions, ratio, g): + """Like _topk_local, and the newest visible compressed position is always among the picks.""" + visible = (positions.long() + 1) // ratio + nmax = max(int(visible.max()), 1) + scores = torch.rand(positions.shape[0], nmax, generator=g, device=DEV) + col = torch.arange(nmax, device=DEV) + scores = scores.masked_fill(col[None, :] == (visible - 1)[:, None], float("inf")) + scores = scores.masked_fill(col[None, :] >= visible[:, None], float("-inf")) + k = min(TOPK, nmax) + idx = scores.topk(k, dim=1).indices + out = torch.full((positions.shape[0], TOPK), -1, dtype=torch.int32, device=DEV) + rank = torch.arange(k, device=DEV) + out[:, :k] = torch.where(rank[None, :] < visible[:, None], idx, -1).to(torch.int32) + return out + + +def _engine_step(ratio, seqs, seed): + """One step's inputs on engine-size tables. seqs: (cached, new) per request in batch order. + Request b's newest SWA page and newest compressed page are the pool's last pages + (page P - 1 - b), so the pool's last rows are produced.""" + g = torch.Generator(device=DEV).manual_seed(seed) + rows = torch.randperm(ENGINE_ROWS, generator=g, device=DEV)[: len(seqs)].to(torch.int32) + req_id = torch.cat([rows[b].repeat(new) for b, (_, new) in enumerate(seqs)]).contiguous() + positions = torch.cat( + [torch.arange(c, c + new, device=DEV, dtype=torch.int32) for c, new in seqs] + ) + swa_table = _engine_table(ENGINE_SWA_PAGES, g) + for b, (c, new) in enumerate(seqs): + swa_table[rows[b], (c + new - 1) // SWA_TPB] = ENGINE_SWA_PAGES - 1 - b + swa_local = _window(positions) + if ratio == 0: + return req_id, swa_table, swa_local, None, None, 0, positions + tpb_c = SWA_TPB // ratio + pages = ENGINE_COMP_PAGES[ratio] + comp_table = _engine_table(pages, g) + for b, (c, new) in enumerate(seqs): + vis = (c + new) // ratio + if vis: + comp_table[rows[b], (vis - 1) // tpb_c] = pages - 1 - b + comp_local = _topk_local_with_newest(positions, ratio, g) + return req_id, swa_table, swa_local, comp_table, comp_local, tpb_c, positions + + +def _step_reference(ratio, step): + req_id, swa_table, swa_local, comp_table, comp_local, tpb_c, _ = step + ref = _reference(req_id, swa_table, swa_local, SWA_TPB, 0) + if ratio: + ref = torch.cat([ref, _reference(req_id, comp_table, comp_local, tpb_c, 0)], dim=1) + return ref + + +def _target_call(ratio, req_id, swa_table, swa_local, comp_table, comp_local, tpb_c): + """The target's call forms (offsets 0).""" + if ratio == 0: + return op(req_id, swa_table, swa_local, SWA_TPB, 0) + return op(req_id, swa_table, swa_local, SWA_TPB, 0, comp_table, comp_local, tpb_c, 0) + + +def _compiled_variants(): + """Number of compiled variants of the op's Triton kernel held by this process.""" + kern = _l2g_kernels._deepseek_v4_local_to_global_kernel + caches = getattr(kern, "device_caches", None) + if caches is not None: + return sum(len(v[0] if isinstance(v, tuple) else v) for v in caches.values()) + return sum(len(v) for v in kern.cache.values()) + + +# One mixed step of 8192 tokens: contexts (fresh, chunked on a cached prefix, 1-token, up to +# position 16383) ahead of 30 decode rows at positions 0 .. 16383. +_CORE_STEP_DECODE = [0, 1, 126, 127, 128, 129, 16383] + list(range(300, 16000, 600))[:23] +CORE_STEP = [ + (0, 3000), + (8192, 2048), + (0, 1), + (0, 129), + (130, 5), + (2, 1570), + (4000, 430), + (15405, 979), +] + [(c, 1) for c in _CORE_STEP_DECODE] + + +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_rows_prefix_bitwise(ratio): + """Every T = 1..8192, not a sample: the 8192-row call equals the reference (engine-size tables, + the pools' last rows produced), every call on the first T rows (prefix views of the same + buffers, as the target slices its metadata) is bitwise equal to the first T rows of it, and the + sweep adds no compiled variant of the Triton kernel. A counter control shows a new constexpr + combination does add one.""" + step = _engine_step(ratio, CORE_STEP, seed=9000 + ratio) + req_id, swa_table, swa_local, comp_table, comp_local, tpb_c, positions = step + n = req_id.numel() + assert n == 8192 + full = _target_call(ratio, *step[:6]) + ref = _step_reference(ratio, step) + assert torch.equal(full, ref) + assert int(full[:, :WINDOW].max()) >= (ENGINE_SWA_PAGES - 1) * SWA_TPB + if ratio: + assert int(full[:, WINDOW:].max()) >= (ENGINE_COMP_PAGES[ratio] - 1) * tpb_c + before = _compiled_variants() + assert before > 0 + bad = torch.zeros((), dtype=torch.int64, device=DEV) + for t in range(1, n + 1): + out = _target_call( + ratio, + req_id[:t], + swa_table, + swa_local[:t], + comp_table, + None if comp_local is None else comp_local[:t], + tpb_c, + ) + bad += (out != full[:t]).any(1).sum() + torch.cuda.synchronize() + after = _compiled_variants() + print( + f"ratio={ratio}: prefixes T=1..{n}: rows differing {int(bad)}; compiled variants {before} -> {after}" + ) + assert int(bad) == 0 + assert after == before + # Counter control: a page size never used elsewhere is a new constexpr, hence a new variant. + op(req_id[:4], swa_table, swa_local[:4], 1000 + ratio, 0) + torch.cuda.synchronize() + assert _compiled_variants() == after + 1 + + +# Captured decode steps (the union of the gates' capture sets per rank): the graph batches +# G of the boot gates (every G in 1..8: max_batch_size 8, padding off) and of the accuracy gates +# ({1, 2, 4, 8, 16, 24, 32}: padding on), each at next_n 1 (T = G rows) and at next_n 6 (the DSpark +# target-verify step, T = 6 G rows: the DSpark boot gates capture every G 1..8 at next_n 6, so +# T = 18, 30, 36 and 42 are in the union too). +CAPTURE_UNION_BATCHES = (1, 2, 3, 4, 5, 6, 7, 8, 16, 24, 32) +CAPTURE_UNION = sorted( + {(b, 1) for b in CAPTURE_UNION_BATCHES} | {(6 * b, 6) for b in CAPTURE_UNION_BATCHES} +) + + +def _decode_seqs(rows, next_n, g, step): + """rows // next_n requests, each verifying next_n positions at random positions; on even steps + request 0 ends at position 16383 and request 1 (when present) starts at position 0.""" + batch = rows // next_n + cached = torch.randint(0, 16384 - next_n + 1, (batch,), generator=g).tolist() + if step % 2 == 0: + cached[0] = 16384 - next_n + if batch > 1: + cached[1] = 0 + return [(c, next_n) for c in cached] + + +@pytest.mark.parametrize("rows,next_n", CAPTURE_UNION, ids=[f"T{r}_n{n}" for r, n in CAPTURE_UNION]) +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_cuda_graph_capture_sets(ratio, rows, next_n): + """Captured as the target captures it: prefix views [:T] of max_num_tokens-row metadata + buffers, static tables, a static [T, 512] top-k buffer, after an eager warm-up at that shape. + Three replays over new requests, positions, tables and top-k: each is bitwise equal to an eager + call on the same inputs and to the reference, and differs from the previous step's reference.""" + g = torch.Generator(device="cpu").manual_seed(1000 * ratio + rows) + max_tokens = 8192 + req_buf = torch.zeros(max_tokens, dtype=torch.int32, device=DEV) + swa_buf = torch.full((max_tokens, WINDOW), -1, dtype=torch.int32, device=DEV) + first = _engine_step(ratio, _decode_seqs(rows, next_n, g, 0), seed=11000 + 100 * ratio + rows) + swa_table = first[1].clone() + comp_table = None if ratio == 0 else first[3].clone() + comp_buf = None if ratio == 0 else torch.empty((rows, TOPK), dtype=torch.int32, device=DEV) + tpb_c = first[5] + + def load(step): + req_buf[:rows].copy_(step[0]) + swa_buf[:rows].copy_(step[2]) + swa_table.copy_(step[1]) + if ratio: + comp_table.copy_(step[3]) + comp_buf.copy_(step[4]) + return _step_reference(ratio, step) + + def call(): + return _target_call( + ratio, req_buf[:rows], swa_table, swa_buf[:rows], comp_table, comp_buf, tpb_c + ) + + prev = load(first) + assert torch.equal(call(), prev) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = call() + for it in range(1, 4): + ref = load( + _engine_step( + ratio, _decode_seqs(rows, next_n, g, it), seed=11000 + 100 * ratio + rows + 7 * it + ) + ) + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, ref), f"replay {it}" + assert torch.equal(out, call()), f"replay {it} vs eager" + assert not torch.equal(out, prev), f"replay {it} equals the previous step" + prev = ref + + +@pytest.mark.parametrize("ratio", [1, 2]) +def test_results_up_to_int32_max_are_exact(ratio): + """offset + page * tpb + l % tpb reaching exactly 2**31 - 1 at a pool's last page: the kernel's + int32 arithmetic is exact up to the top of the int32 range.""" + tpb_c = SWA_TPB // ratio + comp_pages = ENGINE_COMP_PAGES[ratio] + req_id = torch.zeros(1, dtype=torch.int32, device=DEV) + swa_table = torch.tensor([[ENGINE_SWA_PAGES - 1]], dtype=torch.int32, device=DEV) + comp_table = torch.tensor([[comp_pages - 1]], dtype=torch.int32, device=DEV) + swa_local = torch.arange(WINDOW, dtype=torch.int32, device=DEV)[None, :] + comp_local = torch.full((1, TOPK), -1, dtype=torch.int32, device=DEV) + comp_local[0, :tpb_c] = torch.arange(tpb_c, dtype=torch.int32, device=DEV) + swa_off = 2**31 - ENGINE_SWA_PAGES * SWA_TPB + comp_off = 2**31 - comp_pages * tpb_c + out = op( + req_id, swa_table, swa_local, SWA_TPB, swa_off, comp_table, comp_local, tpb_c, comp_off + ) + ref = torch.cat( + [ + _reference(req_id, swa_table, swa_local, SWA_TPB, swa_off), + _reference(req_id, comp_table, comp_local, tpb_c, comp_off), + ], + dim=1, + ) + assert int(out[0, WINDOW - 1]) == 2**31 - 1 and int(out[0, WINDOW + tpb_c - 1]) == 2**31 - 1 + assert torch.equal(out, ref) + + +# --------------------------------------------------------------------------------------------- +# The multi-request engine envelope. One call covers ALL rows of a step, context rows +# first and then generation rows (T <= 8192, <= 32 requests, per-request length <= 16384), with +# req_idx_per_token naming each token's block-table row, the SWA window (positions +# max(0, p-127) .. p, then -1) and, at ratios 1 and 2, the layer's index top-k (valid compressed +# positions < (p+1)//ratio packed first, then -1). Block tables are engine-size: ENGINE_ROWS x +# ENGINE_BLOCKS page ids drawn as a random permutation of the whole engine pool (SWA 188160 pages; +# COMPRESS 453120 x 64 rows at ratio 2, 151040 x 128 at ratio 1). +import base64 # noqa: E402 +import random # noqa: E402 + +from torch.profiler import ProfilerActivity, profile # noqa: E402 + +KERNEL_NAME = "_deepseek_v4_local_to_global_kernel" +MAX_REQUEST_LEN = 16384 + + +def _multi_request_tables(ratio, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + n = ENGINE_ROWS * ENGINE_BLOCKS + swa = torch.randperm(ENGINE_SWA_PAGES, generator=g, device=DEV)[:n].view( + ENGINE_ROWS, ENGINE_BLOCKS + ) + comp = None + if ratio: + comp = torch.randperm(ENGINE_COMP_PAGES[ratio], generator=g, device=DEV)[:n].view( + ENGINE_ROWS, ENGINE_BLOCKS + ) + comp = comp.to(torch.int32) + return swa.to(torch.int32), comp + + +def _multi_request_step(ratio, ctx, gen, next_n, seed): + """Per-token inputs of one step: ctx = [(cached, new)] context requests, then gen = [cached] + generation requests each verifying next_n positions cached .. cached + next_n - 1. + Returns (req_id, swa_local, comp_local, first row of each request, table row of each).""" + g = torch.Generator(device=DEV).manual_seed(seed) + seqs = list(ctx) + [(c, next_n) for c in gen] + rows = torch.randperm(ENGINE_ROWS, generator=g, device=DEV)[: len(seqs)].to(torch.int32) + req_id = torch.cat([rows[b].repeat(new) for b, (_, new) in enumerate(seqs)]).contiguous() + positions = torch.cat( + [torch.arange(c, c + new, device=DEV, dtype=torch.int32) for c, new in seqs] + ) + assert int(positions.max()) < MAX_REQUEST_LEN + swa_local = _window(positions) + comp_local = _topk_local(positions, ratio, g) if ratio else None + starts = [0] + for _, new in seqs[:-1]: + starts.append(starts[-1] + new) + return req_id, swa_local, comp_local, starts, rows + + +def _multi_request_reference(ratio, req_id, swa, swa_local, comp, comp_local): + ref = _reference(req_id, swa, swa_local, SWA_TPB, 0) + if ratio: + ref = torch.cat([ref, _reference(req_id, comp, comp_local, SWA_TPB // ratio, 0)], dim=1) + return ref + + +def _multi_request_random_mixed(rng): + """k in [1, 32] requests, a random number of them context requests (random cached prefix and + new length, T - generation rows split at random) and the rest generation requests with next_n + 1 or 6 (DSpark verify), T in [1, 8192].""" + k = rng.randint(1, 32) + next_n = rng.choice((1, 6)) + n_gen = rng.randint(0, k) + n_ctx = k - n_gen + gen_rows = n_gen * next_n + ctx = [] + if n_ctx: + ctx_rows = rng.randint(gen_rows + n_ctx, 8192) - gen_rows + cuts = sorted(rng.sample(range(1, ctx_rows), n_ctx - 1)) + news = [b - a for a, b in zip([0] + cuts, cuts + [ctx_rows])] + ctx = [(rng.randint(0, MAX_REQUEST_LEN - n), n) for n in news] + gen = [rng.randint(0, MAX_REQUEST_LEN - next_n) for _ in range(n_gen)] + return ctx, gen, next_n + + +def _profiled_calls(prof): + calls = [e for e in prof.events() if e.name == "trtllm::dsv4_local_to_global_indices"] + return len(calls), {kern.name for e in calls for kern in e.kernels} + + +def _check_variants(prof, n_calls, before): + seen, names = _profiled_calls(prof) + after = _compiled_variants() + print( + f"kernel-variant map over {seen} calls: {sorted(names)}; compiled variants {before} -> {after}" + ) + assert seen == n_calls and names == {KERNEL_NAME}, (seen, names) + assert after == before + + +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_random_mixed_partitions(ratio): + """200 seeded random mixed steps per ratio on engine-size tables, each bitwise equal to the + native-torch table lookup, under torch.profiler: every call launches the one Triton kernel, + and no compiled variant is added after one warm-up call of the ratio's form. Every 20th step + with >= 2 requests also runs a control with req_idx_per_token off by one request at a request + boundary (the next request's first token given the previous request's table row); the same + exact check must see it.""" + rng = random.Random(23100 + ratio) + swa, comp = _multi_request_tables(ratio, seed=23200 + ratio) + tpb_c = SWA_TPB // ratio if ratio else 0 + steps = [_multi_request_random_mixed(rng) for _ in range(200)] + warm = _multi_request_step(ratio, [(0, 4)], [], 1, seed=1) + _target_call(ratio, warm[0], swa, warm[1], comp, warm[2], tpb_c) + torch.cuda.synchronize() + before = _compiled_variants() + controls, rows_total, kinds = [], 0, set() + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for i, (ctx, gen, next_n) in enumerate(steps): + req_id, swa_local, comp_local, starts, rows = _multi_request_step( + ratio, ctx, gen, next_n, seed=ratio * 1000 + i + ) + out = _target_call(ratio, req_id, swa, swa_local, comp, comp_local, tpb_c) + ref = _multi_request_reference(ratio, req_id, swa, swa_local, comp, comp_local) + assert torch.equal(out, ref), (i, ctx, gen, next_n) + rows_total += req_id.numel() + kinds.add((bool(ctx), bool(gen), next_n if gen else 0)) + if i % 20 == 0 and len(starts) >= 2: + b = rng.randint(1, len(starts) - 1) + bad_id = req_id.clone() + bad_id[starts[b]] = rows[b - 1] + wrong = _target_call(ratio, bad_id, swa, swa_local, comp, comp_local, tpb_c) + controls.append(int((wrong != ref).sum())) + torch.cuda.synchronize() + print( + f"ratio={ratio}: {len(steps)} steps, {rows_total} rows, step kinds (ctx, gen, next_n) " + f"{sorted(kinds)}; control (req_id off by one request at a boundary) elements differing {controls}" + ) + assert len(controls) >= 8 and all(c > 0 for c in controls), controls + assert { + (True, True, 1), + (True, True, 6), + (False, True, 1), + (False, True, 6), + (True, False, 0), + } <= kinds + _check_variants(prof, len(steps) + len(controls), before) + + +# Every (T, ratio) cell of dsv4_local_to_global_indices in the core runtime-value dumps (boot, +# GSM8K, 24K passkey, gate and accuracy-preview logs). The dump records only T, +# which counts context AND generation rows, so each cell is given one consistent step, as +# "T,ratio,tc,nc,max_q_len,G,next_n" records (space-separated, zlib-compressed, base64): +# - 1577 cells: a trtllm_sparse_mla_dsv4.varlen record at tc == T (nc context requests, one with +# max_q_len new tokens), no generation rows; +# - 25 cells: a deepgemm_fp8_fp4_paged_mqa_logits record with G * next_n == T, generation only; +# - 1094 cells: the largest varlen tc < T whose remainder T - tc is a dumped decode batch G x +# next_n with nc + G <= 32 (context rows first, then generation rows). +# 7 dumped records are garbles and dropped: T 155512 and 673966 (above max_num_tokens 8192), and +# five whose fields are not the call's (e.g. (1849, 128128, 512, 128), (8192, 384, 0, 576)). +DUMPED_CELLS = ( + "eNp13VuS47gOBNCt9ALmw0i9yP1vbNqZhGgS2XEjqowzbd6yBMtSSqbiv89/wf/9/f0n/oulwk+Fvz+/NUYVS4Wf6vj7U/87" + "WMVS4ac63395soqlwk91vf/yYhVLhZ/qfv/lzSqWCj/V8/7Lh1UsFX6q9v7LxiqWCj9VfN4yuAy/LzG+r0yP/uTjX8AK31cQ" + "3z9cj/7k41/AApgLEqfq2Gr81gfmKoLq2Gr81ud3AZzf181HhNgBC1z8m87+/Uu+P4/EcIiK1/r0GBgOUfD4/r9f6H/x+/Ma" + "Fsaw281Wur7L9/vzHBbGsFvjcx+Ox5+yMIZi33Xw/Rn/6bEsjKHYzefyb2n3a2EMxZ7luTEsjGG3/n2zf3/Gf3osC2PY7P5w" + "WXUug65l8LUwhmLf/rw/32Wgn7IwhmInn3vQjtfCGIp1PrfT+ngdWa2G3eKzPDeGhTEU4+sNvsLAGC+r1VDs4XMf2vNaGMNu" + "CD6Xf0t8X5MsjKEYlzO4/DCWHzcA1VDs4nMv2jX+lqxWQ7Gbz71p92thDMUan8tlwJ+yKMY+4HZDj2VhDLsd7IODa56PZWEM" + "xdgH3F7rsSyMoRjXB7eCeiwLYyjGHjq4DI7ntTCG3U6+3pOv4/y8FsZQ7OBz2Rvn8VoYQzH20MnXwZ+yMIZifL0ne+O8Xwtj" + "KNaW52JYGMNuF5fVxWVwjWXwVquhGNfvxddxna+FMRTje4afZXosC2Mo9izPjWFhDMXa8lwMC2PYjZ/TNz+Z9VgWxlCM27Wb" + "W7I7XgtjKMb34M331o3XwhiKsZ9v9ul9vBbGUIz9fLOD7+u1MIZiXL83l999vxbGUIzr6OZaudtrYQy7PVzOD5c9f8rCGIpx" + "OT9cfs9Yfm+1Goody3NjWBhDMb6PHr4/nvH+eKvVUIzr6OGyf67XwhiK8X30cNk/4/3xPG/1a9iN+5039zT1WBbGUIzriPuT" + "99ifnNVqKMZ11Lj8Gl4LYyjGddS4VtrxWhhDMa6jxmXfztfCGIpxHTUu+3a9FsZQjOuIe8d6LAtj2K1zOXcuvx6vhTEU43Lu" + "XH4dr4UxFONy5r73nfveb7UainE58+BKj2VhDMW4zens036/FsZQjMu5c/n157UwhmLcXnVuh3p7LYyhGPfreJShx7Iwhs2e" + "z3ddPTzK0GNZGEOxg889aMdrYQzFLj73pGn/5WthDMVuPvem3WO8rFZDsWd5bgwLY9iNx0IPj370WBbGUIzLObj8Il4LYygG" + "Phc0vBbGUIzrKLjs43gtjKHYyedy2Y+cZ1aroRjXb3CNxvVaGEOxe3luDAtjKNb43EZrY7ysVkOxzud2Wn8tjGE3sDfAdY7P" + "a2EMxdgbPELVY1kYQzH2BrjOgdfCGIqxN3j0qseyMIZi7A0eGeuxLIyhGHuDx7/POP6d1Wooxvcvj2j1WBbGUIy9wSNjPZaF" + "MRRjb/CY+EF/LYxhNx4HPzwmfsYx8axWQ7FYnhvDwhiKsTd4PP2M4+lZrYZi7I2D6/w4XgtjKMbe4LH4M47FZ7UairE3Dq7z" + "43otjKEYe4PH8c84jp/VaijG9XtwvR39tTCG3U6uo5Nr5YzXwhiKcR2dXPYnXgtjKMZ1xPzgGfnBrFZDsXN5bgwLYyjGdcRU" + "Qo9lYQzFuG1nfqDHsjCGYlxHJ5f92V8LY9iNecPvv5aFMRTj+r243q6x3t5qNRTj+r243i68FsZQjOv34rK/jtfCGIpx/TIv" + "eUZeMqvVUIzr9+J6u67XwhiKcf0y83hG1jKr1VCM71+mMHosC2Moxm07Uxg9loUxFGNfXVznV38tjGE3ZjMPc5pn5DSzWg3F" + "2FfMaZ6R08xqNRRjbzBreUbWMqvVUIy9cXOd3+drYQzF2BvMaZ6R08xqNRRjbzCneUZOM6vVUIy9cXOd389rYQzF2BvMeJ6R" + "8cxqNRRjb9xc53d/LYxht4e9wWRIj2VhDMXYG0yd9FgWxlCM2xwmQ8/Ilma1Goqxr5g66bEsjKEY+4rZ0jOypVmthmLsDSZD" + "eiwLYyjG3mDqpMeyMIZi7I2H6/xpr4UxFGNv8EyoHsvCGHZjjvUw03pGpjWr1VCM65e51DNyqVmthmJcv8ylnpFLzWo1FOP6" + "ZS71jFxqVquhGLcbzKWekUvNajUUY2/w7O0zzv/OajUUY28w03pGpjWr1VCMvdG4zlt7LYyhWF+eG8PCGHbr7I3Odd7HOn+r" + "1VCM2w3mcM/I4Wa1Goqxr5jDPSOHm9VqKMa+Yg73jBxuVquhGPuKOdwzcrhZrYZi7KvOfunXa2EMxdhXzPCekeHNajUUY18x" + "w3tGhjer1VCMfcUM7xkZ3qxWQzH2FTO8JzO8t1oNmzVeu9B4tYIey8IYigWfG7R4LYyhGPhc0PBaGEOxg889aMdrYQzFTj73" + "pJ2vhTEUu/jci3a9FsZQ7OZzb9r9WhhDsYfPfWjPa2EMxRqf22jttTCGYp3P7bT+WhjDbsw4G/PONvLOWa2GYuwr5p1t5J2z" + "Wg3F2FfMO9vIO2e1Goqxr5h3tpF3zmo1FGNfMe9sI++c1Wooxr5iWtlG3jmr1VCMfcUktI2L6Ga1Goqxr3idTBvXycxqNRRj" + "XzErbSMrndVqKMa+Ylbaor8WxrAbM87GrFQ/ZWEMxdgbzDvbyDtntRqKsTeYd7aRd85qNRRjbzDvbCPvnNVqKMbeYN7ZRt45" + "q9VQjL3B64LauC5oVquhGHuDWWkbWemsVkMx9gaz0jay0lmthmLsDWalDf21MIbdmI82Jp1tZKWzWg3F2FdMUfVYFsZQjH3F" + "rLSNrHRWq6EY+4pZaRtZ6axWQzH2FbPSNrLSWa2GYuwrZqVtZKWzWg3F2FcH++W4XwtjKMa+Ys7aRs46q9VQjH11sF+O9loY" + "QzH2FTPadvTXwhh24zVcjddztXE916xWQzH2FfPdNvLdWa2GYlieG8PCGIqxr5jOtpHvzmo1FGNfMfnVY1kYQzH2FfPdNvLd" + "Wa2GYuwr5rtt5LuzWg3F2Fe8sk2PZWEMxdhXvLKtjau+Z7UairGvmAq3s78WxrAb8+fGxLmNa99mtRqKsa+YK7eRK89qNRRj" + "bzAbbiMbntVqKMbeYDbcRjY8q9VQjL3BbLiNbHhWq6EYe4PZcBvZ8KxWQzH2BrPhNrLhWa2GYuwNZsNtZMOzWg3F2BvMhtvV" + "Xwtj2I15cGM23EY2PKvVUIy9wWy4jWx4VquhGLc5vNavjWv9ZrUairGvmCu3kSvPajUUY18xV24jV57Vaih2Lc+NYWEMxdhX" + "zJXbyJVntRqKsa+YK7eRK89qNRRjXzFXbiNXntVqKMa+Yq7c7v5aGMNuzJIbc+U2cuVZrYZi7Cvmym3kyrNaDcXYV8yV28iV" + "Z7UairE3mA23kQ3PajUUY2/w+sQ2rk+c1Wooxt5grtxGrjyr1VCMvcFcuY1ceVaroRh7g7lyG7nyrFZDMfYGc+X29NfCGHZj" + "ltyYK7eRK89qNRRjb/CayjauqZzVaijG3mAm3UYmPavVUIzbHGbSbWTSs1oNxdhXzKTbyKRntRqKsa+YSbeRSc9qNRRjXzGT" + "biOTntVqKMa+YibdRiY9q9VQjH3FRLmNTHpWq6EY+4pptR7Lwhh2Yw7dmEm3kUnPajUUY18xk24jk57VaijGvmIm3UYmPavV" + "UIx9xUy6jUx6VquhGPuKmXQbmfSsVkMx9hUz6TYy6VmthmLsDebKbeTKs1oNxdgbzJXbyJVntRqKsTeYK7fMld9qNWzWmSV3" + "5sp95MqzWg3Fgs8NWrwWxlAMfC5oeC2ModjB5x6047UwhmInn3vSztfCGIpdy3NjWBhDsZvPvWn3GC+r1VDs4XMf2vNaGEOx" + "xuc2WnstjKFY53M7rb8WxrAbs+TOXLmPXHlWq6EY+4q5ch+58qxWQzH2FXPlPnLlWa2GYuwr5sp95MqzWg3F2FfMlfvIlWe1" + "Goqxr5gr95Erz2o1FGNfMVfuI1ee1Wooxr5irtxHrjyr1VCMfcVcuY9ceVaroRj7irlyj/5aGMNuvD63M1Hu41rdWa2GYuwr" + "ptV9XKs7q9VQjH3FTLqPTHpWq6EY+4qZdB+Z9KxWQzH2FTPpPjLpWa2GYuwrZtJ9ZNKzWg3F2FfMpPvIpGe1Goqxr5hJ95FJ" + "z2o1FGNfMZPuI5Oe1Wooxr5iJt3RXwtj2I05dGcm3UcmPavVUIx9xUy6j0x6VquhGPuKmXQfmfSsVkMx9hUz6T4y6VmthmLs" + "K2bSfWTSs1oNxdhXzKT7yKRntRqKsTeYK/eRK89qNRRjbzBX7iNXntVqKMbeYK7cj/5aGMNuzJI7c+U+cuVZrYZi7A2mwn3k" + "yrNaDcXYG0yc+7i+eFaroRh7g7lyH7nyrFZDMfYGc+U+cuVZrYZi7A3myn3kyrNaDcW4zWGu3EeuPKvVUIx9xVy5j1x5Vquh" + "WFueG8PCGIqxr5gr9zP7JavVsBuz5M5cuY9ceVaroRj7irlyH7nyrFZDMfYVr2vu47rmWa2GYuwrZtJ9ZNKzWg3F2FfMpPvI" + "pGe1Goqxr5hJ95FJz2o1FGNfMZPuI5Oe1Wooxr5iJt1HJj2r1VCMfcVMuo9MelaroRj7ipl0v/prYQy7MYfuzKT7yKRntRqK" + "sa+YSfeRSc9qNRRjXzGT7iOTntVqKMa+YibdRyY9q9VQjH3FRLmPTHpWq6EY+4ppdR/XRM9qNRRjXzGT7iOTntVqKMa+Yibd" + "RyY9q9VQjH3FTLqPTHpWq6EY+4qZdL/7a2EMuzGH7syk+8ikZ7UaisXy3BgWxlCMfcVMuo9MelaroRj7itdE93FN9KxWQzH2" + "FfPsPvLsWa2GYuwr5tl95NmzWg3F2FfMs/vIs2e1Goqxr5hn95Fnz2o1FGvLc2NYGMNuzKE7M+k+MulZrYZisTw3hoUxFGNv" + "MJPuI5Oe1Wooxt5gJt1HJj2r1VCMvcFMuo9Melarodi1PDeGhTEUY28wk+4jk57VaijG3mAm3UcmPavVUIy9wUy6j0x6Vquh" + "GLc5zKT7yKRntRp2Yw7dmUn3kUnPajUUY18xk+4jk57VaijGvmIm3UcmPavVUIx9xUy6j0x6VquhGPuKmXQfmfSsVkOxa3lu" + "DAtjKMa+4vXUfVxPPavVUIx9xTy7jzx7VquhGPuKeXYfefasVkMx9hXz7J559luths3ioxD7w0Q7qz+/9c5wHNsgkRye4fjQ" + "IIf4yLHfemc4PjXIKT4nh2c4vjTIJb4mh2c4vrdBIjk8w/GjQR7xk2O/9c5w3DRIE7fJ4RmG47MNguTwDMfqk1BfjKz6t94Z" + "jtUnoTUfx+TwDMfqE80T+smZQn/qneFYfaL5HUb157feGY61ijX96PglDs9wrFUcWt7xrrS33hmOuwbp4j45PMOwQuaPrnUe" + "1Z/femc4Vp9Aax4xOTzDMTQIxJgcnuFYzaZpFEf157feGY7PbZBIDs9wrGaD2gfZPrPeGY7VbFD7INtHszA6hmM1G9Q+yPaZ" + "9c5w3LdBIjk8w/Dx2QZBcniGYzXbofY5sn1mvTMcq08OtY9+icMzHKtPDvXFkZuZWe8Mx9c2SCSHZzjWh9ehNT8umv6td4Zj" + "9cmh5X20yeEZjtUnh/ri6JPDMwwrDf6cWt4nJodnONYq1kwTo/rzW+8Mx+c2SCSHZzjWKj61Ss98c896ZzjWKtYkm6P681vv" + "DMfPNkgkh2c41io+tebPNjk8w7AC4c+lt/6VnySz3hmOYxskksMzHGMbBMnhGY6PbZAjOTzD8bkNciaHZzhWn1xqnyv3TDXh" + "qGM4Vp9c6osr90xnvTMcq08urfnrXfP9p14YhhW0fm6tjTvf3LPeGY6PbZBIDs9wrHfxrUV15/ty1jvDsRbsrUV1P5PDMxxr" + "wd5aVPe7qPpPvTAMK+H8KKwc1Z/femc4jm2QSA7PcKy1o+lHxy9xeIbjcxsEyeEZjvXe0aSio/rzW+8Mx882SCSHZzjWKtbU" + "+KP681vvDMO6UPbTtDZaTA7PcKxlognvR/Xnt94ZjrVr0bQM9EscnuFYy6TpVbb3Vb71zjCswOrT1Vz9mBye4Vibgq43d78m" + "h2c4vrdBIjk8w/GzDYLk8AzHWjtda6Pnrtysd4ZjrR3doODz5lWz3hmVQ8FUKIga1Z/femc4PrZBkBye4fjRILd4vC9DWY1j" + "GFYIE7qJwfglDs9wrGWiWCUyVvmpd4ZjbINEcniG42MbBMnhGY4vDXKJr/wD33pnOL63QSI5PMPxsw2C5PAMx02DNHHLP/Ct" + "d4ZhBSmhWCUyVvmpd4bj2AaJ5PAMx+oTxSqRscpPvTMcH9sgkRye4fjcBkFyeIZjNZtilchY5afeGY7VbMpTRvXnt94ZjrsG" + "0ZsbfXJ4hmHlJ6G8ZFR/fuud4Ti2QSI5PMMxtkGQHJ7hWGtHUceo/vzWO8Ox3sWH3rXjmrvfemcY1mVvceoPPD+TwzMca5ko" + "h4jMIX7qneFYzaaD7cgD9596ZzhWs+mIfFR/fuudYVhH2aHjy8hj7p96ZziObZBIDs9wfGyDHMnhGY61YHUwHTp4FodnOFaz" + "6Ug38kj3p94Zjvs2CJLDMwzrwqe49SrvXN6z3hmOYxskksMzHN/bIHdyeIbjZxvkSQ7PcNy2QVpyeIZhHduGjk4jj3R/6p3h" + "WAtWx8Gj+vNb7wzH2p7oIHVUf37rneH42AaJ5PAMx+c2CJLDMxzf2yBncniGY/WJDqZH9ee33hmG21jF6otnnFiI9sl6ZziO" + "bZAjOTzDsT4Bmz7TWn6mzXpnONYy0XHyqP781jvDcdsGieTwDMM9tn99JqMydMAGHXkhD99+6p3hGNsgkRye4fjYBkFyeIbj" + "pkFO8ZmvsmW9Mxz3bZArOTzDsI7ToCMv5JEXdJzmGI7vbZArOTzDsPZcoZ3h768jOTzDsHZYf/hKDs9wjG2QOzk8w/G1DdKT" + "wzMMa1cP/OaE7ow3NKzCaFtHeFLDKqpqX+Rn3Ds5PMOwPrOhvYXvr5fDMxy3bZAnOTzDsD4yf1j3nsUT2yjT4VxXwmFMWPsZ" + "K6xr1VRF0UPXJB3aNnx/HcnhGY6vbZAzOTzDsDYCh27GMi4bZxlWUVVv3qla7Yfa0DEM6333wz05PMOwTscex5E3Sfskh2cY" + "1kmgg9+0+N7vqKWGVVTVbtChOafHlekswyqMXusIkRpWUVWf+vMvOz7J4RmVT3XpqZspnWNH/lSTGkVVRVenpuH8O3xLDauo" + "qrjq1A1wlAFJwyqMclGcui1MJhOn0hujMPqsIxypYRVVdR3JqVuX5M7UqctIjMLos45wpoZVGG3rCFdqWEVVRRjnMW84MzSs" + "wii3FefZ3lsxDA2rqKqD/KlXaliF0XMd4U4Nq6iqE8bnjTGp/Cc1rMLosY4QqWEVRrUk75zgB6lhFVU1p8+pKer7OIY7dehh" + "FFX1OXhqApsj37H6z0ZhVOtiTFUyjlR0k3KnMHprhHtMCPdqWIVRvbM0bUlm7adOjRlF0UsHF5c+Wq88XLh0cOEYjo9tECSH" + "Zzh+NMgzppn8pIZVVNWpoEs3hRtfYWEZVmFUS0O3hhtfCmEZVmH01Ag5AeSrYRVVdVbgGlPeXlPDKozeGuF6vyQ6NKzC6LOO" + "gNSwCqNNI7Rxw41PalhFVe1WTT1TwyqMcoFcmqj1zI7XoYxRGNVyONY+0yeXURjVcjj0l11Twyqq6pPrOuZUiEPDKozGOgJS" + "wyqMal3otlp5KvvSAZ1RGNUmQpOq5vbr0u6tURg91xEiNazCqNbFOSdNGRpWYbSvIyA1rKKqThtcmrg0T/aqrKplpolK81KG" + "S/sLRmFUy2xMOZpbOS0QozB6rSNEaliF0baOcKaGVVTVHsd1z8l8h4ZVGL3WEY7UsAqj2hLoS795zcGlIMAojGo5aNrKdx3r" + "qMcojKr7NHlla1PDKqoqM7+eOVX30LAKo8c6AlLDKozq0+m53imDh4ZVGH3WEZAaVlFVH/iXpmt8tzBKx43CaKwjRGpYhVG9" + "N7VXmSfRL32b1Ciq6huRlyb4y6+KXUqojcIo1hEiNazCqNam9kvzrNulfVijMPqsI0RqWIXRto6A1LAKo30d4UgNqyh664uH" + "U8/UsAqjXJv3mJJvfG7e2ks2CqPnOsKRGlZh9NIIercgx700QlEY7RpBE7COrwfcCveNoqquKLtHljKuR7x1QZlRGD3WEY7U" + "sAqjWpLaM8/Lj27txRuFUS1J3Zjk6lPDKqoqPbo15VleGH0rPTIKo8c6AlLDKoxqOej4YkyQwTKswqiWg27CcfepYRVGb41w" + "j2njXg2rqKqzK7dubtHGXvGtjzSjMHqtI1ypYRVG9Sp044eWnarvJxmF0aYRnneKw6FhFVV1JHFrMqz8yuitIwmjMKp+0O0N" + "3jWkfXCjMHqtI0RqWIVRLUnd6uDJpaPrg4yi6jWWQ38n4x4aRmP9t0gNqzCqZXblNICf1LAKo1pmmkDp/tGwCqNaZroFQH7f" + "6lYwaRRGn3WESA2rMNrXEY7UsIqq2mWYeqaGVRjVctC0P/kNuFvX8RiFUS2HsW//mRpWYVTLQXvm/Z4aVlFV1+DcT078+2pY" + "hdFYR4jUsAqj0Aj4Pb64dSRhFEbPdQSkhlUY1drUJPPtfW1am1Vh9FlHiNSwCqNam5p4J/eKb31NxSiq6jjg1nQ0ketYxwFG" + "YVTbB01Kk197v3W+zCiMaklqaprMEW99xcUoqupY5NaELTnPw61jEaMwinWESA2rMKr9nTF5Sy4dfbPFKIxqXSinzxTwVqZv" + "FEa1LpTAxLtrru+jOIZjrY7xhZhPmxye4fjZBonk8AzHbRsEyeEZjvUe6e/0DJ/k8IzKj7Zdj+62mR9sj7ZHRmEUGgE5TcvL" + "4RmOLw1yjRtbvBpWYfTWCPeYHvDVsAqjTSM87w19hoZVVNWH0TO+jpMX3z86yHIMx1qeY77uz9SwCqOHRjjGzOqvhlUYPTXC" + "Oe4x9GpYhdFrHSFSwyqMaoWOebyR42qFVoXRZx0hUsMqjLZ1BKSGVRjtGqHnhD2f5PAMwzpmfcYU32Mf+tEDozCqltCE2z0b" + "RcesRmFULaFptz8/GlZhVC2hM1p51PDomNUojKoldBya51ofHbMahVGtDt0IMr8LOS5GMIqq+jbMVKSGVRjV2tTk2O8WUMfN" + "RmEU6wiRGlZh9FhHQGpYhVH1g6bLzvMaj+YVMQqj6gdNmp3fBn109G8URu91hEgNqzCqTcSYfPv9y7SJqAqj6ihNwZ0p9KN5" + "SYyiqvKD5/ysH4kKEBzDsVpCM2O/i1inEo3CqFpC82O/n+46lWgURtUS5zvX2svhGY7vbRAkh2c41koZc1Dn54FOKBpFVS2W" + "RzNR56mBR2GGURjV8rxyzsNXwyqMankqtsjTVY/6xiiM6i2m+xvmIeKjiMMojGp1jJmiY2pYhVF9hl7vl79fDs8wrO8bPZrA" + "ueW7V5fhGoVRLUxNiXzlatZd+YzC6LWOEKkwqu2KUo7jmhpWYbStIyA1rKKq8oznyUmvP6lhFUa11MfsGz0mh2c41pbpQc7m" + "83J4hmO9mcYUHjn14qOzG47h+NwGieTwDMdas5qK+DqmhlUYbesISA2rMNrXEY7UsIqqimoeTU78HnoofjEKo1qtmmY4M5On" + "aa1WhVGt03bkF2tfDs9wfG6DRHJ4hmO9s5XeZD78KOkxCqP3OkKkhlUYfdYRkBpWYbStIxypYRVG+zrCmRpWUVWHC48mJ86z" + "VU9XX1WFUa3QfubsNi+HZzi+t0GQHJ7hWEtDiUm8iYnyFceo3Ea+MjmSwzMcc1vadFY5D2+bzmEbhVFoBOS3vl8Oz3B8aJAl" + "3Gs6pjAKo5dGuH73EJq2rUZh9F5HiNSwiqqKSdqYuDbP6jaFNI7hOLZBIjk8w7FWitKb3ONqCh+MwuixjhCpYRVGtUZ0l7Z2" + "TA2rqKrzwQ3rAUnTxCCO4VjLc0zwmh/7TfmFYzjW0lBw8ba+Qg6jMHpqhHPcc+PVsAqjj0a411Wta3wdw3HXINoJzgtymmIK" + "xzB8xDbIkRye4VgtOqbhyNN1TVGDYzg+tkEiOTzD8bkNguTwDMfXNsiRHJ7huGmQlpMRf5LDMwzroP+HkRye4RjbIGdyeIZj" + "LdicgfTKP1AL1jAcP9sgR3J4huO2DXImh2c47tsgV3J4hmEd/bcxNUlObDIuvnUMx7ENEsnhGY713tGdqXI65KZjCaMweq4j" + "IDWswqjeNVfeueWTGlZhVJ/aY26UnO5pnDhxDMdqD914qr8vT91RFVX1La2mWznllKVN39IyCqN6IffYVRy73k2XOjiG42cb" + "5EgOz3CspXHnLdo+qWEVVXVM3XLmjntyeIbjcxskksMzDOtoVLeX7zlRzbgYwCiMas025KQqn+TwDMfqcx3S5AGJbqXuFEbv" + "dYRIDaswqt7QHWnezasOA43CaF9HQGpYRVVFN21MvpGzfOte55bhWN31c7uYoWEVRs91BKSGVRjVe2RMo/nJD3BtDxzDcd8G" + "OZPDMyqPu4D3zzuR9Cc5PMPxoUHwM8mpODzD8bUNciSHZzi+t0HO5PAMx20b5E4OzzCsXfdxvvWTX5cfZ3Adw7GWybj5B8bG" + "XjeetgzHzzbIkRyeYVj7h7oh85yrZpzHdQzH2AaJ5PAMx1omY2bD/IKGbppsGYZ1DNLHrRTygEB3BLYMx9gGieTwDMd6OeMu" + "CHmZwojbHcOwdtp1y9yIvGhPN8C1DMexDYLk8AzH2AY5ksMzHB/bIGdyeIbjcxvkSg7PMKxNnO5VO2fe0x1gLcNxbIOcyeEZ" + "jvVyxhT5+cE+TtY6huNrGySSwzMcaxurs2n5hU/drNUpquqk2dQrNazCqF7I2PfMS3N0V1DLcPxsgxzJ4RmOuwbpeVOWT3J4" + "hmGdu+ljcrfMaMctHR3DsHY1dX/FvH2DODzD8bENEsnhGY7Vqzqp8em50dW5D8dwfG+DIDk8w7BG62PS8zynMe615xiOtT8w" + "bnyXi1s7a0ZhtK8jRGpYxa74KCefeqaGVRg91hGu1LAKo983HfKWdef7p33fdJbh+N4GQXJ4huNHgzy/x3zQbecsw3HXID1n" + "8n05PMMwQ/tfjuTwDMfYBjmSwzMcnxxk3G1Ol/2KwzMca8GOmcHHtS/QreAsw3HbBonk8AzDzMwxbs92RB+vkqG5ZTjWy+G8" + "e+MciDSswmhbR0BqWEXVQ+tl3CvrzMY5tF4Mw/GzDXIkh2c4btsgZ3J4huO+DXIlh2cYPo9tkJ4cnuFYW5Jxc6lxdA7d1Mky" + "HD/bIEgOz3CsBXtqQY7TPr/1znDct0EiOTzD8PXZBkFyeIZjbeev77nzu71DazNfFUbvdQSkhlVUvc/tb/gkh2cYfvRCxpUy" + "CujE4RmG+c1/5F1t/u4Fn8nhGY61hR/B3jitxzo8w/G5DYLk8IzK8dFfol2d77TTLTk8w/GxDdKTwzMc8603bjSSx7LQ/Ugs" + "w7A+On/4Sg7PcHxugzzJ4RmGob9knPttYwOj+zdYhuNzG+RIDs8wzEgB4+4C10hfoNsMWIbjZxsEyeEZhrV1HXP6Q9c/i8Mz" + "DPPw+Zev5PAMw7eWiY4EMS7cZx2eYZiB/C9fyeEZjm8NMm5lMeaehU46eIdzbQ3i6RlvX8nhGYb7eEVHXhqJJz3+4bDODUL0" + "I2cn/ySHZzg+t0EiOTzD8bUNguTwjMqas3oZ+06Pfzicayf2Ox/1oe/SfpLDMwwzUsN3AueD0xFODs8wzA3m3/+Lll10vB7/" + "cDhndPP31zgbrIwZmnfYMgyr6354DK6usw7jh5bMd1rc72vVbNTUsIqq9zYCXg7PMMyrXv/+1+9RB09GfZLDMwzzUrQf1hvr" + "4JWBluG4bYP05PCMyiePbXGOSdS/z33S4x8O56fG4dLqf5fWq2EVRo91BC2qU4cJjuG46eXwVZ650k7t9DqGYfXJeT1cVGPr" + "qWlUnaIq46y/vz7jr7tTwyqqavN4Nv5pbXycaAJSpzB6rSPcqWEVRS9tGK8PJ7L/HnUfyeEZhpUPfOfiPPid9Cc1rKKqjusv" + "jvjNQc/UsIqq2kea4+Ll8AzD6vA5tBpfczVahmF+jeOH1XEXv7FhGYa1Z3Jx+rr+bpku7Zk4huGmZXIxv/h+0eHl8IzKt+KO" + "75z6ByfwvVLDKqrq8PM7KcrB6xcjNazCaFtHOFPDKqryuyG/47bk8AzHxzZITw7PMNy2QfByeEZlfWUZ3+8zH/+Nc3HSsIqq" + "Orr5fon3eO8TJg7PcPxsg9zJ4RmG9aE//z599OkbhZbh+NRfwrM4vDH5nR7/cDhX3vL9jt7xsxOqr+NZhuNnG6Qlh2cY5pd3" + "8PDDnOlFTw7PMKz9h6fH2Lo8qWEVRXXRNRo/P/uYTwu65topqmr9Nu4F625UPT3+4bDe1r9kcniG4aaXc5/retclYN7hXMdH" + "3+vuuLyv0Ye67Mg7rF/7OD09/uEw3pXifmdmOMZ5ldfjHw7nnHQE/eb7grv7r8c/HNWP4LUHxzdauf7jHBxPcniGYV5gdsTf" + "/a8vxw+HZ1QGP78OfJjcMetp6fEPh3O+H/nr+m/MzSINqyiqe6Ycf9fh35F5DeadHJ5hmCd0flh/3cETOpZhmOHJ38MedhGO" + "sddBj384jN/8+sBx64YunyvH122AvKN6+/BaxMZ7sn9nmj30YpuuJfMO4zpj1Zh9fsbbIz3+4XDOEHJxHns0xZD/+A/w/6Fx" + "qGBImuflmm6vaxmG+VXvxq9pffip9ySHZxg+9aL+vtf6N7U9xtinXpJhOD62QZAcnmGYH7L89XcpHGOax6bbYVqG4Ucv53sO" + "IhSM3OnxD4f1Yx/nSY9/OKyzh3ljyr+Otye5V+Ad1tUwD+/oo5hzcHiG474NEsnhGYa7lnDXAtX0Xn9+652x8/+cVuke" +) + + +def _dumped_cell_step(t, tc, nc, max_q_len, g_reqs, next_n, rng): + news = [] + if nc: + rest = tc - max_q_len + news = [max_q_len] + [ + rest // (nc - 1) + (1 if i < rest % (nc - 1) else 0) for i in range(nc - 1) + ] + assert sum(news) + g_reqs * next_n == t and all(n >= 1 for n in news) + ctx = [(rng.randint(0, MAX_REQUEST_LEN - n), n) for n in news] + gen = [rng.randint(0, MAX_REQUEST_LEN - next_n) for _ in range(g_reqs)] + return ctx, gen, next_n + + +def test_dumped_partitions(): + """Every dumped (T, ratio) cell as a consistent step, bitwise equal to the reference, under + torch.profiler (the one Triton kernel, no new compiled variant).""" + cells = [ + tuple(map(int, c.split(","))) + for c in zlib.decompress(base64.b64decode("".join(DUMPED_CELLS))).decode().split() + ] + assert len(cells) == 2696 + rng = random.Random(23300) + tables = {r: _multi_request_tables(r, seed=23400 + r) for r in (0, 1, 2)} + for r in (0, 1, 2): + warm = _multi_request_step(r, [(0, 4)], [], 1, seed=2) + _target_call( + r, warm[0], tables[r][0], warm[1], tables[r][1], warm[2], SWA_TPB // r if r else 0 + ) + torch.cuda.synchronize() + before = _compiled_variants() + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for i, (t, ratio, tc, nc, mq, g_reqs, next_n) in enumerate(cells): + ctx, gen, nn = _dumped_cell_step(t, tc, nc, mq, g_reqs, next_n, rng) + swa, comp = tables[ratio] + req_id, swa_local, comp_local, _, _ = _multi_request_step( + ratio, ctx, gen, nn, seed=60000 + i + ) + assert req_id.numel() == t + tpb_c = SWA_TPB // ratio if ratio else 0 + out = _target_call(ratio, req_id, swa, swa_local, comp, comp_local, tpb_c) + assert torch.equal( + out, _multi_request_reference(ratio, req_id, swa, swa_local, comp, comp_local) + ), cells[i] + torch.cuda.synchronize() + print( + f"{len(cells)} dumped cells exact; ratios {sorted({c[1] for c in cells})}, " + f"T {min(c[0] for c in cells)}..{max(c[0] for c in cells)}" + ) + _check_variants(prof, len(cells), before) + + +# --------------------------------------------------------------------------------------------- +# Pool bound: page ids up to each pool's B_max, for every role the op +# maps. The op takes no pool tensor. A pool's size reaches it only as the magnitude of the page ids +# in the block tables, and every result, page * tpb + slot, is formed in int32 arithmetic. So the +# pool-size axis is exactly the addressing axis, and it is run here at every named bound up to the +# edge of the int32 result range, against an int64 reference (no cast that could hide a wrap). +# Page counts are pages of the role's pool view; the ratio-2 COMPRESS view holds 3 pages per block, +# so its count is 3 x the block count. +# - SWA (128-row pages, every ratio): 188,160 (the engine the dumps came from); 199,920 / 238,560 / +# 360,440 / 468,720 (batch 8 at max_num_tokens 1024 / 2048 / 4096 / 8192, quota +# 148,510,474,240 B); 261,600 (accuracy gate); 491,000 (DSpark included). +# - COMPRESS ratio 1 (128-row pages): 151,040; 175,786 / 231,680 / 268,288 / 300,032 (batch 8 at +# max_num_tokens 8192 / 4096 / 2048 / 1024); 287,573 (the SWA 491,000 configuration); 333,824 +# (batch 32 at 2048). +# - COMPRESS ratio 2 (64-row pages): 453,120; 804,864; 862,719; 1,001,472. +# - The GB300 block ceiling, 2,615,990 blocks (7,847,970 ratio-2 pages), for every role. +# - The int32 edge: 2^24 pages of 128 rows and 2^25 pages of 64 rows, whose last row is 2^31 - 1. +# A pool past it would need 2^31 rows in 298,013,687,808 B of device memory, under 139 B a row; +# a V4.1 KV row holds 512 elements. +# Block tables are the engine's width, max_blocks_per_seq = 132 at max_seq_len 16384 (the SWA +# table a [layer, type] view of a 4-D buffer, as the metadata holds it). +import itertools # noqa: E402 + +POOL_BLOCK_CEILING = 2_615_990 +POOL_INT32_EDGE = {128: 2**24, 64: 2**25} +POOL_BOUNDS = { + "swa": [ + 188_160, + 199_920, + 238_560, + 261_600, + 360_440, + 468_720, + 491_000, + POOL_BLOCK_CEILING, + POOL_INT32_EDGE[128], + ], + 1: [ + 151_040, + 175_786, + 231_680, + 268_288, + 287_573, + 300_032, + 333_824, + POOL_BLOCK_CEILING, + POOL_INT32_EDGE[128], + ], + 2: [ + 453_120, + 804_864, + 862_719, + 1_001_472, + POOL_BLOCK_CEILING, + 3 * POOL_BLOCK_CEILING, + POOL_INT32_EDGE[64], + ], +} +POOL_BOUND_WIDTH = 132 +POOL_STEPS_PER_BOUND = 25 + + +def _reference64(req_id, table, local, tpb): + """One segment's pool rows in int64 (offset 0), never cast: a result at or past 2**31 + stays visible. This is `op.reference` on that segment alone.""" + ref = op.reference(req_id, table, local, tpb, 0) + assert ref.dtype == torch.int64 + return ref + + +def _step_reference64(ratio, req_id, swa, swa_local, comp, comp_local): + ref = _reference64(req_id, swa, swa_local, SWA_TPB) + if ratio: + ref = torch.cat([ref, _reference64(req_id, comp, comp_local, SWA_TPB // ratio)], dim=1) + return ref + + +def _pool_bound_pairs(ratio): + """(SWA bound, COMPRESS bound) per sweep point: every bound of each role the ratio maps.""" + swa = POOL_BOUNDS["swa"] + if ratio == 0: + return [(b, None) for b in swa] + comp = POOL_BOUNDS[ratio] + return list(itertools.zip_longest(swa, comp, fillvalue=None)) + + +def _pool_bound_fill(pairs): + """Pad a zip_longest column with the role's last (largest) bound.""" + out, last = [], None + for b in pairs: + last = b if b is not None else last + out.append(last) + return out + + +def _pool_bound_table(pages, g, layered): + """[ENGINE_ROWS, POOL_BOUND_WIDTH] page ids drawn uniformly from [0, pages). layered: the view + [1, 0] of a [2, 2, ENGINE_ROWS, POOL_BOUND_WIDTH] buffer (the metadata's sliding_block_tables).""" + ids = torch.randint( + 0, pages, (ENGINE_ROWS, POOL_BOUND_WIDTH), generator=g, device=DEV, dtype=torch.int64 + ) + ids = ids.to(torch.int32) + if not layered: + return ids + buf = torch.full((2, 2, ENGINE_ROWS, POOL_BOUND_WIDTH), -1, dtype=torch.int32, device=DEV) + buf[1, 0] = ids + return buf[1, 0] + + +def _pool_bound_far_newest(table, rows, starts, local, tpb, pages, total): + """Each request's newest page (the page of the largest local index its last token names) is + pages - 1 - (b % 4), so the pool's last pages appear in every step.""" + ends = starts[1:] + [total] + for b, end in enumerate(ends): + newest = int(local[end - 1].max()) + if newest >= 0: + table[rows[b], newest // tpb] = pages - 1 - (b % 4) + + +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_page_ids_up_to_every_pool_bound(ratio): + """POOL_STEPS_PER_BOUND seeded random mixed steps (multi-request partitions: k 1..32, T 1..8192, context + then generation rows, next_n 1 or 6) at every (SWA, COMPRESS) bound of the ratio, on width-132 + tables whose page ids are drawn from the whole of [0, B) with each request's newest page among + the pool's last four. Every call equals the int64 reference, so no result wrapped. Pool + translation: step 0's page ids folded into the smallest pool (mod 188160 / the role's smallest + COMPRESS pool) and moved up by delta = B - B_min give outputs that differ by exactly + delta * tpb at every valid entry, with the same -1s. Under torch.profiler: only the one + kernel, and no compiled variant added by the sweep (page magnitude is data). Control: page ids + truncated to their low 18 bits must fail.""" + rng = random.Random(24100 + ratio) + g = torch.Generator(device=DEV).manual_seed(24200 + ratio) + tpb_c = SWA_TPB // ratio if ratio else 0 + pairs = _pool_bound_pairs(ratio) + swa_bounds = _pool_bound_fill([p[0] for p in pairs]) + comp_bounds = _pool_bound_fill([p[1] for p in pairs]) if ratio else [None] * len(pairs) + warm_swa = _pool_bound_table(1024, g, layered=True) + warm_comp = _pool_bound_table(1024, g, layered=False) if ratio else None + warm = _multi_request_step(ratio, [(0, 4)], [], 1, seed=3) + _target_call(ratio, warm[0], warm_swa, warm[1], warm_comp, warm[2], tpb_c) + torch.cuda.synchronize() + before = _compiled_variants() + n_calls, rows_total, kinds, controls, top = 0, 0, set(), [], [] + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for bi, (b_swa, b_comp) in enumerate(zip(swa_bounds, comp_bounds)): + swa = _pool_bound_table(b_swa, g, layered=True) + comp = _pool_bound_table(b_comp, g, layered=False) if ratio else None + max_row = [0, 0] + for s in range(POOL_STEPS_PER_BOUND): + ctx, gen, next_n = _multi_request_random_mixed(rng) + req_id, swa_local, comp_local, starts, rows = _multi_request_step( + ratio, ctx, gen, next_n, seed=24300 + 1000 * ratio + 100 * bi + s + ) + T = req_id.numel() + _pool_bound_far_newest(swa, rows, starts, swa_local, SWA_TPB, b_swa, T) + if ratio: + _pool_bound_far_newest(comp, rows, starts, comp_local, tpb_c, b_comp, T) + out = _target_call(ratio, req_id, swa, swa_local, comp, comp_local, tpb_c) + ref = _step_reference64(ratio, req_id, swa, swa_local, comp, comp_local) + n_calls += 1 + assert int(ref.max()) < 2**31 + assert torch.equal(out.long(), ref), (ratio, b_swa, b_comp, s) + max_row[0] = max(max_row[0], int(ref[:, :WINDOW].max())) + if ratio: + max_row[1] = max(max_row[1], int(ref[:, WINDOW:].max())) + rows_total += T + kinds.add((bool(ctx), bool(gen), next_n if gen else 0)) + if s == 0: + # Pool translation: the step's page ids folded into the smallest pool + # (page mod B_min) and the same ids moved to the top of this bound's pool + # (+ delta, the largest landing on B - 1): both equal the reference, and they + # differ by exactly delta * tpb at every valid entry, with the same -1s. + d_swa = b_swa - ENGINE_SWA_PAGES + lo_swa = torch.remainder(swa, ENGINE_SWA_PAGES) + lo_comp, hi_comp, d_comp = None, None, 0 + if ratio: + base = POOL_BOUNDS[ratio][0] + d_comp = b_comp - base + lo_comp = torch.remainder(comp, base) + hi_comp = lo_comp + d_comp + hi_swa = lo_swa + d_swa + lo = _target_call(ratio, req_id, lo_swa, swa_local, lo_comp, comp_local, tpb_c) + hi = _target_call(ratio, req_id, hi_swa, swa_local, hi_comp, comp_local, tpb_c) + n_calls += 2 + ref_hi = _step_reference64( + ratio, req_id, hi_swa, swa_local, hi_comp, comp_local + ) + assert torch.equal( + lo.long(), + _step_reference64(ratio, req_id, lo_swa, swa_local, lo_comp, comp_local), + ) + assert torch.equal(hi.long(), ref_hi) + shift = torch.full_like(ref_hi, d_swa * SWA_TPB) + if ratio: + shift[:, WINDOW:] = d_comp * tpb_c + assert torch.equal(torch.where(lo >= 0, lo.long() + shift, -1), hi.long()), ( + b_swa, + b_comp, + ) + # Control: page ids truncated to their low 18 bits (a narrow page-id path). + if b_swa > 2**18 or (ratio and b_comp > 2**18): + cut_swa = swa & (2**18 - 1) + cut_comp = comp & (2**18 - 1) if ratio else None + wrong = _target_call( + ratio, req_id, cut_swa, swa_local, cut_comp, comp_local, tpb_c + ) + n_calls += 1 + controls.append(int((wrong.long() != ref).sum())) + top.append((b_swa, b_comp, max_row[0], max_row[1] if ratio else None)) + torch.cuda.synchronize() + for b_swa, b_comp, m_swa, m_comp in top: + print( + f"ratio={ratio} SWA B {b_swa}: largest row {m_swa} (last {b_swa * SWA_TPB - 1})" + + ( + f"; COMPRESS B {b_comp}: largest row {m_comp} (last {b_comp * tpb_c - 1})" + if ratio + else "" + ) + ) + print( + f"ratio={ratio}: {n_calls} calls, {rows_total} rows, kinds {sorted(kinds)}; " + f"control (page ids cut to 18 bits) elements differing {controls}" + ) + assert len(controls) >= 4 and all(c > 0 for c in controls), controls + assert { + (True, True, 1), + (True, True, 6), + (False, True, 1), + (False, True, 6), + (True, False, 0), + } <= kinds + _check_variants(prof, n_calls, before) + + +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_far_end_rows_of_every_bound(ratio): + """At every bound of each role the ratio maps: one token whose table row names the pool's last + pages B-1, B-2, ... and whose local indices enumerate every slot of them (window: positions + 0..127 on page B-1; compressed: 0..511 over pages B-1 .. B-512/tpb). Exact vs int64, and the + pool's last row B * tpb - 1 is produced. At the int32 edge that row is 2**31 - 1.""" + tpb_c = SWA_TPB // ratio if ratio else 0 + req_id = torch.zeros(1, dtype=torch.int32, device=DEV) + swa_local = torch.arange(WINDOW, dtype=torch.int32, device=DEV)[None, :] + comp_local = torch.arange(TOPK, dtype=torch.int32, device=DEV)[None, :] if ratio else None + cols = torch.arange(POOL_BOUND_WIDTH, device=DEV) + pairs = _pool_bound_pairs(ratio) + for b_swa, b_comp in pairs: + for role, B in (("swa", b_swa), ("comp", b_comp)): + if B is None: + continue + swa_B = B if role == "swa" else ENGINE_SWA_PAGES + comp_B = B if role == "comp" else (POOL_BOUNDS[ratio][0] if ratio else None) + swa = (swa_B - 1 - cols).to(torch.int32)[None, :] + comp = (comp_B - 1 - cols).to(torch.int32)[None, :] if ratio else None + out = _target_call(ratio, req_id, swa, swa_local, comp, comp_local, tpb_c) + ref = _step_reference64(ratio, req_id, swa, swa_local, comp, comp_local) + assert torch.equal(out.long(), ref), (ratio, role, B) + tpb = SWA_TPB if role == "swa" else tpb_c + seg = out[0, :WINDOW] if role == "swa" else out[0, WINDOW:] + assert int(seg.max()) == B * tpb - 1, (role, B, int(seg.max())) + swa = torch.tensor([[POOL_INT32_EDGE[SWA_TPB] - 1]], dtype=torch.int32, device=DEV) + comp = ( + torch.tensor([[POOL_INT32_EDGE[tpb_c] - 1] + [0] * 7], dtype=torch.int32, device=DEV) + if ratio + else None + ) + edge_local = torch.where(comp_local < tpb_c, comp_local, -1).to(torch.int32) if ratio else None + out = _target_call(ratio, req_id, swa, swa_local, comp, edge_local, tpb_c) + assert torch.equal( + out.long(), _step_reference64(ratio, req_id, swa, swa_local, comp, edge_local) + ) + assert int(out[0, WINDOW - 1]) == 2**31 - 1 + if ratio: + assert int(out[0, WINDOW + tpb_c - 1]) == 2**31 - 1 + + +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_control_one_page_past_the_int32_edge_wraps(ratio): + """Wrong control: page ids one past the int32 edge (2**24 at 128 rows, 2**25 at 64 rows) give + results of 2**31 and up, which the op does not check (Precondition 5): they come out wrapped + negative, and the same int64 comparison the bound tests use rejects them.""" + tpb_c = SWA_TPB // ratio if ratio else 0 + req_id = torch.zeros(1, dtype=torch.int32, device=DEV) + swa_local = torch.arange(WINDOW, dtype=torch.int32, device=DEV)[None, :] + swa = torch.tensor([[POOL_INT32_EDGE[SWA_TPB]]], dtype=torch.int32, device=DEV) + comp = ( + torch.tensor([[POOL_INT32_EDGE[tpb_c]]], dtype=torch.int32, device=DEV) if ratio else None + ) + comp_local = None + if ratio: + comp_local = torch.full((1, TOPK), -1, dtype=torch.int32, device=DEV) + comp_local[0, :tpb_c] = torch.arange(tpb_c, dtype=torch.int32, device=DEV) + out = _target_call(ratio, req_id, swa, swa_local, comp, comp_local, tpb_c) + ref = _step_reference64(ratio, req_id, swa, swa_local, comp, comp_local) + assert int(ref.max()) >= 2**31 + assert not torch.equal(out.long(), ref) + assert int(out[0, 0]) == -(2**31) # 2**31 wrapped + print(f"ratio={ratio}: one page past the edge -> {int(out[0, 0])} .. {int(out[0, WINDOW - 1])}") + + +@pytest.mark.parametrize("rows,next_n", [(32, 1), (192, 6)], ids=["T32_n1", "T192_n6"]) +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_captured_call_over_every_pool_bound(ratio, rows, next_n): + """The largest graph batches (decode 32; DSpark verify 32 x 6), captured once on width-132 + tables as the target captures the call, then replayed with tables drawn from [0, B) at every + bound of each role (newest pages at the far end): each replay equals the int64 reference and an + eager call on the same inputs.""" + g = torch.Generator(device=DEV).manual_seed(24500 + 10 * ratio + next_n) + gh = torch.Generator(device="cpu").manual_seed(24600 + 10 * ratio + next_n) + tpb_c = SWA_TPB // ratio if ratio else 0 + max_tokens = 8192 + req_buf = torch.zeros(max_tokens, dtype=torch.int32, device=DEV) + swa_buf = torch.full((max_tokens, WINDOW), -1, dtype=torch.int32, device=DEV) + comp_buf = torch.full((rows, TOPK), -1, dtype=torch.int32, device=DEV) if ratio else None + swa_table = _pool_bound_table(1024, g, layered=True) + comp_table = _pool_bound_table(1024, g, layered=False) if ratio else None + + def load(seqs, b_swa, b_comp, seed): + req_id, swa_local, comp_local, starts, rrows = _multi_request_step( + ratio, [], [c for c, _ in seqs], next_n, seed + ) + swa_table.copy_(_pool_bound_table(b_swa, g, layered=True)) + _pool_bound_far_newest(swa_table, rrows, starts, swa_local, SWA_TPB, b_swa, rows) + req_buf[:rows].copy_(req_id) + swa_buf[:rows].copy_(swa_local) + if ratio: + comp_table.copy_(_pool_bound_table(b_comp, g, layered=False)) + _pool_bound_far_newest(comp_table, rrows, starts, comp_local, tpb_c, b_comp, rows) + comp_buf.copy_(comp_local) + return _step_reference64(ratio, req_id, swa_table, swa_local, comp_table, comp_local) + + def call(): + return _target_call( + ratio, req_buf[:rows], swa_table, swa_buf[:rows], comp_table, comp_buf, tpb_c + ) + + pairs = _pool_bound_pairs(ratio) + swa_bounds = _pool_bound_fill([p[0] for p in pairs]) + comp_bounds = _pool_bound_fill([p[1] for p in pairs]) if ratio else [None] * len(pairs) + ref = load(_decode_seqs(rows, next_n, gh, 0), 1024, 1024 if ratio else None, seed=24700) + assert torch.equal(call().long(), ref) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = call() + for i, (b_swa, b_comp) in enumerate(zip(swa_bounds, comp_bounds)): + ref = load(_decode_seqs(rows, next_n, gh, i + 1), b_swa, b_comp, seed=24800 + i) + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.long(), ref), (b_swa, b_comp) + assert torch.equal(out, call()), (b_swa, b_comp) + + +# --------------------------------------------------------------------------------------------- +# The DSpark envelope. With DSpark (max_draft_len 5) and the overlap scheduler +# on (every gate), the executor widens the engine's max_seq_len to 16384 + 5 + 4 + 5 = 16398 +# (py_executor_creator: tokens_per_gen_step - 1 for the overlap scheduler, the one-engine extra +# KV tokens max_draft_len - 1, and tokens_per_gen_step - 1). A request's committed tokens stay +# <= 16384 (the executor keeps the net max_seq_len, and a prompt must leave max_tokens > 0), so +# context rows stay below position 16384. A verify step at the end of a full-length request writes +# generation positions up to 16397: SWA logical page 128 (table column 128), and compressed +# positions up to (16397 + 1) // r - 1 = 16397 (ratio 1) / 8198 (ratio 2), also column 128. The +# manager's tables stay max_blocks_per_seq = 132 columns wide (16398 + 10 -> 129 -> 132). +DSPARK_MAX_POS = 16397 +DSPARK_BOUNDS = { + # engine of the dumps, native-storage cross-product maximum, the GB300 block + # ceiling, the int32 edge + "swa": [188_160, 854_440, POOL_BLOCK_CEILING, POOL_INT32_EDGE[128]], + 1: [151_040, 868_673, POOL_BLOCK_CEILING, POOL_INT32_EDGE[128]], + 2: [453_120, 2_606_019, 3 * POOL_BLOCK_CEILING, POOL_INT32_EDGE[64]], +} +# Contexts ahead of the generation rows: a fresh prompt, a chunk on an 8192-token prefix, 1- and +# 5-token contexts, and chunks ending at positions 16382 and 16383 (the last context position). +DSPARK_CTX = [(0, 3000), (8192, 2048), (0, 1), (130, 5), (15405, 978), (16000, 384)] +# Generation requests at the DSpark tail, as `cached` (first position of the step): +# next_n 1: decode rows at every position 16383 .. 16397; +# next_n 6: verify runs ending at 16397, 16396, 16393, 16392, 16389, 16388, 16384 and 16383. +DSPARK_TAIL = { + 1: list(range(16383, DSPARK_MAX_POS + 1)), + 6: [e - 5 for e in (16397, 16396, 16393, 16392, 16389, 16388, 16384, 16383)], +} + + +def _dspark_bound_pairs(ratio): + comp = DSPARK_BOUNDS[ratio] if ratio else [None] * len(DSPARK_BOUNDS["swa"]) + return list(zip(DSPARK_BOUNDS["swa"], comp)) + + +def _dspark_step(ratio, ctx, gen, next_n, seed): + """Like _multi_request_step, with generation positions up to DSPARK_MAX_POS (contexts below 16384), and the + newest visible compressed position always among a row's top-k picks (so a tail row names + compressed table column 128).""" + g = torch.Generator(device=DEV).manual_seed(seed) + seqs = list(ctx) + [(c, next_n) for c in gen] + assert len(seqs) <= 32 + assert all(c + n <= MAX_REQUEST_LEN for c, n in ctx) + rows = torch.randperm(ENGINE_ROWS, generator=g, device=DEV)[: len(seqs)].to(torch.int32) + req_id = torch.cat([rows[b].repeat(new) for b, (_, new) in enumerate(seqs)]).contiguous() + positions = torch.cat( + [torch.arange(c, c + new, device=DEV, dtype=torch.int32) for c, new in seqs] + ) + assert req_id.numel() <= 8192 and int(positions.max()) <= DSPARK_MAX_POS + swa_local = _window(positions) + comp_local = _topk_local_with_newest(positions, ratio, g) if ratio else None + starts = [0] + for _, new in seqs[:-1]: + starts.append(starts[-1] + new) + return req_id, swa_local, comp_local, starts, rows + + +def _dspark_col128(local, tpb): + """Valid local indices this call resolves through table column 128.""" + return int(((local >= 0) & (torch.div(local, tpb, rounding_mode="floor") == 128)).sum()) + + +@pytest.mark.parametrize("next_n", [1, 6]) +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_dspark_tail_mixed_steps(ratio, next_n): + """Mixed steps at the DSpark tail: 6 contexts (positions < 16384) ahead of 26 generation + requests, the tail cases of DSPARK_TAIL (decode rows at every position 16383..16397, or verify + runs ending at 16383..16397, both parities) plus random others up to position 16397, on + 132-column engine tables at every DSpark pool bound (each request's newest page at the pool's far + end). Every call equals the int64 reference, and table column 128 is read in the SWA and the + compressed segment. Control: the same step on the tables cut to 128 columns (a pre-DSpark + width) differs from the reference exactly at the column-128 entries, which come back -1.""" + g = torch.Generator(device=DEV).manual_seed(25100 + 10 * ratio + next_n) + rng = random.Random(25200 + 10 * ratio + next_n) + tpb_c = SWA_TPB // ratio if ratio else 0 + tail = DSPARK_TAIL[next_n] + gen = list(tail) + [ + rng.randint(0, DSPARK_MAX_POS + 1 - next_n) for _ in range(32 - len(DSPARK_CTX) - len(tail)) + ] + rng.shuffle(gen) + req_id, swa_local, comp_local, starts, rows = _dspark_step( + ratio, DSPARK_CTX, gen, next_n, seed=25300 + 10 * ratio + next_n + ) + T = req_id.numel() + col128 = [_dspark_col128(swa_local, SWA_TPB)] + ( + [_dspark_col128(comp_local, tpb_c)] if ratio else [] + ) + assert all(c > 0 for c in col128), col128 + for b_swa, b_comp in _dspark_bound_pairs(ratio): + swa = _pool_bound_table(b_swa, g, layered=True) + comp = _pool_bound_table(b_comp, g, layered=False) if ratio else None + _pool_bound_far_newest(swa, rows, starts, swa_local, SWA_TPB, b_swa, T) + if ratio: + _pool_bound_far_newest(comp, rows, starts, comp_local, tpb_c, b_comp, T) + out = _target_call(ratio, req_id, swa, swa_local, comp, comp_local, tpb_c) + ref = _step_reference64(ratio, req_id, swa, swa_local, comp, comp_local) + assert int(ref.max()) < 2**31 + assert torch.equal(out.long(), ref), (ratio, next_n, b_swa, b_comp) + # Control: tables one page short of column 128. + cut = _target_call( + ratio, + req_id, + swa[:, :128], + swa_local, + None if comp is None else comp[:, :128], + comp_local, + tpb_c, + ) + diff = cut.long() != ref + assert int(diff.sum()) == sum(col128), (int(diff.sum()), col128) + assert bool((cut[diff] == -1).all()) + print( + f"ratio={ratio} next_n={next_n}: T {T}, {len(gen)} generation requests, column-128 entries " + f"(SWA, compressed) {col128} at bounds {_dspark_bound_pairs(ratio)}" + ) + + +def _dspark_random_mixed(rng): + """_multi_request_random_mixed with generation requests allowed up to position 16397: each generation + request's first position is drawn from the DSpark tail [16384 - next_n, 16398 - next_n] with + probability 1/4 and from [0, 16398 - next_n] otherwise. Contexts stay below 16384.""" + k = rng.randint(1, 32) + next_n = rng.choice((1, 6)) + n_gen = rng.randint(0, k) + n_ctx = k - n_gen + gen_rows = n_gen * next_n + ctx = [] + if n_ctx: + ctx_rows = rng.randint(gen_rows + n_ctx, 8192) - gen_rows + cuts = sorted(rng.sample(range(1, ctx_rows), n_ctx - 1)) + news = [b - a for a, b in zip([0] + cuts, cuts + [ctx_rows])] + ctx = [(rng.randint(0, MAX_REQUEST_LEN - n), n) for n in news] + top = DSPARK_MAX_POS + 1 - next_n + gen = [ + rng.randint(MAX_REQUEST_LEN - next_n, top) if rng.random() < 0.25 else rng.randint(0, top) + for _ in range(n_gen) + ] + return ctx, gen, next_n + + +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_random_mixed_partitions_with_dspark_tail(ratio): + """200 seeded multi-request mixed steps (k 1..32, T 1..8192, contexts then generation rows, next_n 1 or + 6) with generation positions up to 16397, on 132-column tables drawn from [0, B) at the DSpark pool + bounds in turn (newest pages at the far end): every call equals the int64 reference. Under + torch.profiler only the one kernel launches, and the sweep adds no compiled variant after one + warm-up call in the 132-column form.""" + rng = random.Random(25400 + ratio) + g = torch.Generator(device=DEV).manual_seed(25500 + ratio) + tpb_c = SWA_TPB // ratio if ratio else 0 + pairs = _dspark_bound_pairs(ratio) + tables = [ + ( + b_swa, + b_comp, + _pool_bound_table(b_swa, g, layered=True), + _pool_bound_table(b_comp, g, layered=False) if ratio else None, + ) + for b_swa, b_comp in pairs + ] + warm = _dspark_step(ratio, [(0, 4)], [], 1, seed=5) + _target_call(ratio, warm[0], tables[0][2], warm[1], tables[0][3], warm[2], tpb_c) + torch.cuda.synchronize() + before = _compiled_variants() + n_calls, rows_total, tail_rows, col128, kinds = 0, 0, 0, [0, 0], set() + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for i in range(200): + ctx, gen, next_n = _dspark_random_mixed(rng) + b_swa, b_comp, swa, comp = tables[i % len(tables)] + req_id, swa_local, comp_local, starts, rows = _dspark_step( + ratio, ctx, gen, next_n, seed=25600 + 1000 * ratio + i + ) + T = req_id.numel() + _pool_bound_far_newest(swa, rows, starts, swa_local, SWA_TPB, b_swa, T) + if ratio: + _pool_bound_far_newest(comp, rows, starts, comp_local, tpb_c, b_comp, T) + out = _target_call(ratio, req_id, swa, swa_local, comp, comp_local, tpb_c) + ref = _step_reference64(ratio, req_id, swa, swa_local, comp, comp_local) + n_calls += 1 + assert int(ref.max()) < 2**31 + assert torch.equal(out.long(), ref), (ratio, i, ctx, gen, next_n) + rows_total += T + tail_rows += sum(next_n for c in gen if c + next_n - 1 >= MAX_REQUEST_LEN) + col128[0] += _dspark_col128(swa_local, SWA_TPB) + if ratio: + col128[1] += _dspark_col128(comp_local, tpb_c) + kinds.add((bool(ctx), bool(gen), next_n if gen else 0)) + torch.cuda.synchronize() + print( + f"ratio={ratio}: {n_calls} calls, {rows_total} rows, {tail_rows} rows of requests reaching " + f"positions >= 16384, column-128 entries (SWA, compressed) {col128}, kinds {sorted(kinds)}" + ) + assert tail_rows > 0 and col128[0] > 0 and (col128[1] > 0 or not ratio) + assert { + (True, True, 1), + (True, True, 6), + (False, True, 1), + (False, True, 6), + (True, False, 0), + } <= kinds + _check_variants(prof, n_calls, before) + + +def _dspark_decode_seqs(rows, next_n, rng, step): + """rows // next_n generation requests up to position 16397; request 0 ends at 16397 and + request 1 at 16392 (even steps) or 16396 / 16383 (odd steps), the rest random.""" + batch = rows // next_n + top = DSPARK_MAX_POS + 1 - next_n + cached = [rng.randint(0, top) for _ in range(batch)] + ends = (DSPARK_MAX_POS, 16392) if step % 2 == 0 else (16396, 16383) + cached[0] = ends[0] - next_n + 1 + if batch > 1: + cached[1] = ends[1] - next_n + 1 + return cached + + +@pytest.mark.parametrize("rows,next_n", [(32, 1), (192, 6)], ids=["T32_n1", "T192_n6"]) +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_dspark_tail_captured(ratio, rows, next_n): + """The largest graph batches (decode 32; DSpark verify 32 x 6) captured once on 132-column + tables, as the target captures the call (prefix views of max_num_tokens-row buffers, static + tables and top-k), then replayed over DSpark-tail steps (request 0 ending at 16397) with tables + drawn at every DSpark pool bound: each replay equals the int64 reference and an eager call on the same + inputs, differs from the previous step's reference, and reads table column 128.""" + g = torch.Generator(device=DEV).manual_seed(25700 + 10 * ratio + next_n) + rng = random.Random(25800 + 10 * ratio + next_n) + tpb_c = SWA_TPB // ratio if ratio else 0 + req_buf = torch.zeros(8192, dtype=torch.int32, device=DEV) + swa_buf = torch.full((8192, WINDOW), -1, dtype=torch.int32, device=DEV) + comp_buf = torch.full((rows, TOPK), -1, dtype=torch.int32, device=DEV) if ratio else None + swa_table = _pool_bound_table(1024, g, layered=True) + comp_table = _pool_bound_table(1024, g, layered=False) if ratio else None + + def load(step_i, b_swa, b_comp): + gen = _dspark_decode_seqs(rows, next_n, rng, step_i) + req_id, swa_local, comp_local, starts, rrows = _dspark_step( + ratio, [], gen, next_n, seed=25900 + 10 * ratio + step_i + ) + assert _dspark_col128(swa_local, SWA_TPB) > 0 and ( + not ratio or _dspark_col128(comp_local, tpb_c) > 0 + ) + swa_table.copy_(_pool_bound_table(b_swa, g, layered=True)) + _pool_bound_far_newest(swa_table, rrows, starts, swa_local, SWA_TPB, b_swa, rows) + req_buf[:rows].copy_(req_id) + swa_buf[:rows].copy_(swa_local) + if ratio: + comp_table.copy_(_pool_bound_table(b_comp, g, layered=False)) + _pool_bound_far_newest(comp_table, rrows, starts, comp_local, tpb_c, b_comp, rows) + comp_buf.copy_(comp_local) + return _step_reference64(ratio, req_id, swa_table, swa_local, comp_table, comp_local) + + def call(): + return _target_call( + ratio, req_buf[:rows], swa_table, swa_buf[:rows], comp_table, comp_buf, tpb_c + ) + + pairs = _dspark_bound_pairs(ratio) + prev = load(0, *pairs[0]) + assert torch.equal(call().long(), prev) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = call() + for i, (b_swa, b_comp) in enumerate(pairs + pairs[:1]): + ref = load(i + 1, b_swa, b_comp) + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.long(), ref), (i, b_swa, b_comp) + assert torch.equal(out, call()), (i, b_swa, b_comp) + assert not torch.equal(out.long(), prev), i + + +# --------------------------------------------------------------------------------------------- +# The native-storage compact read. Under native KV storage the +# target dequantizes, every step, every compressed entry of every request of the step into ONE +# persistent bf16 buffer: request b's entry e (every e < kv_len // r) at row b * L_r + e, where +# L_r = 132 * 128 / r (16,896 at ratio 1, 8,448 at ratio 2) and b is the request's position +# 0..B-1 in the step (the row of md.compress_block_tables[r] the dequantizer reads for it). The +# attention's compressed columns are then built by this op with a SYNTHETIC page table in place of +# md.compress_block_tables[r]: +# comp_tbl = torch.arange(max_batch_size * 132, dtype=torch.int32).view(max_batch_size, 132) +# allocated once at load and shared by BOTH ratios (one tensor), in the call +# op(req_idx_per_token[:T], swa_table, swa_local[:T], 128, 0, +# comp_tbl, topk, 128 // r, 0) +# so entry e of request b maps to (b * 132 + e // tpb) * tpb + e % tpb = b * L_r + e. The reference +# below forms b * L_r + e directly, never through a table, so what is tested is that the table +# implies the buffer's numbering. The window half is the ordinary call on the real SWA table: the +# [layer, SWA] view [max_batch_size, 132] of a 4-D buffer, as the metadata holds it, with page ids +# up to the native SWA pool's largest size over the max_batch_size axis. +# +# max_batch_size is an axis here: the table is [max_batch_size, 132] for EVERY +# max_batch_size 1..32 (the metadata's compress_block_tables are [max_batch_size, 132] at every +# max_batch_size 1..32 x max_num_tokens 2048 / 8192 x DSpark off / on, max_seq_len 16384 / 16398, +# and the CUDA-graph metadata copy keeps that shape; measured in the dequantizer entries' native-storage +# tests). So the compact buffer the table numbers holds max_batch_size * L_r rows, and the largest +# row a step can produce is (max_batch_size - 1) * L_r + 16398 // r - 1. The witness +# configuration (max_num_tokens 2048, max_batch_size 1) runs as its own configuration. +KVQ_COLS = 132 +KVQ_L = {1: KVQ_COLS * 128, 2: KVQ_COLS * 64} +KVQ_MBS = tuple(range(1, 33)) # every max_batch_size the engine configuration permits +KVQ_LISTED = (1, 2, 3, 4, 8, 16, 32) +KVQ_CFGS = [(8192, s) for s in KVQ_MBS] + [ + (2048, 1) +] # (max_num_tokens, max_batch_size); (2048, 1): witness +KVQ_CFG_IDS = [f"T{t}_mbs{s}" for t, s in KVQ_CFGS] +# Native SWA pages over the max_batch_size axis at the gate memory quota: 941,800 at (max_num_tokens +# 8192, max_batch_size 1, no DSpark), above the 854,440 pages over max_batch_size 8 / 32 only. Page +# ids up to 2**24 are certified above (the pool bound); the native-storage cases draw from this. +KVQ_SWA_PAGES = 941_800 + + +def _kvq_table(seqs): + """The target's synthetic compressed page table for max_batch_size = seqs.""" + return torch.arange(seqs * KVQ_COLS, dtype=torch.int32, device=DEV).view(seqs, KVQ_COLS) + + +def _kvq_swa(seqs, g): + """The real SWA table at max_batch_size = seqs: [seqs, 132] page ids drawn uniformly from the native + SWA pool (KVQ_SWA_PAGES), as the [layer, SWA] view [1, 0] of a [2, 2, seqs, 132] buffer (the + metadata's sliding_block_tables at max_num_sequences = max_batch_size).""" + buf = torch.full((2, 2, seqs, KVQ_COLS), -1, dtype=torch.int32, device=DEV) + buf[1, 0] = torch.randint( + 0, KVQ_SWA_PAGES, (seqs, KVQ_COLS), generator=g, device=DEV, dtype=torch.int64 + ).to(torch.int32) + return buf[1, 0] + + +def _kvq_reference(ratio, req_id, swa, swa_local, comp_local): + """int64: the window half through the real SWA table; the compressed half b * L_r + e for a valid + local e (e >= 0 and within the table's 132 columns), -1 otherwise -- formed without a table.""" + tpb = SWA_TPB // ratio + win = _reference64(req_id, swa, swa_local, SWA_TPB) + e = comp_local.long() + valid = (e >= 0) & (torch.div(e, tpb, rounding_mode="floor") < KVQ_COLS) + comp = torch.where(valid, req_id.long()[:, None] * KVQ_L[ratio] + e, -1) + return torch.cat([win, comp], 1) + + +def _kvq_step(ratio, ctx, gen, next_n, seed, max_t=8192): + """Per-token inputs of a step as the metadata builds them: the requests at step positions + 0..B-1 in order (context requests first, then generation requests each verifying next_n + positions), req_id = the token's request position, the causal window, and the layer's top-k + (valid compressed positions < (p+1)//ratio first, the newest always among them, then -1). + Generation positions reach DSPARK_MAX_POS (DSpark); contexts stay below 16384.""" + g = torch.Generator(device=DEV).manual_seed(seed) + seqs = list(ctx) + [(c, next_n) for c in gen] + assert all(c + n <= MAX_REQUEST_LEN for c, n in ctx) + req_id = torch.cat( + [torch.full((n,), b, dtype=torch.int32, device=DEV) for b, (_, n) in enumerate(seqs)] + ) + positions = torch.cat([torch.arange(c, c + n, device=DEV, dtype=torch.int32) for c, n in seqs]) + assert req_id.numel() <= max_t and int(positions.max()) <= DSPARK_MAX_POS + return ( + req_id.contiguous(), + _window(positions), + _topk_local_with_newest(positions, ratio, g), + len(seqs), + ) + + +def _kvq_call(ratio, req_id, swa, swa_local, tbl, comp_local, tpb=None): + return op( + req_id, + swa, + swa_local, + SWA_TPB, + 0, + tbl, + comp_local, + SWA_TPB // ratio if tpb is None else tpb, + 0, + ) + + +@pytest.mark.parametrize("seqs", KVQ_MBS) +@pytest.mark.parametrize("ratio", [1, 2]) +def test_kvq_synthetic_table_maps_every_row_of_the_buffer(ratio, seqs): + """At max_batch_size = seqs (every value 1..32): every request position b in 0..seqs-1 with every + compressed position e in 0..L_r-1 -- every row of the ratio's view of the compressed read buffer -- + plus positions past the table's 132 columns, through ONE synthetic table used at both ratios, on + the real [seqs, 132] SWA table: the op gives b * L_r + e, each of the seqs * L_r buffer rows + exactly once, and -1 at and past column 132.""" + g = torch.Generator(device=DEV).manual_seed(26000 + 100 * ratio + seqs) + tbl = _kvq_table(seqs) + L = KVQ_L[ratio] + per_req = -(-(L + TOPK) // TOPK) # 512-wide rows of positions 0 .. >= L + 511 + req_id = ( + torch.arange(seqs, dtype=torch.int32, device=DEV).repeat_interleave(per_req).contiguous() + ) + j = torch.arange(per_req, device=DEV).repeat(seqs) + comp_local = (j[:, None] * TOPK + torch.arange(TOPK, device=DEV)[None]).to(torch.int32) + positions = torch.randint( + 0, DSPARK_MAX_POS + 1, (req_id.numel(),), generator=g, device=DEV, dtype=torch.int32 + ) + swa_local = _window(positions) + swa = _kvq_swa(seqs, g) + out = _kvq_call(ratio, req_id, swa, swa_local, tbl, comp_local) + ref = _kvq_reference(ratio, req_id, swa, swa_local, comp_local) + assert torch.equal(out.long(), ref) + comp = out[:, WINDOW:] + rows = comp[comp >= 0].long().sort().values + assert torch.equal(rows, torch.arange(seqs * L, device=DEV)), "every buffer row exactly once" + assert int((comp < 0).sum()) == seqs * (per_req * TOPK - L) + # The other ratio's call on the same table object gives that ratio's numbering. + other = 3 - ratio + out2 = _kvq_call(other, req_id, swa, swa_local, tbl, comp_local) + assert torch.equal(out2.long(), _kvq_reference(other, req_id, swa, swa_local, comp_local)) + print( + f"ratio={ratio} max_batch_size={seqs}: rows 0..{seqs * L - 1} each produced once over {req_id.numel()} " + f"tokens; {int((comp < 0).sum())} entries past column 132 -> -1" + ) + + +def _kvq_random_mixed(rng, seqs, max_t=8192): + """_dspark_random_mixed with at most `seqs` requests (max_batch_size) and T <= max_t.""" + k = rng.randint(1, seqs) + next_n = rng.choice((1, 6)) + n_gen = rng.randint(0, k) + n_ctx = k - n_gen + gen_rows = n_gen * next_n + ctx = [] + if n_ctx: + ctx_rows = rng.randint(gen_rows + n_ctx, max_t) - gen_rows + cuts = sorted(rng.sample(range(1, ctx_rows), n_ctx - 1)) + news = [b - a for a, b in zip([0] + cuts, cuts + [ctx_rows])] + ctx = [(rng.randint(0, MAX_REQUEST_LEN - n), n) for n in news] + top = DSPARK_MAX_POS + 1 - next_n + gen = [ + rng.randint(MAX_REQUEST_LEN - next_n, top) if rng.random() < 0.25 else rng.randint(0, top) + for _ in range(n_gen) + ] + return ctx, gen, next_n + + +def _kvq_far_step(seqs): + """seqs requests (the full batch) whose last is one DSpark verify ending at position 16397, so the + largest row the target can produce, (seqs - 1) * L_r + 16398 // r - 1, is produced: at seqs 1 the + verify alone; at 2 a 4000-token context ahead of it; from 3 on also seqs - 2 short contexts.""" + ctx = [] if seqs == 1 else [(16384 - 4000, 4000)] + [(97 * i, 3 + i) for i in range(seqs - 2)] + return ctx, [DSPARK_MAX_POS - 5], 6 + + +def _kvq_num_steps(cfg): + """200 random steps at every listed max_batch_size and at the witness configuration, 50 elsewhere.""" + return 200 if cfg[1] in KVQ_LISTED or cfg[0] != 8192 else 50 + + +@pytest.mark.parametrize("cfg", KVQ_CFGS, ids=KVQ_CFG_IDS) +@pytest.mark.parametrize("ratio", [1, 2]) +def test_kvq_synthetic_table_random_mixed_steps(ratio, cfg): + """Random multi-request steps of the native-storage call at one (max_num_tokens, max_batch_size): 200 + seeded steps (50 at a max_batch_size outside KVQ_LISTED) of k <= max_batch_size requests, T + 1..max_num_tokens, contexts then generation rows at next_n 1 or 6, generation positions up to + 16397, plus the far-end step, on the real [max_batch_size, 132] SWA table (each request's newest + window page among the pool's last four) and the synthetic [max_batch_size, 132] table: every + call equals the reference (window half through the SWA table, compressed half b * L_r + e), + every compressed row lies inside the max_batch_size * L_r-row buffer, and the far-end step + produces the buffer's largest reachable row (max_batch_size - 1) * L_r + 16398 // r - 1. Under + torch.profiler only the one kernel launches, and no variant is compiled after one warm-up call + in the target's form at this max_batch_size.""" + max_t, seqs = cfg + rng = random.Random(26100 + 1000 * ratio + 10 * seqs + (max_t == 2048)) + g = torch.Generator(device=DEV).manual_seed(26200 + 1000 * ratio + 10 * seqs + (max_t == 2048)) + tbl = _kvq_table(seqs) + swa = _kvq_swa(seqs, g) + steps = [_kvq_random_mixed(rng, seqs, max_t) for _ in range(_kvq_num_steps(cfg))] + [ + _kvq_far_step(seqs) + ] + warm = _kvq_step(ratio, [(0, 4)], [], 1, seed=7) + _kvq_call(ratio, warm[0], swa, warm[1], tbl, warm[2]) + torch.cuda.synchronize() + before = _compiled_variants() + rows_total, kinds, top, ks = 0, set(), -1, set() + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for i, (ctx, gen, next_n) in enumerate(steps): + req_id, swa_local, comp_local, B = _kvq_step( + ratio, + ctx, + gen, + next_n, + seed=26300 + 100_000 * ratio + 1000 * seqs + i + 500 * (max_t == 2048), + max_t=max_t, + ) + assert B <= seqs + T = req_id.numel() + lens = [q for _, q in ctx] + [next_n] * len(gen) + starts = [sum(lens[:b]) for b in range(B)] + _pool_bound_far_newest( + swa, list(range(B)), starts, swa_local, SWA_TPB, KVQ_SWA_PAGES, T + ) + out = _kvq_call(ratio, req_id, swa, swa_local, tbl, comp_local) + ref = _kvq_reference(ratio, req_id, swa, swa_local, comp_local) + assert torch.equal(out.long(), ref), (ratio, cfg, i, ctx, gen, next_n) + comp = ref[:, WINDOW:] + assert int(comp.max()) < seqs * KVQ_L[ratio] + top = max(top, int(comp.max())) + rows_total += T + kinds.add((bool(ctx), bool(gen), next_n if gen else 0)) + ks.add(B) + torch.cuda.synchronize() + far = (seqs - 1) * KVQ_L[ratio] + (DSPARK_MAX_POS + 1) // ratio - 1 + print( + f"ratio={ratio} cfg={cfg}: {len(steps)} calls, {rows_total} rows, requests {min(ks)}..{max(ks)}, kinds " + f"{sorted(kinds)}; largest compressed row {top} (far end {far})" + ) + assert top == far, (top, far) + if len(steps) > 100: # the 200-step configurations reach every step kind + assert {(False, True, 1), (False, True, 6), (True, False, 0)} <= kinds, kinds + if seqs >= 2: + assert {(True, True, 1), (True, True, 6)} <= kinds, kinds + _check_variants(prof, len(steps), before) + + +@pytest.mark.parametrize("cfg", KVQ_CFGS, ids=KVQ_CFG_IDS) +@pytest.mark.parametrize("ratio", [1, 2]) +def test_kvq_synthetic_table_cuda_graph_capture_set(ratio, cfg): + """The capture set at one (max_num_tokens, max_batch_size): decode G in EVERY 1..max_batch_size + (padding off captures all of them; padding on a subset plus max_batch_size) at next_n 1 and 6 + (DSpark verify), each captured as the target captures the call (prefix views of max_num_tokens-row + req / window buffers, the static [max_batch_size, 132] SWA and synthetic tables, a static [T, 512] + top-k) after one eager call at that shape. Three replays over new steps (request 0 ending at 16397 + or 16392, a fresh SWA table): each equals the reference and an eager call bitwise, and differs from + the previous step's reference.""" + max_t, seqs = cfg + g = torch.Generator(device=DEV).manual_seed(26400 + 1000 * ratio + 10 * seqs + (max_t == 2048)) + rng = random.Random(26500 + 1000 * ratio + 10 * seqs + (max_t == 2048)) + req_buf = torch.zeros(max_t, dtype=torch.int32, device=DEV) + swa_buf = torch.full((max_t, WINDOW), -1, dtype=torch.int32, device=DEV) + swa = _kvq_swa(seqs, g) + tbl = _kvq_table(seqs) + n_graphs = 0 + for next_n in (1, 6): + for G in range(1, seqs + 1): + rows = G * next_n + top_buf = torch.full((rows, TOPK), -1, dtype=torch.int32, device=DEV) + + def load(step_i): + gen = [rng.randint(0, DSPARK_MAX_POS + 1 - next_n) for _ in range(G)] + gen[0] = (DSPARK_MAX_POS if step_i % 2 == 0 else 16392) - next_n + 1 + req_id, swa_local, comp_local, _ = _kvq_step( + ratio, + [], + gen, + next_n, + seed=26600 + 97 * step_i + 1000 * rows + seqs, + max_t=max_t, + ) + swa.copy_(_kvq_swa(seqs, g)) + req_buf[:rows].copy_(req_id) + swa_buf[:rows].copy_(swa_local) + top_buf.copy_(comp_local) + return _kvq_reference(ratio, req_id, swa, swa_local, comp_local) + + def call(): + return _kvq_call(ratio, req_buf[:rows], swa, swa_buf[:rows], tbl, top_buf) + + prev = load(0) + assert torch.equal(call().long(), prev) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = call() + for it in range(1, 4): + ref = load(it) + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.long(), ref), (cfg, G, next_n, it) + assert torch.equal(out, call()), (cfg, G, next_n, it) + assert not torch.equal(out.long(), prev), (cfg, G, next_n, it) + prev = ref + del graph + n_graphs += 1 + print( + f"ratio={ratio} cfg={cfg}: {n_graphs} captures (G 1..{seqs} x next_n 1, 6) x 3 replays, " + f"replay = eager = reference" + ) + assert n_graphs == 2 * seqs + + +@pytest.mark.parametrize("seqs", [1, 2, 32]) +@pytest.mark.parametrize("ratio", [1, 2]) +def test_kvq_synthetic_table_wrong_numberings_are_rejected(ratio, seqs): + """The exact check sees a wrong compact numbering, on 20 random steps at max_batch_size `seqs` + (1, 2, 32): the table off by one (arange + 1, every valid entry one page late), the other ratio's + page size on the shared table, and req_idx_per_token off by one request at a request boundary. + + At max_batch_size 1 two of them are not wrong numberings, and the test pins that instead of + counting them: there is no request boundary, and request 0's row of the table is arange(132), so the + other ratio's page size gives the same row e wherever e < 132 * that page size (and -1 past it: + ratio 1 entries e >= 8,448 under the 64-row page size); ratio 2's entries e <= 8,198 all map to e.""" + rng = random.Random(26700 + 10 * ratio + seqs) + g = torch.Generator(device=DEV).manual_seed(26800 + 10 * ratio + seqs) + swa = _kvq_swa(seqs, g) + tbl = _kvq_table(seqs) + counts = {"table_off_by_one": [], "other_ratio_page_size": [], "req_off_by_one_request": []} + same_at_b0 = [] + n = 0 + while n < 20: + ctx, gen, next_n = _kvq_random_mixed(rng, seqs) + req_id, swa_local, comp_local, B = _kvq_step( + ratio, ctx, gen, next_n, seed=26900 + 1000 * ratio + 10 * seqs + n + ) + if not bool(((comp_local >= 0).any(1) & (req_id >= (1 if seqs > 1 else 0))).any()): + continue # the controls need a valid compressed entry (of a request at position >= 1 when there is one) + ref = _kvq_reference(ratio, req_id, swa, swa_local, comp_local) + assert torch.equal(_kvq_call(ratio, req_id, swa, swa_local, tbl, comp_local).long(), ref) + wrong = {"table_off_by_one": _kvq_call(ratio, req_id, swa, swa_local, tbl + 1, comp_local)} + other = _kvq_call( + ratio, req_id, swa, swa_local, tbl, comp_local, tpb=SWA_TPB // (3 - ratio) + ) + if seqs == 1: + # Not a wrong numbering at b = 0: exactly e where e < 132 * other page size, -1 past it. + e = comp_local.long() + tpb_o = SWA_TPB // (3 - ratio) + exp = torch.where( + (e >= 0) & (torch.div(e, tpb_o, rounding_mode="floor") < KVQ_COLS), e, -1 + ) + assert torch.equal(other[:, WINDOW:].long(), exp) and torch.equal( + other[:, :WINDOW], ref[:, :WINDOW] + ) + same_at_b0.append(int((other.long() != ref).sum())) + else: + wrong["other_ratio_page_size"] = other + starts = torch.nonzero(req_id[1:] != req_id[:-1]).flatten() + 1 + bad = req_id.clone() + bad[int(starts[rng.randrange(starts.numel())])] -= 1 + wrong["req_off_by_one_request"] = _kvq_call(ratio, bad, swa, swa_local, tbl, comp_local) + for k, w in wrong.items(): + counts[k].append(int((w.long() != ref).sum())) + n += 1 + print( + f"ratio={ratio} max_batch_size={seqs}: elements differing from the reference per control over 20 steps: " + f"{counts}" + + ( + f"; other ratio's page size at b = 0 (not a wrong numbering) differs in {same_at_b0}" + if seqs == 1 + else "" + ) + ) + for k, c in counts.items(): + if seqs == 1 and k != "table_off_by_one": + assert c == [], (k, c) # not applicable at max_batch_size 1 + else: + assert len(c) == 20 and all(x > 0 for x in c), (k, c) + if seqs == 1 and ratio == 2: + assert all(x == 0 for x in same_at_b0), same_at_b0 + + +# --------------------------------------------------------------------------------------------- +# The capture union (CAPTURE_UNION) at ratio 0 in the native-storage form. The native-storage target's ratio-0 layers +# make the same window-only call as every storage mode does, on the real SWA table as the metadata +# holds it: the [layer, SWA] view [max_batch_size, 132] of a 4-D buffer, page ids from the native SWA +# pool (up to KVQ_SWA_PAGES), with req_idx_per_token the token's request position in the step. (The +# native-storage ratio-1/2 call is test_kvq_synthetic_table_cuda_graph_capture_set, which captures every G +# in 1..max_batch_size at next_n 1 and 6.) +@pytest.mark.parametrize("rows,next_n", CAPTURE_UNION, ids=[f"T{r}_n{n}" for r, n in CAPTURE_UNION]) +def test_capture_union_ratio0_on_native_swa_tables(rows, next_n): + """Every (T, next_n) of the capture union, captured as the target captures the ratio-0 call + (prefix views [:T] of 8192-row req / window buffers, the static SWA table) after one eager call at + that shape, at max_batch_size 8 (the boot gates; G <= 8 only) and 32 (the accuracy gates). Three + replays over new steps: G generation requests at positions 0, 1, ... (request position b), each + verifying next_n positions at random up to 16397, request 0 ending at 16397 or 16392 with its + newest SWA page the native pool's last page, and a fresh SWA table. Each replay equals the int64 + reference and an eager call bitwise, produces the pool's last row, and differs from the previous + step's reference.""" + G = rows // next_n + n_graphs = 0 + for seqs in (8, 32): + if G > seqs: + continue + g = torch.Generator(device=DEV).manual_seed(27000 + 100 * seqs + rows + next_n) + rng = random.Random(27100 + 100 * seqs + rows + next_n) + req_buf = torch.zeros(8192, dtype=torch.int32, device=DEV) + swa_buf = torch.full((8192, WINDOW), -1, dtype=torch.int32, device=DEV) + swa = _kvq_swa(seqs, g) + + def load(step_i): + end0 = DSPARK_MAX_POS if step_i % 2 == 0 else 16392 + starts = [rng.randint(0, DSPARK_MAX_POS + 1 - next_n) for _ in range(G)] + starts[0] = end0 - next_n + 1 + positions = torch.cat( + [torch.arange(c, c + next_n, device=DEV, dtype=torch.int32) for c in starts] + ) + req_id = ( + torch.arange(G, device=DEV, dtype=torch.int32) + .repeat_interleave(next_n) + .contiguous() + ) + swa_local = _window(positions) + swa.copy_(_kvq_swa(seqs, g)) + swa[0, end0 // SWA_TPB] = KVQ_SWA_PAGES - 1 + req_buf[:rows].copy_(req_id) + swa_buf[:rows].copy_(swa_local) + ref = _reference64(req_id, swa, swa_local, SWA_TPB) + # Request 0's last token attends its own position, on the pool's last page. + assert ( + int(ref[next_n - 1, WINDOW - 1]) == (KVQ_SWA_PAGES - 1) * SWA_TPB + end0 % SWA_TPB + ) + return ref + + def call(): + return op(req_buf[:rows], swa, swa_buf[:rows], SWA_TPB, 0) + + prev = load(0) + assert torch.equal(call().long(), prev) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = call() + for it in range(1, 4): + ref = load(it) + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.long(), ref), (seqs, it) + assert torch.equal(out, call()), (seqs, it) + assert not torch.equal(out.long(), prev), (seqs, it) + prev = ref + del graph + n_graphs += 1 + print( + f"ratio=0 native SWA table T={rows} next_n={next_n}: {n_graphs} captures (max_batch_size " + f"{[s for s in (8, 32) if G <= s]}) x 3 replays, replay = eager = reference" + ) + assert n_graphs == (2 if G <= 8 else 1) + + +# --------------------------------------------------------------------------------------------- +# The engine envelope (per rank): max_num_tokens up to 16384, max_batch_size up to 256, +# max_seq_len L up to 1,048,576 with DSpark on and off, and the CUDA-graph batches GSET (every G in +# 1..32 plus TRT-LLM's padded list at cuda_graph_config max_batch_size 256), each captured at +# next_n 1 (T = G) and 6 (the DSpark verify step, T = 6 G). The metadata's block tables are +# [max_batch_size, W] with W = ceil(Lm / 128) + 3 columns (Lm = L + 14 with DSpark, L without); +# generation positions reach Lm + 5 (the graph metadata's max_draft_tokens) and context positions +# stay below L. Under decoder bounded replay a context step's decoder pass holds at most window = +# 128 rows per finishing request (its stash included), so one call's rows reach 128 * 256 = 32768. +# Page ids are drawn from [0, B) at the DSpark pool bounds (the engine of the dumps, the native +# cross-product maximum, the GB300 block ceiling and the int32 edge), each request's newest page +# among the pool's last four; every comparison is against the int64 reference. +GSET = tuple(range(1, 33)) + (40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256) +ENVELOPE_CAPTURE = [(G, n) for n in (1, 6) for G in GSET] +ENVELOPE_MAX_TOKENS = 16384 +ENVELOPE_MAX_BATCH = 256 +ENVELOPE_MAX_ROWS = WINDOW * ENVELOPE_MAX_BATCH +ENVELOPE_LS = (16384, 32768, 65536, 131072, 262144, 1048576) +ENVELOPE_CFGS = [(L, ds) for L in ENVELOPE_LS for ds in (True, False)] +ENVELOPE_CFG_IDS = [f"L{L}_{'dspark' if ds else 'nodspark'}" for L, ds in ENVELOPE_CFGS] + + +def _envelope_cfg(L, dspark): + lm = L + 14 if dspark else L + return {"L": L, "Lm": lm, "W": -(-lm // SWA_TPB) + 3, "max_gen": lm + 5} + + +ENVELOPE_WIDTHS = {cfg: _envelope_cfg(*cfg)["W"] for cfg in ENVELOPE_CFGS} +assert sorted(ENVELOPE_WIDTHS.values()) == [ + 131, + 132, + 259, + 260, + 515, + 516, + 1027, + 1028, + 2051, + 2052, + 8195, + 8196, +] + + +def _envelope_table(pages, g, width, rows=ENVELOPE_MAX_BATCH, layered=False): + """[rows, width] page ids drawn uniformly from [0, pages). layered: the [layer, SWA] view [1, 0] + of a [2, 2, rows, width] buffer (the metadata's sliding_block_tables).""" + ids = torch.randint(0, pages, (rows, width), generator=g, device=DEV, dtype=torch.int64).to( + torch.int32 + ) + if not layered: + return ids + buf = torch.full((2, 2, rows, width), -1, dtype=torch.int32, device=DEV) + buf[1, 0] = ids + return buf[1, 0] + + +def _envelope_topk(positions, ratio, g): + """The V4.1 index top-k shape without a [T, visible] score matrix (visible reaches ~1M here): + min(512, v) valid compressed positions in [0, v), v = (p + 1) // ratio, the newest v - 1 always + among them at a random column, the others uniform in [0, v) (the op maps element by element, so + a repeat changes nothing), then -1.""" + t = positions.shape[0] + vis = (positions.long() + 1) // ratio + k = vis.clamp(max=TOPK) + col = torch.arange(TOPK, device=DEV) + draw = (torch.rand(t, TOPK, generator=g, device=DEV, dtype=torch.float64) * vis[:, None]).long() + draw = torch.minimum(draw, (vis - 1).clamp(min=0)[:, None]) + newest = (torch.rand(t, generator=g, device=DEV, dtype=torch.float64) * k).long() + draw = torch.where(col[None] == newest[:, None], (vis - 1)[:, None], draw) + return torch.where(col[None] < k[:, None], draw, -1).to(torch.int32) + + +def _envelope_step(ratio, seqs, g, rows=None, tail_t0=None): + """Per-token inputs of requests seqs = [(cached, new)] in step order. rows: each request's + block-table row (default: distinct rows drawn from ENVELOPE_MAX_BATCH). tail_t0: per request, None or + the bounded-replay window start t0 -- that request's rows then carry the decoder pass's window, + positions t0 .. p in columns 0 .. p - t0, then -1. Returns req_id, positions, the window, the + top-k (ratio > 0), each request's first and last row, and the table rows.""" + n = len(seqs) + if rows is None: + rows = torch.randperm(ENVELOPE_MAX_BATCH, generator=g, device=DEV)[:n] + rows = rows.to(torch.int32) + total = sum(q for _, q in seqs) + cached = torch.tensor([c for c, _ in seqs], device=DEV, dtype=torch.int64) + lens = torch.tensor([q for _, q in seqs], device=DEV, dtype=torch.int64) + ends = torch.cumsum(lens, 0) + starts = ends - lens + req_id = torch.repeat_interleave(rows, lens, output_size=total).contiguous() + tok = torch.arange(total, device=DEV) + positions = (torch.repeat_interleave(cached - starts, lens, output_size=total) + tok).to( + torch.int32 + ) + swa_local = _window(positions) + if tail_t0 is not None and any(t is not None for t in tail_t0): + t0 = torch.tensor([-1 if t is None else t for t in tail_t0], device=DEV, dtype=torch.int64) + t0r = torch.repeat_interleave(t0, lens, output_size=total) + cols = torch.arange(WINDOW, device=DEV) + tail_win = torch.where( + cols[None] <= (positions.long() - t0r)[:, None], t0r[:, None] + cols[None], -1 + ) + swa_local = torch.where((t0r >= 0)[:, None], tail_win, swa_local.long()).to(torch.int32) + comp_local = _envelope_topk(positions, ratio, g) if ratio else None + return req_id, positions, swa_local, comp_local, starts, ends - 1, rows + + +def _envelope_far_newest(table, rows, last, local, tpb, pages): + """Each request's newest page (the page of the largest local index its last row names) is one of + the pool's last four: pages - 1 - (b % 4) for request b.""" + newest = local[last].max(1).values.long() + b = torch.arange(newest.numel(), device=DEV) + ok = newest >= 0 + table[rows.long()[ok], newest[ok] // tpb] = (pages - 1 - b[ok] % 4).to(torch.int32) + + +def _envelope_tables(ratio, W, b_swa, b_comp, g, step): + """Engine tables [ENVELOPE_MAX_BATCH, W] at one bound pair, with the step's newest pages at the far end.""" + _, _, swa_local, comp_local, _, last, rows = step + swa = _envelope_table(b_swa, g, W, layered=True) + _envelope_far_newest(swa, rows, last, swa_local, SWA_TPB, b_swa) + comp = None + if ratio: + comp = _envelope_table(b_comp, g, W) + _envelope_far_newest(comp, rows, last, comp_local, SWA_TPB // ratio, b_comp) + return swa, comp + + +def _envelope_random_mixed(rng, c, max_t=ENVELOPE_MAX_TOKENS, max_k=ENVELOPE_MAX_BATCH): + """A multi-request step at the engine envelope: k in [1, max_k] requests, context rows first, then + generation rows at next_n 1 or 6, T in [1, max_t]. One step in ten is generation-only and at least + one in ten context-only; the rest mix both. Each context sits at a random cached prefix below L; + each generation request starts in the tail [L - next_n, max_gen + 1 - next_n] with probability 1/4 + and anywhere up to max_gen + 1 - next_n otherwise.""" + L, top = c["L"], c["max_gen"] + k = rng.randint(1, max_k) + next_n = rng.choice((1, 6)) + u = rng.random() + if u < 0.1: + n_gen = k + elif u < 0.2 or k == 1: + n_gen = 0 + else: + n_gen = rng.randint(1, k - 1) + n_ctx = k - n_gen + gen_rows = n_gen * next_n + ctx = [] + if n_ctx: + ctx_rows = rng.randint(gen_rows + n_ctx, max_t) - gen_rows + cuts = sorted(rng.sample(range(1, ctx_rows), n_ctx - 1)) + news = [b - a for a, b in zip([0] + cuts, cuts + [ctx_rows])] + ctx = [(rng.randint(0, L - n), n) for n in news] + g_top = top + 1 - next_n + gen = [ + rng.randint(L - next_n, g_top) if rng.random() < 0.25 else rng.randint(0, g_top) + for _ in range(n_gen) + ] + return ctx, gen, next_n + + +def _envelope_col_count(local, tpb, col): + """Valid local indices this call resolves through table column `col`.""" + return int(((local >= 0) & (torch.div(local, tpb, rounding_mode="floor") == col)).sum()) + + +def _envelope_big_step(c, rng): + """One 32768-row step (the decoder-pass maximum, 128 rows x 256 requests, context rows first): + 8 contexts (a fresh prompt, a chunk on an 8192-token prefix, 1-, 5- and 129-token contexts, a + chunk ending at L - 1), 190 decoder-tail requests of 128 rows at random positions below L with the + bounded-replay window, then 47 DSpark verify runs (the first ending at max_gen) and 4 decode rows + (at L - 1, max_gen and two random positions): 249 requests.""" + L, top = c["L"], c["max_gen"] + ctx = [ + (0, 3000), + (8192, 2048), + (0, 1), + (0, 129), + (130, 5), + (2, 1570), + (4000, 430), + (L - 979, 979), + ] + tails = [(rng.randint(0, L - 128), 128) for _ in range(190)] + verify = [(top - 5, 6)] + [(rng.randint(0, top - 5), 6) for _ in range(46)] + decode = [(L - 1, 1), (top, 1), (rng.randint(0, top), 1), (rng.randint(0, top), 1)] + seqs = ctx + tails + verify + decode + t0 = [None] * len(ctx) + [c0 for c0, _ in tails] + [None] * (len(verify) + len(decode)) + assert sum(q for _, q in seqs) == ENVELOPE_MAX_ROWS and len(seqs) <= ENVELOPE_MAX_BATCH + return seqs, t0 + + +@pytest.mark.parametrize( + "cfg", [(16384, True), (1048576, True)], ids=["L16384_dspark", "L1048576_dspark"] +) +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_envelope_rows_prefix_bitwise(ratio, cfg): + """Every T = 1..32768, not a sample (max_num_tokens 16384; decoder bounded replay's 128 x 256 rows): + the 32768-row call equals the int64 reference on [256, W] tables at the GB300 block ceiling with + the pools' last rows produced, every call on the first T rows (prefix views of the same buffers) + is bitwise equal to the first T rows of it, and the sweep adds no compiled variant. Counter + control: a new constexpr (an unused page size) adds exactly one.""" + c = _envelope_cfg(*cfg) + W = c["W"] + rng = random.Random(28000 + 10 * ratio + (cfg[0] > 16384)) + g = torch.Generator(device=DEV).manual_seed(28050 + 10 * ratio + (cfg[0] > 16384)) + tpb_c = SWA_TPB // ratio if ratio else 0 + seqs, t0 = _envelope_big_step(c, rng) + step = _envelope_step(ratio, seqs, g, tail_t0=t0) + req_id, positions, swa_local, comp_local = step[:4] + b_swa, b_comp = _dspark_bound_pairs(ratio)[2] + swa, comp = _envelope_tables(ratio, W, b_swa, b_comp, g, step) + n = req_id.numel() + assert n == ENVELOPE_MAX_ROWS and int(positions.max()) == c["max_gen"] + full = _target_call(ratio, req_id, swa, swa_local, comp, comp_local, tpb_c) + ref = _step_reference64(ratio, req_id, swa, swa_local, comp, comp_local) + assert torch.equal(full.long(), ref) + assert int(ref[:, :WINDOW].max()) >= (b_swa - 4) * SWA_TPB + if ratio: + assert int(ref[:, WINDOW:].max()) >= (b_comp - 4) * tpb_c + torch.cuda.synchronize() + before = _compiled_variants() + bad = torch.zeros((), dtype=torch.int64, device=DEV) + for t in range(1, n + 1): + out = _target_call( + ratio, + req_id[:t], + swa, + swa_local[:t], + comp, + None if comp_local is None else comp_local[:t], + tpb_c, + ) + bad += (out != full[:t]).any(1).sum() + torch.cuda.synchronize() + after = _compiled_variants() + print( + f"ratio={ratio} cfg={cfg} W={W}: prefixes T=1..{n}: rows differing {int(bad)}; compiled variants " + f"{before} -> {after}; positions up to {int(positions.max())}" + ) + assert int(bad) == 0 + assert after == before + op(req_id[:4], swa, swa_local[:4], 3000 + 10 * ratio + (cfg[0] > 16384), 0) + torch.cuda.synchronize() + assert _compiled_variants() == after + 1 + + +ENVELOPE_RANDOM_STEPS = 240 + + +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_envelope_random_mixed_partitions(ratio): + """Random multi-request steps at the engine envelope: 240 seeded steps per ratio, k in [1, 256] requests, + T in [1, 16384], contexts at random cached prefixes below L then generation rows (next_n 1 or 6) + up to max_gen, cycling through the 12 (L, DSpark) table widths (20 steps each) and the DSpark pool + bound pairs, on [256, W] tables at random rows; each call bitwise equal to the int64 reference. Every + 20th step with >= 2 requests also runs req_idx_per_token off by one request at a request boundary, + which the same check must reject. Every step also runs on the tables cut to the rows it names + (fresh, aligned [m, W] buffers, m = the largest row + 1: a smaller max_batch_size's tables), + bitwise equal. Under torch.profiler only the one kernel launches, and no compiled variant is added + after one warm-up call per width.""" + rng = random.Random(28100 + ratio) + g = torch.Generator(device=DEV).manual_seed(28200 + ratio) + tpb_c = SWA_TPB // ratio if ratio else 0 + bounds = _dspark_bound_pairs(ratio) + before_warm = _compiled_variants() + for cfg in ENVELOPE_CFGS: + W = ENVELOPE_WIDTHS[cfg] + warm = _envelope_step(ratio, [(0, 4)], g) + swa, comp = _envelope_tables(ratio, W, 1024, 1024, g, warm) + _target_call(ratio, warm[0], swa, warm[2], comp, warm[3], tpb_c) + torch.cuda.synchronize() + before = _compiled_variants() + n_calls, rows_total, kinds, controls, ks, ts, tail_rows = 0, 0, set(), [], [], [], 0 + far = {} + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for i in range(ENVELOPE_RANDOM_STEPS): + cfg = ENVELOPE_CFGS[i % len(ENVELOPE_CFGS)] + c = _envelope_cfg(*cfg) + b_swa, b_comp = bounds[(i // len(ENVELOPE_CFGS)) % len(bounds)] + ctx, gen, next_n = _envelope_random_mixed(rng, c) + seqs = ctx + [(p, next_n) for p in gen] + step = _envelope_step(ratio, seqs, g) + req_id, positions, swa_local, comp_local, starts, _, rows = step + swa, comp = _envelope_tables(ratio, c["W"], b_swa, b_comp, g, step) + out = _target_call(ratio, req_id, swa, swa_local, comp, comp_local, tpb_c) + ref = _step_reference64(ratio, req_id, swa, swa_local, comp, comp_local) + n_calls += 1 + assert int(ref.max()) < 2**31 + assert torch.equal(out.long(), ref), (ratio, i, cfg, next_n, len(ctx), len(gen)) + # Row-count invariance: the same tables cut to the rows the step names (a smaller + # max_batch_size's tables, fresh 16-byte-aligned buffers) give the same output. + m = int(rows.max()) + 1 + swa_m = swa[:m].clone() + comp_m = comp[:m].clone() if ratio else None + assert torch.equal( + _target_call(ratio, req_id, swa_m, swa_local, comp_m, comp_local, tpb_c), out + ) + n_calls += 1 + T = req_id.numel() + rows_total += T + ks.append(len(seqs)) + ts.append(T) + tail_rows += sum(next_n for p in gen if p + next_n - 1 >= c["L"]) + far[cfg] = max(far.get(cfg, 0), int(positions.max())) + kinds.add((bool(ctx), bool(gen), next_n if gen else 0)) + if i % 20 == 0 and len(seqs) >= 2: + b = rng.randint(1, len(seqs) - 1) + bad_id = req_id.clone() + bad_id[int(starts[b])] = rows[b - 1] + wrong = _target_call(ratio, bad_id, swa, swa_local, comp, comp_local, tpb_c) + n_calls += 1 + controls.append(int((wrong.long() != ref).sum())) + torch.cuda.synchronize() + print( + f"ratio={ratio}: {n_calls} calls, {rows_total} rows, requests {min(ks)}..{max(ks)}, T {min(ts)}..{max(ts)}, " + f"{tail_rows} rows at positions >= L, kinds {sorted(kinds)}; largest position per (L, DSpark) {far}; " + f"width warm-ups added {before - before_warm} compiled variants; control (req_id off by one request) " + f"elements differing {controls}" + ) + assert len(controls) >= 8 and all(x > 0 for x in controls), controls + assert { + (True, True, 1), + (True, True, 6), + (False, True, 1), + (False, True, 6), + (True, False, 0), + } <= kinds + assert max(ks) > 200 and max(ts) > 15000 and tail_rows > 0 + _check_variants(prof, n_calls, before) + + +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_envelope_unaligned_table_views(ratio): + """The metadata's sliding_block_tables is one [layers, types, max_num_sequences, W] buffer, so the + [layer, SWA] view the op reads starts (layer * types + SWA) * max_num_sequences * W * 4 bytes in. + When max_num_sequences * W is not a multiple of 4 (an odd W with max_batch_size not a multiple + of 4), some of those views are not 16-byte aligned, and Triton specializes pointer arguments on + 16-byte alignment. Here: the views [0, 1], [1, 0] and [1, 1] of a [2, 2, 7, W] buffer at W 131 and + 8195 (offsets 4, 8 and 12 mod 16 bytes), each on 5 random steps of up to 7 requests: every call + equals the int64 reference and the same call on an aligned copy; the first unaligned call adds + exactly one compiled variant and no later call adds any (no earlier case in this file hands the + op an unaligned SWA table). Captured on the [1, 0] view after an eager + call (T = 7 and 42), three replays over new tables and steps each equal the reference and an eager + call.""" + rng = random.Random(30100 + ratio) + g = torch.Generator(device=DEV).manual_seed(30000 + ratio) + tpb_c = SWA_TPB // ratio if ratio else 0 + mbs = 7 + added, offsets = [], {} + for cfg in ((16384, False), (1048576, False)): + c = _envelope_cfg(*cfg) + W = c["W"] + buf = torch.randint( + 0, POOL_INT32_EDGE[128], (2, 2, mbs, W), generator=g, device=DEV, dtype=torch.int64 + ).to(torch.int32) + comp = _envelope_table(DSPARK_BOUNDS[ratio][2], g, W, rows=mbs) if ratio else None + views = {k: buf[k] for k in ((0, 1), (1, 0), (1, 1))} + for k, v in views.items(): + offsets[(W, k)] = v.data_ptr() % 16 + assert v.data_ptr() % 16 != 0, (W, k) + for s_i in range(5): + ctx, gen, next_n = _envelope_random_mixed(rng, c, max_k=mbs) + seqs = ctx + [(p, next_n) for p in gen] + step = _envelope_step(ratio, seqs, g, rows=torch.arange(len(seqs), device=DEV)) + req_id, _, swa_local, comp_local = step[:4] + for k, v in views.items(): + before = _compiled_variants() + out = _target_call(ratio, req_id, v, swa_local, comp, comp_local, tpb_c) + torch.cuda.synchronize() + added.append(_compiled_variants() - before) + assert torch.equal( + out.long(), _step_reference64(ratio, req_id, v, swa_local, comp, comp_local) + ), (W, k, s_i) + assert torch.equal( + _target_call(ratio, req_id, v.clone(), swa_local, comp, comp_local, tpb_c), out + ) + assert added[0] == 1 and all(a == 0 for a in added[1:]), added + c = _envelope_cfg(16384, False) + W = c["W"] + buf = torch.randint( + 0, POOL_INT32_EDGE[128], (2, 2, mbs, W), generator=g, device=DEV, dtype=torch.int64 + ).to(torch.int32) + view = buf[1, 0] + comp = _envelope_table(DSPARK_BOUNDS[ratio][2], g, W, rows=mbs) if ratio else None + n_graphs = 0 + for next_n in (1, 6): + rows = mbs * next_n + req_buf = torch.zeros(ENVELOPE_MAX_TOKENS, dtype=torch.int32, device=DEV) + swa_buf = torch.full((ENVELOPE_MAX_TOKENS, WINDOW), -1, dtype=torch.int32, device=DEV) + top_buf = torch.full((rows, TOPK), -1, dtype=torch.int32, device=DEV) if ratio else None + + def load(it): + starts = [rng.randint(0, c["max_gen"] + 1 - next_n) for _ in range(mbs)] + starts[0] = c["max_gen"] - next_n + 1 - (it % 2) + step = _envelope_step( + ratio, [(p, next_n) for p in starts], g, rows=torch.arange(mbs, device=DEV) + ) + req_id, _, swa_local, comp_local = step[:4] + view.copy_( + torch.randint( + 0, POOL_INT32_EDGE[128], (mbs, W), generator=g, device=DEV, dtype=torch.int64 + ).to(torch.int32) + ) + req_buf[:rows].copy_(req_id) + swa_buf[:rows].copy_(swa_local) + if ratio: + comp.copy_(_envelope_table(DSPARK_BOUNDS[ratio][2], g, W, rows=mbs)) + top_buf.copy_(comp_local) + return _step_reference64(ratio, req_id, view, swa_local, comp, comp_local) + + def call(): + return _target_call(ratio, req_buf[:rows], view, swa_buf[:rows], comp, top_buf, tpb_c) + + prev = load(0) + assert torch.equal(call().long(), prev) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = call() + for it in range(1, 4): + ref = load(it) + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.long(), ref), (next_n, it) + assert torch.equal(out, call()), (next_n, it) + assert not torch.equal(out.long(), prev), (next_n, it) + prev = ref + del graph + n_graphs += 1 + print( + f"ratio={ratio}: unaligned [layer, SWA] views, byte offsets mod 16 {offsets}; compiled variants added per " + f"call {added} (first call {added[0]}); {n_graphs} captures on the [1, 0] view x 3 replays exact" + ) + + +# Every (T, ratio) cell of dsv4_local_to_global_indices in the runtime-value dumps of the envelope +# engines (max_batch_size 256 with graphs over the padded list to 256; max_num_tokens 16384 / 8192 at +# L 32768 with DSpark): "T:ratios". +ENVELOPE_DUMPED_CELLS = ( + "1:012 2:012 5:012 6:012 7:012 12:012 16:012 24:012 34:012 48:012 96:012 128:01 134:1 144:012 192:012 " + "240:012 288:012 336:012 384:012 432:012 480:012 528:012 576:012 597:012 624:012 672:012 720:012 768:012 " + "1152:012 1536:012 2252:012 4559:012 8147:012 8192:012 8467:012 9193:02 9240:012 9524:012 10441:012 " + "10444:012 10447:012 10451:012 16298:012 16327:012 16328:012 16378:012 16384:012" +) + + +def _envelope_partition_exact(rng, T, c): + """A random step of exactly T rows: n_gen generation requests at next_n 1 or 6 and the rest of + the rows over n_ctx context requests, n_ctx + n_gen <= 256.""" + L, top = c["L"], c["max_gen"] + next_n = rng.choice((1, 6)) + n_gen = rng.randint(0, min(ENVELOPE_MAX_BATCH - 1, T // next_n)) + ctx_rows = T - n_gen * next_n + ctx = [] + if ctx_rows: + n_ctx = rng.randint(1, min(ctx_rows, ENVELOPE_MAX_BATCH - n_gen)) + cuts = sorted(rng.sample(range(1, ctx_rows), n_ctx - 1)) + news = [b - a for a, b in zip([0] + cuts, cuts + [ctx_rows])] + ctx = [(rng.randint(0, L - n), n) for n in news] + gen = [rng.randint(0, top + 1 - next_n) for _ in range(n_gen)] + return ctx, gen, next_n + + +def test_envelope_dumped_cells(): + """Every dumped (T, ratio) cell of the envelope engines as a consistent random step of exactly + T rows (<= 256 requests) at a (L, DSpark) width cycling over the 12, on [256, W] tables at the DSpark + pool bounds: bitwise equal to the int64 reference.""" + rng = random.Random(28300) + g = torch.Generator(device=DEV).manual_seed(28400) + cells = [ + (int(t), int(r)) + for tok in ENVELOPE_DUMPED_CELLS.split() + for t, rs in [tok.split(":")] + for r in rs + ] + assert len(cells) == 137 + for i, (T, ratio) in enumerate(cells): + c = _envelope_cfg(*ENVELOPE_CFGS[i % len(ENVELOPE_CFGS)]) + b_swa, b_comp = _dspark_bound_pairs(ratio)[i % 4] + ctx, gen, next_n = _envelope_partition_exact(rng, T, c) + step = _envelope_step(ratio, ctx + [(p, next_n) for p in gen], g) + req_id, _, swa_local, comp_local = step[:4] + assert req_id.numel() == T + swa, comp = _envelope_tables(ratio, c["W"], b_swa, b_comp, g, step) + tpb_c = SWA_TPB // ratio if ratio else 0 + out = _target_call(ratio, req_id, swa, swa_local, comp, comp_local, tpb_c) + assert torch.equal( + out.long(), _step_reference64(ratio, req_id, swa, swa_local, comp, comp_local) + ), (T, ratio) + print( + f"{len(cells)} dumped cells exact; T {min(t for t, _ in cells)}..{max(t for t, _ in cells)}" + ) + + +@pytest.mark.parametrize("cfg", ENVELOPE_CFGS, ids=ENVELOPE_CFG_IDS) +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_envelope_widths_far_end(ratio, cfg): + """Every (L, DSpark) table width with rows at its far end, at every DSpark pool bound pair. + (a) Tail steps (next_n 1 and 6): 6 contexts (two ending at L - 1) ahead of 250 generation + requests, which include decode rows at every position L - 1 .. max_gen (next_n 1) or verify runs + ending at every one of them (next_n 6); every call equals the int64 reference. Control: the same + call on tables cut just before the farthest column the step reads differs from the reference + exactly at the entries that read that column, which come back -1. + (b) Far-end enumeration: table rows 0 and 255 name pages B - W .. B - 1, so the last column is the + pool's last page; one token per row reads every slot of the last column (window) and the last + 512 compressed entries of the table's range, and one reads positions at column W (past the table): + exact, the pool's last row B * tpb - 1 is produced, and every past-the-table entry is -1.""" + c = _envelope_cfg(*cfg) + L, W, top = c["L"], c["W"], c["max_gen"] + rng = random.Random(28500 + 100 * ratio + ENVELOPE_CFGS.index(cfg)) + g = torch.Generator(device=DEV).manual_seed(28600 + 100 * ratio + ENVELOPE_CFGS.index(cfg)) + tpb_c = SWA_TPB // ratio if ratio else 0 + bounds = _dspark_bound_pairs(ratio) + ctx = [(0, 3000), (L - 7048, 2048), (0, 1), (130, 5), (L - 979, 979), (L - 384, 384)] + tail_ends = list(range(L - 1, top + 1)) + col_s = top // SWA_TPB + col_c = ((top + 1) // ratio - 1) // tpb_c if ratio else None + counts = [] + for next_n in (1, 6): + starts = [e - next_n + 1 for e in tail_ends] + starts += [ + rng.randint(0, top + 1 - next_n) + for _ in range(ENVELOPE_MAX_BATCH - len(ctx) - len(starts)) + ] + rng.shuffle(starts) + step = _envelope_step(ratio, ctx + [(s, next_n) for s in starts], g) + req_id, positions, swa_local, comp_local = step[:4] + assert req_id.numel() <= ENVELOPE_MAX_TOKENS and int(positions.max()) == top + n_s = _envelope_col_count(swa_local, SWA_TPB, col_s) + n_c = _envelope_col_count(comp_local, tpb_c, col_c) if ratio else 0 + assert n_s > 0 and (n_c > 0 or not ratio) + counts.append((next_n, n_s, n_c)) + for b_swa, b_comp in bounds: + swa, comp = _envelope_tables(ratio, W, b_swa, b_comp, g, step) + out = _target_call(ratio, req_id, swa, swa_local, comp, comp_local, tpb_c) + ref = _step_reference64(ratio, req_id, swa, swa_local, comp, comp_local) + assert int(ref.max()) < 2**31 + assert torch.equal(out.long(), ref), (ratio, cfg, next_n, b_swa, b_comp) + cut = _target_call( + ratio, + req_id, + swa[:, :col_s], + swa_local, + None if comp is None else comp[:, :col_c], + comp_local, + tpb_c, + ) + diff = cut.long() != ref + assert int(diff.sum()) == n_s + n_c, (int(diff.sum()), n_s, n_c) + assert bool((cut[diff] == -1).all()) + req_id = torch.tensor( + [0, 0, ENVELOPE_MAX_BATCH - 1, ENVELOPE_MAX_BATCH - 1], dtype=torch.int32, device=DEV + ) + ar = torch.arange(WINDOW, dtype=torch.int32, device=DEV) + swa_local = torch.stack([(W - 1) * SWA_TPB + ar, W * SWA_TPB + ar] * 2) + comp_local = None + if ratio: + ak = torch.arange(TOPK, dtype=torch.int32, device=DEV) + comp_local = torch.stack([W * tpb_c - TOPK + ak, W * tpb_c + ak] * 2) + for b_swa, b_comp in bounds: + swa = _envelope_table(b_swa, g, W, layered=True) + swa[0] = b_swa - W + torch.arange(W, device=DEV, dtype=torch.int32) + swa[ENVELOPE_MAX_BATCH - 1] = swa[0] + comp = None + if ratio: + comp = _envelope_table(b_comp, g, W) + comp[0] = b_comp - W + torch.arange(W, device=DEV, dtype=torch.int32) + comp[ENVELOPE_MAX_BATCH - 1] = comp[0] + out = _target_call(ratio, req_id, swa, swa_local, comp, comp_local, tpb_c) + ref = _step_reference64(ratio, req_id, swa, swa_local, comp, comp_local) + assert torch.equal(out.long(), ref), (ratio, cfg, b_swa, b_comp) + assert ( + int(out[0, WINDOW - 1]) == b_swa * SWA_TPB - 1 + and int(out[2, WINDOW - 1]) == b_swa * SWA_TPB - 1 + ) + assert bool((out[1] == -1).all()) and bool((out[3] == -1).all()) + if ratio: + assert int(out[0, -1]) == b_comp * tpb_c - 1 + print( + f"ratio={ratio} cfg={cfg} W={W}: tail positions {L - 1}..{top}; far column (SWA {col_s}, compressed " + f"{col_c}) entries per step (next_n, SWA, compressed) {counts}; bounds {bounds}" + ) + + +@pytest.mark.parametrize("cfg", ENVELOPE_CFGS, ids=ENVELOPE_CFG_IDS) +def test_envelope_tail_slots_form(cfg): + """The decoder-replay slot call (the target's _TailSlots): req_id = each tail row's request + position, the layer's SWA table, and a [tcd, 128] local table of [row position, -1 x 127]. Each + finishing request contributes at most window = 128 rows, so tcd reaches 128 x 256 = 32768. 20 + random steps (1..256 requests, 1..128 tail rows each, tails ending anywhere up to L - 1) and the + full 32768-row step (every tail ending at L - 1 .. L - 7), on [256, W] SWA tables at the DSpark + SWA pool bounds: each call equals the int64 reference, so column 0 is the row's own slot and every other + column is -1.""" + c = _envelope_cfg(*cfg) + L, W = c["L"], c["W"] + idx = ENVELOPE_CFGS.index(cfg) + rng = random.Random(28700 + idx) + g = torch.Generator(device=DEV).manual_seed(28800 + idx) + steps = [] + for _ in range(20): + nreq = rng.randint(1, ENVELOPE_MAX_BATCH) + steps.append( + [ + (n, rng.randint(n, L)) + for n in ( + WINDOW if rng.random() < 0.5 else rng.randint(1, WINDOW) for _ in range(nreq) + ) + ] + ) + steps.append([(WINDOW, L - b % 7) for b in range(ENVELOPE_MAX_BATCH)]) + sizes = [] + for i, tails in enumerate(steps): + b_swa = DSPARK_BOUNDS["swa"][i % 4] + seqs = [(N - n, n) for n, N in tails] + step = _envelope_step(0, seqs, g, rows=torch.arange(len(seqs), device=DEV)) + req_id, positions = step[0], step[1] + self_local = torch.full((req_id.numel(), WINDOW), -1, dtype=torch.int32, device=DEV) + self_local[:, 0] = positions + swa = _envelope_table(b_swa, g, W, layered=True) + _envelope_far_newest(swa, step[6], step[5], self_local, SWA_TPB, b_swa) + out = op(req_id, swa, self_local, SWA_TPB, 0) + ref = _reference64(req_id, swa, self_local, SWA_TPB) + assert torch.equal(out.long(), ref), (cfg, i) + assert bool((out[:, 1:] == -1).all()) and bool((out[:, 0] >= 0).all()) + sizes.append(req_id.numel()) + assert sizes[-1] == ENVELOPE_MAX_ROWS + print(f"cfg={cfg} W={W}: tail-slot calls of {sizes} rows exact") + + +@pytest.mark.parametrize("cfg", ENVELOPE_CFGS, ids=ENVELOPE_CFG_IDS) +@pytest.mark.parametrize("ratio", [0, 1, 2]) +def test_envelope_capture_gset_runtime_tables(ratio, cfg): + """CUDA-graph capture over GSET x next_n {1, 6} (T = G and 6 G, 1..1536) in the target's + runtime-table forms (ratio 0: window only; ratios 1 and 2: the runtime compress table), at every + (L, DSpark) width: prefix views [:T] of 16384-row req / window buffers, static [256, W] SWA (the + [layer, SWA] view of a 4-D buffer) and compress tables, a static [T, 512] top-k, after one eager + call at that shape. Three replays over new steps (request 0 ending at max_gen or L - 1, table rows + random or the request positions, tables at a new DSpark pool bound pair): each replay equals the int64 + reference and an eager call bitwise, differs from the previous step's reference, and request 0's + last row reads its own position on the pool's last page.""" + c = _envelope_cfg(*cfg) + L, W, top = c["L"], c["W"], c["max_gen"] + idx = ENVELOPE_CFGS.index(cfg) + rng = random.Random(29000 + 100 * ratio + idx) + g = torch.Generator(device=DEV).manual_seed(29100 + 100 * ratio + idx) + tpb_c = SWA_TPB // ratio if ratio else 0 + bounds = _dspark_bound_pairs(ratio) + req_buf = torch.zeros(ENVELOPE_MAX_TOKENS, dtype=torch.int32, device=DEV) + swa_buf = torch.full((ENVELOPE_MAX_TOKENS, WINDOW), -1, dtype=torch.int32, device=DEV) + swa_table = _envelope_table(1024, g, W, layered=True) + comp_table = _envelope_table(1024, g, W) if ratio else None + n_graphs = 0 + for G, next_n in ENVELOPE_CAPTURE: + rows = G * next_n + top_buf = torch.full((rows, TOPK), -1, dtype=torch.int32, device=DEV) if ratio else None + + def load(it): + b_swa, b_comp = bounds[(n_graphs + it) % len(bounds)] + starts = [rng.randint(0, top + 1 - next_n) for _ in range(G)] + end0 = top if it % 2 == 0 else L - 1 + starts[0] = end0 - next_n + 1 + trows = torch.arange(G, device=DEV) if it % 2 else None + step = _envelope_step(ratio, [(s, next_n) for s in starts], g, rows=trows) + req_id, _, swa_local, comp_local, _, last, trow = step + swa_table.copy_(_envelope_table(b_swa, g, W, layered=True)) + _envelope_far_newest(swa_table, trow, last, swa_local, SWA_TPB, b_swa) + req_buf[:rows].copy_(req_id) + swa_buf[:rows].copy_(swa_local) + if ratio: + comp_table.copy_(_envelope_table(b_comp, g, W)) + _envelope_far_newest(comp_table, trow, last, comp_local, tpb_c, b_comp) + top_buf.copy_(comp_local) + ref = _step_reference64(ratio, req_id, swa_table, swa_local, comp_table, comp_local) + assert int(ref[next_n - 1, WINDOW - 1]) == (b_swa - 1) * SWA_TPB + end0 % SWA_TPB + return ref + + def call(): + return _target_call( + ratio, req_buf[:rows], swa_table, swa_buf[:rows], comp_table, top_buf, tpb_c + ) + + prev = load(0) + assert torch.equal(call().long(), prev) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = call() + for it in range(1, 4): + ref = load(it) + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.long(), ref), (cfg, G, next_n, it) + assert torch.equal(out, call()), (cfg, G, next_n, it) + assert not torch.equal(out.long(), prev), (cfg, G, next_n, it) + prev = ref + del graph + n_graphs += 1 + print( + f"ratio={ratio} cfg={cfg} W={W}: {n_graphs} captures (G in GSET x next_n 1, 6; T 1..1536) x 3 replays, " + f"replay = eager = reference" + ) + assert n_graphs == 2 * len(GSET) + + +# The native-storage synthetic table at the engine envelope: [max_batch_size, W] = arange(mbs * W) +# for max_batch_size up to 256 and every (L, DSpark) width W, numbering request b's compressed entry e +# as row b * L_r + e of the compact read buffer, L_r = W * 128 / r. The window half is the real SWA +# table [mbs, W] (the [layer, SWA] view of a 4-D buffer) with page ids drawn from [0, 2**24). +def _envelope_kvq_reference(ratio, W, req_id, swa, swa_local, comp_local): + tpb = SWA_TPB // ratio + win = _reference64(req_id, swa, swa_local, SWA_TPB) + e = comp_local.long() + valid = (e >= 0) & (torch.div(e, tpb, rounding_mode="floor") < W) + return torch.cat([win, torch.where(valid, req_id.long()[:, None] * (W * tpb) + e, -1)], 1) + + +def _envelope_kvq_table(mbs, W): + return torch.arange(mbs * W, dtype=torch.int32, device=DEV).view(mbs, W) + + +@pytest.mark.parametrize("cfg", ENVELOPE_CFGS, ids=ENVELOPE_CFG_IDS) +@pytest.mark.parametrize("ratio", [1, 2]) +def test_envelope_kvq_maps_every_row_of_the_buffer(ratio, cfg): + """max_batch_size 256 at every (L, DSpark) width: every request position b < 256 with every + compressed position e in 0 .. L_r + 511 (in calls of <= 16384 rows), through the synthetic table: + each call equals the reference, each of the 256 * L_r buffer rows comes out exactly once, and every + position at or past column W gives -1.""" + c = _envelope_cfg(*cfg) + W = c["W"] + idx = ENVELOPE_CFGS.index(cfg) + g = torch.Generator(device=DEV).manual_seed(29300 + 100 * ratio + idx) + mbs = ENVELOPE_MAX_BATCH + tpb = SWA_TPB // ratio + L_r = W * tpb + tbl = _envelope_kvq_table(mbs, W) + swa = _envelope_table(POOL_INT32_EDGE[128], g, W, rows=mbs, layered=True) + per_req = -(-(L_r + TOPK) // TOPK) + total = mbs * per_req + counts = torch.zeros(mbs * L_r, dtype=torch.int32, device=DEV) + n_minus, calls = 0, 0 + ak = torch.arange(TOPK, device=DEV) + for t0 in range(0, total, ENVELOPE_MAX_TOKENS): + tok = torch.arange(t0, min(t0 + ENVELOPE_MAX_TOKENS, total), device=DEV) + req_id = (tok // per_req).to(torch.int32) + comp_local = ((tok % per_req)[:, None] * TOPK + ak[None]).to(torch.int32) + positions = torch.randint( + 0, c["max_gen"] + 1, (tok.numel(),), generator=g, device=DEV, dtype=torch.int32 + ) + swa_local = _window(positions) + out = _kvq_call(ratio, req_id, swa, swa_local, tbl, comp_local) + assert torch.equal( + out.long(), _envelope_kvq_reference(ratio, W, req_id, swa, swa_local, comp_local) + ), (cfg, t0) + comp = out[:, WINDOW:] + vals = comp[comp >= 0].long() + counts.index_add_(0, vals, torch.ones_like(vals, dtype=torch.int32)) + n_minus += int((comp < 0).sum()) + calls += 1 + assert bool((counts == 1).all()), "every buffer row exactly once" + assert n_minus == mbs * (per_req * TOPK - L_r) + print( + f"ratio={ratio} cfg={cfg} W={W}: rows 0..{mbs * L_r - 1} each produced once over {total} tokens in " + f"{calls} calls; {n_minus} entries past column {W} -> -1" + ) + + +def _envelope_kvq_far_step(c): + """256 requests (the full batch) whose last is one DSpark verify ending at max_gen, so the largest + row a step can produce, 255 * L_r + (max_gen + 1) // r - 1, is produced.""" + L = c["L"] + ctx = [(L - 4000, 4000)] + [ + ((97 * i) % (L - 64), 3 + i % 50) for i in range(ENVELOPE_MAX_BATCH - 2) + ] + return ctx, [c["max_gen"] - 5], 6 + + +@pytest.mark.parametrize("cfg", ENVELOPE_CFGS, ids=ENVELOPE_CFG_IDS) +@pytest.mark.parametrize("ratio", [1, 2]) +def test_envelope_kvq_random_mixed_steps(ratio, cfg): + """Random multi-request steps of the native-storage call at max_batch_size 256 and every (L, DSpark) + width: 200 seeded steps (k in [1, 256] requests at positions 0..k-1, T in [1, 16384], contexts + below L then generation rows at next_n 1 or 6 up to max_gen) plus the far-end step; every call + equals the reference (window half through the SWA table, compressed half b * L_r + e), every + compressed row lies inside the 256 * L_r-row buffer, and the far-end step produces the largest + reachable row. Row-count invariance: each step also runs with the synthetic table at + max_batch_size k and at a random max_batch_size in [k, 256] (fresh tensors), bitwise equal. + Under torch.profiler only the one kernel launches, and no variant is compiled after one warm-up + call in this form.""" + c = _envelope_cfg(*cfg) + W = c["W"] + idx = ENVELOPE_CFGS.index(cfg) + rng = random.Random(29500 + 100 * ratio + idx) + g = torch.Generator(device=DEV).manual_seed(29600 + 100 * ratio + idx) + mbs = ENVELOPE_MAX_BATCH + tpb = SWA_TPB // ratio + L_r = W * tpb + tbl = _envelope_kvq_table(mbs, W) + swa = _envelope_table(POOL_INT32_EDGE[128], g, W, rows=mbs, layered=True) + steps = [_envelope_random_mixed(rng, c) for _ in range(200)] + [_envelope_kvq_far_step(c)] + warm = _envelope_step(ratio, [(0, 4)], g, rows=torch.arange(1, device=DEV)) + _kvq_call(ratio, warm[0], swa, warm[2], tbl, warm[3]) + torch.cuda.synchronize() + before = _compiled_variants() + rows_total, kinds, top, ks, n_calls = 0, set(), -1, [], 0 + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for i, (ctx, gen, next_n) in enumerate(steps): + seqs = ctx + [(p, next_n) for p in gen] + k = len(seqs) + step = _envelope_step(ratio, seqs, g, rows=torch.arange(k, device=DEV)) + req_id, _, swa_local, comp_local, _, last, trow = step + _envelope_far_newest(swa, trow, last, swa_local, SWA_TPB, POOL_INT32_EDGE[128]) + out = _kvq_call(ratio, req_id, swa, swa_local, tbl, comp_local) + ref = _envelope_kvq_reference(ratio, W, req_id, swa, swa_local, comp_local) + assert torch.equal(out.long(), ref), (ratio, cfg, i, len(ctx), len(gen), next_n) + m = rng.randint(k, mbs) + assert torch.equal( + _kvq_call(ratio, req_id, swa, swa_local, _envelope_kvq_table(k, W), comp_local), out + ) + assert torch.equal( + _kvq_call(ratio, req_id, swa, swa_local, _envelope_kvq_table(m, W), comp_local), out + ) + n_calls += 3 + cmax = int(ref[:, WINDOW:].max()) + assert cmax < mbs * L_r + top = max(top, cmax) + rows_total += req_id.numel() + kinds.add((bool(ctx), bool(gen), next_n if gen else 0)) + ks.append(k) + torch.cuda.synchronize() + far = (mbs - 1) * L_r + (c["max_gen"] + 1) // ratio - 1 + print( + f"ratio={ratio} cfg={cfg} W={W}: {len(steps)} steps ({n_calls} calls), {rows_total} rows, requests " + f"{min(ks)}..{max(ks)}, kinds {sorted(kinds)}; largest compressed row {top} (far end {far})" + ) + assert top == far, (top, far) + assert { + (True, True, 1), + (True, True, 6), + (False, True, 1), + (False, True, 6), + (True, False, 0), + } <= kinds + _check_variants(prof, n_calls, before) + + +# Capture of the synthetic-table call. The target builds the per-request read buffer (and with it the +# synthetic table) only when it holds no more rows than the per-query one: max_batch_size * W * 128 <= +# max_num_tokens * 512, i.e. max_batch_size * W <= 65536 at max_num_tokens 16384; above that it passes +# the runtime compress table (test_envelope_capture_gset_runtime_tables). So at each width: every +# max_batch_size in GSET with max_batch_size * W <= 65536, plus 256 (a superset, the numbering at the +# largest batch), each capturing every G in GSET up to max_batch_size at next_n 1 and 6. At W 132 the +# max_batch_size values 1..32 are test_kvq_synthetic_table_cuda_graph_capture_set. +ENVELOPE_KVQ_PER_REQUEST_MAX = 4 * ENVELOPE_MAX_TOKENS + + +def _envelope_kvq_capture_mbs(cfg): + W = ENVELOPE_WIDTHS[cfg] + out = [m for m in GSET if m * W <= ENVELOPE_KVQ_PER_REQUEST_MAX or m == ENVELOPE_MAX_BATCH] + return [m for m in out if m > 32] if cfg == (16384, True) else out + + +@pytest.mark.parametrize("cfg", ENVELOPE_CFGS, ids=ENVELOPE_CFG_IDS) +@pytest.mark.parametrize("ratio", [1, 2]) +def test_envelope_kvq_capture_gset(ratio, cfg): + """The native-storage call captured as the target captures it, at every max_batch_size of + _envelope_kvq_capture_mbs(cfg): every G in GSET up to max_batch_size at next_n 1 and 6, prefix views of + 16384-row req / window buffers, the static [mbs, W] SWA and synthetic tables and a static + [T, 512] top-k, after one eager call at that shape. Three replays over new steps (request 0 + ending at max_gen or L - 1, a fresh SWA table): each equals the reference and an eager call + bitwise, and differs from the previous step's.""" + c = _envelope_cfg(*cfg) + L, W, top = c["L"], c["W"], c["max_gen"] + rng = random.Random(29800 + 1000 * ratio + 10 * ENVELOPE_CFGS.index(cfg)) + g = torch.Generator(device=DEV).manual_seed( + 29900 + 1000 * ratio + 10 * ENVELOPE_CFGS.index(cfg) + ) + req_buf = torch.zeros(ENVELOPE_MAX_TOKENS, dtype=torch.int32, device=DEV) + swa_buf = torch.full((ENVELOPE_MAX_TOKENS, WINDOW), -1, dtype=torch.int32, device=DEV) + mbs_list = _envelope_kvq_capture_mbs(cfg) + n_graphs = 0 + for mbs, (G, next_n) in ((m, gn) for m in mbs_list for gn in ENVELOPE_CAPTURE): + if G > mbs: + continue + if G == 1 and next_n == 1: + tbl = _envelope_kvq_table(mbs, W) + swa = _envelope_table(POOL_INT32_EDGE[128], g, W, rows=mbs, layered=True) + rows = G * next_n + top_buf = torch.full((rows, TOPK), -1, dtype=torch.int32, device=DEV) + + def load(it): + starts = [rng.randint(0, top + 1 - next_n) for _ in range(G)] + starts[0] = (top if it % 2 == 0 else L - 1) - next_n + 1 + step = _envelope_step( + ratio, [(s, next_n) for s in starts], g, rows=torch.arange(G, device=DEV) + ) + req_id, _, swa_local, comp_local, _, last, trow = step + swa.copy_(_envelope_table(POOL_INT32_EDGE[128], g, W, rows=mbs, layered=True)) + _envelope_far_newest(swa, trow, last, swa_local, SWA_TPB, POOL_INT32_EDGE[128]) + req_buf[:rows].copy_(req_id) + swa_buf[:rows].copy_(swa_local) + top_buf.copy_(comp_local) + return _envelope_kvq_reference(ratio, W, req_id, swa, swa_local, comp_local) + + def call(): + return _kvq_call(ratio, req_buf[:rows], swa, swa_buf[:rows], tbl, top_buf) + + prev = load(0) + assert torch.equal(call().long(), prev) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = call() + for it in range(1, 4): + ref = load(it) + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.long(), ref), (cfg, mbs, G, next_n, it) + assert torch.equal(out, call()), (cfg, mbs, G, next_n, it) + assert not torch.equal(out.long(), prev), (cfg, mbs, G, next_n, it) + prev = ref + del graph + n_graphs += 1 + print( + f"ratio={ratio} cfg={cfg} W={W}: max_batch_size {mbs_list}: {n_graphs} captures x 3 replays, " + f"replay = eager = reference" + ) + assert n_graphs == sum(2 * sum(1 for G in GSET if G <= m) for m in mbs_list) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_flashinfer_top_k_ragged_deterministic.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_flashinfer_top_k_ragged_deterministic.py new file mode 100644 index 000000000000..91410a8c04b6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_flashinfer_top_k_ragged_deterministic.py @@ -0,0 +1,2383 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the flashinfer_top_k_ragged_deterministic entry (FlashInfer ragged top-k in its +deterministic mode, ties at the k-th value to the smaller column), at the DeepSeek-V4.1 target's two +index top-k calls: the context rows (per-request column ranges ``[ks, ke)`` over the prefill logits) +and the generation rows (per-row visible lengths over the paged logits, next_n 1 and 6), over the +envelope of up to 16384 tokens and 256 requests per step, max_seq_len up to 1048576 (DSpark on and +off), and CUDA-graph batches 1..32, 40..128 step 8, 192 and 256. + +The entry owns what is certified: `CELLS` is the list, `reference` and `compare` are the gate. The +reference ranks each row's own range by the op's order (ties go to the smaller column), takes its +first 512 columns sorted ascending, then -1; the entry must equal it exactly, element for element, +including the order and the -1 padding. Because that answer is a function of the scores alone, the +same test also shows the output is bitwise identical across calls, across fresh processes and +across GPUs. This file turns a cell's spec into inputs and keeps the checks a cell cannot carry: +the CUDA-graph capture set, the mixed-step views, the row-prefix sweep of the 16384-row step, the +seeded random and dumped steps, the cross-process and cross-GPU digests, the first call inside a +capture, the wrong controls, the guards and the rejections. + +The cross-process checks are plain two-device checks, not a collective: fresh processes, each on one +GPU, recompute digests that this process compares. Run as a script (``python --digests +``) it writes the digests the cross-process test compares; ``--capture-digests +`` writes the replay digests of the CUDA-graph capture cases the cross-process capture +test compares; ``--first-call-in-capture `` makes the process's first call inside a CUDA +graph capture. +""" + +import base64 +import hashlib +import json +import os +import random +import subprocess +import sys +import tempfile +import zlib + +import pytest +import torch +from flashinfer import top_k_ragged_transform +from flashinfer.topk import TopKTieBreak +from torch.profiler import ProfilerActivity, profile + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.flashinfer_top_k_ragged_deterministic import ( # noqa: E501 + flashinfer_top_k_ragged_deterministic as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "flashinfer_top_k_ragged_deterministic requires a CUDA device" + +DEV = "cuda" +K = 512 +# The two kernels one call launches on the pinned path (deterministic, ties SMALL, dsa_graph_safe): +# FlashInfer's one-block-per-row filtered top-k with scalar loads, then the per-row ascending sort. +KERNELS = ( + "flashinfer::sampling::FilteredTopKUnifiedKernel K).nonzero().flatten() + col = torch.arange(w, device=DEV)[None, :] + step = _row_step(w) + for a in range(0, rows.numel(), step): + rr = rows[a : a + step] + x = torch.where(col < L[rr][:, None], win[rr].float(), float("-inf")) + if large: + x = x.flip(1) + order = torch.sort(x, dim=1, descending=True, stable=True).indices[:, :K] + if large: + order = w - 1 - order + ref[rr] = order.sort(1).values + del x, order + return ref + + +def _kth_ties(win, lens, ref): + """Per row: (in-range columns equal to the K-th selected score, selected columns equal to it).""" + L = lens.long().clamp(min=0) + n, w = win.shape + col = torch.arange(w, device=DEV)[None, :] + eq_in = torch.zeros(n, dtype=torch.int64, device=DEV) + eq_sel = torch.zeros(n, dtype=torch.int64, device=DEV) + step = _row_step(w) + for a in range(0, n, step): + rs = slice(a, a + step) + vals = torch.gather(win[rs].float(), 1, ref[rs].clamp(min=0)) + kth = torch.where(ref[rs] >= 0, vals, float("inf")).min(1).values + eq_in[rs] = ((win[rs] == kth[:, None]) & (col < L[rs, None])).sum(1) + eq_sel[rs] = ((vals == kth[:, None]) & (ref[rs] >= 0)).sum(1) + return eq_in, eq_sel + + +def _straddling_rows(win, lens, ref): + """Rows where the K-th selected score is tied with an in-range score that was left out.""" + sel = lens.long() > K + if not bool(sel.any()): + return 0 + eq_in, eq_sel = _kth_ties(win, lens, ref) + return int(((eq_in > eq_sel) & sel).sum()) + + +def _max_tied(win, lens, ref): + """The largest number of in-range columns tied with a selecting row's K-th selected score.""" + eq_in, _ = _kth_ties(win, lens, ref) + return int(torch.where(lens.long() > K, eq_in, 0).max()) + + +def _gen(seed): + return torch.Generator(device=DEV).manual_seed(seed) + + +def _distinct(n, w, seed): + """fp32 [n, w], each row a random permutation of distinct values of both signs.""" + perm = torch.argsort(torch.rand(n, w, generator=_gen(seed), device=DEV), dim=1, stable=True) + return (perm.float() - w // 2) * 0.125 + + +def _tie_at_kth(n, w, lens, seed): + """fp32 [n, w]. In every row with lens > K, in random in-range columns: `above` in 0..511 + distinct scores above 5, then `m` exact copies of 5.0 with above < K < above + m (so the tie at + the K-th value straddles the cut), the rest distinct and below 5.""" + g = _gen(seed) + L = lens.long().clamp(min=0) + col = torch.arange(w, device=DEV)[None, :] + u = torch.where(col < L[:, None], torch.rand(n, w, generator=g, device=DEV), 2.0) + rank = torch.argsort(torch.argsort(u, dim=1, stable=True), dim=1, stable=True) + above = torch.randint(0, K, (n,), generator=g, device=DEV) + extra = torch.randint(1, 65, (n,), generator=g, device=DEV) + m = torch.minimum(K - above + extra, L - above).clamp(min=0) + return torch.where( + rank < above[:, None], + 100.0 + rank.float(), + torch.where(rank < (above + m)[:, None], 5.0, -1.0 - u), + ) + + +KINDS = ("distinct", "bf16", "levels", "zero_plateau", "tie_at_kth") + + +def _scores(kind, n, w, lens, seed): + """fp32 [n, w] scores of one kind: distinct; bf16-rounded normals and 8 integer levels (heavy + exact ties everywhere); an exact-zero plateau (relu(x - 3): ~0.1 % positive, so the K-th value of + every selecting row is 0.0, tied with thousands of columns); an exact tie constructed to + straddle the K-th value.""" + if kind == "distinct": + return _distinct(n, w, seed) + if kind == "tie_at_kth": + return _tie_at_kth(n, w, lens, seed) + g = _gen(seed) + if kind == "levels": + return torch.randint(0, 8, (n, w), generator=g, device=DEV).float() + x = torch.randn(n, w, generator=g, device=DEV) + if kind == "bf16": + return x.to(torch.bfloat16).float() + assert kind == "zero_plateau" + return torch.relu(x - 3.0) + + +def _mixed(n, w, lens, seed): + """Rows cycling through the five kinds.""" + out = torch.empty(n, w, device=DEV) + for j, kind in enumerate(KINDS[:n]): + rows = torch.arange(j, n, len(KINDS), device=DEV) + out[rows] = _scores(kind, rows.numel(), w, lens[rows], seed * 8 + j) + return out + + +def _zeros(n): + return torch.zeros(n, dtype=torch.int32, device=DEV) + + +def _call(scores, lens, row_starts=None, offsets=None): + return op(scores, _zeros(scores.shape[0]) if offsets is None else offsets, lens, row_starts, K) + + +def _flashinfer_kernels(prof): + """Names of the FlashInfer kernels launched inside a profiler session (the test's own + reference math launches only torch kernels).""" + return [ + e.name for e in prof.events() if e.device_type.name == "CUDA" and "flashinfer" in e.name + ] + + +def _profiled(fn, names): + """fn() in its own profiler session, appending the FlashInfer kernels it launched to `names` + (for loops whose own torch work would make one session too large to parse).""" + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + out = fn() + torch.cuda.synchronize() + names.extend(_flashinfer_kernels(prof)) + return out + + +def _assert_variants(prof, n_calls, what): + """`prof` is a profiler session or the kernel names `_profiled` collected.""" + names = prof if isinstance(prof, list) else _flashinfer_kernels(prof) + kinds = sorted(set(names)) + print(f"[{what}] kernel-variant map over {n_calls} calls: {len(names)} launches of {kinds}") + assert len(kinds) == 2 and all(any(k in name for name in kinds) for k in KERNELS), kinds + assert len(names) == 2 * n_calls, (len(names), n_calls) + + +def _sha(t): + return hashlib.sha256( + t.detach().contiguous().view(torch.uint8).reshape(-1).cpu().numpy().tobytes() + ).hexdigest() + + +# ─── the call's semantics ────────────────────────────────────────────────────────────────────── + + +def _drive_semantics(): + """Short rows (length -3, 0, 1, 511, 512) are 0 .. len-1 then -1; long rows are the reference + exactly; offsets are added to valid entries only; row_starts makes the columns window-local. + The result is a fresh contiguous int32 [N, 512].""" + n, w = 96, 4096 + g = torch.Generator().manual_seed(1) + rs = torch.randint(0, 2048, (n,), generator=g) + lens = torch.minimum(torch.randint(-3, 4097, (n,), generator=g), w - rs) + lens[:8] = torch.tensor([-3, 0, 1, 511, 512, 513, 1000, 2048]) + rs[:8] = torch.tensor([0, 5, 7, 100, 3, 0, 9, 2048]) + rs, lens = rs.int().to(DEV), lens.int().to(DEV) + win = _mixed(n, w, lens, seed=11) # row i's range in window coordinates + x = torch.stack([win[i].roll(int(rs[i])) for i in range(n)]) # x[i, rs + j] = win[i, j] + off = torch.randint(0, 1 << 20, (n,), generator=g).int().to(DEV) + out = _call(x, lens, row_starts=rs, offsets=off) + assert out.dtype == torch.int32 and out.shape == (n, K) and out.is_contiguous() + assert out.data_ptr() != x.data_ptr() + op.compare(out, op.reference(x, off, lens, rs, K)) + # offsets are added to valid entries only: the window-local answer plus each row's offset + ref = _reference(win, lens) + assert torch.equal(out.long(), torch.where(ref >= 0, ref + off[:, None].long(), -1)) + out0 = _call(x, lens.clamp(max=w)) # row_starts None: ranges start at column 0 + op.compare(out0, op.reference(x, _zeros(n), lens.clamp(max=w), None, K)) + print(f"straddling tie rows among {n}: {_straddling_rows(win, lens, ref)}") + + +# ─── generation rows: the paged-logits call at every engine width and row count ─────────────── +# The target calls the entry once per index-source layer in every step with generation rows: +# scores = the paged-logits output [rows, W] (row stride ceil256(W), unspecified past each row's +# length), or at W > 16384 on re-index layers the contiguous torch.where(candidate_mask, logits, +# -inf) copy; lengths = compressed_visible_lens_cuda[r][tc:T] ((p + 1) // r per row); row_starts +# None; offsets = zeros. W is the metadata's static max_compressed_kv_lens[r]. Rows: at most 256 +# requests x next_n 6 = 1536. +FORMS = [ + (c.spec["width"], c.spec["ratio"], c.spec["layout"]) + for c in op.CELLS + if c.spec["kind"] == "gen" +] +ENGINE_WIDTHS = list(dict.fromkeys((w, r) for w, r, _ in FORMS)) +MAX_SEQ_LENS = [16384, 32768, 65536, 131072, 262144, 1048576] +MAX_ROWS = 1536 +FORM_IDS = [f"w{w}_r{r}_{f}" for w, r, f in FORMS] +assert len(ENGINE_WIDTHS) == 40 and len(FORMS) == 59 + + +def _dspark_vis(batch, ratio, width, g, padding): + """Per-row visible lengths (p + j + 1) // ratio of a verify step (next_n 6): request 0 at the cap + (its last row sees exactly `width`), request 1 at position 0 or 1, the rest random, and with + `padding` a graph padding request (all 0) last.""" + cap = width * ratio - 6 + p = torch.randint(0, cap + 1, (batch,), generator=g) + p[0] = cap + if batch > 1: + p[1] = int(torch.randint(0, 2, (1,), generator=g)) + r = torch.arange(batch * 6) + vis = (p[r // 6] + r % 6 + 1) // ratio + if padding and batch > 2: + vis[-6:] = 0 + return vis.to(torch.int32).to(DEV) + + +def _decode_vis(batch, ratio, width, g, padding): + """Per-row visible lengths (p + 1) // ratio of a next_n 1 step: request 0 at the cap, request 1 + at position 0 or 1, one straddling K, the rest random, and with `padding` a 0 row last.""" + p = torch.randint(0, width * ratio, (batch,), generator=g) + p[0] = width * ratio - 1 + if batch > 1: + p[1] = int(torch.randint(0, 2, (1,), generator=g)) + if batch > 2: + p[2] = K * ratio + vis = (p + 1) // ratio + if padding and batch > 3: + vis[-1] = 0 + return vis.to(torch.int32).to(DEV) + + +def _candidate_keep(vis, width, g): + """A re-index layer's keep mask [rows, width]: a row with more than 2048 blocks of 8 keeps its + newest block plus 2047 older ones, a shorter row keeps everything visible; False past the row.""" + rows = vis.numel() + nb_all = -(-width // 8) + L = vis.long() + nb = (L + 7) // 8 + blk = torch.arange(nb_all, device=DEV)[None, :] + older = blk < (nb - 1)[:, None] + s = torch.where(older, torch.rand(rows, nb_all, generator=g, device=DEV), -1.0) + keep = torch.zeros(rows, nb_all, dtype=torch.bool, device=DEV) + keep.scatter_(1, s.topk(min(2047, nb_all), dim=1).indices, True) + keep = (keep & older) | (blk == (nb - 1)[:, None]) + keep = torch.where((nb <= 2048)[:, None], blk < nb[:, None], keep) + col = torch.arange(width, device=DEV)[None, :] + return keep.repeat_interleave(8, 1)[:, :width] & (col < L[:, None]) + + +def _decode_inputs(width, form, vis, seed, out=None, kind="mixed"): + """The target's generation-row scores for lengths `vis`: the paged-logits view (NaN past every + row's length and in the stride padding) or the candidate-masked contiguous copy (-inf outside + the kept blocks). Returns (the entry's scores, the scores the reference ranks). With `out` + (a (paged buffer, masked buffer) pair) it writes in place.""" + rows = vis.numel() + base = ( + _mixed(rows, width, vis, seed) if kind == "mixed" else _scores(kind, rows, width, vis, seed) + ) + col = torch.arange(width, device=DEV)[None, :] + stride = -(-width // 256) * 256 + buf = torch.full((rows, stride), float("nan"), device=DEV) if out is None else out[0] + if out is not None: + buf.fill_(float("nan")) + buf[:, :width] = base.masked_fill(col >= vis[:, None].long(), float("nan")) + if form == "paged_view": + return buf[:, :width], base + masked = torch.where(_candidate_keep(vis, width, _gen(seed)), base, float("-inf")) + assert masked.stride() == (width, 1) + if out is not None: + out[1].copy_(masked) + masked = out[1] + return masked, masked + + +def _drive_gen(width, ratio, form): + """Every row count of the generation rows: the 1536-row call (256 DSpark verify requests, row 0 at + the full width, mixed tie kinds) equals the reference exactly, and a call on the first r rows + equals the first r rows of it bitwise for every r in 1..1536 (one block per row: rows only size + the grid). A next_n 1 step (1536 one-row requests) is exact too. Every call launches the same two + kernels. Control: lengths one past each row's own read a NaN column and are seen.""" + seed = width + 13 * ratio + len(form) + g = torch.Generator().manual_seed(seed) + vis = _dspark_vis(MAX_ROWS // 6, ratio, width, g, padding=False) + vis[0] = width + scores, ranked = _decode_inputs(width, form, vis, seed) + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + full = _call(scores, vis) + prefix = [_call(scores[:r], vis[:r]) for r in range(1, MAX_ROWS + 1)] + torch.cuda.synchronize() + ref = op.reference(scores, _zeros(MAX_ROWS), vis, None, K) + op.compare(full, ref) + ref = ref.long() + bad = [r for r, out in enumerate(prefix, 1) if not torch.equal(out, full[:r])] + assert not bad, f"prefixes differing from the {MAX_ROWS}-row call: {bad[:10]}" + del prefix + _assert_variants(prof, MAX_ROWS + 1, f"sweep w{width} r{ratio} {form}") + vis1 = _decode_vis(MAX_ROWS, ratio, width, g, padding=True) + s1, _ = _decode_inputs(width, form, vis1, seed + 1) + op.compare(_call(s1, vis1), op.reference(s1, _zeros(MAX_ROWS), vis1, None, K)) + print( + f"[sweep w{width} r{ratio} {form}] stride {scores.stride()}; straddling tie rows " + f"{_straddling_rows(ranked, vis, ref)} of {MAX_ROWS}" + ) + if form == "paged_view": + over = _call(scores, (vis + 1).clamp(max=width)) + wrong = int((over.long() != ref).any(1).sum()) + print(f"[sweep w{width}] control lengths + 1: rows differing {wrong} of {MAX_ROWS}") + assert wrong > 0 + + +# Every CUDA-graph batch of the envelope, per rank: Gset = every batch 1..32 plus +# TRT-LLM's padded capture list at cuda_graph_config max_batch_size 256 (1, 2, 4, 8..128 step 8, +# 192, 256; the padding-off list 1..32, 64, 128, 256 is a subset), each as decode (next_n 1: G rows) +# and as DSpark verify (next_n 6: 6G rows), at every form above (all 40 static widths, both layouts). +GSET = list(range(1, 33)) + [40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256] +assert len(GSET) == 46 +GRAPH_CASES = [(w, r, f, n, b) for w, r, f in FORMS for n in (1, 6) for b in GSET] +GRAPH_IDS = [f"w{w}_r{r}_{f}_nextn{n}_G{b}" for w, r, f, n, b in GRAPH_CASES] + + +_FP_WEIGHTS = {} + + +def _fingerprint(t): + """Digest of a 2-D tensor's bit patterns, reduced on the GPU (the capture inputs reach 6.4 GB): + per row, the sum of its 32-bit words times seeded random odd 64-bit weights, wrapping modulo + 2**64 (so the sum's order does not matter), then SHA-256 over the row values.""" + rows, cols = t.shape + words = t.view(torch.int32) + if cols not in _FP_WEIGHTS: + g = torch.Generator(device=DEV).manual_seed(20260926) + _FP_WEIGHTS[cols] = torch.randint(-(2**62), 2**62, (cols,), generator=g, device=DEV) | 1 + w = _FP_WEIGHTS[cols] + fp = torch.empty(rows, dtype=torch.int64, device=DEV) + step = _row_step(cols) + for a in range(0, rows, step): + fp[a : a + step] = (words[a : a + step].to(torch.int64) * w).sum(1) + return _sha(fp) + + +def _capture_case(width, ratio, form, next_n, batch, digests=True): + """One case of the capture set: an eager call at the shape, the call captured, then 3 replays on + new lengths (0/1, the width, a padding request) and new mixed-tie scores written in place. + Returns (whether the eager call equals the reference, one record per replay: with `digests`, + digests of the inputs the op reads (fingerprint of the scores buffer, SHA-256 of the lengths) + and of the replay's output (SHA-256); whether the replay equals the reference and an eager + call on the same inputs bitwise, and the rows where it differs from the previous step's + reference).""" + rows = batch * next_n + seed = ( + width + + 7 * ratio + + 1000 * batch + + 100000 * next_n + + (1 if form == "candidate_masked" else 0) + ) + g = torch.Generator().manual_seed(seed) + stride = -(-width // 256) * 256 + bufs = (torch.empty(rows, stride, device=DEV), torch.empty(rows, width, device=DEV)) + static, read = (bufs[0][:, :width], bufs[0]) if form == "paged_view" else (bufs[1], bufs[1]) + lens = torch.empty(rows, dtype=torch.int32, device=DEV) + off = _zeros(rows) + make_vis = _dspark_vis if next_n == 6 else _decode_vis + + def load(it): + vis = make_vis(batch, ratio, width, g, padding=it % 2 == 1) + _, ranked = _decode_inputs(width, form, vis, seed * 10 + it, out=bufs) + lens.copy_(vis) + return _reference(ranked, vis) + + prev = load(0) + assert int(lens.max()) == width + eager0 = bool(torch.equal(op(static, off, lens, None, K).long(), prev)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(static, off, lens, None, K) + replays = [] + for it in range(1, 4): + ref = load(it) + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + got = out.clone() + replays.append( + { + **( + { + "in": hashlib.sha256( + (_fingerprint(read) + _sha(lens)).encode() + ).hexdigest(), + "out": _sha(got), + } + if digests + else {} + ), + "exact": bool(torch.equal(got.long(), ref)), + "eager": bool(torch.equal(got, op(static, off, lens, None, K))), + "stale": int((got.long() != prev).any(1).sum()), + } + ) + prev = ref + del graph + return eager0, replays + + +@pytest.mark.parametrize("width,ratio,form,next_n,batch", GRAPH_CASES, ids=GRAPH_IDS) +def test_cuda_graph_engine_capture_set(width, ratio, form, next_n, batch): + """The capture set at every form: the call captured after one eager call at its shape, decode + (next_n 1, G rows) and DSpark verify (next_n 6, 6G rows) at every G in Gset. Then 3 replays + on new lengths (0/1, the width, a padding request) and new mixed-tie scores written in place. + Each replay equals the reference exactly and an eager call on the same inputs bitwise, and + differs from the previous step's reference (a replay re-reading the captured step is seen).""" + rows = batch * next_n + eager0, replays = _capture_case(width, ratio, form, next_n, batch, digests=False) + assert eager0 + for it, rec in enumerate(replays, 1): + print( + f"[graph {form} w{width} r{ratio} next_n{next_n} rows{rows} replay {it}] rows differing " + f"from the previous step's reference: {rec['stale']} of {rows}" + ) + assert rec["exact"], f"replay {it}" + assert rec["eager"], f"replay {it} != eager" + assert rec["stale"] > 0 + + +# The cross-process capture digests: every form at G 1, 7, 40 and 256 (decode and verify; up to +# 1536 rows), plus the whole 8192-token-envelope set (G 1..8, 16, 24, 32) at its 16 forms. +_OLD_WIDTHS = set(ENGINE_WIDTHS[:11]) +DIGEST_GRAPH = [ + (gid, c) + for gid, c in zip(GRAPH_IDS, GRAPH_CASES) + if c[4] in (1, 7, 40, 256) + or ((c[0], c[1]) in _OLD_WIDTHS and c[4] in (1, 2, 3, 4, 5, 6, 7, 8, 16, 24, 32)) +] + + +def _capture_digests(): + """{case id: [eager call exact, replay records]} over the digest capture cases (see _capture_case).""" + res = {} + for gid, case in DIGEST_GRAPH: + eager0, replays = _capture_case(*case) + res[gid] = [eager0, replays] + torch.cuda.empty_cache() + return res + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs") +def test_capture_set_replays_identical_across_processes_and_gpus(): + """The digest capture cases (all 59 forms at G 1, 7, 40, 256 and the 8192-token-envelope set at + its 16 forms; 760 captures) are run again here and in 3 fresh processes on each of 2 GPUs: the + replays' output digests are identical in all 7 processes (and every replay equals the reference + and the eager call on the same inputs), and so are the input digests, so no difference could + hide behind a changed input.""" + torch.cuda.empty_cache() # the fresh processes share this GPU + local = _capture_digests() + assert len(local) == 760 + bad = [ + k + for k, (e0, recs) in local.items() + if not (e0 and all(r["exact"] and r["eager"] for r in recs)) + ] + assert not bad, bad[:10] + devs = _visible_devices() + assert len(devs) >= 2, f"this test needs 2 GPUs, got CUDA_VISIBLE_DEVICES={devs}" + tmp = tempfile.mkdtemp(prefix="fitopk_") + procs = [] + for d in devs[:2]: + for j in range(3): + path = os.path.join(tmp, f"capture_{d}_{j}.json") + env = _script_env(d) + p = subprocess.Popen( + [sys.executable, __file__, "--capture-digests", path], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + procs.append((d, j, path, p)) + uuids = {} + for d, j, path, p in procs: + log, _ = p.communicate() + assert p.returncode == 0, f"capture-digest process {d}/{j} failed:\n{log[-3000:]}" + with open(path) as f: + got = json.load(f) + uuids.setdefault(d, set()).add(got["uuid"]) + differing = sorted(k for k in local if got["cases"].get(k) != local[k]) + print( + f"process {j} on GPU {d} ({got['name']} {got['uuid']}): {len(got['cases'])} capture cases, " + f"differing from this process: {differing[:10]}" + ) + assert not differing and set(got["cases"]) == set(local) + assert len({u for s in uuids.values() for u in s}) == 2, uuids + n_replays = sum(len(recs) for _, recs in local.values()) + print( + f"{len(local)} capture cases ({n_replays} replays) bitwise identical here and in 6 fresh processes " + f"on GPUs {uuids}" + ) + + +# ─── mixed steps: the target's views ────────────────────────────────────────────────────────── +MAX_TOKENS = 16384 # the envelope's max_num_tokens: the metadata's per-token buffers hold this many + + +@pytest.mark.parametrize("tc_at", ["1", "2", "573", "7674", "8193", "max-1", "max"]) +@pytest.mark.parametrize("next_n", [1, 6]) +@pytest.mark.parametrize("gen_reqs", [32, 256]) +def test_mixed_step_views(tc_at, next_n, gen_reqs): + """In a step with tc context tokens and gen = G x next_n generation rows (G 32 or 256; T = tc + + gen <= 16384) the target passes lengths as views of the metadata's 16384-entry + compressed_visible_lens_cuda[r] ([:tc] for the context call, [tc:T] for the generation call, whose + base is only 4-byte aligned at odd tc) and offsets as [:rows] views of one 16384-entry zero + buffer. Both calls equal the same calls on contiguous copies bitwise, and neither buffer is + written.""" + gen = gen_reqs * next_n + tc = {"max": MAX_TOKENS - gen, "max-1": MAX_TOKENS - 1 - gen}.get(tc_at) or int(tc_at) + T = tc + gen + width, ratio = 16398, 1 + g = torch.Generator().manual_seed(tc * 7 + next_n + gen_reqs) + vis_all = torch.randint(0, 8192, (MAX_TOKENS,), generator=g).int().to(DEV) + vis_all[:tc] = ( + torch.arange(tc, device=DEV) + 1 + ).int() # one context request at positions 0 .. tc-1 + vis_all[tc:T] = ( + _dspark_vis(gen_reqs, ratio, width, g, padding=False) + if next_n == 6 + else _decode_vis(gen_reqs, ratio, width, g, padding=False) + ) + zeros_all = _zeros(MAX_TOKENS) + keep_v, keep_z = vis_all.clone(), zeros_all.clone() + s_gen, _ = _decode_inputs(width, "paged_view", vis_all[tc:T].clone(), seed=tc) + got = op(s_gen, zeros_all[:gen], vis_all[tc:T], None, K) + assert torch.equal(got, op(s_gen, _zeros(gen), vis_all[tc:T].clone(), None, K)) + ks = _zeros(tc) + buf = torch.full((tc, -(-(tc + 256) // 256) * 256), float("nan"), device=DEV) + buf[:, :tc] = torch.where( + torch.arange(tc, device=DEV)[None, :] < vis_all[:tc, None].long(), + _mixed(tc, tc, vis_all[:tc], seed=tc + 1), + float("nan"), + ) + ctx = op(buf[:, :tc], zeros_all[:tc], vis_all[:tc], ks, K) + assert torch.equal(ctx, op(buf[:, :tc], _zeros(tc), vis_all[:tc].clone(), ks.clone(), K)) + assert torch.equal(ctx.long(), _reference(buf[:, :tc], vis_all[:tc])) + assert torch.equal(vis_all, keep_v) and torch.equal(zeros_all, keep_z) + + +# ─── context rows: per-request ranges over the prefill logits ───────────────────────────────── +# The target calls the entry once per index-source layer in every step with context requests (never +# captured: a step with a context request is not CUDA-graph eligible) with scores = the prefill +# logits [tc, S] (a view with row stride align(S + 256, 256), unspecified outside each row's range), +# or at re-index layers when S > 16384 the contiguous torch.where(candidate_mask, logits, -inf) +# copy; row_starts = ks (the request's first column); lengths = compressed_visible_lens_cuda[r][:tc] +# = ke - ks = (p + 1) // r; offsets = zeros. S = sum of kv_len // r over the step's requests. + + +def _packed_rows(ratio, reqs): + """The target's rows of one context step, requests as (cached, new) in batch order: request b's + compressed entries are columns [cu_b, cu_b + kv_b // ratio) and its token at position p scores + [ks, ke) = [cu_b, cu_b + (p + 1) // ratio). Returns (ks, ke, S).""" + ks, ke, col0 = [], [], 0 + for cached, new in reqs: + kv = cached + new + p = torch.arange(cached, kv, device=DEV) + ks.append(torch.full((new,), col0, dtype=torch.int32, device=DEV)) + ke.append((col0 + torch.div(p + 1, ratio, rounding_mode="floor")).to(torch.int32)) + col0 += kv // ratio + return torch.cat(ks), torch.cat(ke), col0 + + +def _packed_logits(reqs, ratio, contiguous, seed, kind="mixed"): + """[tc, S] scores of a packed step: each row's [ks, ke) holds scores of `kind`, every other + column NaN. contiguous=False gives the prefill-logits layout (row stride align(S + 256, 256)). + A request whose rows x width exceed one chunk is filled a row chunk at a time.""" + ks, ke, s = _packed_rows(ratio, reqs) + tc = ks.numel() + stride = s if contiguous else -(-(s + 256) // 256) * 256 + buf = torch.full((tc, max(stride, 1)), float("nan"), device=DEV) + row0 = 0 + for b, (cached, new) in enumerate(reqs): + width = (cached + new) // ratio + if width: + c0 = int(ks[row0]) + col = torch.arange(width, device=DEV)[None, :] + step = new if new * width <= _CHUNK else _row_step(width) + for a in range(0, new, step): + e = min(new, a + step) + rs = slice(row0 + a, row0 + e) + lens = ke[rs] - ks[rs] + sd = seed * 64 + b if step == new else (seed * 64 + b) * 8192 + a // step + vals = ( + _mixed(e - a, width, lens, sd) + if kind == "mixed" + else _scores(kind, e - a, width, lens, sd) + ) + buf[rs, c0 : c0 + width] = torch.where( + col < lens[:, None].long(), vals, float("nan") + ) + del vals + row0 += new + return buf[:, :s], ks, ke + + +def _packed_bad_rows(logits, ks, ke, reqs, ratio, out, ref=None): + """Rows whose output differs from the reference over each row's own [ks, ke) span (the call's + ``op.reference``, or ``ref`` when given), and the step's straddling rows.""" + if ref is None: + ref = op.reference(logits, _zeros(ks.numel()), ke - ks, ks, K) + bad = int((out != ref).any(1).sum()) + straddle, row0 = 0, 0 + for cached, new in reqs: + width = (cached + new) // ratio + rs = slice(row0, row0 + new) + if width > K: + c0 = int(ks[row0]) + win = logits[rs, c0 : c0 + width] + straddle += _straddling_rows(win, ke[rs] - ks[rs], ref[rs].long()) + row0 += new + return bad, straddle + + +def _packed_check(logits, ks, ke, reqs, ratio, out): + """The call's output gated by the entry against its reference; returns the straddling rows.""" + ref = op.reference(logits, _zeros(ks.numel()), ke - ks, ks, K) + op.compare(out, ref) + return _packed_bad_rows(logits, ks, ke, reqs, ratio, out, ref)[1] + + +def _run_packed(logits, ks, ke): + return _call(logits, ke - ks, row_starts=ks) + + +# 16384 context tokens of a mixed chunked-prefill step (max_num_tokens 16384). The first 8192: a +# fresh 3000-token prompt, a chunk on an 8192-token cached prefix, 1- and 5-token contexts, contexts +# whose rows all see <= 512, and a chunk ending at position 16383. Then a fresh 4096-token prompt, +# chunks ending at positions 32767, 131071 and 65549 (long-context max_seq_len values; 65550 is the +# 64K DSpark manager length) and a 1-token context. Columns start at unaligned offsets. +FULL_CONTEXT_STEP = [ + (0, 3000), + (8192, 2048), + (0, 1), + (0, 129), + (130, 5), + (2, 1600), + (4000, 430), + (15405, 979), + (0, 4096), + (32768 - 1000, 1000), + (131072 - 2048, 2048), + (3, 1), + (65550 - 1047, 1047), +] + + +@pytest.mark.parametrize("ratio", [1, 2]) +def test_rows_prefix_bitwise(ratio): + """Every row count of the context rows 1..16384: the 16384-row call equals the reference exactly on + every row, and a call on the first r rows (same scores, ks, lengths) equals the first r rows of + that call bitwise, for every r in 1..16384. A one-row-shifted call fails the same comparison.""" + ks, ke, s = _packed_rows(ratio, FULL_CONTEXT_STEP) + assert ks.numel() == MAX_TOKENS + logits, ks, ke = _packed_logits(FULL_CONTEXT_STEP, ratio, contiguous=False, seed=5000 + ratio) + lens = ke - ks + full = _call(logits, lens, row_starts=ks) + bad, straddle = _packed_bad_rows(logits, ks, ke, FULL_CONTEXT_STEP, ratio, full) + assert bad == 0 + assert int(lens.min()) <= 1 and int((lens <= K).sum()) > 0 and int((lens > K).sum()) > 0 + differing = torch.zeros((), dtype=torch.int64, device=DEV) + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for r in range(1, MAX_TOKENS + 1): + differing += (_call(logits[:r], lens[:r], row_starts=ks[:r]) != full[:r]).any(1).sum() + torch.cuda.synchronize() + print( + f"ratio={ratio} S={s} lengths {int(lens.min())}..{int(lens.max())}: straddling tie rows {straddle}; " + f"prefixes r=1..{MAX_TOKENS}: rows differing from the {MAX_TOKENS}-row call {int(differing)}" + ) + assert int(differing) == 0 + _assert_variants(prof, MAX_TOKENS, f"prefix sweep ratio {ratio}") + for r in (1, 2, 3, 100, 513, 4096, 8191, 8192, 12000, 16383): + shifted = _call(logits[1 : r + 1], lens[1 : r + 1], row_starts=ks[1 : r + 1]) + assert not torch.equal(shifted, full[:r]), f"shifted call at r={r} not seen" + + +# Every per-row length: the 16384-token chunks of one request at the 256K DSpark manager length +# (262158 tokens; the last chunk has 14), so ratio-1 rows see every length 1..262158 and ratio-2 rows +# every length 0..131079. +EVERY_LEN_LM = 262158 +EVERY_LEN_CHUNKS = -(-EVERY_LEN_LM // MAX_TOKENS) + + +def _drive_every_len(ratio, chunk): + """The 17 chunks of one 262158-token request: ratio-1 rows see every length 1..262158 and ratio-2 + rows every length 0..131079, exact.""" + start = chunk * MAX_TOKENS + reqs = [(start, min(MAX_TOKENS, EVERY_LEN_LM - start))] + logits, ks, ke = _packed_logits(reqs, ratio, contiguous=False, seed=6000 + 100 * ratio + chunk) + out = _run_packed(logits, ks, ke) + straddle = _packed_check(logits, ks, ke, reqs, ratio, out) + lens = ke - ks + print( + f"ratio={ratio} chunk={chunk}: lengths {int(lens.min())}..{int(lens.max())}, straddling tie rows {straddle}" + ) + assert int(lens.max()) == (start + reqs[0][1]) // ratio + + +def _drive_widest(ratio): + """The widest step the gate configuration can build: 32 context requests (max_batch_size), each + a 256-token chunk ending at max_seq_len 16384, so 8192 rows over S = 32 * 16384 // ratio columns + in the prefill-logits layout (8192 x 524544 fp32 = 17.2 GB at ratio 1). The last row starts past + element 2**32 at ratio 1 and past 2**31 at ratio 2. Exact.""" + reqs = [(16384 - 256, 256)] * 32 + logits, ks, ke = _packed_logits(reqs, ratio, contiguous=False, seed=7000 + ratio) + stride = logits.stride(0) + assert (8192 - 1) * stride > (2**32 if ratio == 1 else 2**31) + out = _run_packed(logits, ks, ke) + straddle = _packed_check(logits, ks, ke, reqs, ratio, out) + print( + f"widest step ratio={ratio}: S={logits.shape[1]} row stride={stride}, straddling tie rows {straddle}" + ) + + +def _ctx_candidate_keep(ks, ke, s, seed): + """Keep mask [tc, S] of a re-index layer's context rows: a row seeing more than 16384 positions + keeps its newest 8-wide block plus 2047 older ones, a shorter row keeps its whole range.""" + lens = (ke - ks).long() + keep = torch.zeros(ks.numel(), s, dtype=torch.bool, device=DEV) + col = torch.arange(s, device=DEV)[None, :] + inside = (col >= ks[:, None].long()) & (col < ke[:, None].long()) + for c0 in ks.unique().tolist(): + rows = (ks == c0).nonzero().flatten() + w = int(lens[rows].max()) + local = _candidate_keep(lens[rows].int(), w, _gen(seed + c0)) + keep[rows, c0 : c0 + w] = local + return keep & inside + + +def _masked_copy(logits, ks, ke, seed): + """The re-index layers' contiguous candidate-masked copy of a one-request step's logits (every + row starts at column 0), built a row chunk at a time: a row seeing more than 16384 positions + keeps its newest 8-wide block plus 2047 older ones, -inf elsewhere and past the row.""" + tc, s = logits.shape + assert int(ks.max()) == 0 + lens = (ke - ks).int() + out = torch.empty(tc, s, device=DEV) + step = _row_step(s) + for a in range(0, tc, step): + e = min(tc, a + step) + out[a:e] = torch.where( + _candidate_keep(lens[a:e], s, _gen(seed + a)), logits[a:e], float("-inf") + ) + return out + + +def _drive_long_ctx(form): + """The 24K passkey step from the runtime dump, packed: 7674 rows x 24058 columns, a 4001-token + chunk on a 16384-token cached prefix (every row sees more than 16384 positions), then a fresh + 3673-token prompt. In the prefill-logits layout, and as the re-index layers' contiguous + candidate-masked copy (-inf outside the kept blocks). Exact.""" + reqs = [(16384, 4001), (0, 3673)] + logits, ks, ke = _packed_logits(reqs, 1, contiguous=False, seed=8100) + assert logits.shape == (7674, 24058) + if form == "candidate_masked": + keep = _ctx_candidate_keep(ks, ke, logits.shape[1], seed=8200) + logits = torch.where(keep, logits, float("-inf")) + assert logits.stride() == (24058, 1) + out = _run_packed(logits, ks, ke) + straddle = _packed_check(logits, ks, ke, reqs, 1, out) + print(f"24K step {form}: straddling tie rows {straddle}") + + +# The max_seq_len axis's far end: the last 16384-token chunk of one request at each manager length Lm +# of the envelope (L + 14 with DSpark, L without), so its rows see up to Lm // r compressed positions, +# in the prefill-logits layout and, at ratio 1 above 16384 positions, as the re-index layers' +# candidate-masked copy. +# Context row counts in the max_num_tokens-16384 engine dumps (prefill top-k records). +DUMPED_CTX_ROWS = [ + 1, + 2, + 5, + 12, + 16, + 28, + 128, + 597, + 2252, + 4553, + 8147, + 8192, + 8317, + 9055, + 9108, + 9392, + 10441, + 10444, + 10451, + 16292, + 16327, + 16328, + 16378, + 16384, +] + + +def _relaxed_prefix_rows(n): + """The row counts swept where a call takes over 1 ms, in place of every row count: every r in + 1..256, every r = {0, 1, -1} mod the SM count (the kernels run one block per row, so the SM count + sets the wave boundaries) and mod 256, up to n, and the dumped row counts.""" + sms = torch.cuda.get_device_properties(0).multi_processor_count + rows = set(range(1, 257)) | set(DUMPED_CTX_ROWS) + for q in (sms, 256): + for m in range(q, n + q, q): + rows |= {m - 1, m, m + 1} + return sorted(r for r in rows if 1 <= r <= n) + + +def _drive_far_end(lm, ratio, form): + """The last 16384-token chunk of a request at manager length Lm: 16384 rows seeing + (Lm - 16383) // r .. Lm // r compressed positions over S = Lm // r columns (16384 x 1049088 + fp32 = 64 GiB at Lm 1048590), exact, one launch of each kernel. At the widest (Lm 1048590, ratio + 1, a call of about 0.1 s) the row-prefix set for calls over 1 ms (_relaxed_prefix_rows) is bitwise + equal to the leading rows of the 16384-row call.""" + reqs = [(lm - MAX_TOKENS, MAX_TOKENS)] + logits, ks, ke = _packed_logits(reqs, ratio, contiguous=False, seed=9000 + lm + ratio) + lens = ke - ks + assert int(lens.max()) == lm // ratio and logits.shape == (MAX_TOKENS, lm // ratio) + if form == "candidate_masked": + logits = _masked_copy(logits, ks, ke, seed=9100 + lm) + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + full = _call(logits, lens, row_starts=ks) + torch.cuda.synchronize() + _assert_variants(prof, 1, f"far end Lm {lm} ratio {ratio} {form}") + straddle = _packed_check(logits, ks, ke, reqs, ratio, full) + print( + f"far end Lm={lm} r={ratio} {form}: S={logits.shape[1]} row stride={logits.stride(0)} " + f"({logits.shape[0] * logits.stride(0) * 4 / 2**30:.1f} GiB), lengths {int(lens.min())}..{int(lens.max())}, " + f"straddling tie rows {straddle}" + ) + if (lm, ratio, form) == (1048590, 1, "mqa_view"): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + _call(logits, lens, row_starts=ks) + end.record() + torch.cuda.synchronize() + rows = _relaxed_prefix_rows(MAX_TOKENS) + differing = torch.zeros((), dtype=torch.int64, device=DEV) + for r in rows: + differing += (_call(logits[:r], lens[:r], row_starts=ks[:r]) != full[:r]).any(1).sum() + print( + f"far end Lm={lm}: the 16384-row call takes {start.elapsed_time(end):.1f} ms; {len(rows)} prefixes " + f"(r {rows[0]}..{rows[-1]}): rows differing from the 16384-row call {int(differing)}" + ) + assert int(differing) == 0 + del logits, full + torch.cuda.empty_cache() + + +# ─── context rows in row tiles ───────────────────────────────────────────────────────────────── +# At long contexts the target runs a step's context rows in tiles (`_index_row_tiles` in the +# target's modeling.py): tile [a, b) gets its own fresh prefill-logits buffer [b - a, S] (row stride +# align(S + 256, 256)), zeros[:b - a], and the step's lengths[a:b] and ks[a:b] (views at row offset +# a), and at re-index layers the tile's candidate-masked copy. A tile holds at most 2**28 fp32 +# logits: the whole step when it fits, otherwise floor(2**28 / row stride) rows, rounded down to a +# multiple of 64 when that is at least 64. +TILE_ELEMS = 1 << 28 + + +def _row_tiles(rows, s_len): + stride = -(-(s_len + 256) // 256) * 256 + per = max(1, TILE_ELEMS // stride) + if per >= rows: + return [(0, rows)] + if per >= 64: + per -= per % 64 + return [(a, min(rows, a + per)) for a in range(0, rows, per)] + + +def _tile_logits(reqs, ratio, ks, ke, a, b, seed, kind="mixed"): + """Rows [a, b) of a packed step as a fresh prefill-logits buffer [b - a, S] (row stride + align(S + 256, 256)): each row's [ks, ke) holds scores of `kind`, every other column NaN.""" + s = sum((c + n) // ratio for c, n in reqs) + buf = torch.full((b - a, -(-(s + 256) // 256) * 256), float("nan"), device=DEV) + row0 = 0 + for j, (cached, new) in enumerate(reqs): + lo, hi, width = max(a, row0), min(b, row0 + new), (cached + new) // ratio + if lo < hi and width: + c0 = int(ks[lo]) + col = torch.arange(width, device=DEV)[None, :] + step = _row_step(width) + for x in range(lo, hi, step): + y = min(hi, x + step) + lens = ke[x:y] - ks[x:y] + sd = (seed * 1024 + j) * 4096 + (x - lo) // step + # a one-row chunk (one-row tiles) takes the kinds in turn by row, not always the first + k = KINDS[x % len(KINDS)] if kind == "mixed" and y - x == 1 else kind + vals = ( + _mixed(y - x, width, lens, sd) + if k == "mixed" + else _scores(k, y - x, width, lens, sd) + ) + buf[x - a : y - a, c0 : c0 + width] = torch.where( + col < lens[:, None].long(), vals, float("nan") + ) + del vals + row0 += new + return buf[:, :s] + + +def _tile_bad_rows(logits, ks, lens, out, ref=None): + """Rows of one call whose output differs from the reference over the row's own range (the + call's ``op.reference``, or ``ref`` when given), and its straddling rows; rows are grouped by + their column start (one request's rows each).""" + n = ks.numel() + if ref is None: + ref = op.reference(logits, _zeros(n), lens, ks, K) + bad = int((out != ref).any(1).sum()) + cut = ((ks[1:] != ks[:-1]).nonzero().flatten() + 1).tolist() + straddle = 0 + for p, q in zip([0] + cut, cut + [n]): + w, c0 = int(lens[p:q].max()), int(ks[p]) + if w > K: + win = logits[p:q, c0 : c0 + w] + straddle += _straddling_rows(win, lens[p:q], ref[p:q].long()) + return bad, straddle + + +def _run_tiles(reqs, ratio, masked, seed, only_reqs=None): + """Every tile of a packed step (with `only_reqs`, the tiles holding those requests' rows), each + called as the target calls it and gated by the entry; returns (tiles, rows, rows wrong, + straddling rows, S).""" + ks, ke, s = _packed_rows(ratio, reqs) + lens = ke - ks + zeros = _zeros(MAX_TOKENS) + tiles = _row_tiles(ks.numel(), s) + if only_reqs is not None: + starts = [sum(n for _, n in reqs[:j]) for j in range(len(reqs))] + keep = [(starts[j], starts[j] + reqs[j][1]) for j in only_reqs] + tiles = [(a, b) for a, b in tiles if any(a < hi and lo < b for lo, hi in keep)] + rows = bad = straddle = 0 + for t, (a, b) in enumerate(tiles): + logits = _tile_logits(reqs, ratio, ks, ke, a, b, seed * 100003 + t) + if masked: + logits = torch.where( + _ctx_candidate_keep(ks[a:b], ke[a:b], s, seed + a), logits, float("-inf") + ) + out = op(logits, zeros[: b - a], lens[a:b], ks[a:b], K) + ref = op.reference(logits, zeros[: b - a], lens[a:b], ks[a:b], K) + op.compare(out, ref) + bd, st = _tile_bad_rows(logits, ks[a:b], lens[a:b], out, ref) + rows, bad, straddle = rows + b - a, bad + bd, straddle + st + del logits, out + return tiles, rows, bad, straddle, s + + +# Tiled steps (the cells' `reqs`, as (cached, new) requests): the last 16384-token chunk at the 1M +# manager lengths (192-row tiles at ratio 1, 448 at ratio 2); an 8192-token chunk ending at 1048219, +# the 1M engine dump's longest context (256-row tiles); a 16384-token prompt (S 16384: tiles of +# 16128 and 256 rows); tiles crossing request boundaries at mixed long contexts; two 8192-token +# chunks ending at 1048590 (S 2,097,180: 64-row tiles); and 256 requests of 64 tokens ending at +# 1048590 (S 268,439,040 at ratio 1 and 134,219,520 at ratio 2, both above 2**28 / 2 columns: +# one-row tiles, of which those holding the rows of requests 0, 128 and 255 run). +def _drive_tiles(name, reqs, ratio, form): + """Each tile of a long-context step called as the target calls it (its own padded logits + buffer, or the candidate-masked copy, zeros[:rows], lengths[a:b] and ks[a:b] views at row + offset a): every row of every tile equals the reference exactly.""" + reqs = [tuple(r) for r in reqs] + only = (0, 128, 255) if name == "b256_1M" else None + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + tiles, rows, bad, straddle, s = _run_tiles( + reqs, ratio, form == "candidate_masked", 9500 + ratio, only + ) + torch.cuda.synchronize() + sizes = sorted({b - a for a, b in tiles}) + print( + f"[tiles {name} r{ratio} {form}] S={s}: {len(tiles)} tiles of {sizes} rows ({rows} rows), " + f"straddling tie rows {straddle}" + ) + _assert_variants(prof, len(tiles), f"tiles {name} r{ratio} {form}") + assert all((b - a) * (-(-(s + 256) // 256) * 256) <= TILE_ELEMS or b - a == 1 for a, b in tiles) + assert bad == 0 + + +# The context top-k's (rows, S, ratio) cells in the long-context engine dumps (tep4, DSpark, +# max_num_tokens 8192, max_batch_size 256, max_seq_len 131072, 262144 and 1048576; the target's +# row-tiled prefill), the ratio from the prefill-logits record of the same (rows, S); a cell seen at +# both ratios runs at both. Each runs as one tile: `rows` consecutive positions ending a request of +# S * ratio tokens. +DUMPED_TILE_CELLS = [ + (1, 1, 1), + (2, 1, 2), + (2, 2, 1), + (64, 32768, 1), + (64, 32768, 2), + (128, 8192, 1), + (128, 65536, 1), + (128, 65536, 2), + (128, 98304, 1), + (128, 98304, 2), + (128, 130773, 1), + (128, 192512, 2), + (128, 196608, 1), + (128, 196608, 2), + (128, 221184, 1), + (128, 221184, 2), + (128, 225280, 2), + (128, 229376, 1), + (128, 229376, 2), + (128, 261860, 1), + (128, 282624, 2), + (128, 286720, 1), + (128, 286720, 2), + (128, 290816, 2), + (128, 294912, 1), + (128, 294912, 2), + (128, 299008, 2), + (128, 421888, 2), + (128, 425984, 1), + (128, 425984, 2), + (128, 430080, 2), + (128, 434176, 1), + (128, 434176, 2), + (128, 438272, 2), + (128, 442368, 1), + (128, 442368, 2), + (128, 446464, 2), + (128, 450560, 1), + (128, 450560, 2), + (128, 454656, 2), + (128, 458752, 1), + (128, 458752, 2), + (128, 462848, 2), + (128, 523951, 1), + (128, 524288, 1), + (128, 532480, 1), + (128, 540672, 1), + (128, 548864, 1), + (128, 557056, 1), + (128, 565248, 1), + (128, 573440, 1), + (128, 581632, 1), + (128, 589824, 1), + (128, 598016, 1), + (128, 606208, 1), + (128, 614400, 1), + (128, 622592, 1), + (128, 630784, 1), + (128, 638976, 1), + (128, 647168, 1), + (128, 655360, 1), + (128, 663552, 1), + (128, 671744, 1), + (128, 679936, 1), + (128, 688128, 1), + (128, 696320, 1), + (128, 1048219, 1), + (155, 1048219, 1), + (175, 261975, 2), + (175, 523951, 1), + (192, 163840, 1), + (192, 163840, 2), + (192, 704512, 1), + (192, 712704, 1), + (192, 720896, 1), + (192, 729088, 1), + (192, 737280, 1), + (192, 745472, 1), + (192, 753664, 1), + (192, 761856, 1), + (192, 770048, 1), + (192, 778240, 1), + (192, 786432, 1), + (192, 794624, 1), + (192, 802816, 1), + (192, 811008, 1), + (192, 819200, 1), + (192, 827392, 1), + (192, 835584, 1), + (219, 524109, 2), + (256, 131072, 1), + (256, 131072, 2), + (256, 843776, 1), + (256, 851968, 1), + (256, 860160, 1), + (256, 868352, 1), + (256, 876544, 1), + (256, 884736, 1), + (256, 892928, 1), + (256, 901120, 1), + (256, 909312, 1), + (256, 917504, 1), + (256, 925696, 1), + (256, 933888, 1), + (256, 942080, 1), + (256, 950272, 1), + (256, 958464, 1), + (256, 966656, 1), + (256, 974848, 1), + (256, 983040, 1), + (256, 991232, 1), + (256, 999424, 1), + (256, 1007616, 1), + (256, 1015808, 1), + (256, 1024000, 1), + (256, 1032192, 1), + (256, 1040384, 1), + (256, 1048219, 1), + (320, 704512, 1), + (320, 712704, 1), + (320, 720896, 1), + (320, 729088, 1), + (320, 737280, 1), + (320, 745472, 1), + (320, 753664, 1), + (320, 761856, 1), + (320, 770048, 1), + (320, 778240, 1), + (320, 786432, 1), + (320, 794624, 1), + (320, 802816, 1), + (320, 811008, 1), + (320, 819200, 1), + (320, 827392, 1), + (320, 835584, 1), + (384, 606208, 1), + (384, 614400, 1), + (384, 622592, 1), + (384, 630784, 1), + (384, 638976, 1), + (384, 647168, 1), + (384, 655360, 1), + (384, 663552, 1), + (384, 671744, 1), + (384, 679936, 1), + (384, 688128, 1), + (384, 696320, 1), + (448, 352256, 1), + (448, 352256, 2), + (448, 356352, 2), + (448, 360448, 1), + (448, 360448, 2), + (448, 364544, 2), + (448, 368640, 1), + (448, 368640, 2), + (448, 372736, 2), + (448, 376832, 1), + (448, 376832, 2), + (448, 380928, 2), + (448, 524109, 2), + (448, 524288, 1), + (448, 532480, 1), + (448, 540672, 1), + (448, 548864, 1), + (448, 557056, 1), + (448, 565248, 1), + (448, 573440, 1), + (448, 581632, 1), + (448, 589824, 1), + (448, 598016, 1), + (512, 69632, 2), + (512, 102400, 2), + (512, 135168, 2), + (512, 139264, 1), + (512, 139264, 2), + (512, 167936, 2), + (512, 172032, 1), + (512, 172032, 2), + (512, 200704, 2), + (512, 204800, 1), + (512, 204800, 2), + (512, 208896, 2), + (512, 262144, 1), + (512, 262144, 2), + (512, 266240, 2), + (512, 270336, 1), + (512, 270336, 2), + (512, 274432, 2), + (512, 278528, 1), + (512, 278528, 2), + (512, 323584, 2), + (512, 327680, 1), + (512, 327680, 2), + (512, 331776, 2), + (512, 335872, 1), + (512, 335872, 2), + (512, 339968, 2), + (512, 344064, 1), + (512, 344064, 2), + (512, 348160, 2), + (512, 385024, 1), + (512, 385024, 2), + (512, 389120, 2), + (512, 393216, 1), + (512, 393216, 2), + (512, 397312, 2), + (512, 401408, 1), + (512, 401408, 2), + (512, 405504, 2), + (512, 409600, 1), + (512, 409600, 2), + (512, 413696, 2), + (512, 417792, 1), + (512, 417792, 2), + (512, 466944, 1), + (512, 466944, 2), + (512, 471040, 2), + (512, 475136, 1), + (512, 475136, 2), + (512, 479232, 2), + (512, 483328, 1), + (512, 483328, 2), + (512, 487424, 2), + (512, 491520, 1), + (512, 491520, 2), + (512, 495616, 2), + (512, 499712, 1), + (512, 499712, 2), + (512, 503808, 2), + (512, 507904, 1), + (512, 507904, 2), + (512, 512000, 2), + (512, 516096, 1), + (512, 516096, 2), + (512, 520192, 2), + (512, 523951, 1), + (576, 233472, 2), + (576, 237568, 1), + (576, 237568, 2), + (576, 241664, 2), + (576, 245760, 1), + (576, 245760, 2), + (576, 421888, 2), + (576, 425984, 1), + (576, 425984, 2), + (576, 430080, 2), + (576, 434176, 1), + (576, 434176, 2), + (576, 438272, 2), + (576, 442368, 1), + (576, 442368, 2), + (576, 446464, 2), + (576, 450560, 1), + (576, 450560, 2), + (576, 454656, 2), + (576, 458752, 1), + (576, 458752, 2), + (576, 462848, 2), + (640, 385024, 1), + (640, 385024, 2), + (640, 389120, 2), + (640, 393216, 1), + (640, 393216, 2), + (640, 397312, 2), + (640, 401408, 1), + (640, 401408, 2), + (640, 405504, 2), + (640, 409600, 1), + (640, 409600, 2), + (640, 413696, 2), + (640, 417792, 1), + (640, 417792, 2), + (704, 106496, 1), + (704, 106496, 2), + (704, 303104, 1), + (704, 303104, 2), + (704, 307200, 2), + (704, 311296, 1), + (704, 311296, 2), + (704, 315392, 2), + (704, 319488, 1), + (704, 319488, 2), + (704, 352256, 1), + (704, 352256, 2), + (704, 356352, 2), + (704, 360448, 1), + (704, 360448, 2), + (704, 364544, 2), + (704, 368640, 1), + (704, 368640, 2), + (704, 372736, 2), + (704, 376832, 1), + (704, 376832, 2), + (704, 380928, 2), + (740, 261860, 1), + (768, 143360, 2), + (768, 323584, 2), + (768, 327680, 1), + (768, 327680, 2), + (768, 331776, 2), + (768, 335872, 1), + (768, 335872, 2), + (768, 339968, 2), + (768, 344064, 1), + (768, 344064, 2), + (768, 348160, 2), + (832, 176128, 2), + (832, 180224, 1), + (832, 180224, 2), + (832, 303104, 1), + (832, 303104, 2), + (832, 307200, 2), + (832, 311296, 1), + (832, 311296, 2), + (832, 315392, 2), + (832, 319488, 1), + (832, 319488, 2), + (896, 212992, 1), + (896, 212992, 2), + (896, 217088, 2), + (896, 282624, 2), + (896, 286720, 1), + (896, 286720, 2), + (896, 290816, 2), + (896, 294912, 1), + (896, 294912, 2), + (896, 299008, 2), + (960, 261975, 2), + (960, 262144, 1), + (960, 262144, 2), + (960, 266240, 2), + (960, 270336, 1), + (960, 270336, 2), + (960, 274432, 2), + (960, 278528, 1), + (960, 278528, 2), + (1024, 36864, 2), + (1024, 73728, 1), + (1024, 73728, 2), + (1024, 147456, 1), + (1024, 147456, 2), + (1024, 249856, 2), + (1024, 253952, 1), + (1024, 253952, 2), + (1024, 258048, 2), + (1024, 261860, 1), + (1088, 110592, 2), + (1088, 233472, 2), + (1088, 237568, 1), + (1088, 237568, 2), + (1088, 241664, 2), + (1088, 245760, 1), + (1088, 245760, 2), + (1152, 184320, 2), + (1152, 188416, 1), + (1152, 188416, 2), + (1152, 221184, 1), + (1152, 221184, 2), + (1152, 225280, 2), + (1152, 229376, 1), + (1152, 229376, 2), + (1216, 212992, 1), + (1216, 212992, 2), + (1216, 217088, 2), + (1280, 114688, 1), + (1280, 114688, 2), + (1280, 151552, 2), + (1280, 200704, 2), + (1280, 204800, 1), + (1280, 204800, 2), + (1280, 208896, 2), + (1344, 192512, 2), + (1344, 196608, 1), + (1344, 196608, 2), + (1408, 77824, 2), + (1408, 184320, 2), + (1408, 188416, 1), + (1408, 188416, 2), + (1472, 118784, 2), + (1472, 176128, 2), + (1472, 180224, 1), + (1472, 180224, 2), + (1536, 155648, 1), + (1536, 155648, 2), + (1536, 159744, 2), + (1536, 167936, 2), + (1536, 172032, 1), + (1536, 172032, 2), + (1600, 163840, 1), + (1600, 163840, 2), + (1664, 81920, 1), + (1664, 81920, 2), + (1664, 122880, 1), + (1664, 122880, 2), + (1664, 155648, 1), + (1664, 155648, 2), + (1664, 159744, 2), + (1728, 40960, 1), + (1728, 40960, 2), + (1728, 151552, 2), + (1749, 130773, 1), + (1792, 147456, 1), + (1792, 147456, 2), + (1856, 143360, 2), + (1920, 135168, 2), + (1920, 139264, 1), + (1920, 139264, 2), + (1956, 130930, 2), + (1984, 130930, 2), + (1984, 131072, 1), + (1984, 131072, 2), + (2048, 86016, 2), + (2048, 126976, 2), + (2048, 130773, 1), + (2176, 122880, 1), + (2176, 122880, 2), + (2240, 118784, 2), + (2304, 45056, 2), + (2304, 90112, 1), + (2304, 90112, 2), + (2304, 114688, 1), + (2304, 114688, 2), + (2368, 110592, 2), + (2496, 106496, 1), + (2496, 106496, 2), + (2560, 94208, 2), + (2560, 102400, 2), + (2688, 98304, 1), + (2688, 98304, 2), + (2816, 49152, 1), + (2816, 49152, 2), + (2816, 94208, 2), + (2944, 90112, 1), + (2944, 90112, 2), + (3072, 86016, 2), + (3200, 53248, 2), + (3264, 81920, 1), + (3264, 81920, 2), + (3392, 77824, 2), + (3584, 57344, 1), + (3584, 57344, 2), + (3584, 73728, 1), + (3584, 73728, 2), + (3840, 69632, 2), + (3861, 65386, 2), + (3904, 61440, 2), + (4032, 65386, 2), + (4032, 65536, 1), + (4032, 65536, 2), + (4288, 61440, 2), + (4608, 57344, 1), + (4608, 57344, 2), + (4992, 53248, 2), + (5376, 49152, 1), + (5376, 49152, 2), + (5888, 45056, 2), + (6464, 40960, 1), + (6464, 40960, 2), + (7168, 36864, 2), + (8128, 32768, 1), + (8128, 32768, 2), + (8192, 4096, 2), + (8192, 8192, 1), + (8192, 8192, 2), + (8192, 12288, 2), + (8192, 16384, 1), + (8192, 16384, 2), + (8192, 20480, 2), + (8192, 24576, 1), + (8192, 24576, 2), + (8192, 28672, 2), +] + + +def test_dumped_tile_cells(): + """Every dumped tile cell (471), exact; every call launches the same two kernels.""" + assert len(DUMPED_TILE_CELLS) == 471 + zeros = _zeros(MAX_TOKENS) + straddle = 0 + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for i, (rows, s, ratio) in enumerate(DUMPED_TILE_CELLS): + reqs = [(s * ratio - rows, rows)] + ks, ke, s2 = _packed_rows(ratio, reqs) + assert s2 == s and _row_tiles(rows, s) == [(0, rows)] + logits = _tile_logits(reqs, ratio, ks, ke, 0, rows, 60000 + i) + out = op(logits, zeros[:rows], ke - ks, ks, K) + bad, st = _tile_bad_rows(logits, ks, ke - ks, out) + assert bad == 0, (rows, s, ratio) + straddle += st + del logits, out + torch.cuda.synchronize() + print(f"{len(DUMPED_TILE_CELLS)} dumped tile cells: all exact; straddling tie rows {straddle}") + _assert_variants(prof, len(DUMPED_TILE_CELLS), "dumped tile cells") + + +# Random multi-request steps over the envelope: T <= 16384 context tokens in k <= 256 requests, each +# request's context up to the step's manager length Lm (a MAX_SEQ_LENS value, DSpark on or off). The +# test holds at most 2**32 fp32 logits (16 GiB) per random step: a draw whose tc x S exceeds that has +# its cached prefixes scaled down to fit, so large S comes with small tc. +MANAGER_LENS = sorted({L + d for L in MAX_SEQ_LENS for d in (0, 14)}) +MULTI_REQUEST_MAX_REQS = 256 +MULTI_REQUEST_CAP = 2**32 + + +def _multi_request_random_step(rng, ratio): + """A random context step: T in [1, 16384] rows split into k in [1, 256] requests in random order, + each with a random new length and a random cached prefix (kv <= Lm, Lm drawn from MANAGER_LENS), + scaled down if T x S > 2**32. Returns (requests, Lm).""" + t = rng.randint(1, MAX_TOKENS) + k = rng.randint(1, min(MULTI_REQUEST_MAX_REQS, t)) + lm = rng.choice(MANAGER_LENS) + s_max = MULTI_REQUEST_CAP // t + while True: + cuts = sorted(rng.sample(range(1, t), k - 1)) + news = [b - a for a, b in zip([0] + cuts, cuts + [t])] + cached = [rng.randint(0, lm - n) for n in news] + if sum((c + n) // ratio for c, n in zip(cached, news)) > s_max: + f = (s_max * ratio - t) / sum(cached) + cached = [int(c * f) for c in cached] + reqs = list(zip(cached, news)) + s = sum((c + n) // ratio for c, n in reqs) + if s: # the target skips the call when S == 0 + assert s <= s_max + return reqs, lm + + +@pytest.mark.parametrize("ratio", [1, 2]) +def test_multi_request_random_partitions(ratio): + """Random multi-request steps over the context rows: 200 seeded random packed steps per ratio (T + up to 16384, up to 256 requests, contexts up to the envelope's manager lengths), alternating the + two layouts, mixed tie kinds, against the reference exactly; every call launches the same two + kernels. Every 20th step also runs a control with one request's ks shifted by one column, which + the same check must see.""" + rng = random.Random(23000 + ratio) + steps = [_multi_request_random_step(rng, ratio) for _ in range(200)] + controls, rows, straddle, reqs_seen, calls, names = [], 0, 0, [], 0, [] + max_s, max_t, max_len = 0, 0, 0 + for i, (reqs, lm) in enumerate(steps): + logits, ks, ke = _packed_logits(reqs, ratio, contiguous=bool(i % 2), seed=ratio * 1000 + i) + out = _profiled(lambda logits=logits: _run_packed(logits, ks, ke), names) + calls += 1 + bad, st = _packed_bad_rows(logits, ks, ke, reqs, ratio, out) + assert bad == 0, (i, lm, len(reqs), bad) + rows, straddle = rows + ks.numel(), straddle + st + max_s, max_t = max(max_s, logits.shape[1]), max(max_t, ks.numel()) + max_len = max(max_len, int((ke - ks).max())) + reqs_seen.append(len(reqs)) + if i % 20 == 0: + lens = (ke - ks).long() + row0, pick = 0, None + for cached, new in reqs: + if int(lens[row0 : row0 + new].max()) >= 2: + pick = (row0, new) + row0 += new + if pick is not None: + ks_bad = ks.clone() + a, b = pick[0], pick[0] + pick[1] + ks_bad[a:b] += (lens[a:b] >= 1).int() # keeps ks <= ke + wrong = _profiled( + lambda logits=logits: _call(logits, ke - ks_bad, row_starts=ks_bad), names + ) + calls += 1 + controls.append(_packed_bad_rows(logits, ks, ke, reqs, ratio, wrong)[0]) + del logits, out + torch.cuda.empty_cache() + print( + f"ratio={ratio}: {len(steps)} steps, {rows} rows, T up to {max_t}, request counts " + f"{min(reqs_seen)}..{max(reqs_seen)} ({sum(k > 32 for k in reqs_seen)} steps above 32), max S {max_s}, " + f"longest row {max_len}, straddling tie rows {straddle}; control (one request's ks + 1) failing rows " + f"{controls}" + ) + assert max(reqs_seen) > 200 and max_t > 16000 and max_s > 2**22 + assert len(controls) >= 10 and all(c > 0 for c in controls), controls + _assert_variants(names, calls, f"random steps ratio {ratio}") + + +# Every (tc, S) cell of the context top-k in the core runtime-value dumps (boot, GSM8K, 24K passkey, +# gate and accuracy-preview logs), joined with the ratio of the prefill-logits record at the same +# (tc, S) and with a sparse-MLA varlen record (nc requests, max_q_len) at the same tc when one is +# consistent; otherwise the smallest (nc, max_q_len) that is (9 cells). As "tc,S,ratio,nc,max_q_len" +# records, space-separated, zlib-compressed, base64. 1933 cells; 3 dumped cells are garbles outside +# any consistent step and are dropped: (48, 0) and (848, 0) (the target never calls the top-k with +# S = 0) and (81, 6) (81 rows cannot see only 6 compressed columns in <= 32 requests at either ratio). +DUMPED_CELLS = ( + "eNpdnW267CgIhLcyCzg/4iey/40NVIHS95kJN/VqdyfGGNHgaX+N//3XzWKzvQ7W/2vz71BM32+beU3sTNi+nwn7v2l4ImUe" + "F/Z/C6Hjb0x+t+1D77OYDD3/xtiRPqG39Eif/63PPi+KdNuHlsWvh+76NxY/b/vQ0vh56GmfFR6b7UOPqT/at1b01vP0+v5m" + "HJ/tQ584Pur1jm8taFnxfdD7b+wT6RtahkS66W1lc3gNbB/at3a1/s3eI12hz5yRblqsnD5hujTo0yId2srvjEgf0LLj96Ht" + "2E4cvxxo39rVXr75/QotbUW6aT92/Zh+GrRv7WorO43fPxt6nP6jfWtFy55P+3frYn79oGeTH+1bK/qMVbSd/47fs/rmWnqm" + "u7b6J/l9E1pW/D70susf56cL+uw4P2j7rMb1VYH2rV1td8MX9UcP9Gkj0l1b2WuUvyq0by30/tqtP7ZPLetHy2D5Xb2+ov38" + "NfIPaBka6a7tenwj0je0b+1qu/7fjHSF9q2ltnttxPnZPvSO+4/ayzc+b02Ia4n7jdp/P87PWhLXOxoWarHP5/cLtKwR6a79" + "+OL8mkL71lJ3O//G+8v2oX1rV9v923qkL+gd9YPajq9F+fQN7Vu7Wkq6QPvWrrbja3H+XaHn13+0b63o882nre6NFtd7fNC+" + "tatbSW/QvrWr7fq0KL/RoX1rV9v1ifvJ9qGzvqf2rRWd9Z96lc8v6B33H7WXT1zfIdA72j/oaefX4/fmB+1bu9rO78TxzwYt" + "EtcP2s7vxO/PDi0Sxwtt17/H9Z8D2rd2tZ1/tK+2D71nft61Xf8e13du6JnnE9q3VvTJ84O26x/PRtuH3ln/XS8//7h+9jxx" + "vfP+hW4lvUH71q628+9RPqtTR3ub2rdWtBwpepTPD+id9xe0l0+U75rQO9pHai+fKA97nrn2raXedn5Zn/YH7Vu7upX0Rq37" + "R++sb6FFR9F2viOurz0vXc+sT6F3PE9Tn9mKHuXzA9q3drXV7xH1Yy/oGc/D1L61orVVbeUz4nrsDe1bu1pKukDPbI9C73he" + "pz7ZPkEf+3zUr32gfWupxcs/6o980DOvZ2jfWtEnry90v/0z24fO/hm199/y+yb0zvsPepX0BX3PL7RvrejTZ9FeflG/ZEP7" + "1q62+2tE+2/9Ede+tdTHzn9GfTgftG/t6lbSG/SO/gu11acsv9Op568+kvlD61e093+jfp0B7Vu7et7+se1D+9auXiV9QWf/" + "mdrqz4z6cwR65/0NbfVjRv04B/odP7VvreijrWgtn1do31pq9fKN77P+luvZ9UfvLI/QJ/rv1K18vkHv/H1oq38z6p92aN/a" + "1cPqUxyf9edcn57f79rLN+qf9edc+9autvo1o37phvatXS0lXaB39O+oz/UnbB86/QdqLekKvaN/51o+K7/F+mf70L61q1tJ" + "b9C+tautfFaL9A7tW7t6lPQB7Vu72spn9Uif0BL9D2ornzUifUNL9A+prT7dzx9o/Vakm25+fry+tg/tW7u6lfQG7Vu7ul//" + "zvahZzyPU6e/l/pMKXqUzw9oif439bz+je1DS/hP1KukL+j0f6jt+i+N9APtW7taS7pCS/RPoe1ZNXZc3/5BTzk/2rdWtMbz" + "lbqVzzdo39rVVn47rn/v0L61q0dJH9AS/WPqef0n24ee0d9KLdFep9Zoz6lX+fyCTv+L2v3juJ7Wn3YtPc/PtZR0gZboX1Nb" + "+e8Z6Qda4v6n1pKu0BL9bWjvX++of9afdi3Rv6b28ov6Y/1l1xL9Z+pR0gf0jP5Xat9a0Rr+LrWXb5Sn9addz3V+tET/O7W2" + "XfQqn1/QEv1v6jf+YPvQ6S9SS0kX6HnPhzrHK1LrPX/X5/qXtg/tW7taS7pCp/8J7f1/ifpr/X3XEv1/6lbSG7RE/43aro9E" + "/bX+vmvf2tWjpA9oif4/tY/fRP2cC3rO70f71oo+5yvaylfielt/37WEP0AtJV2gJdsn6HP9Y9uHlhg+o9aSrtDpP0OvNx5g" + "+9AS/gF1v/6z7UNL+APUo6QP6PSvqd/4iO1D5/gI9SrpC9q3drWXT9Qv8wdc3/sjtGR7Flqj/0It5fMC7Vu7+tzxN9uHzvFB" + "ai3pCp3jc9Dun5y43uaPuJbwT6hbSW/QM/oDqSX8k9Qn+gfU7p9G/TN/w7WEP0I9SvqAlvA/qOcdL7N96Bwfo14lfUFL+CPU" + "+47/2T60hP9BLSVdoHN8kNrK90T9M3/CtYR/Qa0lXaF9a6nd/zhRP82fcC3hb1C3kt6gJfrT1P36/7YPPeVXS/gjqXVVPcrn" + "B3SOH1DPO35g+9A5fkC9SvqClny+QXv5Rn0wf8S1hH9CLSVdoH1rV3v5Rv2UAy0x/kqtJV2hJfwbaPdvYvzV9qHv/RVasr0K" + "fZ8/0K18vkFL+D/UPp4Q9dP8Gde+tatHSR/QEv4NtZWvRv00f8b1zP5QaAl/J7Vm/wl6lc8vaAn/h3rf8Q7bh5Z7fq6lpAt0" + "jn9QW/lr1F/zd1zLyd837f5NjD/bPrSEP0PdSnqDlvBfqPsdz7B9aAn/hXqU9AGd4xvUXn5R/8xfcS3hv1Cvkr6gfWtXe/lE" + "/TN/xrWEf0MtJV2gJdsvaC+fuF56oPe3frSEv5O69ezAA2j5AoWWcIBcH3N25scKaPvQJxwg6lbSG/QJB4jaGoSvRXqnll99" + "wiFKrdEBoB7l8wP6fJnuet4BW9uHzgFa6lXSF/QJB4jaHngxwG770CccKGop6QLtW7vaJxBmpB9o39rVWtIV+sQAPXTz8uUF" + "s33oEw4YdSvpDfrE9aPudwLD9qnjhk99wgFLrafqUT4/qKM+pM4JkdQa9YN63gku24fOCS7qVdIX9AkHjdrL/0T6hj7RgaaW" + "ki7QvrWrvfzzeA/0iQaKWku6Qp+Wx2/aHboW9dccNtcnHDpqK8+YoLB96BMOGvUo6QP6hINGPe+Amu1D5wAa9SrpizoaxNQ5" + "oJg6HxDUPuCd+Tf06ZnuWkq6QJ9w0KjPHaCxfegcoKHWkq7QJxw06OHlF/XTHDLXJxw06lbSG/SJCRBqL9+of+awuT7hwFCP" + "kj6gTzhs1PNOONo+9AmHjHqV9AWdE5DUb8LA9qFPOGDUUtIFejX90TmhkLp94UEReAHnBw70CQ+MWku6Qp/7edPzuyPmtg99" + "wgOjbiW9QecIOrUVYMyw2D70yQYWet4ZbNuHzhls6lXSF/QJj4t63xEf24c+4WFRS0kX6BwBorbyiRkY24c+4WFRa0lX6BMe" + "FrS/ARAzMLYPfWKGhrqV9AZ9wgOj9hGrqGDmcbk+0UOlHiV9QK/w6FOf8MhSt5Y1DGDeIX3bh17ZYqQOFyZ1DvGnbl82KQna" + "/gpY5RcW9Fl5BK79EkUdNqfN9QmnjlpKukCfcOKozx00tn3oNb4fnYPGqdt3D9CBli9Q6BNeHvT+7qym7UOf/Dx0K+kNOmc5" + "qX0WIIrUvDbXK+ts6BNeXer2ZSUGGOULBvQJN4/ap13zCyf0CTeP2ko4polsH/qEG0ctJV2gT7hx1FbCMU1k+9An3DhqLekK" + "fcKNg/Zpo5gGsX3oE24cdSvpDTqnRaj7fY3G9qHzNRrqUdIH9Ak3rehZgBdYVEnz01wfyXTXq6Qv6BN+G7UXaFRJ89Ncn/Db" + "qKWkC/QJv43aCzRqnPlprk/4bdRa0hX6hN8G7X5ajDPZPnSOK1G3kt6gT/hl1D4PFRXK/DDXJ7td0KOkD+gTfhm1lV/MK9k+" + "9Ak/jHqV9AV9wu+itvKLeSfbhz7ZbYCWki7QJ/wuap+HigpnfpbrE35X6jZj4DBB11FzaPkGhT4xMwXtM1Exs2T70Dkzl/r9" + "QoD3CwCtfEODzrkxar8GUWnNl3N9wrejHiV9QJ/w7aj9GkQdNV/O9QnfjnqV9AV9wrej3ncsyvahT/h21FLSBTrHpqj9GmQB" + "HOijef6utaQr9MmuuWl11y3mpmwfWsN1o24lvUFruG7UVn4xN2X70BquGvUo6QNawzWjnnfuyfahNVwzaiufmLuyfWgN14ta" + "SrpAa7he1D63NSP9QGu4XtRa0hVaw/WCdtcr5rZsH1rD9aJuJb1Ba7he1F4+O9I7tIarRT1K+oDWcKWo531XwfahNVwn6lXS" + "F/S630ed7zKkbu0egIN9JytsHzonJ6ilpAu0hm9Ffe5gkO1D5+APtZZ0hdbwraDdt4rJL9uH1vCtqFtJb9ArbqjUGpNlqdsX" + "dyCBXYGYHbN9aA3njHqU9AGt0ZOitisQs1+2D609v9/1KukLWsMZo7YCjtkv24fWcL6opaQLtIbzRW0FHLNftg+t4XxRa0lX" + "aA3nC9qdr5gds31oDeeLupX0Bq3hfFH3O5tk+9A5e0Q9SvqA1nC+qK284qlv/5r28ooLaM6Va42uMLWXV1Q4c65cazhb1FLS" + "BVrDuaL28ooKZ76T63wiUGtJV2gNXwrafamYrbJ9aA1firqV9Aat4UtRv8Er24fWmM2iHiV9QOdgFrXVt5jNsn1oDV+LepX0" + "Ba3RTaW28ovZKtuH1vC1qKWkC7SGr0Vt5ROzVbYPreFLQbsvFbNVtg+t4UtRt5LeoDV8Keo3+m/70DnaTz1K+oDW8J2ovXyi" + "/phj41rDFaJeJX1Bazg61PvOzts+dM7GU0tJF2gNR4fa6lfMRtk+tIZjQ60lXaE1/Bho92PClbN9aA0/hrqV9Aat4cdQW/nF" + "bJPtQ2v4LdSjpA9oDTeF2mf/o/6YW+Jaw02hXiV9QWvMRlFb+cVsk+1Da7gx1FLSBVrDjaH22IKoX+a2uNZwY6i1pCu0hhsD" + "7W5MzDbZPrSGG0Pt5RP1x9wS1xpuCvUo6QN65f0YWsNtSd1a3sAA7/UJ24fWcFuoV0lf0Pk6BbUXYFQwc1Nca7gt1FLSBVrD" + "baH2AooKJgqt4ZZAY/ooKpC5IYrpoigg6Df6a/vQOdpLPUr6gNZwO6jt/GM6x/ahNdwO6jc9bPvQORpJLSVdoHO6mPp42Eik" + "H2gNt4JaS7pCazgN0O40xHSO7UNrOA3UraQ3aA2XgNrLJyqIuQCuVzZooTVchNStZQsIMMoXDGgNn4HaCzAqiPkIrjV7LNBe" + "gFEBzAdwreETUEtJF2gNn4D63OF524fW8AmotaQrdA7Xu26fOQAr5mtckNAUZMUSky4uSGgKGjXTIKEpaFqmnpkmCU1Bq2Za" + "JDQFiUcUZSYhoSnoWKaZmQ4JzUPNY5NWZLKOPAhNQV4EOzP1IOf7h9D8orbbDxr1mwYJTUFeTpKZJglNQatmWiQ0BW3LdDLT" + "JqH5RdlTeui614m8NG+mQ0JTkNZMSkLzkPX6V8uy6x8JTUGtZmokNAX1O3rqgoSmoFEzDZI7iJrISrNlreuLhKYgK82Wta5v" + "EpqCrJxa1rp+SGgeGn52WetGI6EpyM8u68roJDQF+QByVoMxSWgK8gPPajA2CU1B505muCChKUhrJiW5cxyBrBe9el67+ZHQ" + "FOTDsnlZZiehKWjUTIOEpqB1B3hdkNwh3URWBD2v3dwkNAVZEfS8dvOQ0Dy0/Ozy2q2PhKYgP6a8LGuR0BRkrVjPy7KEhKYg" + "P6Ys8XVIaB7yMe6R5bQ7CU1BVlVGltOeJDQFrZppkdAUZIU5sjD3JqEpSGomIaEpyGcPssT3IaF5yHp9a2SJy0dCU5AXQd4t" + "0kloCvIDzxKXTUJTkNXxkSUuSkLzkDX+a2YdPx8JTUGtZmokNAXZgc+8dqeT0BRk127mZTmThKagVTMtEpqCrAhmXpazSWgK" + "kppJSGgK8nDcvHZHSWge8pjPmddOPxKaglrN1EhoCvJyygusnYSmoFEzDRKaguYdWnJBcgeTEq2aaZHQFHQQkRyZDglNQVoz" + "KQnNRc27YjH+6oKEpqBWMzUSmoLs7GKk1QUJTUFWC2I41QUJTUF24DFm6oKE5qGG4OnI5L0lJzQF+WilZKZFQlOQFWaMILog" + "oSlIayYloXnIexg7y8m7E05oCvLo7ZaZOglNQQsh3pFpkdAUZIW5szC9E+CEpiA7u52F6Z0AJzQFac2kJDQPDT+7LPHxkdAU" + "tO7MtAuSOxedyA/8ZKZNQlOQ1ExCQlOQH3helqEkNA95v0DysngnwAlNQa1maiQ0BfXry7sgud57olEzDRKagqycJC+w9xSc" + "0BS0PQA/M20SmoKsCCSv3VQSmoe88yB57bynAIcyvcpEdgdLXrs1SWgKEgT/RyYhoXnIboLrppggue8LBPLH68ljkiA0BbWa" + "qZHQFNSxQEBk6iQ0BdmBnzxwERKaguxuOVmf5JDQPORPTs0L7I9JJzQF2QXWvMD+BHRCU5CvBJAX2B9uTmge8sZe85i8ZXdC" + "c1G3BnunS9290XZCU1CrmRoJTUHTMvXMNEloHmrtvuLpguS+1JlILNPJTEJC85AHy6dz073FckLz0PBM+U0jCE1BUjMJCc1D" + "02PmvQg6YucBaC7ZCJuP79mNhOYhuy13dvtMkNAUtBDIjh9TIaC5xN+FRgPW8a41QFtwfpL467T5HO+42HiBNq8/0PgQC85M" + "w6+jE5qH2kDgNX7Me90OaC4ZA7HHOGS/AR3QXDK5/ArP6hDQPLJuFo9wAaC5xN8PGl/5FvdRGMadxOOEj8ThBqC5xENTcQN1" + "hHoA0FyCaNHFLP6IckDzyMLrWsxyCBpDKC8RxGx6Fn/lFYAmyUTYZIuT/ghoHhG8uY9vcQfKQVuYHbjkILjyHosDmkv89dxv" + "x7csgrbbjG8BUfuWqFY+0OCA5hKP4mtxuO6tO6B5ZCJw8hadA5pL/FXJj6WrODgfJtiYwb7EfyhqnTdoDmgeYcQjD1cJaC7x" + "oKce5bI2Ac0lHjcz4oxkEtA8Ml+WNQjOyIPrBDQ/hPHsl2y80I4L4D0xB20pjyWIIKYPWfzJ4oDmEg9LyZ/2joMDmkfGy8LP" + "eNFJFB2JIJgOWbQT0FxyECdXvsUdQZhHOkLmWAyHgOaRebP4S6sANJd44MGMiqkBaB7ZeDWbdXcStN2j6EgE4W281Q4BzSPn" + "ZvEIFwCaJCsWuUGDuLz1c0JT0EKkG37Kr6wDmke4mAxPWwloHjkvyxQCmks8pHmtqDEfAc0j7WZBu++A5hHU9MgyCWgeGSXL" + "IaB5ZOLlWx7uTKC/oHGFgB+SFZzkIEiOl2AT0DyieM2cVXMQtJ23LEhnKB0PdxPQPDIQPccL0AloHpkvi999DmgeWYipQzu1" + "OsGJ47+A5odItJlBDiLveNKNgOYSb+32Lpdx4IE+ymX0+NstUd6HgOYRQQjdq3Xei4F55NwsHmXGTsPtOZDo7Vz4K64ANJdg" + "LZy4ABqA5pGBcLt7yzqgeWQjwg5Z9kjAgvLXqABoHpGXxXuCDs5pv4Dmh0g06CTrQyweG4tFQPMIwvBeq+SA5pGJiLxXCgBx" + "idw3ckBzyUZ0HLOgMfG+y4oOTJD+sswAJ8spAc0PkehMBdkIl3vlj6fN3qX8/WlzsiwbAc0j530LisXHOmEeUcTZsTU8BDSX" + "eGhmZlmNgOYRRN1F1Q1A88hEAN7tTBFEcXvPwwHNI+tlwS3i3egl9abxuMcTTZJuAppLDkPzXln6wCnMIwOvTcMDwjSGk8aA" + "/ocmQvhehQGIK9s3Ac0j62VBebsvCHMJIgOj7NzPckDzyEBwHm+1TkDzyHxZ0A4ASFSPQ0DzyLpZ/GVtAJpHSq9nB6BJshF4" + "98VtrwQ0jyzE2vE6dgKaR3bJIgmyek8CmkfOzeLvfgPQXIJuc9bdALerGmQisO4+ptGpvj1rkvWyeAvqgOaRjfescbh+B2Ox" + "OkZcJ+mMsOOt9hHQPNJuFn+vBYDmkYm4Ol6RTUDzyCpZlIDmkY3oOvwQTrGjSxalSyIvy2gENI/o69j5Excr1nHZuiQelnWy" + "t6UEt1yCbETWMcsgoHlEEFzHyhCA5hHG1/HZPgloLpkMsXtVykcDuarcJe1m8fgHAJpH+s3CNwSc0BTk0XatXAJv3Xc28SQb" + "AX2v9vqTcsfjMoggpo8lIwQ0j7z+lgdUANBc4nE9PVqCPgloHlkvyw5A88i+WTz8CoDmkt2ue88ZPrj26d8n6jeTv5oIQPPI" + "QKgeW5QTIBq3C2h+CAOoL5kI6OM59QCqv4DmhzDK+JJ1v4WluTFksUr5esxN/hCP30smntVBBMGDrHkfAc0j52bx8LIAM59J" + "m4SmIL2Z/NV+AJpLhCufsvA+AppH+s2C7gfBfr0CB5r3RYLrID9y8k4hWQhRZMXbBDSPbAQlvjbRp0e5kNslB3GJPP5JQHOJ" + "B8lkFu8nA4T3ewHNI1zwj1k+AppHxsuC28SDYWZcaax1fOA05TMVZJUsSkDzyH5Z0Eb6fCjMI3LdNX9fnI5ZemdBDl7hvb0y" + "rNbGQJZLuCYwK/wioLlEETIQhSsE7aRHRDIQV4ksqFDe79jR+QgyEVp5+2lYSW3nGADJulk4medEo449QvOLTox9JNqI04z2" + "95DQFKTXveWbQFhcjSusJfL11HII0IPpAGgeaS+LP4kwQJijhEH6zcJ3z7DOGhdbe2gg+PNWC6y1JjFOGGTiHXi2NAHauS0N" + "yEKIKLNsAppH9s3ibwYD0DwiCBxlzdkENI8oQkHxLf6YcUBzSftuFs4Vkqx/CE1B/WbCiBjBfk0jVm2T1l/T6Ouy3SzuRmKh" + "NmnjOZa+NNtZ8gZZsVab3OFRkP2y8Jc3okf5y52A5hG9AwZ80Q9DAzk+EMj7dTmq4K+KYS22Lb+A5ofoWJUMhJbyNvsIaB6Z" + "iCa9LTkWWJPo+gVZL4t3IbCmmuT9TLIRYnodAyyrJtn1Izk3C3xlrJwmOWAH4uH0O58qg4DmkXmdKsW1B4hycXcagZuM3rxk" + "vSw+H4A1yyQLikQQonobbixDJtkXJDmIUr19Qaw8JjFMQjLfNYpbdJZLkghRrK91xApjMttrHX1NsczisXsBWjnJCee+l5P0" + "SPjMglbGx1UkBleCzDtowHdVMTyQYwSJBMGwzOTvrWAFMYn+YaJzM8EZwypiMs9zz3zdsBPjKR5NAkBzCZa3n6VKA6z3bMfi" + "4rK+92z3QdbMwrd4SG6T/5HQFDRuJkWJeqAG1x67ZL4s7kMTyC+geWTdLOidYDUyWev1V3wM9X7LagnOL6B5RF6WswloHjkv" + "Sx8JslVBnfY37Lg82UNaMvFjflVWNj1Am9HB/PVDQPPIfI6pj2dilTDJXi9J8V19TB8LgUl2YEkUYcCsIpOA5hLvZ8YYkMec" + "ANA80m6WqJ0+RCTSaoX1rubJE/e+A9bx+mLaqiCVLLBE/Yv2MtFA8PFz1bDEl2R/ItBCBDIrJZ7E3kMVyX4I0b6ZPGAIgOYS" + "RGprae68uynZ5yRBAHHplWBJLTmt9kqwUnB7nUwsqyXZMSUZNwvficHSWlxf66GFqOPbe8XyWJI9URCEPucl2wQ0j7SXZQWg" + "ecRDEVa5+bxvqPktJBMRxjxx9FO8cyiazRmRIMyYZ4V7wOd/JCaBEul14uLSeJ9OsmMHdBBQ/L0nOxaGOjkkRdIRc8zv8clp" + "LAZ1stcWaNxM6OhiAajzla7vQWxxf2MKWAOKC0Fdsl4WbzOwDNT51mtFDiKIsyUXJaF5yMOEv3QIfSwI6zWd7F8FQo/otQhY" + "g+lkd4pkliyDgOYRQURzNN1KQlPQKZnmIKEpiLHP7JAIAc0l3l36tDyKQdr3D6EpqJVMfrs7oSloIHT61QRvIE72qkjmy7IH" + "Ac0j63lJ3sGAg3S9JJKNeJdycbxHpdfhIDoItI5MHwnNQwM3TdbMQUJT0EZEdPS1OwlNQXIzeRAZAM0j530P4hywRNAZeURA" + "Hk0ag1kezQlA88i8WfgSPJbyObe/AuRRnT1dBh/jwHo5Z6WDQITuQzRbh4DmkYVoZp7XIqB5ZJcskuCnwPwxftZPgXlYZmRC" + "XxOL0Zx8mJMwEjqq5kdC89BGGHTc8YuA5pGBiOh3Dj6AdXJIi+Q5hNYCTBKagg7DWd7F8b7xl31UEkW09KsH/tw++fAGEQRK" + "Zz3YJDQFjeKxS5DinhMthE6Xghb48KsWtIcnZia8gYY1Vk48TxMxyJqlKAQ0lxyMrecRfSQ0BS0ES9+nHNYrOfkgJBHES7MQ" + "FwHNJYpQ6HzmbhKagjriolmBOgEie37JyYHKJP3LnjPJQvh0af8wDqLVOcGqGXfcwf0KrKGhX14yol4yeR8cC2noHZwA8gdL" + "uvIIQiO51eGQ0BTUbiYPogWgeQR9hGzwDglNQaNk8h4zFqvQlr00olky+RgZVqzQljcDEBZBz5vTAx6w6oPe5wYRYqJL3cPy" + "DprPjUATAdDP/8LyDNqr/4VVGTITXpTFGg3as6oT7fJNPm2IhRq053AV0EAcdHmUYrUEHa08SrHoQXi1DVM0WAJBs60OtIpf" + "6+MecGmfX0skCI9mD6AT0Dyi78fgl2E9Ax1aXDWE+acDiteJEfSv6YASeYMuu4zXITJfVy/jdQjIv5n8NVmE5+u6dQBoImg5" + "atMmoSlo30wIPEUcvqbPBuIN6+m1jLxN1duwEh1EKD9fhWT+Q2gK0pJpBKF5CC7OLt1hpT/TS3cYMeOZCfEIiBr/suo+pLcB" + "TtQZWfDQQjR0nPAioSloIyS6XCdvjlXquCrixjMT+tqIIleR0v1G8Pj14Nw3Qii5SpYBECLI7wjtCnK9I29OEWGu5ysTHggs" + "v5n4RR2h1r+EpqCJ+Oq4dpOEpiBBTPXrxCCEXI+UTgwix28mNJY+8K4nawEQXKLbWk0SmoIEIdRxvp2EpqBzM7Hy+gtZGm9l" + "OekR2Hw7X4HICxLGSN+nNxB5QYdB0nEjfkTkBWnJ5TEgQOQPIdr4y1JQISIvaDNO+raiQOQFSc31BSIvCMeVrZY9EoDIH0IE" + "cOYaTYnIC9qMle5/Ay+skDS4lYUQ/xC+W58EIbItb775EZE/hCjSfitoJyIvCJ09xn24umj9g3rGkVYUgR+BVns9E4R/ApEX" + "NBl03P8m/u4PCf/yTCHEP4QTIEmEN/DrqQGRFzRKLlY1b82+O2BDJAwVvs8BIPKHENh4+4atE5EXtBgH7NdUUGt9fIphjEnQ" + "h5v7DUwAkRc0GdDFXFuIep9Zjx01rC+wMpd1o4DIC0Jo7nhPKCDyh3BPrLwnvBo5In8Ioez7dgE7EflDYzFcFrn8PVkg8odm" + "hMKWo/dZ1ja/evQIds5ugYdpAJE/hBhlPKgHxklAiC/ZH2Nho51RIvKCBuNc74MZiLygWXJ5JA4QeUGL4bA8dhlE5A9JxJ5e" + "ZwCI/CFGg946L0TkF3VvDRtayYE5HBDiSxC910YpB0fkD43OyBe0T2ORdL4WngQhWem/esAZEPlFiA7pmK4cmIID6QsPlCT+" + "dmvbO35tkxBfgqC0mDZALA4Q+UP+GkTHZO3g2g5AfWLs5CH/LvSsMPhJQpxkMvwl6pTfu1jBnX83NEljPBPz+PMAq2gDP+Ln" + "n08uj7UBIi9IPReP21ccBiG+xIujt7i6IwjxJT6+gOuJ8whCfInXP1xOHOMgIX4E4VGjfI9fHuBLvDPV41UJFj/WhQK/B+Rv" + "CPQZRX0GCfElCEeKi+YLg4IQJ/HFRxr/iCcXwSDqfJ3qIlyQFb/mR4QVNFZeEH8Y483vvqNC+m2D9RX4QngSBHXJ94oake4r" + "fozE3YCeLkSXRUTuPy/WFuK9xPFFWY9DQnyJP/T592EHFyoAIi9IGWqFU5MgxCB22Pijo1lIvoIrCDG+SP1jiLbK8TX0TxyR" + "I5c3E3ili3+0E8tzkRBf4p1Z/mHAwe4joq34akkSb6D4x/UGBj9AiB/xmKt58qYNRH4R5uj5Z9biu50QP+LFveKmnZ2E+BJ/" + "WPHPmY14rjkiL8gDQkaLY/LDdNQ5qnaRNz5c3HBEG+yoc/3Ci7xdGdHcsN+Bt+plZ6MExLCy+eolXhDm9EgQ/BGDGQ08+zl4" + "rfJEISRSRnzhm46SEF/iXTMuKc4nHhH5Qwhm+ikFR+QF+UI8IyvwCUR+EQZ85spScD8CQ+4cCLrIPWAuz8bblIg80WgMMPEW" + "Y+FPEIKMgZYnCV6Exw2z8GcMQcbEcy6It8uNY7ve47J+AtA40QMNtPBu7UAu74iDEP9hwN+IDyUKynwhMhNkTPSfkuCVqRnf" + "s4MQJxl4X3VLnP8IRH4R3ogWjcecuXBE5H7c3uE+ny+YHUuOd/6hgQ+zplzvJJGvANy4Ns6HFT2BTme8eaCGP4GyIheuKNhp" + "PbIFw5+6wPCT0isDOuz8PeSr/nKxhw+3JdCJlXMT+UqcsfgN1nAk0o+h6RftmstX+3ekH9dsTDTxN+yFx9X2IDoNg+0X+dqR" + "7G7CrfI/tYGl81scqrWXh+xMrlvwGP4GCKLf468yODrsx13kiyu+XDhLn3eP2O+LfLl/XpL4CxyOzmR0fyJfHvzDqM7h+gKJ" + "5B90GH/uaPtB+AJyn+bfS/FVpB3BxJ8U6f8DshTQjg==" +) + + +# The context top-k's (tc, S) cells in the max_num_tokens-16384 engine dumps of this envelope (tep4 +# DSpark max_batch_size 8 and dep4 DSpark max_batch_size 256, max_seq_len 32768; and the same tep4 +# engine at max_num_tokens 8192), with nc and max_q_len from the sparse-MLA varlen record at the +# same tc, as (tc, S, ratio, nc, max_q_len). +DUMPED_CELLS_16K = [ + (1, 1, 1, 1, 1), + (2, 2, 1, 1, 2), + (2, 1, 2, 1, 2), + (5, 5, 1, 1, 5), + (5, 2, 2, 1, 5), + (12, 12, 1, 1, 12), + (12, 6, 2, 1, 12), + (16, 16, 1, 1, 16), + (16, 8, 2, 1, 16), + (28, 28, 1, 2, 16), + (28, 14, 2, 2, 16), + (128, 8192, 1, 1, 128), + (128, 10441, 1, 1, 128), + (128, 10444, 1, 1, 128), + (128, 10451, 1, 1, 128), + (128, 16339, 1, 1, 128), + (128, 16384, 1, 1, 128), + (597, 610, 1, 1, 597), + (597, 305, 2, 1, 597), + (2252, 10444, 1, 1, 2252), + (2252, 5222, 2, 1, 2252), + (4553, 10441, 1, 1, 4553), + (4553, 5220, 2, 1, 4553), + (8147, 16339, 1, 2, 5888), + (8147, 8169, 2, 2, 5888), + (8192, 8192, 1, 1, 8192), + (8192, 4096, 2, 1, 8192), + (8317, 8759, 1, 25, 905), + (8317, 4372, 2, 25, 905), + (9055, 9754, 1, 27, 893), + (9055, 4870, 2, 27, 893), + (9108, 9959, 1, 28, 894), + (9108, 4972, 2, 28, 894), + (9392, 10234, 1, 28, 890), + (9392, 5108, 2, 28, 890), + (10441, 10441, 1, 1, 10441), + (10441, 5220, 2, 1, 10441), + (10444, 10444, 1, 1, 10444), + (10444, 5222, 2, 1, 10444), + (10451, 10451, 1, 1, 10451), + (10451, 5225, 2, 1, 10451), + (16292, 16625, 1, 24, 903), + (16292, 8307, 2, 24, 903), + (16327, 16327, 1, 23, 944), + (16327, 8159, 2, 23, 944), + (16328, 16328, 1, 23, 933), + (16328, 8158, 2, 23, 933), + (16378, 16695, 1, 24, 934), + (16378, 8344, 2, 24, 934), + (16384, 16384, 1, 1, 16384), + (16384, 8192, 2, 1, 16384), +] + + +def _dumped_cell_reqs(tc, s, ratio, nc, max_q_len, max_len=16384): + """One step consistent with a dumped cell: nc requests, the first with max_q_len new tokens, the + rest splitting tc - max_q_len evenly (where even lengths alone would give more than S compressed + entries, pairs of them become odd, keeping the total and the longest); cached prefixes (multiples + of ratio, filled in request order up to max_len) chosen so the sum of kv // ratio is exactly S.""" + rest = tc - max_q_len + news = [max_q_len] + [ + rest // (nc - 1) + (1 if i < rest % (nc - 1) else 0) for i in range(nc - 1) + ] + if s < sum(n // ratio for n in news): + even = [i for i in range(1, nc) if news[i] % 2 == 0] + for i, j in zip(even[0::2], even[1::2]): + if news[i] < max_q_len: + news[i] += 1 + news[j] -= 1 + extra = s - sum(n // ratio for n in news) + assert extra >= 0 and min(news) >= 1 and sum(news) == tc + reqs = [] + for n in news: + e = min(extra, (max_len - n) // ratio) + reqs.append((ratio * e, n)) + extra -= e + assert extra == 0 + return reqs + + +def test_multi_request_dumped_partitions(): + """Every dumped (tc, S) cell as a consistent packed step, both layouts alternating, mixed tie + kinds, against the reference exactly: the 1933 cells of the 8192-token gate dumps, and the 51 + cells of the 16384-token envelope dumps (kv up to the 32K DSpark manager length 32782); every call + launches the same two kernels.""" + blob = zlib.decompress(base64.b64decode("".join(DUMPED_CELLS))).decode() + cells = [tuple(map(int, c.split(","))) + (16384,) for c in blob.split()] + assert len(cells) == 1933 + cells += [c + (32782,) for c in DUMPED_CELLS_16K] + straddle = 0 + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for i, (tc, s, ratio, nc, mq, max_len) in enumerate(cells): + reqs = _dumped_cell_reqs(tc, s, ratio, nc, mq, max_len) + logits, ks, ke = _packed_logits(reqs, ratio, contiguous=bool(i % 2), seed=50000 + i) + assert logits.shape == (tc, s) + out = _run_packed(logits, ks, ke) + bad, st = _packed_bad_rows(logits, ks, ke, reqs, ratio, out) + assert bad == 0, (tc, s, ratio, nc, mq) + straddle += st + torch.cuda.synchronize() + print( + f"{len(cells)} dumped cells: all exact; request counts {sorted({c[3] for c in cells})}; " + f"straddling tie rows {straddle}" + ) + _assert_variants(prof, len(cells), "dumped cells") + + +# ─── ties at the K-th value ─────────────────────────────────────────────────────────────────── +def _tie_case(form, width, ratio, layout, kind, seed): + """(scores, lengths, row_starts, ranked windows as (rows, win, lens) triples) of one tie case.""" + if form == "decode": + g = torch.Generator().manual_seed(seed) + vis = _dspark_vis(32, ratio, width, g, padding=False) + vis[0] = width + scores, ranked = _decode_inputs(width, layout, vis, seed, kind=kind) + return scores, vis, None, [(slice(0, vis.numel()), ranked, vis)] + reqs = [(width * ratio - 2048, 2048)] + logits, ks, ke = _packed_logits(reqs, ratio, contiguous=False, seed=seed, kind=kind) + assert logits.shape[1] == width + return logits, ke - ks, ks, [(slice(0, 2048), logits, ke - ks)] + + +def _drive_ties(form, width, ratio, layout, kind): + """Scores whose K-th value is tied with scores left out (bf16-rounded normals, 8 integer levels, an + exact-zero plateau, and a tie constructed at the K-th value): the selected set is the reference's + exactly, ties going to the smaller column, and repeat calls are bitwise identical.""" + scores, lens, rs, ranked = _tie_case( + form, width, ratio, layout, kind, seed=width + ratio + len(kind) + ) + out = _call(scores, lens, row_starts=rs) + ref_all = op.reference(scores, _zeros(scores.shape[0]), lens, rs, K) + op.compare(out, ref_all) + straddle, tied = 0, 0 + for rows, win, L in ranked: + ref = ref_all[rows].long() + straddle += _straddling_rows(win, L, ref) + tied = max(tied, _max_tied(win, L, ref)) + selecting = int((lens > K).sum()) + print( + f"[ties {form} w{width} r{ratio} {layout} {kind}] straddling rows {straddle} of {selecting} selecting; " + f"most columns tied with a row's K-th score {tied}" + ) + assert straddle > 0 + if (width, layout, kind) == (32768, "paged_view", "zero_plateau"): + # more candidates at the cut than the kernel's 16K-entry candidate buffer holds + assert tied > 16384 + if (width, layout, kind) == (1048595, "paged_view", "zero_plateau"): + assert tied > 200_000 + if kind == "tie_at_kth" and layout != "candidate_masked": + # (the candidate mask sets some of the constructed scores to -inf, which can move the cut + # off the tie) + assert straddle == selecting + for _ in range(3): + assert torch.equal(_call(scores, lens, row_starts=rs), out) + + +def _drive_signed_zero(): + """Scores are ranked by their fp32 bit pattern, so -0.0 ranks below +0.0 and the two are not a + tie: with 300 positive scores, 1500 columns of -0.0 first and +0.0 after them, the 212 remaining + picks are the smallest +0.0 columns, not the -0.0 columns a value-equality ranking would take.""" + n, w = 4, 4096 + x = torch.zeros(n, w, device=DEV) + x[:, :1500] = -0.0 + x[:, 3000:3300] = torch.arange(1, 301, device=DEV).float() + lens = torch.full((n,), w, dtype=torch.int32, device=DEV) + out = _call(x, lens) + op.compare(out, op.reference(x, _zeros(n), lens, None, K)) + expected = torch.cat([torch.arange(1500, 1712), torch.arange(3000, 3300)]).to(DEV) + assert all(torch.equal(out[i].long(), expected) for i in range(n)) + # control: a value-equality ranking (-0.0 == +0.0, ties to the smaller column) + value_ref = _control_reference(x, lens) + assert not torch.equal(out.long(), value_ref) and int(value_ref[0, 0]) == 0 + + +def _drive_few_finite(): + """A selecting row with fewer than K finite scores in range (the rest -inf) gets every finite + column plus the smallest -inf columns: the reference's answer, since -inf entries tie with each + other.""" + n, w = 16, 4096 + g = torch.Generator().manual_seed(3) + x = torch.full((n, w), float("-inf"), device=DEV) + finite = [0, 1, 7, 300, 511, 512, 513, 2000] * 2 + for i, f in enumerate(finite): + cols = torch.randperm(w, generator=g)[:f].to(DEV) + x[i, cols] = _distinct(1, max(f, 1), seed=i)[0, :f] + lens = torch.full((n,), w, dtype=torch.int32, device=DEV) + out = _call(x, lens) + op.compare(out, op.reference(x, _zeros(n), lens, None, K)) + + +# ─── determinism across calls, processes and GPUs ───────────────────────────────────────────── +DIGEST_DECODE = [ + (16384, 1, "paged_view"), + (16403, 1, "candidate_masked"), + (32768, 1, "paged_view"), + (32773, 1, "candidate_masked"), + (8201, 2, "paged_view"), + (1048595, 1, "paged_view"), + (1048595, 1, "candidate_masked"), + (524297, 2, "paged_view"), +] +DIGEST_PREFILL = [(8192, 1), (4096, 2), (131072, 1)] + + +def _digest_cases(repeats): + """{case: {"in": input digest (fingerprint of the scores, SHA-256 of the int32 vectors), "out": + SHA-256 of the output, "exact": equal to the reference}} over the generation and context forms and + every score kind (including the ties), plus a CUDA-graph replay; every case is called `repeats` + times and must be bitwise identical each time.""" + res = {} + + def record(name, scores, lens, rs, ranked): + out = _call(scores, lens, row_starts=rs) + for _ in range(repeats - 1): + assert torch.equal(_call(scores, lens, row_starts=rs), out), f"{name}: repeat differs" + exact = all(torch.equal(out[rows].long(), _reference(win, L)) for rows, win, L in ranked) + ins = _fingerprint(scores) + _sha(lens) + ("" if rs is None else _sha(rs)) + res[name] = { + "in": hashlib.sha256(ins.encode()).hexdigest(), + "out": _sha(out), + "exact": exact, + } + return out + + for kind in KINDS: + for width, ratio, layout in DIGEST_DECODE: + record( + f"decode.w{width}.r{ratio}.{layout}.{kind}", + *_tie_case( + "decode", width, ratio, layout, kind, seed=width + 3 * ratio + len(kind) + ), + ) + for width, ratio in DIGEST_PREFILL: + record( + f"prefill.w{width}.r{ratio}.{kind}", + *_tie_case( + "prefill", width, ratio, "mqa_view", kind, seed=width + 5 * ratio + len(kind) + ), + ) + logits, ks, ke = _packed_logits(FULL_CONTEXT_STEP, 1, contiguous=False, seed=5001) + record( + "prefill.full_context_step.r1.mixed", + logits, + ke - ks, + ks, + [ + ( + slice(a, a + n), + logits[a : a + n, int(ks[a]) : int(ks[a]) + (c + n)], + (ke - ks)[a : a + n], + ) + for a, (c, n) in zip( + [sum(q[1] for q in FULL_CONTEXT_STEP[:i]) for i in range(len(FULL_CONTEXT_STEP))], + FULL_CONTEXT_STEP, + ) + ], + ) + # a captured call replayed on the same inputs + scores, lens, _, ranked = _tie_case( + "decode", 16403, 1, "candidate_masked", "tie_at_kth", seed=99 + ) + off = _zeros(scores.shape[0]) + eager = op(scores, off, lens, None, K) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(scores, off, lens, None, K) + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + res["decode.w16403.graph_replay.tie_at_kth"] = { + "in": _fingerprint(scores), + "out": _sha(out), + "exact": bool(torch.equal(out, eager)) + and bool(torch.equal(out.long(), _reference(ranked[0][1], lens))), + } + return res + + +def _script_env(device=None): + """The environment of a fresh process running this file as a script: the one device it may see, + and the directory ``__extra_import_path__`` names (where ``_validating`` lives), which pytest + supplies to a collected file but a script has to be given.""" + env = dict(os.environ) + if device is not None: + env["CUDA_VISIBLE_DEVICES"] = device + here = os.path.dirname(os.path.abspath(__file__)) + tests_dir = os.path.normpath(os.path.join(here, *__extra_import_path__[0].split("/"))) + env["PYTHONPATH"] = os.pathsep.join(p for p in (tests_dir, env.get("PYTHONPATH", "")) if p) + return env + + +def _visible_devices(): + env = os.environ.get("CUDA_VISIBLE_DEVICES") + if env: + return [d.strip() for d in env.split(",") if d.strip()] + return [str(i) for i in range(torch.cuda.device_count())] + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs") +def test_bitwise_identical_across_calls_processes_and_gpus(): + """Every digest case (generation rows at eight engine widths up to 1048595 in both layouts, + context rows at both ratios up to S 131072 and the 16384-row FULL_CONTEXT_STEP, each with distinct, + bf16-rounded, 8-level, zero-plateau and constructed K-th-value ties, plus a graph replay; 57 + cases) is called 5 times here, and again in 3 fresh + processes on each of 2 GPUs. Every output is bitwise identical everywhere and equals the + reference; the inputs' digests are identical too, so no difference could hide behind one.""" + torch.cuda.empty_cache() # the fresh processes share this GPU + local = _digest_cases(repeats=5) + assert all(c["exact"] for c in local.values()), [k for k, c in local.items() if not c["exact"]] + assert len(local) == 57 + devs = _visible_devices() + assert len(devs) >= 2, f"this test needs 2 GPUs, got CUDA_VISIBLE_DEVICES={devs}" + tmp = tempfile.mkdtemp(prefix="fitopk_") + procs = [] + for d in devs[:2]: + for j in range(3): + path = os.path.join(tmp, f"digests_{d}_{j}.json") + env = _script_env(d) + p = subprocess.Popen( + [sys.executable, __file__, "--digests", path], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + procs.append((d, j, path, p)) + uuids = {} + for d, j, path, p in procs: + log, _ = p.communicate() + assert p.returncode == 0, f"digest process {d}/{j} failed:\n{log[-3000:]}" + with open(path) as f: + got = json.load(f) + uuids.setdefault(d, set()).add(got["uuid"]) + differing = sorted(k for k in local if got["cases"].get(k) != local[k]) + print( + f"process {j} on GPU {d} ({got['name']} {got['uuid']}): {len(got['cases'])} cases, " + f"differing from this process: {differing}" + ) + assert not differing and set(got["cases"]) == set(local) + assert len({u for s in uuids.values() for u in s}) == 2, uuids + print( + f"{len(local)} cases bitwise identical over 5 calls here and 6 fresh processes on GPUs {uuids}" + ) + + +def _first_call_in_capture(): + """The process's first call of the op is inside a CUDA graph capture; the replay is exact.""" + scores, lens, _, ranked = _tie_case( + "decode", 16403, 1, "candidate_masked", "tie_at_kth", seed=77 + ) + off = _zeros(scores.shape[0]) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(scores, off, lens, None, K) + graph.replay() + torch.cuda.synchronize() + return { + "exact": bool(torch.equal(out.long(), _reference(ranked[0][1], lens))), + "out": _sha(out), + } + + +def test_first_call_inside_graph_capture(): + """A fresh process whose first call of the op (FlashInfer loads its top-k module then) is inside + a CUDA graph capture: the capture succeeds and the replay equals the reference.""" + path = os.path.join(tempfile.mkdtemp(prefix="fitopk_"), "first.json") + p = subprocess.run( + [sys.executable, __file__, "--first-call-in-capture", path], + env=_script_env(), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + assert p.returncode == 0, p.stdout[-3000:] + with open(path) as f: + got = json.load(f) + print(f"first call inside capture: {got}") + assert got["exact"] + + +# ─── wrong controls ─────────────────────────────────────────────────────────────────────────── +def _raw(scores, lens, rs, deterministic, tie_break): + """FlashInfer's call as the entry makes it, with the two determinism knobs free.""" + x = scores + if not scores.is_contiguous(): + x = scores.as_strided((scores.size(0), scores.stride(0)), (scores.stride(0), 1)) + return top_k_ragged_transform( + x, + _zeros(scores.shape[0]), + lens, + K, + deterministic=deterministic, + tie_break=tie_break, + dsa_graph_safe=True, + row_starts=rs, + ) + + +CONTROL_FORMS = [("decode", 16384, 1, "paged_view"), ("prefill", 8192, 1, "mqa_view")] + + +@pytest.mark.parametrize("kind", ["bf16", "zero_plateau", "tie_at_kth"]) +@pytest.mark.parametrize( + "form,width,ratio,layout", CONTROL_FORMS, ids=[f"{f}_w{w}" for f, w, r, lay in CONTROL_FORMS] +) +def test_control_nondeterministic_mode_varies(form, width, ratio, layout, kind): + """The same call with deterministic=False, tie_break NONE (the same kernel family otherwise) is + seen to vary: over 6 calls on the same inputs its raw rows differ from the first call's, and on + tied rows its selected sets differ from the entry's. The entry's 6 calls are identical.""" + scores, lens, rs, _ = _tie_case(form, width, ratio, layout, kind, seed=width + 17 + len(kind)) + entry = [_call(scores, lens, row_starts=rs) for _ in range(6)] + assert all(torch.equal(e, entry[0]) for e in entry) + nd = [_raw(scores, lens, rs, False, TopKTieBreak.NONE) for _ in range(6)] + raw = sum(int((o != nd[0]).any(1).sum()) for o in nd[1:]) + sets_vs_entry = sum(int((o.sort(1).values != entry[0]).any(1).sum()) for o in nd) + sets_between = sum(int((o.sort(1).values != nd[0].sort(1).values).any(1).sum()) for o in nd[1:]) + rows = scores.shape[0] + print( + f"[control nondet {form} w{width} {kind}] over 6 calls x {rows} rows: raw rows differing from call 0 " + f"{raw}, set rows differing from call 0 {sets_between}, set rows differing from the entry {sets_vs_entry}" + ) + assert raw > 0 and sets_vs_entry > 0 + + +@pytest.mark.parametrize("kind", ["distinct", "bf16", "zero_plateau", "tie_at_kth"]) +@pytest.mark.parametrize( + "form,width,ratio,layout", CONTROL_FORMS, ids=[f"{f}_w{w}" for f, w, r, lay in CONTROL_FORMS] +) +def test_control_tie_break_large_differs_on_tie_rows(form, width, ratio, layout, kind): + """The same deterministic call with ties to the LARGER column differs from the entry on exactly + the rows whose K-th value straddles the cut (all selecting rows of the constructed tie, none on + distinct scores), and equals the larger-column reference.""" + scores, lens, rs, ranked = _tie_case( + form, width, ratio, layout, kind, seed=width + 29 + len(kind) + ) + entry = _call(scores, lens, row_starts=rs) + large = _raw(scores, lens, rs, True, TopKTieBreak.LARGE) + rows, win, L = ranked[0] + straddle = _straddling_rows(win, L, _reference(win, L)) + differing = int((large != entry).any(1).sum()) + print( + f"[control large {form} w{width} {kind}] rows differing from the entry {differing}, straddling {straddle}" + ) + assert torch.equal(large[rows].long(), _control_reference(win, L, large=True)) + assert differing == straddle + assert (differing > 0) == (kind != "distinct") + + +# ─── what the op accepts silently, and what it rejects ──────────────────────────────────────── +def test_is_valid_guards_what_the_op_accepts_silently(): + """The op reads offsets, lengths and row_starts through an int32 pointer and never compares the + lengths of offsets / lengths with the row count. Each case below is accepted by the op itself + (called outside ``validating``) without an error and gives wrong rows; the entry's guard + rejects each.""" + n, w = 8, 4096 + x = _distinct(n, w, seed=4) + lens = torch.full((n,), 4000, dtype=torch.int32, device=DEV) + rs = torch.full((n,), 5, dtype=torch.int32, device=DEV) + off = torch.full((n,), 7, dtype=torch.int32, device=DEV) + good = _call(x, lens, row_starts=rs, offsets=off) + + def unguarded(o, L, r): + return op(x, o, L, r, K) + + lens8 = torch.tensor([4000] * 4 + [0] * 4, dtype=torch.int32, device=DEV) + off8 = torch.tensor([7] * 4 + [1000] * 4, dtype=torch.int32, device=DEV) + silent = { + "int64 lengths": unguarded(off, lens.long(), rs), + "int64 offsets": unguarded(off.long(), lens, rs), + "int64 row_starts": unguarded(off, lens, rs.long()), + "short lengths (the 4 after it read)": unguarded(off, lens8[:4], rs), + "short offsets (the 4 after it read)": unguarded(off8[:4], lens, rs), + } + for name, got in silent.items(): + wrong = int((got != good).any(1).sum()) + print(f"op itself, {name}: no error, rows wrong {wrong} of {n}") + assert wrong > 0, name + for match, args in ( + ("offsets and lengths must be int32", (x, off, lens.long(), rs, K)), + ("offsets and lengths must be int32", (x, off.long(), lens, rs, K)), + ("row_starts must be int32", (x, off, lens, rs.long(), K)), + ("one entry per row", (x, off, lens8[:4], rs, K)), + ("one entry per row", (x, off8[:4], lens, rs, K)), + ): + with pytest.raises(AssertionError, match=match): + with validating(op): + op(*args) + with validating(op): # the guard admits the well-formed call + assert torch.equal(op(x, off, lens, rs, K), good) + + +def _past_storage_view(): + """A [4, 1000] view with row stride 1280 whose storage ends at the last row's column 1000.""" + store = torch.randn(3 * 1280 + 1000, device=DEV) + return store.as_strided((4, 1000), (1280, 1)) + + +@pytest.mark.parametrize( + "case,match", + [ + ("transposed", "input must be contiguous"), + ("k 4096", "operation not supported"), + ("fp64", "Unsupported dtype torch.float64"), + ("short row_starts", r"row_starts must have shape \(num_rows,\)"), + ("strided lengths", "lengths must be contiguous"), + ("padding past the storage", "out of bounds for storage"), + ], +) +def test_rejections(case, match): + """Layouts and arguments the call refuses, loudly (host-side checks, before any launch).""" + lens = torch.full((2,), 4000, dtype=torch.int32, device=DEV) + if case == "transposed": + args = (torch.randn(4096, 2, device=DEV).t(), _zeros(2), lens, None, K) + elif case == "k 4096": + args = ( + torch.randn(2, 8192, device=DEV), + _zeros(2), + torch.full((2,), 8192, dtype=torch.int32, device=DEV), + None, + 4096, + ) + elif case == "fp64": + args = (torch.randn(2, 4096, device=DEV, dtype=torch.float64), _zeros(2), lens, None, K) + elif case == "short row_starts": + args = (torch.randn(2, 4096, device=DEV), _zeros(2), lens, _zeros(1), K) + elif case == "strided lengths": + args = ( + torch.randn(2, 4096, device=DEV), + _zeros(2), + torch.full((4,), 4000, dtype=torch.int32, device=DEV)[::2], + None, + K, + ) + else: + args = ( + _past_storage_view(), + _zeros(4), + torch.full((4,), 1000, dtype=torch.int32, device=DEV), + None, + K, + ) + with pytest.raises((RuntimeError, AssertionError), match=match): + op(*args) + torch.cuda.synchronize() + + +# ─── the certified cells ────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Each cell drives the entry the way the target calls it, and every call it makes is gated by + the entry's own reference and compare (the drivers above).""" + certifying(op) + spec = cell.spec + kind = spec["kind"] + if kind == "semantics": + _drive_semantics() + elif kind == "gen": + _drive_gen(spec["width"], spec["ratio"], spec["layout"]) + elif kind == "every_len": + _drive_every_len(spec["ratio"], spec["chunk"]) + elif kind == "widest": + _drive_widest(spec["ratio"]) + elif kind == "long_ctx": + _drive_long_ctx(spec["layout"]) + elif kind == "far_end": + _drive_far_end(spec["lm"], spec["ratio"], spec["layout"]) + elif kind == "tiles": + _drive_tiles(spec["name"], spec["reqs"], spec["ratio"], spec["layout"]) + elif kind == "ties": + _drive_ties(spec["rows"], spec["width"], spec["ratio"], spec["layout"], spec["scores"]) + elif kind == "signed_zero": + _drive_signed_zero() + else: + assert kind == "few_finite", kind + _drive_few_finite() + + +if __name__ == "__main__": + mode, path = sys.argv[1], sys.argv[2] + props = torch.cuda.get_device_properties(0) + if mode == "--digests": + payload = {"uuid": str(props.uuid), "name": props.name, "cases": _digest_cases(repeats=2)} + elif mode == "--capture-digests": + payload = {"uuid": str(props.uuid), "name": props.name, "cases": _capture_digests()} + else: + assert mode == "--first-call-in-capture" + payload = {"uuid": str(props.uuid), "name": props.name, **_first_call_in_capture()} + with open(path, "w") as f: + json.dump(payload, f) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_indexer_k_cache_gather.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_indexer_k_cache_gather.py new file mode 100644 index 000000000000..94051e72857c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_indexer_k_cache_gather.py @@ -0,0 +1,1167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the indexer_k_cache_gather catalog entry. + +The entry owns what is certified: `CELLS` is the list, `reference` and `compare` are the gate. +This file turns a cell's spec into tensors and keeps the checks a cell cannot carry: the consumer +views and the decoded value set, the DeepGEMM consumer, skipped rows, the op's rejections and the +guard, the wrong-addressing controls, the engine-size pools, every row count up to the per-step +maximum, and pools up to the largest any DeepSeek-V4.1 cache manager can have. + +The DeepSeek-V4.1 prefill index-K read path over the native MXFP4 pool: index K is quantized with +``torch.ops.trtllm.fp4_quantize(x, None, 32, True, False)`` (MXFP4, UE8M0, linear scale layout), +stored in a paged ``[pages, block_kv, 1, 68]`` uint8 pool whose page holds ``block_kv x 64`` code +bytes followed by ``block_kv x 4`` scale bytes, and gathered back to contiguous rows for +``tensorrt_llm.deep_gemm.fp8_fp4_mqa_logits``. + +The pool is written here with native torch indexing (not the scatter op), the slot mappings are +computed here from a block table with native torch arithmetic, and beside the entry's reference +the gathered rows must equal the quantizer's own bytes. Deliberately wrong addressing schemes are +run as controls and must fail. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm import deep_gemm +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.indexer_k_cache_gather import ( # noqa: E501 + indexer_k_cache_gather as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "indexer_k_cache_gather requires a CUDA device" + +DEV = "cuda" +DIM = 128 # index_head_dim +CODE_BYTES = 64 # packed E2M1 bytes per key (the op's head_dim argument) +SCALE_BYTES = 4 # four UE8M0 bytes per key +KEY_BYTES = CODE_BYTES + SCALE_BYTES +# The pool's page count is the cache manager's number, never the batch's: every case allocates +# POOL_FACTOR x the pages its batch needs (+ POOL_EXTRA) and hands out the batch's pages through a +# random permutation, so "the page the block table names" and "the page this request happens to +# own" are different numbers. Chosen here, not taken from a runtime dump. +POOL_FACTOR = 3 +POOL_EXTRA = 8 + +_E2M1 = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float64, + device=DEV, +) +_MIDPOINTS = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], device=DEV) + + +# ─── inputs ─────────────────────────────────────────────────────────────────────────────────── + + +def _index_k(S: int, seed: int) -> torch.Tensor: + """bf16 [S, 128] with per-row magnitudes spread over 2^-12..2^12 (varied UE8M0 bytes).""" + g = torch.Generator(device=DEV).manual_seed(seed) + x = torch.randn(S, DIM, generator=g, device=DEV) + x = x * torch.exp2(torch.randint(-12, 13, (S, 1), generator=g, device=DEV).float()) + if S > 3: + x[3, 32:64] = 0.0 # an all-zero 32-block + return x.to(torch.bfloat16) + + +def _quantize(x: torch.Tensor): + """The target's index-K quantizer: -> codes uint8 [S, 64], scale bytes uint8 [S, 4].""" + codes, sf = torch.ops.trtllm.fp4_quantize(x, None, 32, True, False) + S = x.shape[0] + assert codes.shape == (S, CODE_BYTES) and codes.dtype == torch.uint8 + assert sf.shape == (S * SCALE_BYTES,) and sf.dtype == torch.uint8 + return codes, sf.view(S, SCALE_BYTES) + + +def _reference_values(x: torch.Tensor) -> torch.Tensor: + """Checkpoint fp4_act_quant(k, 32, True) value set, native torch: fp64 [S, 128]. + + Per 32-block: amax clamped to 6*2^-126, scale = 2^ceil(log2(amax * (1/6))) taken from the + fp32 bits, y = E2M1_rne(clamp(x / scale, -6, 6)) * scale. + """ + xf = x.float().unflatten(-1, (4, 32)) + amax = xf.abs().amax(-1).clamp_min(6 * 2.0**-126) + bits = (amax * (1.0 / 6.0)).contiguous().view(torch.int32) + e = ((bits >> 23) & 0xFF) - 127 + ((bits & 0x7FFFFF) != 0).to(torch.int32) + scale = ((e + 127) << 23).view(torch.float32) + y = (xf / scale.unsqueeze(-1)).clamp(-6.0, 6.0) + a = y.abs() + code = (a.unsqueeze(-1) > _MIDPOINTS).sum(-1) + code = code + ((a == 0.75) | (a == 1.75) | (a == 3.5)).to(code.dtype) # ties to even code + mag = _E2M1[code] + return (torch.where(y < 0, -mag, mag) * scale.double().unsqueeze(-1)).flatten(-2) + + +def _decode(k_int8: torch.Tensor, k_sf_int32: torch.Tensor) -> torch.Tensor: + """int8 [S, 64] + int32 [S] -> fp64 [S, 128]: element 2j in the low nibble of byte j, byte b of + the little-endian int32 is the UE8M0 exponent of elements 32b..32b+31.""" + p = k_int8.view(torch.uint8).long() + v = torch.stack([_E2M1[p & 0xF], _E2M1[p >> 4]], dim=-1).flatten(-2) + w = k_sf_int32.long() & 0xFFFFFFFF + sb = torch.stack([(w >> (8 * b)) & 0xFF for b in range(4)], dim=-1) + scale = ((sb + (1023 - 127)) << 52).view(torch.float64) # 2^(byte - 127), exact + return (v.unflatten(-1, (4, 32)) * scale.unsqueeze(-1)).flatten(-2) + + +# ─── pool, block table, slot mappings (native torch; the caller's arithmetic) ─────────────────── + + +def _block_tables(lens, block_kv: int, seed: int): + """Per-request page lists from a random permutation of a pool 3x too large -> (table, pages).""" + need = [-(-n // block_kv) for n in lens] + num_pages = POOL_FACTOR * sum(need) + POOL_EXTRA + g = torch.Generator(device="cpu").manual_seed(seed) + perm = torch.randperm(num_pages, generator=g) + table = torch.full((len(lens), max(max(need), 1)), -1, dtype=torch.int32) + off = 0 + for r, n in enumerate(need): + table[r, :n] = perm[off : off + n].to(torch.int32) + off += n + return table.to(DEV), num_pages + + +def _positions(lens, table: torch.Tensor, block_kv: int): + """Full slot-mapping order (request by request, position 0..len-1) -> (page, slot) int64.""" + req = torch.repeat_interleave( + torch.arange(len(lens), device=DEV), torch.tensor(lens, device=DEV) + ) + starts = torch.tensor([0] + list(lens[:-1]), device=DEV).cumsum(0) + pos = torch.arange(sum(lens), device=DEV) - starts[req] + page = table[req, pos // block_kv].long() + return page, pos % block_kv + + +def _slot_mappings(page: torch.Tensor, slot: torch.Tensor, block_kv: int): + """Flat element indices into the logical [pages, block_kv, 1, 68] view.""" + page_elems = block_kv * KEY_BYTES + data = page * page_elems + slot * CODE_BYTES + scale = page * page_elems + block_kv * CODE_BYTES + slot * SCALE_BYTES + return data.contiguous(), scale.contiguous() + + +def _new_pool(num_pages: int, block_kv: int, seed: int, layers: int = 1, layer: int = 0): + """Random background bytes; returns the owner layer's [pages, block_kv, 1, 68] view. + + layers > 1 builds a layer-interleaved buffer [pages, layers, block_kv, 68] whose owner view has + page stride layers * block_kv * 68 (not contiguous).""" + g = torch.Generator(device=DEV).manual_seed(seed) + buf = torch.randint(0, 256, (num_pages, layers, block_kv, KEY_BYTES), generator=g, device=DEV) + buf = buf.to(torch.uint8) + return buf, buf[:, layer].unsqueeze(2) + + +def _write(pool: torch.Tensor, page, slot, codes, sf4, block_kv: int): + """Store rows in the kernel layout: a page's block_kv x 64 code bytes, then its block_kv x 4 + scale bytes.""" + flat = pool.view(pool.shape[0], block_kv * KEY_BYTES) # a view, also for the strided pool + cd = slot[:, None] * CODE_BYTES + torch.arange(CODE_BYTES, device=DEV) + cs = block_kv * CODE_BYTES + slot[:, None] * SCALE_BYTES + torch.arange(SCALE_BYTES, device=DEV) + flat[page[:, None], cd] = codes + flat[page[:, None], cs] = sf4 + + +def _build(lens, block_kv: int, seed: int, layers: int = 1): + x = _index_k(sum(lens), seed) + codes, sf4 = _quantize(x) + table, num_pages = _block_tables(lens, block_kv, seed) + buf, pool = _new_pool(num_pages, block_kv, seed + 1, layers, layers // 2) + page, slot = _positions(lens, table, block_kv) + _write(pool, page, slot, codes, sf4, block_kv) + d_idx, s_idx = _slot_mappings(page, slot, block_kv) + return dict( + x=x, + codes=codes, + sf4=sf4, + table=table, + num_pages=num_pages, + buf=buf, + pool=pool, + page=page, + slot=slot, + d_idx=d_idx, + s_idx=s_idx, + ) + + +def _row_mismatch(k, k_scale, codes, sf4) -> float: + """Fraction of rows whose 64 code bytes or 4 scale bytes differ from the expected ones.""" + bad = (k.view(torch.uint8) != codes).any(1) | ( + k_scale.view(torch.uint8).view(-1, 4) != sf4 + ).any(1) + return bad.float().mean().item() + + +def _assert_rows(k, k_scale, codes, sf4): + n = codes.shape[0] + assert k.dtype == torch.float8_e4m3fn and k.shape == (n, CODE_BYTES) and k.is_contiguous() + assert k_scale.dtype == torch.float32 and k_scale.shape == (n, 1) and k_scale.is_contiguous() + assert torch.equal(k.view(torch.uint8), codes), "code bytes differ" + assert torch.equal(k_scale.view(torch.uint8).view(n, SCALE_BYTES), sf4), "scale bytes differ" + + +# ─── certified behaviour ───────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + block_kv, layers = spec["block_kv"], spec.get("layers", 1) + b = _build(spec["lens"], block_kv, seed=zlib.crc32(cell.why.encode()), layers=layers) + if layers > 1: + # the owner layer's page view of a layer-interleaved buffer: slot mappings stay logical + # indices into [pages, block_kv, 1, 68] + assert not b["pool"].is_contiguous() + assert b["pool"].stride(0) == layers * block_kv * KEY_BYTES + before = b["buf"].clone() + for k_start, n in spec["chunks"]: + out = op(b["pool"], b["d_idx"], b["s_idx"], k_start, n, CODE_BYTES) + op.compare(out, op.reference(b["pool"], b["d_idx"], b["s_idx"], k_start, n, CODE_BYTES)) + # the gathered rows are fp4_quantize's own bytes, in fresh contiguous outputs + _assert_rows(*out, b["codes"][k_start : k_start + n], b["sf4"][k_start : k_start + n]) + # a repeat call is bitwise identical + k2, k_scale2 = op(b["pool"], b["d_idx"], b["s_idx"], k_start, n, CODE_BYTES) + assert torch.equal(k2.view(torch.uint8), out[0].view(torch.uint8)) + assert torch.equal(k_scale2.view(torch.uint8), out[1].view(torch.uint8)) + # the gather never writes the pool (nor, for a strided view, the other layers) + assert torch.equal(b["buf"], before), "the gather wrote into the pool" + + +@pytest.mark.parametrize("block_kv", [64, 128]) +def test_consumer_views_and_reference_values(block_kv): + """The int8 / int32 views DeepGEMM takes are no-copy, contiguous, 16-byte aligned, carry the + quantizer's bytes, and decode to the checkpoint's fp4_act_quant value set exactly.""" + lens = [700, 1, 2500, 37] + b = _build(lens, block_kv, seed=block_kv) + S = sum(lens) + k, k_scale = op(b["pool"], b["d_idx"], b["s_idx"], 0, S, CODE_BYTES) + kv = k.view(torch.int8) + kv_sf = k_scale.view(torch.int32).reshape(-1) + assert kv.dtype == torch.int8 and kv.shape == (S, CODE_BYTES) and kv.is_contiguous() + assert kv_sf.dtype == torch.int32 and kv_sf.shape == (S,) and kv_sf.is_contiguous() + assert kv.data_ptr() == k.data_ptr() and kv_sf.data_ptr() == k_scale.data_ptr() + assert kv_sf.data_ptr() % 16 == 0 and kv.data_ptr() % 16 == 0 + assert torch.equal(kv, b["codes"].view(torch.int8)) + assert torch.equal(kv_sf, b["sf4"].contiguous().view(torch.int32).view(-1)) + for byte in range(4): # little-endian: byte b of the word is the UE8M0 of block b + assert torch.equal((kv_sf.long() >> (8 * byte)) & 0xFF, b["sf4"][:, byte].long()) + assert torch.equal(_decode(kv, kv_sf), _reference_values(b["x"])), ( + "value set differs from fp4_act_quant" + ) + + +def test_negative_slot_rows_are_skipped_others_exact(): + """A row whose data or scale slot is negative is skipped; every other row is exact.""" + lens = [900] + block_kv = 64 + b = _build(lens, block_kv, seed=5) + d_idx, s_idx = b["d_idx"].clone(), b["s_idx"].clone() + d_idx[::5] = -1 + s_idx[2::7] = -1 + k, k_scale = op(b["pool"], d_idx, s_idx, 0, 900, CODE_BYTES) + keep = (d_idx >= 0) & (s_idx >= 0) + assert k.shape == (900, CODE_BYTES) and k_scale.shape == (900, 1) + assert torch.equal(k.view(torch.uint8)[keep], b["codes"][keep]) + assert torch.equal(k_scale.view(torch.uint8).view(900, 4)[keep], b["sf4"][keep]) + ref_k, ref_scale = op.reference(b["pool"], d_idx, s_idx, 0, 900, CODE_BYTES) + assert torch.equal(k.view(torch.uint8)[keep], ref_k.view(torch.uint8)[keep]) + assert torch.equal(k_scale.view(torch.uint8)[keep], ref_scale.view(torch.uint8)[keep]) + + +def test_consumer_accepts_gathered_view(): + """The gathered views drive DeepGEMM fp8_fp4_mqa_logits to the same logits, bit for bit, as + the quantizer's own tensors; the op's raw float8 dtype is refused by the consumer.""" + lens = [1200, 800] + b = _build(lens, 64, seed=21) + S = sum(lens) + k, k_scale = op(b["pool"], b["d_idx"], b["s_idx"], 0, S, CODE_BYTES) + T, H = 64, 32 + g = torch.Generator(device=DEV).manual_seed(22) + qx = torch.randn(T * H, DIM, generator=g, device=DEV).to(torch.bfloat16) + q_codes, q_sf = torch.ops.trtllm.fp4_quantize(qx, None, 32, True, False) + q = q_codes.view(torch.int8).view(T, H, CODE_BYTES) + q_sf = q_sf.view(torch.int32).view(T, H) + w = torch.randn(T, H, generator=g, device=DEV) * 2.0**-6 + ks = torch.zeros(T, dtype=torch.int32, device=DEV) + ke = torch.randint(1, S + 1, (T,), generator=g, device=DEV).to(torch.int32) + kv_src = (b["codes"].view(torch.int8), b["sf4"].contiguous().view(torch.int32).view(-1)) + ref = deep_gemm.fp8_fp4_mqa_logits((q, q_sf), kv_src, w, ks, ke, clean_logits=True) + got = deep_gemm.fp8_fp4_mqa_logits( + (q, q_sf), + (k.view(torch.int8), k_scale.view(torch.int32).reshape(-1)), + w, + ks, + ke, + clean_logits=True, + ) + assert torch.equal(got, ref) + with pytest.raises(RuntimeError): + deep_gemm.fp8_fp4_mqa_logits( + (q, q_sf), (k, k_scale.view(torch.int32).reshape(-1)), w, ks, ke, clean_logits=True + ) + + +# ─── rejections (each probed on sm_103) ────────────────────────────────────────────────────── + + +def test_rejections(): + b = _build([37], 64, seed=9) + pool, d, s = b["pool"], b["d_idx"], b["s_idx"] + cases = [ + ( + lambda: op(pool, d, s, 0, 37, 96), + "head_dim must be 128 \\(FP8\\) or 64", + ), + ( + lambda: op(pool.view(pool.shape[0], 64, 68), d, s, 0, 37, 64), + "k_cache must be a 4D Tensor", + ), + ( + lambda: op(pool, d.int(), s.int(), 0, 37, 64), + "slot_mapping_fp8 must be int64", + ), + (lambda: op(pool, d, s, 1, 37, 64), "slot_mapping_fp8 too short"), + ( + lambda: op(pool.view(pool.shape[0], 32, 2, 68), d, s, 0, 37, 64), + "k_cache dimension 2 must be 1", + ), + ( + lambda: op(pool, d.cpu(), s.cpu(), 0, 37, 64), + "All tensors must be CUDA tensors", + ), + ( + lambda: op(pool, d.view(1, 37), s.view(1, 37), 0, 37, 64), + "slot_mapping_fp8 must be a 1D Tensor", + ), + ( + lambda: op(pool, torch.stack([d, d], 1)[:, 0], s, 0, 37, 64), + "slot_mapping_fp8 must be contiguous", + ), + ( + lambda: op(pool.view(torch.int8), d, s, 0, 37, 64), + "expected scalar type Byte but found Char", + ), + ] + # The op's own rejections, outside `validating` (the guard would fire first on head_dim 96). + for fn, msg in cases: + with pytest.raises(RuntimeError, match=msg): + fn() + # The entry's guard: head_dim 128 on the 68-byte pool. The op itself accepts it and returns the + # neighbouring key's bytes (test_controls_fail shows it outside `validating`). + with pytest.raises(AssertionError, match="must equal k_cache.size\\(3\\) - 4"): + with validating(op): + op(pool, d, s, 0, 37, 128) + + +# ─── deliberately wrong addressing schemes: every one must be visibly wrong ────────────────── + + +@pytest.mark.parametrize("block_kv", [64, 128]) +def test_controls_fail(block_kv): + lens = [700, 1, 2500, 37] + b = _build(lens, block_kv, seed=31 + block_kv) + S = sum(lens) + codes, sf4 = b["codes"], b["sf4"] + page, slot = b["page"], b["slot"] + + def run(d_idx, s_idx, head_dim=CODE_BYTES): + k, k_scale = op(b["pool"], d_idx, s_idx, 0, S, head_dim) + return k[:, :CODE_BYTES].contiguous(), k_scale + + good = _row_mismatch(*run(b["d_idx"], b["s_idx"]), codes, sf4) + assert good == 0.0 + + fractions = {} + # 1. wrong block table: the same requests paged through another permutation of the pool (a + # stale table, or another owner layer's) -- every entry is a valid page, just not theirs + wrong_table, _ = _block_tables(lens, block_kv, seed=999 + block_kv) + wrong_page, _ = _positions(lens, wrong_table, block_kv) + fractions["wrong block table"] = _row_mismatch( + *run(*_slot_mappings(wrong_page, slot, block_kv)), codes, sf4 + ) + # 2. off-by-one page offset: page id + 1 (kept inside the pool) + fractions["page id + 1"] = _row_mismatch( + *run(*_slot_mappings((page + 1) % b["num_pages"], slot, block_kv)), codes, sf4 + ) + # 3. off-by-one slot inside the page (wraps inside the same page) + fractions["slot + 1 within the page"] = _row_mismatch( + *run(*_slot_mappings(page, (slot + 1) % block_kv, block_kv)), codes, sf4 + ) + # 4. key-interleaved addressing ([64 codes, 4 scales] per key) over the paged layout + pe = block_kv * KEY_BYTES + fractions["key-interleaved addressing"] = _row_mismatch( + *run(page * pe + slot * KEY_BYTES, page * pe + slot * KEY_BYTES + CODE_BYTES), codes, sf4 + ) + # 5. page stride taken as block_kv x 64 (scale region forgotten) + fractions["page stride without the scale region"] = _row_mismatch( + *run( + page * block_kv * CODE_BYTES + slot * CODE_BYTES, + page * block_kv * CODE_BYTES + block_kv * CODE_BYTES + slot * SCALE_BYTES, + ), + codes, + sf4, + ) + + # 6. the FP8 head_dim (128) through the unguarded op on the 68-byte pool is accepted: bytes 64..127 + # of a row are the next slot's codes (or, for a page's last slot, its scale region) + k128, _ = op(b["pool"], b["d_idx"], b["s_idx"], 0, S, 128) + assert k128.shape == (S, 128) + assert torch.equal(k128.view(torch.uint8)[:, :CODE_BYTES], codes) + nxt = torch.cat([codes[1:], codes[:1]]) + same_page_next = torch.cat( + [(page[1:] == page[:-1]) & (slot[1:] == slot[:-1] + 1), torch.tensor([False], device=DEV)] + ) + upper = k128.view(torch.uint8)[:, CODE_BYTES:] + assert torch.equal(upper[same_page_next], nxt[same_page_next]), "upper half is not the next key" + frac_next = same_page_next.float().mean().item() + + print(f"\n[controls block_kv={block_kv} S={S}] row-mismatch fractions (correct run: {good}):") + for name, frac in fractions.items(): + print(f" {name}: {frac:.4f}") + print( + f" head_dim 128 on the MXFP4 pool: accepted; rows whose bytes 64..127 are the next key's " + f"codes: {frac_next:.4f}" + ) + assert frac_next >= 0.9 + for name, frac in fractions.items(): + assert frac >= 0.9, f"control '{name}' was not caught (mismatch fraction {frac})" + + +# ─── the engine's pools, every row count, one kernel ───────────────────────────────────────── +# +# The core target calls this op only in a step with context tokens (`ModelingV2Core._index`, +# `if tc:` branch), with k_token_start 0 and num_tokens = v41_ctx_index_k_len[r] +# (`_prepare_v41_ctx_index_k` of the V4.1 metadata: S_r = sum over the step's context requests of +# kv_len // r, cached entries included). Context and mixed steps run eagerly, so the call is never +# captured. One request contributes at most Lm // r rows, Lm = max_seq_len (+ 14 with DSpark); the +# longest request of the envelope (max_seq_len 1,048,576, DSpark on) gives 1,048,590 at ratio 1 and +# 524,295 at ratio 2, the N_max of the exhaustive row sweep below. A step with several long context +# requests sums them (up to max_batch_size 256 of them); that axis is swept on the largest pool by +# test_rows_to_2_25_on_b_max_pool and bounded by the defect test after it. Observed in the core +# runtime dumps: up to 2938 (boot, r = 1), 24058 (the 24K passkey run, r = 1) and 12029 (r = 2). +# +# The pools are sized from the same dumps. The cache manager allocates B blocks; the ratio-2 +# INDEXER_COMPRESS allocation holds 3 pages of 64 entries per block, which owners 2, 8 and 14 see +# as contiguous views starting 0, 1 and 2 pages into it ([3B - shift, 64, 1, 68]; their page ids +# are 3 x block), and the ratio-1 pool (owner 20) holds one page of 128 entries per block +# ([B, 128, 1, 68], page id = block). B = 137386 / 142336 (boot dumps) and 151040 (the largest +# pool of the runtime dumps, 453120 x 64 and 151040 x 128). +ENGINE_BLOCKS = (137386, 142336, 151040) +OWNER_SHIFT = {2: 0, 8: 1, 14: 2} +N_MAX = {1: 1_048_590, 2: 524_295} +DUMPED_ROWS = { # s_len values the core runtime dumps hold (boot gates and an accuracy-gate run) + 1: (1, 2, 16, 48, 323, 576, 1938, 2048, 2050, 2938, 4618, 8192, 16384, 24058), + 2: (1, 8, 24, 161, 288, 969, 1024, 1025, 1469, 2309, 4096, 8192, 12029), +} +# One request's whole index K at each certified max_seq_len L: Lm // r with Lm = L + 14 (DSpark on) +# and Lm = L (DSpark off). +L_VALUES = (16384, 32768, 65536, 131072, 262144, 1048576) +ONE_REQUEST_ROWS = { + r: tuple(sorted({(L + d) // r for L in L_VALUES for d in (0, 14)})) for r in (1, 2) +} +_ENGINE_CACHE = {} + + +def _engine_alloc(ratio: int, blocks: int, seed: int = 0) -> torch.Tensor: + """The ratio's whole INDEXER_COMPRESS allocation at engine size, random background bytes. + One per (ratio, blocks) per process; the gather never writes it (checked where used).""" + key = (ratio, blocks) + if key not in _ENGINE_CACHE: + _ENGINE_CACHE.clear() # keep one engine-size allocation alive at a time + torch.cuda.empty_cache() + pages = 3 * blocks if ratio == 2 else blocks + g = torch.Generator(device=DEV).manual_seed(seed + blocks + ratio) + _ENGINE_CACHE[key] = torch.randint( + 0, 256, (pages, 128 // ratio, 1, KEY_BYTES), dtype=torch.uint8, device=DEV, generator=g + ) + return _ENGINE_CACHE[key] + + +def _owner_view(alloc: torch.Tensor, ratio: int, owner: int) -> torch.Tensor: + """What get_indexer_k_cache_buffers(owner) returns: a contiguous view `shift` pages in.""" + view = alloc[OWNER_SHIFT[owner] :] if ratio == 2 else alloc + assert view.is_contiguous() + return view + + +def _engine_table(lens, block_kv: int, blocks: int, ratio: int, seed: int): + """Page ids for the requests, drawn from every block of the pool: the first request's first + page is block 0 and the last request's last page is block `blocks - 1` (the far end); page id = + 3 x block at ratio 2 (the shared-block scale of the three ratio-2 owners), = block at ratio 1.""" + need = [-(-n // block_kv) for n in lens] + total = sum(need) + g = torch.Generator(device="cpu").manual_seed(seed) + mid = torch.randperm(blocks - 2, generator=g)[: max(total - 2, 0)] + 1 + chosen = torch.cat([torch.tensor([0]), mid, torch.tensor([blocks - 1])]) + table = torch.full((len(lens), max(need)), -1, dtype=torch.int64) + off = 0 + for r, n in enumerate(need): + table[r, :n] = chosen[off : off + n] + off += n + return (table * (3 if ratio == 2 else 1)).to(DEV) + + +def _engine_rows(ratio: int, blocks: int, owner: int, lens, seed: int): + """Write fp4_quantize bytes for `lens` requests into the owner's view of the engine-size pool + (native torch indexing) and return the view, the slot mappings and the expected bytes.""" + block_kv = 128 // ratio + alloc = _engine_alloc(ratio, blocks) + view = _owner_view(alloc, ratio, owner) + table = _engine_table(lens, block_kv, blocks, ratio, seed) + page, slot = _positions(lens, table, block_kv) + x = _index_k(sum(lens), seed) + codes, sf4 = _quantize(x) + _write(view, page, slot, codes, sf4, block_kv) + d_idx, s_idx = _slot_mappings(page, slot, block_kv) + return alloc, view, page, slot, d_idx, s_idx, codes, sf4 + + +@pytest.mark.parametrize("blocks", ENGINE_BLOCKS) +@pytest.mark.parametrize("ratio,owner", [(2, 2), (2, 8), (2, 14), (1, 20)]) +def test_engine_pool_far_end(ratio, owner, blocks): + """Every owner's view of the engine-size pool: requests paged over the whole pool, the last one + ending on the last slot of block `blocks - 1` (for owner 14 at ratio 2 and owner 20 that is the + last byte of the allocation). Rows are the quantizer's bytes; the allocation is unchanged.""" + block_kv = 128 // ratio + lens = [3 * block_kv, 5, 2 * block_kv + 7, 4 * block_kv] + alloc, view, page, slot, d_idx, s_idx, codes, sf4 = _engine_rows( + ratio, blocks, owner, lens, seed=blocks + owner + ) + last_byte = OWNER_SHIFT.get(owner, 0) * block_kv * KEY_BYTES + int(s_idx[-1]) + SCALE_BYTES + assert ( + int(page[-1]) == (blocks - 1) * (3 if ratio == 2 else 1) and int(slot[-1]) == block_kv - 1 + ) + before = alloc.clone() + k, k_scale = op(view, d_idx, s_idx, 0, sum(lens), CODE_BYTES) + _assert_rows(k, k_scale, codes, sf4) + assert torch.equal(alloc, before), "the gather wrote into the pool" + print( + f"[engine pool r={ratio} owner={owner}] view {tuple(view.shape)} alloc {alloc.numel()} B; " + f"far row's scale word ends at byte {last_byte} of the allocation; bit-exact" + ) + del before + + +@pytest.mark.parametrize("ratio", [1, 2]) +def test_every_row_count_prefix_bitwise(ratio): + """Every row count on the target's call (k_token_start 0, num_tokens r), over the engine-size pool. + + (a) the N_max call returns the quantizer's bytes for all N_max rows (32 requests paged over the + whole pool, the far end included); (b) every r in 1 .. N_max returns exactly the first r rows + of that call, bitwise. The sweep alternates between two row orders (the second is the first + rolled by N_max // 2 + 1), so a call that left rows unwritten in a recycled allocation would + show the other order's bytes; (c) the dumped s_len values with a fresh packed batch each, and + one request's whole index K at every certified max_seq_len (ONE_REQUEST_ROWS), each a fresh + single-request batch paged over the pool.""" + n_max = N_MAX[ratio] + owner = 14 if ratio == 2 else 20 + blocks = ENGINE_BLOCKS[-1] + lens = [n_max // 32] * 32 + lens[-1] += n_max - sum(lens) + alloc, view, page, slot, d_a, s_a, codes, sf4 = _engine_rows( + ratio, blocks, owner, lens, seed=ratio + ) + before = alloc.clone() + k_full, ks_full = op(view, d_a, s_a, 0, n_max, CODE_BYTES) + _assert_rows(k_full, ks_full, codes, sf4) # (a) + ref = [(k_full.view(torch.uint8), ks_full.view(torch.uint8).view(n_max, SCALE_BYTES))] + shift = n_max // 2 + 1 + d_b, s_b = d_a.roll(shift).contiguous(), s_a.roll(shift).contiguous() + ref.append((ref[0][0].roll(shift, 0).contiguous(), ref[0][1].roll(shift, 0).contiguous())) + maps = [(d_a, s_a), (d_b, s_b)] + assert (ref[1][0] != ref[0][0]).any(1).float().mean().item() > 0.99 + bad = torch.zeros((), dtype=torch.int64, device=DEV) + for r in range(1, n_max + 1): # (b) + v = r & 1 + k, ks = op(view, maps[v][0], maps[v][1], 0, r, CODE_BYTES) + bad += (k.view(torch.uint8) != ref[v][0][:r]).any(1).sum() + bad += (ks.view(torch.uint8).view(r, SCALE_BYTES) != ref[v][1][:r]).any(1).sum() + bad = int(bad) + assert torch.equal(alloc, before), "the gather wrote into the pool" + del before + print( + f"[rows r={ratio}] N_max {n_max} on {tuple(view.shape)}: every r in 1..{n_max} " + f"prefix-bitwise, mismatching rows {bad}" + ) + assert bad == 0 + del k_full, ks_full, ref, maps, d_a, s_a, d_b, s_b, codes, sf4 + for s_len in DUMPED_ROWS[ratio]: # (c) + g = torch.Generator(device="cpu").manual_seed(s_len) + cuts = sorted(torch.randperm(s_len - 1, generator=g)[: min(3, s_len - 1)].add(1).tolist()) + req_lens = [b - a for a, b in zip([0] + cuts, cuts + [s_len])] + _, view, _, _, d, s, c, sf = _engine_rows( + ratio, blocks, owner, req_lens, seed=s_len + ratio + ) + k, ks = op(view, d, s, 0, s_len, CODE_BYTES) + _assert_rows(k, ks, c, sf) + for s_len in ONE_REQUEST_ROWS[ratio]: + _, view, page, _, d, s, c, sf = _engine_rows( + ratio, blocks, owner, [s_len], seed=s_len + 7 * ratio + ) + assert int(page[-1]) == (blocks - 1) * ( + 3 if ratio == 2 else 1 + ) # its last page is the far end + k, ks = op(view, d, s, 0, s_len, CODE_BYTES) + _assert_rows(k, ks, c, sf) + print( + f"[rows r={ratio}] dumped s_len {DUMPED_ROWS[ratio]} and one request's rows at every L " + f"{ONE_REQUEST_ROWS[ratio]}: bit-exact" + ) + + +def test_one_kernel_at_every_row_count(): + """The op is an ahead-of-time compiled kernel whose row count is only its grid: the same single + kernel launches at 1 row, at dumped counts and at N_max (no per-shape variant).""" + from torch.autograd import DeviceType + from torch.profiler import ProfilerActivity, profile + + n_max = N_MAX[1] + _, view, _, _, d, s, _, _ = _engine_rows(1, ENGINE_BLOCKS[-1], 20, [n_max], seed=5) + seen = set() + for r in (1, 2, 1025, 2938, 24058, n_max): + with profile(activities=[ProfilerActivity.CUDA]) as prof: + op(view, d, s, 0, r, CODE_BYTES) + torch.cuda.synchronize() + kernels = [e.name for e in prof.events() if e.device_type == DeviceType.CUDA] + print(f"[kernels r={r}] {kernels}") + assert len(kernels) == 1 and "indexerKCacheGatherUnifiedKernel" in kernels[0], kernels + seen.add(kernels[0]) + assert len(seen) == 1 + + +# Beyond the engine's pools: an allocation past 2^32 bytes, rows on the pages that straddle the +# 2^31- and 2^32-byte offsets and on the last page. The kernel unravels the int64 flat index and +# forms an int64 byte offset; a caller that computed the offsets in 32 bits would wrap them. +MARGIN_PAGES = {64: 1_000_000, 128: 500_000} # 4.352e9 bytes either way + + +@pytest.mark.parametrize("block_kv", [64, 128]) +def test_past_2_to_the_32_bytes_and_wrap_control(block_kv): + page_bytes = block_kv * KEY_BYTES + pages = MARGIN_PAGES[block_kv] + _ENGINE_CACHE.clear() + torch.cuda.empty_cache() + g = torch.Generator(device=DEV).manual_seed(block_kv) + pool = torch.randint( + 0, 256, (pages, block_kv, 1, KEY_BYTES), dtype=torch.uint8, device=DEV, generator=g + ) + assert pool.numel() > 2**32 + far = [2**31 // page_bytes, 2**32 // page_bytes, pages - 1] + for p, b in zip(far[:2], (2**31, 2**32)): + assert p * page_bytes < b < (p + 1) * page_bytes # the page straddles the boundary + lens = [block_kv] * 3 + table = torch.tensor(far, dtype=torch.int64, device=DEV).view(3, 1) + page, slot = _positions(lens, table, block_kv) + x = _index_k(sum(lens), seed=block_kv) + codes, sf4 = _quantize(x) + _write(pool, page, slot, codes, sf4, block_kv) + d_idx, s_idx = _slot_mappings(page, slot, block_kv) + k, k_scale = op(pool, d_idx, s_idx, 0, sum(lens), CODE_BYTES) + _assert_rows(k, k_scale, codes, sf4) + + # control: the same offsets wrapped to signed 32 bits (negative ones are skipped by the op, + # the rest read the wrong bytes) + def wrap(t): + return ((t + 2**31) % 2**32) - 2**31 + + kw, ksw = op(pool, wrap(d_idx), wrap(s_idx), 0, sum(lens), CODE_BYTES) + frac = _row_mismatch(kw, ksw, codes, sf4) + past = (d_idx >= 2**31).float().mean().item() + print( + f"[margin block_kv={block_kv}] pool {pool.numel()} B, pages {far}: bit-exact; int32-wrapped " + f"offsets: {frac:.4f} of rows wrong ({past:.4f} of rows lie past 2^31)" + ) + assert frac >= past > 0.5 + del pool + + +@pytest.mark.parametrize("ratio", [2]) +def test_control_wrong_owner_view(ratio): + """Owner 14's page ids read through owner 2's view (the three ratio-2 views overlap, shifted by + one page each): every row is another owner's bytes.""" + blocks = ENGINE_BLOCKS[-1] + lens = [640, 3, 200] + alloc, view14, page, slot, d, s, codes, sf4 = _engine_rows(ratio, blocks, 14, lens, seed=77) + view2 = _owner_view(alloc, ratio, 2) + k, ks = op(view14, d, s, 0, sum(lens), CODE_BYTES) + _assert_rows(k, ks, codes, sf4) + kw, ksw = op(view2, d, s, 0, sum(lens), CODE_BYTES) + frac = _row_mismatch(kw, ksw, codes, sf4) + print(f"[control] owner 14's ids through owner 2's view: {frac:.4f} of rows wrong") + assert frac >= 0.9 + + +# Pool-size invariance. Every V4.1 persistent role of one cache manager shares one block +# count B, so B is capped by the whole GPU (284208 MiB = 298,013,687,808 B) divided by the manager's +# persistent bytes per block. The cheapest manager that holds an index-K pool is the PP2 stage-1 +# manager in native storage (layers 20-39: 356 B/token, 45,568 B per block, measured); an unsplit +# native manager costs 113,920 B per block and a stage-0 one 68,352. So no V4.1 manager on a GB300 +# holds more than floor(298,013,687,808 / 45,568) = 6,539,977 blocks. The largest pool measured (a +# stage-1 manager built alone at half the GPU) is 3,245,959 blocks, and the unsplit ceiling 2,615,990; both +# are run as intermediate sizes. One B_MAX serves both ratios (the ratio-2 owners live only in +# unsplit and stage-0 managers, whose own ceilings are lower), so the ratio-2 allocation at B_MAX, +# 3 x 6,539,977 pages of 64 entries (85,385,939,712 B), is a superset of any real one. +# The same rows, at the same page ids, are gathered from exact-size allocations (the pool ends +# where the allocation ends) from B_min, the smallest pool that holds their pages, up to that +# bound. Pages no row owns hold POISON, so a read outside the rows' pages would show. +B_MAX = 6_539_977 +INV_BLOCKS = (481, 1493, 151040, 852832, 2_615_990, 3_245_959, B_MAX) +B_OTHER = 852832 # the control's wrong page count +POISON = 0xFF +FAR_BYTES = (2**31, 2**32, 2**33, 2**34, 2**35, 2**36) # those below the allocation's size + + +def _poisoned_alloc(ratio: int, blocks: int) -> torch.Tensor: + """The ratio's whole INDEXER_COMPRESS allocation for `blocks` blocks, every byte POISON.""" + _ENGINE_CACHE.clear() + torch.cuda.empty_cache() + pages = 3 * blocks if ratio == 2 else blocks + return torch.full((pages, 128 // ratio, 1, KEY_BYTES), POISON, dtype=torch.uint8, device=DEV) + + +def _wrap32(t: torch.Tensor) -> torch.Tensor: + return (((t + 2**31) % 2**32) - 2**31).contiguous() + + +def _owned_pages(ratio: int, rows, seed: int): + """The sorted allocation pages the rows live on and their content: random background bytes + (independent of B) with every row's code bytes and scale word placed in the kernel layout. + rows = [(owner, view page ids, slots, codes, scale bytes)].""" + block_kv = 128 // ratio + owned = torch.unique(torch.cat([OWNER_SHIFT.get(o, 0) + p for o, p, *_ in rows])) + g = torch.Generator(device=DEV).manual_seed(seed) + content = torch.randint( + 0, 256, (owned.numel(), block_kv * KEY_BYTES), generator=g, device=DEV + ).to(torch.uint8) + cd = torch.arange(CODE_BYTES, device=DEV) + cs = torch.arange(SCALE_BYTES, device=DEV) + for owner, page, slot, codes, sf4 in rows: + i = torch.searchsorted(owned, OWNER_SHIFT.get(owner, 0) + page) + content[i[:, None], slot[:, None] * CODE_BYTES + cd] = codes + content[i[:, None], block_kv * CODE_BYTES + slot[:, None] * SCALE_BYTES + cs] = sf4 + return owned, content + + +def _native_gather(alloc: torch.Tensor, ratio: int, owner: int, d: torch.Tensor, s: torch.Tensor): + """Native torch: the rows' bytes read from the whole allocation at the view's flat indices.""" + base = OWNER_SHIFT.get(owner, 0) * (128 // ratio) * KEY_BYTES + flat = alloc.view(-1) + return ( + flat[(base + d)[:, None] + torch.arange(CODE_BYTES, device=DEV)], + flat[(base + s)[:, None] + torch.arange(SCALE_BYTES, device=DEV)], + ) + + +def _pool_unchanged(alloc: torch.Tensor, owned: torch.Tensor, content: torch.Tensor) -> bool: + """The owned pages still hold `content` and every other byte of the allocation is POISON.""" + if not torch.equal(alloc[owned].view(owned.numel(), -1), content): + return False + alloc[owned] = POISON + word = torch.full((8,), POISON, dtype=torch.uint8, device=DEV).view(torch.int64) + ok = not bool((alloc.view(-1).view(torch.int64) != word).any()) + alloc[owned] = content.view(owned.numel(), -1, 1, KEY_BYTES) + return ok + + +def _gather_all(alloc, ratio, rows, wrap=False, page_fn=None): + """One call per owner through its view -> [(k bytes [n, 64], scale bytes [n, 4], d, s)].""" + block_kv = 128 // ratio + out = [] + for owner, page, slot, codes, _ in rows: + d, s = _slot_mappings(page if page_fn is None else page_fn(page), slot, block_kv) + if wrap: + d, s = _wrap32(d), _wrap32(s) + k, ks = op(_owner_view(alloc, ratio, owner), d, s, 0, codes.size(0), CODE_BYTES) + out.append((k.view(torch.uint8), ks.view(torch.uint8).view(-1, SCALE_BYTES), d, s)) + torch.cuda.synchronize() + return out + + +def _invariance_rows(ratio: int): + """Five requests (3 pages + 11 entries, 5, 2 pages + 7, 5000, and 4 whole pages ending on the last + slot of block B_min - 1), every compressed position 0 .. len-1 as the target's gather lists them, + over B_min blocks, each used once; page ids 3 x block at ratio 2. Every owner of the ratio holds + them with its own keys. Returns (B_min, rows).""" + block_kv = 128 // ratio + lens = [3 * block_kv + 11, 5, 2 * block_kv + 7, 5000, 4 * block_kv] + b_min = sum(-(-n // block_kv) for n in lens) + table = _engine_table(lens, block_kv, b_min, ratio, seed=60 + ratio) + page, slot = _positions(lens, table, block_kv) + assert int(page[-1]) == (b_min - 1) * (3 if ratio == 2 else 1) and int(slot[-1]) == block_kv - 1 + rows = [] + for owner in (2, 8, 14) if ratio == 2 else (20,): + codes, sf4 = _quantize(_index_k(sum(lens), seed=700 + owner)) + rows.append((owner, page, slot, codes, sf4)) + return b_min, rows + + +@pytest.mark.parametrize("ratio", [1, 2]) +def test_pool_size_invariance(ratio): + """The same rows at the same page ids, gathered through every owner's view of the + manager layout (ratio 2: one 3B-page allocation, owners 2 / 8 / 14 at 0 / 1 / 2 pages in; + ratio 1: [B, 128, 1, 68]) at B = B_min, 481, 1493, 151040, 852832 and B_MAX, each an exact-size + allocation with every unowned page POISON. At every B both outputs are bitwise identical to + B_min's, to a native-torch gather of the allocation's bytes and to the quantizer's bytes, and + the allocation is unchanged. Control: at B_MAX the page ids counted from the end of a + B_OTHER-block pool (a B-dependent addressing) fail the same check.""" + b_min, rows = _invariance_rows(ratio) + owned, content = _owned_pages(ratio, rows, seed=80 + ratio) + sizes = (b_min,) + INV_BLOCKS + assert list(sizes) == sorted(sizes) + first = None + for blocks in sizes: + alloc = _poisoned_alloc(ratio, blocks) + alloc[owned] = content.view(owned.numel(), -1, 1, KEY_BYTES) + out = _gather_all(alloc, ratio, rows) + first = out if first is None else first + same = all(torch.equal(a[0], b[0]) and torch.equal(a[1], b[1]) for a, b in zip(out, first)) + native = True + quant = True + for (owner, _, _, codes, sf4), (k, ks, d, s) in zip(rows, out): + rk, rs = _native_gather(alloc, ratio, owner, d, s) + native &= torch.equal(k, rk) and torch.equal(ks, rs) + quant &= torch.equal(k, codes) and torch.equal(ks, sf4) + kept = _pool_unchanged(alloc, owned, content) + print( + f"[invariance r={ratio}] B={blocks}: allocation {tuple(alloc.shape)} ({alloc.numel()} B), " + f"{sum(r[3].size(0) for r in rows)} rows on {owned.numel()} pages: equal to B_min {same}, " + f"to the native gather {native}, to the quantizer {quant}, pool unchanged {kept}" + ) + assert same and native and quant and kept, f"B={blocks}: result depends on the pool size" + if blocks != sizes[-1]: + del alloc + k3 = 3 if ratio == 2 else 1 + wrong = _gather_all(alloc, ratio, rows, page_fn=lambda p: k3 * (B_OTHER - 1) - p) + frac = min(_row_mismatch(k, ks, c, sf) for (_, _, _, c, sf), (k, ks, _, _) in zip(rows, wrong)) + print( + f"[invariance control r={ratio}] page ids counted from the end of a {B_OTHER}-block pool on " + f"B={B_MAX}: {frac:.4f} of rows wrong (lowest over owners)" + ) + assert frac >= 0.9 + del alloc, first, out, wrong + + +def _far_rows(ratio: int, blocks: int, seed: int): + """Every slot of the allocation pages that straddle each FAR_BYTES offset and of the last page + of block blocks - 1 of every owner, read through the owner whose view holds the page + (allocation page p belongs to owner shift p % 3 at ratio 2, view page id p - shift = 3 x block).""" + block_kv = 128 // ratio + page_bytes = block_kv * KEY_BYTES + alloc_bytes = (3 * blocks if ratio == 2 else blocks) * page_bytes + marks = [x for x in FAR_BYTES if x < alloc_bytes] + pages = [x // page_bytes for x in marks] + for x, p in zip(marks, pages): + assert p * page_bytes < x < (p + 1) * page_bytes + pages += [3 * blocks - 3, 3 * blocks - 2, 3 * blocks - 1] if ratio == 2 else [blocks - 1] + rows = [] + for owner in (2, 8, 14) if ratio == 2 else (20,): + shift = OWNER_SHIFT.get(owner, 0) + mine = [p - shift for p in pages if ratio == 1 or p % 3 == shift] + if not mine: + continue + assert all(m % (3 if ratio == 2 else 1) == 0 for m in mine) + page = torch.tensor(mine, dtype=torch.int64, device=DEV).repeat_interleave(block_kv) + slot = torch.arange(block_kv, device=DEV).repeat(len(mine)) + codes, sf4 = _quantize(_index_k(page.numel(), seed=seed + owner)) + rows.append((owner, page, slot, codes, sf4)) + return rows + + +@pytest.mark.parametrize("ratio", [1, 2]) +def test_far_end_at_b_max(ratio): + """Far-end addressing on the B_MAX allocation (ratio 2: 19,619,931 pages, 85,385,939,712 B; + ratio 1: 6,539,977 pages, 56,923,959,808 B), every other page POISON: every slot of the pages + straddling bytes 2^31 .. 2^35 (and 2^36 at ratio 2) and of the last page of block B_MAX - 1 of every owner + (for owners 14 and 20 the last slot's scale word is the allocation's last 4 bytes). Both whole + outputs equal a native-torch gather of the allocation and the quantizer's bytes. Control: the + same offsets wrapped to signed 32 bits fail.""" + rows = _far_rows(ratio, B_MAX, seed=950 + ratio) + owned, content = _owned_pages(ratio, rows, seed=990 + ratio) + alloc = _poisoned_alloc(ratio, B_MAX) + assert alloc.numel() > 2**35 + alloc[owned] = content.view(owned.numel(), -1, 1, KEY_BYTES) + out = _gather_all(alloc, ratio, rows) + ok = True + refs = [] # kept alive through the control, so no freed block holding correct rows is reused + for (owner, _, _, codes, sf4), (k, ks, d, s) in zip(rows, out): + refs.append(_native_gather(alloc, ratio, owner, d, s)) + rk, rs = refs[-1] + ok &= ( + torch.equal(k, rk) + and torch.equal(ks, rs) + and torch.equal(k, codes) + and torch.equal(ks, sf4) + ) + last = rows[-1] + base = OWNER_SHIFT.get(last[0], 0) * (128 // ratio) * KEY_BYTES + tail = int(base + out[-1][3][-1]) + SCALE_BYTES == alloc.numel() + kept = _pool_unchanged(alloc, owned, content) + print( + f"[far end r={ratio}] B={B_MAX}, allocation {alloc.numel()} B, pages {owned.tolist()}: native " + f"gather and quantizer bytes {ok}, owner {last[0]}'s last scale word ends the allocation {tail}, " + f"pool unchanged {kept}" + ) + assert ok and tail and kept + wrong = _gather_all(alloc, ratio, rows, wrap=True) + frac = sum( + _row_mismatch(k, ks, c, sf) * c.size(0) + for (_, _, _, c, sf), (k, ks, _, _) in zip(rows, wrong) + ) + n = sum(r[3].size(0) for r in rows) + frac /= n + past = sum(int(((d >= 2**31) | (s >= 2**31)).sum()) for _, _, d, s in out) / n + print( + f"[far end control r={ratio}] offsets wrapped to signed 32 bits ({past:.4f} of rows past 2^31): " + f"{frac:.4f} of rows wrong" + ) + assert frac >= past > 0.5 + del alloc, out, wrong, refs + + +# ─── many long context requests in one step: rows past one request, up to the step maximum ─── +# +# S_r sums every context request of the step, so an engine of the envelope (max_batch_size 256, +# max_seq_len 1,048,576) can pass up to 256 x 1,048,590 = 268,439,040 rows at ratio 1 and +# 134,219,520 at ratio 2; a B_MAX-block pool holds 837,117,056 entries at ratio 1, so every such +# row can address a distinct slot. Calls that large run in tens of milliseconds, so above one +# request's N_max a sample of counts replaces the exhaustive sweep: densely up to ROW_BOUND, then +# coarsely to the step maximum. +# +# ROW_BOUND is 2^25 because that is where the kernel's output offset `token_idx * head_dim` used +# to overflow 32 bits. It is only an interval boundary now: the offsets are computed in 64 bits +# since `static_cast(token_idx)` was added to indexerKCacheGather.cu, and the counts on +# both sides of it are certified below. +ROW_BOUND = 2**25 +N_STEP_MAX = {r: 256 * N_MAX[r] for r in (1, 2)} # max_batch_size context requests of Lm // r rows +_BIJ_A = 1_000_003 # odd prime, coprime with every B_MAX x block_kv slot count + + +def _bijection_rows(ratio: int, n: int, c: int, far_ends: bool = True): + """Row i of n reads slot (A * i + c) mod P of the owner's B_MAX-block view, P = B_MAX x block_kv + (a bijection, so rows < P read distinct slots spread over the whole pool); with far_ends, row 0 + is slot 0 of block 0 and row n - 1 the last slot of block B_MAX - 1. -> (view page ids, slots).""" + block_kv = 128 // ratio + k = 3 if ratio == 2 else 1 + sg = (torch.arange(n, dtype=torch.int64, device=DEV) * _BIJ_A + c) % (B_MAX * block_kv) + page, slot = (sg // block_kv) * k, sg % block_kv + if far_ends: + page[0], slot[0] = 0, 0 + page[-1], slot[-1] = (B_MAX - 1) * k, block_kv - 1 + return page, slot + + +def _native_rows(alloc, ratio, owner, d, s, lo, hi): + """Native torch: rows [lo, hi) read from the whole allocation at the owner view's flat indices.""" + base = OWNER_SHIFT.get(owner, 0) * (128 // ratio) * KEY_BYTES + flat = alloc.view(-1) + return ( + flat[(base + d[lo:hi])[:, None] + torch.arange(CODE_BYTES, device=DEV)], + flat[(base + s[lo:hi])[:, None] + torch.arange(SCALE_BYTES, device=DEV)], + ) + + +def _rows_wrong(k8, ks8, alloc, ratio, owner, d, s, lo, hi, chunk=1 << 22): + """(code rows, scale rows) of output rows [lo, hi) that differ from the native gather.""" + bc = torch.zeros((), dtype=torch.int64, device=DEV) + bs = torch.zeros((), dtype=torch.int64, device=DEV) + for a in range(lo, hi, chunk): + b = min(hi, a + chunk) + rc, rs = _native_rows(alloc, ratio, owner, d, s, a, b) + bc += (k8[a:b] != rc).any(1).sum() + bs += (ks8[a:b] != rs).any(1).sum() + return int(bc), int(bs) + + +def _checksum(alloc: torch.Tensor) -> torch.Tensor: + """Wrapping sum of the allocation's int64 words: a changed word changes it.""" + return alloc.view(-1).view(torch.int64).sum() + + +def _sampled_counts(ratio: int): + """Row counts in (N_MAX, ROW_BOUND]: every r = 0, +1, -1 mod 2^16 (one CUDA block per row, so the + grid has no row tile; 2^16 steps cover every power-of-two boundary from 2^16 up), every sum of + nc whole requests at one certified max_seq_len (nc = 2..256), and ROW_BOUND - 1, ROW_BOUND.""" + lo, hi = N_MAX[ratio], ROW_BOUND + counts = {m * 2**16 + e for m in range(lo // 2**16, hi // 2**16 + 1) for e in (-1, 0, 1)} + counts |= {nc * x for x in ONE_REQUEST_ROWS[ratio] for nc in range(2, 257)} + counts |= {hi - 1, hi} + return sorted(c for c in counts if lo < c <= hi) + + +def _sampled_counts_high(ratio: int): + """Row counts in (ROW_BOUND, N_STEP_MAX]: the first rows past the old 32-bit overflow point + (ROW_BOUND + 1, +2, +3, +64, +1023), every multiple of ROW_BOUND and its neighbours (the + overflow would have recurred at each one), the sums of nc = 32 .. 256 whole longest requests, + and the step maximum with its predecessor.""" + lo, hi = ROW_BOUND, N_STEP_MAX[ratio] + counts = {lo + e for e in (1, 2, 3, 64, 1023)} + counts |= {m * lo + e for m in range(2, hi // lo + 1) for e in (-1, 0, 1)} + counts |= {nc * N_MAX[ratio] for nc in (32, 33, 64, 65, 128, 129, 200, 255, 256)} + counts |= {hi - 1, hi} + return sorted(c for c in counts if lo < c <= hi) + + +@pytest.mark.parametrize("ratio", [2, 1]) +def test_rows_to_2_25_on_b_max_pool(ratio): + """Every row count above one request's rows, on the B_MAX allocation (owner 14's view at ratio 2, + owner 20's at ratio 1), random background bytes. (a) The ROW_BOUND = 2^25-row call equals a + native-torch gather of the allocation, every row; the rows read slots spread over the whole + view by a bijection, row 0 the first slot and the last row the last slot of block B_MAX - 1, + whose scale word is the allocation's last 4 bytes. (b) Every sampled count r in + (N_MAX, ROW_BOUND] (_sampled_counts) returns exactly the first r rows of that call, alternating + two row orders as in the exhaustive sweep. (c) The same single kernel launches at 2^25 rows. + The allocation's wrapping int64 word sum is unchanged (the gather only reads).""" + from torch.autograd import DeviceType + from torch.profiler import ProfilerActivity, profile + + owner = 14 if ratio == 2 else 20 + block_kv = 128 // ratio + alloc = _engine_alloc(ratio, B_MAX, seed=31) + view = _owner_view(alloc, ratio, owner) + n = ROW_BOUND + page, slot = _bijection_rows(ratio, n, c=12_345) + d_a, s_a = _slot_mappings(page, slot, block_kv) + del page, slot + base = OWNER_SHIFT.get(owner, 0) * block_kv * KEY_BYTES + assert int(base + s_a[-1]) + SCALE_BYTES == alloc.numel() # the far row ends the allocation + before = _checksum(alloc) + k_full, ks_full = op(view, d_a, s_a, 0, n, CODE_BYTES) + k8, ks8 = k_full.view(torch.uint8), ks_full.view(torch.uint8).view(n, SCALE_BYTES) + wrong = _rows_wrong(k8, ks8, alloc, ratio, owner, d_a, s_a, 0, n) + print( + f"[rows to 2^25 r={ratio}] view {tuple(view.shape)} of {alloc.numel()} B: {n}-row call vs the " + f"native gather, code/scale rows wrong {wrong}" + ) + assert wrong == (0, 0) # (a) + with profile(activities=[ProfilerActivity.CUDA]) as prof: # (c) + op(view, d_a, s_a, 0, n, CODE_BYTES) + torch.cuda.synchronize() + kernels = [e.name for e in prof.events() if e.device_type == DeviceType.CUDA] + assert len(kernels) == 1 and "indexerKCacheGatherUnifiedKernel" in kernels[0], kernels + # (b) row words as int64 [n, 8] / int32 [n] keep the comparisons' temporaries small + shift = n // 2 + 1 + ref = [(k8.view(torch.int64), ks8.view(torch.int32).view(-1))] + ref.append((ref[0][0].roll(shift, 0).contiguous(), ref[0][1].roll(shift, 0).contiguous())) + maps = [(d_a, s_a), (d_a.roll(shift).contiguous(), s_a.roll(shift).contiguous())] + counts = _sampled_counts(ratio) + bad = torch.zeros((), dtype=torch.int64, device=DEV) + for i, r in enumerate(counts): + v = i & 1 + k, ks = op(view, maps[v][0], maps[v][1], 0, r, CODE_BYTES) + bad += (k.view(torch.uint8).view(torch.int64) != ref[v][0][:r]).any(1).sum() + bad += (ks.view(torch.int32).view(-1) != ref[v][1][:r]).sum() + del k, ks + bad = int(bad) + unchanged = bool(_checksum(alloc) == before) + print( + f"[rows to 2^25 r={ratio}] {len(counts)} sampled counts in ({N_MAX[ratio]}, {n}] " + f"(min {counts[0]}, max {counts[-1]}): mismatching rows {bad}; one kernel at 2^25; allocation " + f"checksum unchanged {unchanged}" + ) + assert bad == 0 and unchanged + del k_full, ks_full, k8, ks8, ref, maps, d_a, s_a, before + + +def test_rows_past_2_25_exact_with_a_skipped_band(): + """Rows at index >= 2^25 with a skipped band, the arrangement that measured the pre-fix defect: + n = 2^26 + 2^24 rows over the B_MAX ratio-1 pool with rows [2^25, 2^26) skipped (-1). Every + valid row must be its own; nothing writes outside the outputs. Before the kernel's offsets were + computed in 64 bits, row 2^26 + j landed on output row j (all 2^24 of them) and output rows + [2^26, n) were never written, while every scale word (offset token_idx * 4) was right; if those + counts come back, the 64-bit cast is gone from indexerKCacheGather.cu.""" + ratio, owner, block_kv = 1, 20, 128 + alloc = _engine_alloc(ratio, B_MAX, seed=31) + lo, hi = ROW_BOUND, 2 * ROW_BOUND + n = hi + (1 << 24) + page, slot = _bijection_rows(ratio, n, c=777, far_ends=False) + d, s = _slot_mappings(page, slot, block_kv) + del page, slot + d[lo:hi] = -1 + s[lo:hi] = -1 + before = _checksum(alloc) + k, ks = op(alloc, d, s, 0, n, CODE_BYTES) + k8, ks8 = k.view(torch.uint8), ks.view(torch.uint8).view(n, SCALE_BYTES) + low = _rows_wrong(k8, ks8, alloc, ratio, owner, d, s, 0, 1 << 24) + mid = _rows_wrong(k8, ks8, alloc, ratio, owner, d, s, 1 << 24, lo) + high = _rows_wrong(k8, ks8, alloc, ratio, owner, d, s, hi, n) + # the pre-fix aliasing: output row j < 2^24 holding row 2^26 + j's code bytes + alias = 0 + for a in range(0, 1 << 24, 1 << 22): + rc, _ = _native_rows(alloc, ratio, owner, d, s, hi + a, hi + a + (1 << 22)) + alias += int((k8[a : a + (1 << 22)] == rc).all(1).sum()) + unchanged = bool(_checksum(alloc) == before) + print( + f"[skipped band] n={n}, rows [2^25, 2^26) skipped: code/scale rows wrong in [0, 2^24) {low}, " + f"[2^24, 2^25) {mid}, [2^26, n) {high}; output rows j < 2^24 holding row 2^26 + j's codes " + f"{alias} (pre-fix: {1 << 24}); pool unchanged {unchanged}" + ) + assert low == (0, 0) and mid == (0, 0) and high == (0, 0) + assert alias <= (1 << 24) // 1000 # the pre-fix aliasing wrote every one of them + assert unchanged + del k, ks, k8, ks8, d, s + + +@pytest.mark.parametrize("ratio", [2, 1]) +def test_rows_to_the_step_maximum(ratio): + """Row counts above ROW_BOUND, to the largest row count an engine of the envelope can pass in + one call: N_STEP_MAX = max_batch_size 256 context requests of Lm // r rows each (268,439,040 at + ratio 1, 134,219,520 at ratio 2), on the B_MAX allocation with random background bytes. Rows + read slots spread over the whole view by a bijection, row 0 the first slot and the last row the + last slot of block B_MAX - 1. (a) The N_STEP_MAX call equals a native-torch gather of the + allocation, every row. (b) Every sampled count in (ROW_BOUND, N_STEP_MAX] (_sampled_counts_high) + returns exactly the first rows of that call, alternating two row orders so a call that left + rows unwritten in a recycled allocation would show the other order's bytes. (c) One kernel at + N_STEP_MAX. The allocation's int64 word sum is unchanged.""" + from torch.autograd import DeviceType + from torch.profiler import ProfilerActivity, profile + + owner = 14 if ratio == 2 else 20 + block_kv = 128 // ratio + alloc = _engine_alloc(ratio, B_MAX, seed=31) + view = _owner_view(alloc, ratio, owner) + n = N_STEP_MAX[ratio] + page, slot = _bijection_rows(ratio, n, c=4_242) + d_a, s_a = _slot_mappings(page, slot, block_kv) + del page, slot + base = OWNER_SHIFT.get(owner, 0) * block_kv * KEY_BYTES + assert int(base + s_a[-1]) + SCALE_BYTES == alloc.numel() # the far row ends the allocation + before = _checksum(alloc) + k_full, ks_full = op(view, d_a, s_a, 0, n, CODE_BYTES) + k8, ks8 = k_full.view(torch.uint8), ks_full.view(torch.uint8).view(n, SCALE_BYTES) + wrong = _rows_wrong(k8, ks8, alloc, ratio, owner, d_a, s_a, 0, n) + print( + f"[step max r={ratio}] view {tuple(view.shape)} of {alloc.numel()} B: {n}-row call vs the " + f"native gather, code/scale rows wrong {wrong}" + ) + assert wrong == (0, 0) # (a) + with profile(activities=[ProfilerActivity.CUDA]) as prof: # (c) + op(view, d_a, s_a, 0, n, CODE_BYTES) + torch.cuda.synchronize() + kernels = [e.name for e in prof.events() if e.device_type == DeviceType.CUDA] + assert len(kernels) == 1 and "indexerKCacheGatherUnifiedKernel" in kernels[0], kernels + # (b) row words as int64 [n, 8] / int32 [n] keep the comparisons' temporaries small + shift = n // 2 + 1 + ref = [(k8.view(torch.int64), ks8.view(torch.int32).view(-1))] + ref.append((ref[0][0].roll(shift, 0).contiguous(), ref[0][1].roll(shift, 0).contiguous())) + maps = [(d_a, s_a), (d_a.roll(shift).contiguous(), s_a.roll(shift).contiguous())] + counts = _sampled_counts_high(ratio) + bad = torch.zeros((), dtype=torch.int64, device=DEV) + for i, r in enumerate(counts): + v = i & 1 + k, ks = op(view, maps[v][0], maps[v][1], 0, r, CODE_BYTES) + bad += (k.view(torch.uint8).view(torch.int64) != ref[v][0][:r]).any(1).sum() + bad += (ks.view(torch.int32).view(-1) != ref[v][1][:r]).sum() + del k, ks + bad = int(bad) + unchanged = bool(_checksum(alloc) == before) + print( + f"[step max r={ratio}] {len(counts)} sampled counts in ({ROW_BOUND}, {n}] " + f"(min {counts[0]}, max {counts[-1]}): mismatching rows {bad}; one kernel at {n}; allocation " + f"checksum unchanged {unchanged}" + ) + assert bad == 0 and unchanged + del k_full, ks_full, k8, ks8, ref, maps, d_a, s_a diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_indexer_k_cache_scatter.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_indexer_k_cache_scatter.py new file mode 100644 index 000000000000..555d8e7e00fe --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_indexer_k_cache_scatter.py @@ -0,0 +1,1247 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the indexer_k_cache_scatter catalog entry. + +The entry owns what is certified: `CELLS` is the list, `reference` and `compare` are the gate. +This file turns a cell's spec into tensors and keeps the checks a cell cannot carry: the +reader-layout placement and decode, a live cache manager's pools and offsets, the wrong-layout +controls, the op's rejections and the guards, the engine-size pools, every row count, the CUDA-graph +capture sets, and pools up to the largest any DeepSeek-V4.1 cache manager can have. + +The DeepSeek-V4.1 index-K write path: packed MXFP4 keys go into the owner layer's native +INDEXER_COMPRESS pool, uint8 ``[pages, block_kv, 1, 68]``, whose page holds ``block_kv x 64`` +code bytes followed by ``block_kv x 4`` scale bytes (the layout the DeepGEMM paged logits kernel +reads). The inputs are produced by the target's own quantizer, ``torch.ops.trtllm.fp4_quantize(x, +None, 32, True, False)``, whose linear scale output is viewed as int32 ``[T, 1]``. Beside the +entry's flat-offset reference, the expected pool is also built from page and slot with native torch +indexing (the pre-call snapshot with exactly those bytes placed at the paged offsets), and a +reader-layout decode of the pool must give back the quantizer's values. +""" + +import zlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.indexer_k_cache_scatter import ( # noqa: E501 + indexer_k_cache_scatter as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "indexer_k_cache_scatter requires a CUDA device" + +DEV = "cuda" +DIM = 128 # index_head_dim +CODE_BYTES = 64 # two E2M1 codes per byte +SCALE_BYTES = 4 # one UE8M0 byte per 32 elements +ENTRY_BYTES = CODE_BYTES + SCALE_BYTES +# The pool's page count is the cache manager's number, never the batch's: each case allocates +# POOL_FACTOR x the pages its batch needs (+ POOL_EXTRA) and hands its requests pages drawn by a +# random permutation. Chosen here, not taken from a runtime dump. +POOL_FACTOR = 3 +POOL_EXTRA = 8 + +_E2M1 = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float64, + device=DEV, +) + + +def _keys(n: int, seed: int): + """bf16 [n, 128] with a different power-of-two magnitude per 32-block, quantized by the + target's quantizer -> (codes uint8 [n, 64], scale bytes uint8 [n, 4]).""" + g = torch.Generator(device=DEV).manual_seed(seed) + x = torch.randn(n, DIM, device=DEV, generator=g) + mag = torch.randint(-8, 9, (n, DIM // 32, 1), device=DEV, generator=g).float() + x = (x.view(n, DIM // 32, 32) * torch.exp2(mag)).view(n, DIM).to(torch.bfloat16) + codes, sf = torch.ops.trtllm.fp4_quantize(x, None, 32, True, False) + assert codes.shape == (n, CODE_BYTES) and codes.dtype == torch.uint8 + assert sf.shape == (n * SCALE_BYTES,) and sf.dtype == torch.uint8 + return codes, sf.view(n, SCALE_BYTES) + + +def _as_scale_word(scale_bytes: torch.Tensor) -> torch.Tensor: + """uint8 [n, 4] -> int32 [n, 1], a reinterpretation (byte g of each row = byte g of the word).""" + return scale_bytes.contiguous().view(torch.int32).view(-1, 1) + + +def _decode(codes: torch.Tensor, scale_bytes: torch.Tensor) -> torch.Tensor: + """(uint8 [n, 64], uint8 [n, 4]) -> fp64 [n, 128]; element 2i is the low nibble of byte i, + scale byte g multiplies elements 32g .. 32g+31 by 2^(byte - 127).""" + v = torch.stack([_E2M1[(codes & 0xF).long()], _E2M1[(codes >> 4).long()]], -1).flatten(-2) + sc = torch.exp2(scale_bytes.double() - 127.0) + return (v.unflatten(-1, (4, 32)) * sc.unsqueeze(-1)).flatten(-2) + + +class _Batch: + """Requests with ``cached[b]`` entries already in the pool and ``new[b]`` entries to write.""" + + def __init__(self, block_kv: int, cached, new, seed: int, num_pages=None): + g = torch.Generator().manual_seed(seed) + self.block_kv = block_kv + cached = torch.tensor(cached, dtype=torch.int64) + new = torch.tensor(new, dtype=torch.int64) + pages = ((cached + new + block_kv - 1) // block_kv).clamp_min(1) + need = int(pages.sum()) + self.num_pages = POOL_FACTOR * need + POOL_EXTRA if num_pages is None else num_pages + assert self.num_pages >= need + perm = torch.randperm(self.num_pages, generator=g) + self.table = torch.full((len(cached), int(pages.max())), -1, dtype=torch.int64) + o = 0 + for b, p in enumerate(pages.tolist()): + self.table[b, :p] = perm[o : o + p] + o += p + self.n = int(new.sum()) + self.req = torch.repeat_interleave(torch.arange(len(cached)), new) + start = torch.cumsum(new, 0) - new + self.pos = cached[self.req] + torch.arange(self.n) - start[self.req] + self.table = self.table.to(DEV) + self.req = self.req.to(DEV) + self.pos = self.pos.to(DEV) + + def pages_slots(self, table=None, block_kv=None): + table = self.table if table is None else table + block_kv = self.block_kv if block_kv is None else block_kv + page = table[self.req, self.pos // block_kv] + return page, self.pos % block_kv + + def slot_mappings(self, table=None, block_kv=None): + """The in-tree formula (``_compute_slot_mappings`` in the DSA indexer).""" + block_kv = self.block_kv if block_kv is None else block_kv + page, slot = self.pages_slots(table, block_kv) + base = page * block_kv * ENTRY_BYTES + return base + slot * CODE_BYTES, base + block_kv * CODE_BYTES + slot * SCALE_BYTES + + def pool(self, seed: int, extra_pages: int = 0) -> torch.Tensor: + g = torch.Generator(device=DEV).manual_seed(seed) + return torch.randint( + 0, + 256, + (self.num_pages + extra_pages, self.block_kv, 1, ENTRY_BYTES), + dtype=torch.uint8, + device=DEV, + generator=g, + ) + + def expected(self, before, codes, scale_bytes, rows=None): + """Native-torch placement: the snapshot with exactly these rows' bytes at their offsets.""" + rows = torch.arange(self.n, device=DEV) if rows is None else rows + page, slot = self.pages_slots() + page, slot = page[rows], slot[rows] + exp = before.clone().view(before.size(0), -1) + ar_c = torch.arange(CODE_BYTES, device=DEV) + ar_s = torch.arange(SCALE_BYTES, device=DEV) + exp[page.unsqueeze(1), (slot * CODE_BYTES).unsqueeze(1) + ar_c] = codes[rows] + exp[ + page.unsqueeze(1), (self.block_kv * CODE_BYTES + slot * SCALE_BYTES).unsqueeze(1) + ar_s + ] = scale_bytes[rows] + return exp.view_as(before) + + def read(self, pool: torch.Tensor): + """The DeepGEMM paged reader's view of every token: its page's code bytes, then scales.""" + page, slot = self.pages_slots() + flat = pool.reshape(pool.size(0), -1) if pool.is_contiguous() else None + if flat is None: # page-strided view: address pages through stride(0) + flat = torch.as_strided( + pool, (pool.size(0), self.block_kv * ENTRY_BYTES), (pool.stride(0), 1) + ) + codes = flat[ + page.unsqueeze(1), + (slot * CODE_BYTES).unsqueeze(1) + torch.arange(CODE_BYTES, device=DEV), + ] + sb = flat[ + page.unsqueeze(1), + (self.block_kv * CODE_BYTES + slot * SCALE_BYTES).unsqueeze(1) + + torch.arange(SCALE_BYTES, device=DEV), + ] + return codes, sb + + +def _row_mismatch(batch: _Batch, pool, codes, scale_bytes) -> float: + rc, rs = batch.read(pool) + bad = (_decode(rc, rs) != _decode(codes, scale_bytes)).any(-1) + return bad.float().mean().item() + + +def _cell_batch(spec, seed: int) -> _Batch: + """A cell's requests: explicit (cached, new) lists, or `requests` drawn below `cached_below` / + up to `new_upto` entries (the random decode and DSpark-verify batches).""" + g = torch.Generator().manual_seed(seed) + reqs = spec.get("requests") + cached = ( + spec["cached"] + if "cached" in spec + else torch.randint(0, spec["cached_below"], (reqs,), generator=g).tolist() + ) + new = ( + spec["new"] + if "new" in spec + else torch.randint(0, spec["new_upto"] + 1, (reqs,), generator=g).tolist() + ) + return _Batch(spec["block_kv"], cached, new, seed=seed) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + seed = zlib.crc32(cell.why.encode()) + b = _cell_batch(spec, seed) + bk = b.block_kv + codes, sb = _keys(b.n, seed=seed + 1) + d_map, s_map = b.slot_mappings() + skip = torch.zeros(b.n, dtype=torch.bool, device=DEV) + if spec.get("skip"): + skip[::5] = True + if spec["skip"] in ("both", "data_only"): + d_map = torch.where(skip, -1, d_map) + if spec["skip"] in ("both", "scale_only"): + s_map = torch.where(skip, -1, s_map) + pad = spec.get("page_pad", 0) + page_bytes = bk * ENTRY_BYTES + if pad: + # A view into a wider allocation: page stride block_kv * 68 + pad. + backing = torch.randint( + 0, 256, (b.num_pages, page_bytes + pad), dtype=torch.uint8, device=DEV + ) + pool = backing[:, :page_bytes].view(b.num_pages, bk, 1, ENTRY_BYTES) + assert not pool.is_contiguous() and pool.stride(0) == page_bytes + pad + backing_before = backing.clone() + else: + pool = b.pool(seed=seed + 2) + n = b.n if spec.get("num_tokens") is None else spec["num_tokens"] + word = _as_scale_word(sb) + before = pool.clone(memory_format=torch.contiguous_format) + + out = op(codes, word, pool, d_map, s_map, n) + torch.cuda.synchronize() + assert out is pool, "the op writes the pool in place and returns it" + op.compare(out, op.reference(codes, word, before, d_map, s_map, n)) + + # The paged layout the reader expects: the snapshot with exactly the written rows' bytes placed + # at their page and slot (every other byte of the pool untouched). + rows = (~skip[:n]).nonzero().flatten() + assert torch.equal(pool, b.expected(before, codes, sb, rows=rows)) + if pad: + assert torch.equal(backing[:, page_bytes:], backing_before[:, page_bytes:]) + if rows.numel() == b.n: + # the reader-layout decode of the pool is the quantizer's value set, row for row + rc, rs = b.read(pool) + assert torch.equal(_decode(rc, rs), _decode(codes, sb)) + print(f"[scatter {cell.why}] tokens={b.n} n={n} pages={b.num_pages} bit-exact") + + +def test_cuda_graph_replay(): + """Capture once, then replay with fresh keys and fresh slot offsets written into the static + buffers (the decode / DSpark-verify step is graph-captured); padding rows carry -1.""" + n = 48 + pages = 200 # the pool outlives every batch replayed into it + b0 = _Batch(64, [0] * 16, [3] * 16, seed=12, num_pages=pages) + pool = b0.pool(seed=13) + codes_s = torch.zeros(n, CODE_BYTES, dtype=torch.uint8, device=DEV) + scale_s = torch.zeros(n, 1, dtype=torch.int32, device=DEV) + d_s = torch.full((n,), -1, dtype=torch.int64, device=DEV) + s_s = torch.full((n,), -1, dtype=torch.int64, device=DEV) + op(codes_s, scale_s, pool, d_s, s_s, n) # warm-up (all rows skipped) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(codes_s, scale_s, pool, d_s, s_s, n) + torch.cuda.synchronize() + for trial, (cached, new) in enumerate( + (([5 * t for t in range(16)], [3] * 16), ([64 * t + 63 for t in range(12)], [2] * 12)) + ): + b = _Batch(64, cached, new, seed=20 + trial, num_pages=pages) + codes, sb = _keys(b.n, seed=30 + trial) + d_map, s_map = b.slot_mappings() + codes_s.zero_() + scale_s.zero_() + d_s.fill_(-1) + s_s.fill_(-1) + codes_s[: b.n] = codes + scale_s[: b.n] = _as_scale_word(sb) + d_s[: b.n] = d_map + s_s[: b.n] = s_map + before = pool.clone() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(pool, b.expected(before, codes, sb)), f"replay {trial}" + + +# ─── the real cache manager's pools and the metadata's own offsets ─────────────────────────── + +# (cached, new) per request, context requests first; every request here feeds its step as the +# target sees it. prefill: fresh prompts of odd and even length (1 emits no ratio-2 entry); +# mixed: a fresh and a chunked context (chunk starting mid-page at 513) ahead of four decode rows +# of both parities, one ending on a page edge; decode_b32: the target's widest decode graph batch. +_V41_DECODE32 = torch.randint( + 1, 400, (28,), generator=torch.Generator().manual_seed(77) +).tolist() + [127, 128, 2047, 2048] +_V41_STEPS = { + "prefill": ([(0, 1), (0, 130), (0, 257), (0, 700), (0, 512)], 5), + "mixed": ([(0, 300), (513, 388), (129, 1), (2048, 1), (63, 1), (1000, 1)], 2), + "decode_b32": ([(c, 1) for c in _V41_DECODE32], 0), +} +_V41_MAX_BATCH = 40 +_V41_MAX_TOKENS = 4096 * _V41_MAX_BATCH * 2 + + +@pytest.mark.parametrize("step", list(_V41_STEPS)) +def test_v41_cache_manager_pools(step): + """The target's call on the pools a live DeepseekV4CacheManager allocates (V4.1 mode, + tokens_per_block 128, max_batch_size 40, max_seq_len 4096, max_tokens 327680, pools + uint8 [6720, 64, 1, 68] for ratio 2 and [2240, 128, 1, 68] for ratio 1), with the offsets + of the prepared DeepseekV4TrtllmAttentionMetadata: ``v41_index_k_slot_data[r][:n]`` / + ``v41_index_k_slot_scale[r][:n]`` and ``n = num_total_compressed_tokens[r]``, reserved rows + included. Owners 2, 8 and 14 share page ids and interleave their pages in one allocation + (each owner's view starts one page after the previous one's), so every ratio-2 call is + checked over the whole allocation. The expected bytes are placed from the reader's table, + ``indexer_compress_block_tables[r]``, and the packed-row order (context entries, then + generation entries, in request order) with native torch indexing; every other byte of the + allocation, and the same owner's COMPRESS pool, must be unchanged.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + DeepseekV4CacheManager, + DeepseekV4TrtllmAttentionMetadata, + ) + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + from tensorrt_llm._torch.metadata import KVCacheParams + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from tensorrt_llm.bindings import DataType, SamplingConfig + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig + from tensorrt_llm.mapping import Mapping + + batch, nc = _V41_STEPS[step] + text = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=128, + head_dim=512, + index_head_dim=DIM, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, + ) + config = SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **{k: v for k, v in text.items() if k != "model_type"}, + ) + sparse = DeepSeekV4SparseAttentionConfig( + index_head_dim=DIM, window_size=128, compress_ratios=[1] * 43 + ) + manager = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, max_tokens=_V41_MAX_TOKENS, event_buffer_max_size=0 + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=40, + num_kv_heads=1, + head_dim=512, + tokens_per_block=128, + max_seq_len=4096, + max_batch_size=_V41_MAX_BATCH, + max_input_len=4096, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=8192, + sparse_attn_config=sparse, + pretrained_config=config, + ) + requests = [] + try: + cached = [c for c, _ in batch] + new = [n for _, n in batch] + B = len(batch) + for i, (c, n) in enumerate(batch): + req = LlmRequest( + request_id=i, + max_new_tokens=64, + input_tokens=list(range(c + n)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + assert manager.prepare_context(req), f"manager could not hold request {i}" + assert manager.resize_context(req, req.context_chunk_size) + requests.append(req) + md = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor(new, dtype=torch.int), + request_ids=list(range(B)), + max_num_requests=manager.max_batch_size, + num_contexts=nc, + prompt_lens=[c + n for c, n in batch], + max_num_tokens=8192, + kv_cache_manager=manager, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + sparse_attention_config=sparse, + ) + md.prepare() + torch.cuda.synchronize() + assert (md.num_contexts, md.num_seqs) == (nc, B) + + for r, owners in ((2, (2, 8, 14)), (1, (20,))): + block_kv = 128 // r + n = md.num_total_compressed_tokens[r] + # Packed rows: every (request, entry) the step completes, then reserved rows (-1). + rows = [(b, k) for b, (c, m) in enumerate(batch) for k in range(c // r, (c + m) // r)] + assert 0 < len(rows) <= n, (step, r, len(rows), n) + d_map = md.v41_index_k_slot_data[r][:n] + s_map = md.v41_index_k_slot_scale[r][:n] + assert (d_map[len(rows) :] == -1).all() and (s_map[len(rows) :] == -1).all() + assert (d_map[: len(rows)] >= 0).all() and (s_map[: len(rows)] >= 0).all() + table = md.indexer_compress_block_tables[r][:B].long() + req = torch.tensor([b for b, _ in rows], device=DEV) + ent = torch.tensor([k for _, k in rows], device=DEV) + page = table[req, ent // block_kv] + slot = ent % block_kv + assert (page >= 0).all() + # The whole allocation the ratio's owners interleave in: the first owner's view reaches + # the same end address as every other owner's. + first = manager.get_indexer_k_cache_buffers(owners[0]) + ends = { + manager.get_indexer_k_cache_buffers(o).data_ptr() + + manager.get_indexer_k_cache_buffers(o).numel() + for o in owners + } + assert ends == {first.data_ptr() + first.numel()} + first.copy_( + torch.randint( + 0, + 256, + first.shape, + dtype=torch.uint8, + device=DEV, + generator=torch.Generator(device=DEV).manual_seed(r), + ) + ) + for owner in owners: + pool = manager.get_indexer_k_cache_buffers(owner) + assert pool.dtype == torch.uint8 and tuple(pool.shape[1:]) == ( + block_kv, + 1, + ENTRY_BYTES, + ) + assert pool.is_contiguous() + page_bytes = block_kv * ENTRY_BYTES + shift, rem = divmod(pool.data_ptr() - first.data_ptr(), page_bytes) + assert rem == 0 and 0 <= shift < len(owners), (owner, shift, rem) + comp = manager.get_buffers(owner, T.COMPRESS) + comp_before = comp.clone() + before = first.clone() + # What the target passes: fp4_quantize's data as is, its linear scales as int32 [n, 1]. + codes, sb = _keys(n, seed=1000 * r + owner) + word = sb.reshape(-1).view(torch.int32).reshape(n, 1) + op(codes, word, pool, d_map, s_map, n) + torch.cuda.synchronize() + exp = before.clone().view(first.size(0), -1) + idx = torch.arange(len(rows), device=DEV) + ar_c = torch.arange(CODE_BYTES, device=DEV) + ar_s = torch.arange(SCALE_BYTES, device=DEV) + exp[(page + shift).unsqueeze(1), (slot * CODE_BYTES).unsqueeze(1) + ar_c] = codes[ + idx + ] + exp[ + (page + shift).unsqueeze(1), + (block_kv * CODE_BYTES + slot * SCALE_BYTES).unsqueeze(1) + ar_s, + ] = sb[idx] + assert torch.equal(first.view(first.size(0), -1), exp), ( + f"{step} r={r} owner {owner}: pool bytes" + ) + assert torch.equal(comp, comp_before), ( + f"{step} r={r} owner {owner}: COMPRESS pool changed" + ) + # The reader's view of the owner's pool decodes to the quantizer's values. + flat = pool.view(pool.size(0), -1) + rc = flat[page.unsqueeze(1), (slot * CODE_BYTES).unsqueeze(1) + ar_c] + rs = flat[ + page.unsqueeze(1), + (block_kv * CODE_BYTES + slot * SCALE_BYTES).unsqueeze(1) + ar_s, + ] + assert torch.equal(_decode(rc, rs), _decode(codes[idx], sb[idx])) + print( + f"[v41 {step} r={r} owner {owner}] n={n} entries={len(rows)} reserved={n - len(rows)} " + f"pool={tuple(pool.shape)} page_shift={shift} bit-exact" + ) + finally: + for req in requests: + manager.free_resources(req) + manager.shutdown() + + +# ─── deliberately wrong inputs: the checks above must see each one ─────────────────────────── + + +def _control_batch(): + b = _Batch(64, [0, 100, 250, 7, 64, 500], [300, 90, 120, 33, 64, 200], seed=40) + codes, sb = _keys(b.n, seed=41) + return b, codes, sb + + +def test_control_wrong_slot_mapping(): + """Every request writes with its neighbour's page-table row.""" + b, codes, sb = _control_batch() + d_map, s_map = b.slot_mappings(table=b.table.roll(1, 0)) + pool = b.pool(seed=42) + op(codes, _as_scale_word(sb), pool, d_map, s_map, b.n) + torch.cuda.synchronize() + frac = _row_mismatch(b, pool, codes, sb) + print(f"[control] wrong slot mapping (rolled table): {frac:.4f} of rows wrong") + assert frac > 0.9 + + +def test_control_key_interleaved_offsets(): + """Offsets for a [64 codes, 4 scales]-per-key layout instead of the page-blocked one.""" + b, codes, sb = _control_batch() + page, slot = b.pages_slots() + base = page * 64 * ENTRY_BYTES + slot * ENTRY_BYTES + pool = b.pool(seed=43) + op(codes, _as_scale_word(sb), pool, base, base + CODE_BYTES, b.n) + torch.cuda.synchronize() + frac = _row_mismatch(b, pool, codes, sb) + print(f"[control] key-interleaved offsets: {frac:.4f} of rows wrong") + assert frac > 0.9 + + +def test_control_page_size(): + """A 64-entry pool addressed as 128-entry pages.""" + b, codes, sb = _control_batch() + d_map, s_map = b.slot_mappings(table=b.table // 2, block_kv=128) + pool = b.pool(seed=44, extra_pages=2) # a 128-entry page at the last odd page overruns by < 2 + op(codes, _as_scale_word(sb), pool, d_map, s_map, b.n) + torch.cuda.synchronize() + frac = _row_mismatch(b, pool, codes, sb) + print(f"[control] 64-entry pages addressed as 128: {frac:.4f} of rows wrong") + assert frac > 0.9 + + +@pytest.mark.parametrize("how", ["byte_reversed", "group_major"]) +def test_control_scale_bytes_misinterleaved(how): + """The four UE8M0 bytes packed into the word in the wrong order: big-endian, or the linear + [T*4] scale output read as if it were group-major [4, T].""" + b, codes, sb = _control_batch() + if how == "byte_reversed": + wrong = sb.flip(1) + else: + wrong = sb.reshape(-1).view(SCALE_BYTES, b.n).t() + d_map, s_map = b.slot_mappings() + pool = b.pool(seed=45) + op(codes, _as_scale_word(wrong), pool, d_map, s_map, b.n) + torch.cuda.synchronize() + frac = _row_mismatch(b, pool, codes, sb) + print(f"[control] scale bytes {how}: {frac:.4f} of rows wrong") + assert frac > 0.9 + + +# ─── rejections: the op's own (loud) and the entry's guards (run under `validating`) ────────── + + +def _reject_inputs(): + b = _Batch(64, [0], [8], seed=50) + codes, sb = _keys(8, seed=51) + d_map, s_map = b.slot_mappings() + return b, codes, sb, d_map, s_map + + +_REJECTIONS = { + # the quantizer's raw scale output, not viewed as int32 [T, 1] + "scale_uint8_1d": (RuntimeError, "k_scale must be 2D [num_tokens, scale_elements]"), + "scale_uint8_2d": (RuntimeError, "k_scale must have 4-byte elements (e.g. float32), got %d1"), + "scale_int32_1d": (RuntimeError, "k_scale must be 2D [num_tokens, scale_elements]"), + "slot_int32": (RuntimeError, "slot_mapping_fp8 must be int64"), + "pool_int8": (RuntimeError, "expected scalar type Byte but found Char"), + "pool_3d": ( + RuntimeError, + "k_cache must be 4D [num_blocks, block_size, 1, per_token_size], got %d dims3", + ), + "pool_dim2": (RuntimeError, "k_cache dimension 2 must be 1, got %d2"), + "data_noncontig": (RuntimeError, "k_fp8 must be contiguous"), + "slot_cpu": (RuntimeError, "All tensors must be CUDA tensors"), + # the entry's is_valid guards: accepted silently by the op (FP8-width rows overwrite + # neighbouring entries; num_tokens past the supplied rows / slot entries reads past those + # tensors), so never driven outside `validating` + "data_fp8_width": (AssertionError, "k_data row bytes + 4 scale bytes"), + "num_tokens_past_rows": (AssertionError, "num_tokens must not exceed"), + "num_tokens_past_slots": (AssertionError, "num_tokens must not exceed"), +} + + +@pytest.mark.parametrize("case", list(_REJECTIONS)) +def test_rejections(case): + b, codes, sb, d_map, s_map = _reject_inputs() + word = _as_scale_word(sb) + pool = b.pool(seed=52) + args = dict( + k_data=codes, + k_scale=word, + k_cache=pool, + slot_mapping_data=d_map, + slot_mapping_scale=s_map, + num_tokens=8, + ) + if case == "scale_uint8_1d": + args["k_scale"] = sb.reshape(-1) + elif case == "scale_uint8_2d": + args["k_scale"] = sb + elif case == "scale_int32_1d": + args["k_scale"] = word.view(-1) + elif case == "slot_int32": + args["slot_mapping_data"], args["slot_mapping_scale"] = d_map.int(), s_map.int() + elif case == "pool_int8": + args["k_cache"] = pool.view(torch.int8) + elif case == "pool_3d": + args["k_cache"] = pool.view(pool.size(0), 64, ENTRY_BYTES) + elif case == "pool_dim2": + args["k_cache"] = pool.view(pool.size(0), 32, 2, ENTRY_BYTES) + elif case == "data_noncontig": + args["k_data"] = torch.cat([codes, codes], 1)[:, ::2] + elif case == "slot_cpu": + args["slot_mapping_data"], args["slot_mapping_scale"] = d_map.cpu(), s_map.cpu() + elif case == "data_fp8_width": + args["k_data"] = torch.zeros(8, 128, dtype=torch.uint8, device=DEV) + elif case == "num_tokens_past_rows": + args["k_data"] = codes[:4] + elif case == "num_tokens_past_slots": + args["slot_mapping_data"], args["slot_mapping_scale"] = d_map[:4], s_map[:4] + exc, text = _REJECTIONS[case] + before = pool.clone() + # Every case runs under `validating`: the op's own rejections pass the guard and are raised by + # the op; the guard cases raise AssertionError before anything is launched. + with pytest.raises(exc) as info: + with validating(op): + op(**args) + torch.cuda.synchronize() + assert text in str(info.value) + assert torch.equal(pool, before) + + +# ─── the engine's pools, every row count, the capture sets ─────────────────────────────────── +# +# The core target makes this call once per KV-source layer per step (`ModelingV2Core._compress`) +# with n = num_total_compressed_tokens[r] rows of v41_index_k_slot_data[r] / _scale[r]: every +# entry the step completes, then reserved rows (-1). At ratio 1 n = T (up to max_num_tokens, at +# most 16384 in the envelope); at ratio 2 n <= (T + the step's requests) / 2, below 16384 too, and +# the metadata's slot buffers hold max_num_tokens rows. N_MAX below is 16384 for both. +# +# The pools are sized from the core runtime dumps. The cache manager allocates B blocks; the +# ratio-2 INDEXER_COMPRESS allocation holds 3 pages of 64 entries per block, which owners 2, 8 and +# 14 see as contiguous views starting 0, 1 and 2 pages into it ([3B - shift, 64, 1, 68]; their page +# ids are 3 x block), and the ratio-1 pool (owner 20) holds one page of 128 entries per block +# ([B, 128, 1, 68], page id = block). B = 137386 / 142336 (boot dumps) and 151040 (the largest +# pool of the runtime dumps: 453120 x 64 and 151040 x 128, about 1.97e9 and 1.31e9 bytes). +ENGINE_BLOCKS = (137386, 142336, 151040) +OWNER_SHIFT = {2: 0, 8: 1, 14: 2} +N_MAX = 16384 +# The longest request of the envelope: max_seq_len 1,048,576 with DSpark (manager length L + 14), +# whose last compressed entry is Lm // r - 1. +LM_1M = 1_048_576 + 14 +_ENGINE_CACHE = {} + + +def _engine_alloc(ratio: int, blocks: int) -> torch.Tensor: + """The ratio's whole INDEXER_COMPRESS allocation at engine size, random background bytes. One + at a time per process; tests snapshot it before writing and compare the whole allocation.""" + key = (ratio, blocks) + if key not in _ENGINE_CACHE: + _ENGINE_CACHE.clear() + torch.cuda.empty_cache() + pages = 3 * blocks if ratio == 2 else blocks + g = torch.Generator(device=DEV).manual_seed(blocks + ratio) + _ENGINE_CACHE[key] = torch.randint( + 0, + 256, + (pages, 128 // ratio, 1, ENTRY_BYTES), + dtype=torch.uint8, + device=DEV, + generator=g, + ) + return _ENGINE_CACHE[key] + + +def _owner_view(alloc: torch.Tensor, ratio: int, owner: int) -> torch.Tensor: + """What get_indexer_k_cache_buffers(owner) returns: a contiguous view `shift` pages in.""" + view = alloc[OWNER_SHIFT[owner] :] if ratio == 2 else alloc + assert view.is_contiguous() + return view + + +def _engine_batch(ratio: int, blocks: int, cached, new, seed: int) -> _Batch: + """A step's requests paged over the whole engine pool. Their pages are distinct blocks drawn + at random, except that request 0's first page is block 0 and the last request's last page is + block `blocks - 1` (the far end); page id = 3 x block at ratio 2, = block at ratio 1.""" + block_kv = 128 // ratio + pages = [max(1, -(-(c + m) // block_kv)) for c, m in zip(cached, new)] + total = sum(pages) + b = _Batch(block_kv, cached, new, seed=seed, num_pages=total) + b.num_pages = None # the pool is the engine allocation, not _Batch.pool() + g = torch.Generator(device="cpu").manual_seed(seed) + mid = torch.randperm(blocks - 2, generator=g)[: max(total - 2, 0)] + 1 + chosen = ( + torch.cat([torch.tensor([0]), mid, torch.tensor([blocks - 1])])[:total] + if total > 1 + else torch.tensor([0]) + ) + table = torch.full((len(cached), max(pages)), -1, dtype=torch.int64) + off = 0 + for r, p in enumerate(pages): + table[r, :p] = chosen[off : off + p] + off += p + b.table = (table * (3 if ratio == 2 else 1)).to(DEV) + return b + + +def _place(dst_view: torch.Tensor, b: _Batch, codes, scale_bytes, rows): + """Native-torch placement of `rows` (a device index tensor) into `dst_view` in the reader's + layout (writes through to the allocation the view belongs to).""" + page, slot = b.pages_slots() + page, slot = page[rows], slot[rows] + flat = dst_view.view(dst_view.size(0), -1) + ar_c = torch.arange(CODE_BYTES, device=DEV) + ar_s = torch.arange(SCALE_BYTES, device=DEV) + flat[page.unsqueeze(1), (slot * CODE_BYTES).unsqueeze(1) + ar_c] = codes[rows] + flat[page.unsqueeze(1), (b.block_kv * CODE_BYTES + slot * SCALE_BYTES).unsqueeze(1) + ar_s] = ( + scale_bytes[rows] + ) + + +def _far_step(block_kv: int, nreq: int, per_req: int, seed: int): + """(cached, new) for nreq requests writing per_req entries each: request 0 starts inside its + first page, the last request's last entry is its last page's last slot, the rest random.""" + g = torch.Generator().manual_seed(seed) + cached = torch.randint(0, 40 * block_kv, (nreq,), generator=g).tolist() + cached[0] = ( + int(torch.randint(0, block_kv - per_req % block_kv, (1,), generator=g)) + if per_req < block_kv + else 0 + ) + last_end = (int(torch.randint(1, 40, (1,), generator=g)) + (per_req // block_kv) + 1) * block_kv + cached[-1] = last_end - per_req + return cached, [per_req] * nreq + + +@pytest.mark.parametrize("blocks", ENGINE_BLOCKS) +@pytest.mark.parametrize("ratio,owner", [(2, 2), (2, 8), (2, 14), (1, 20)]) +def test_engine_pool_far_end(ratio, owner, blocks): + """Every owner's view of the engine-size pool. 40 requests write 6 entries each over pages drawn + from the whole pool; the last one's last entry is the last slot of block `blocks - 1` (for owner + 14 at ratio 2 and owner 20, the allocation's last page). The whole allocation equals the + snapshot with exactly those bytes placed; the reader-layout decode is the quantizer's values.""" + block_kv = 128 // ratio + alloc = _engine_alloc(ratio, blocks) + view = _owner_view(alloc, ratio, owner) + cached, new = _far_step(block_kv, 40, 6, seed=blocks + owner) + b = _engine_batch(ratio, blocks, cached, new, seed=blocks + owner) + page, slot = b.pages_slots() + assert ( + int(page[-1]) == (blocks - 1) * (3 if ratio == 2 else 1) and int(slot[-1]) == block_kv - 1 + ) + codes, sb = _keys(b.n, seed=b.n + owner) + d_map, s_map = b.slot_mappings() + before = alloc.clone() + op(codes, _as_scale_word(sb), view, d_map, s_map, b.n) + torch.cuda.synchronize() + _place(_owner_view(before, ratio, owner), b, codes, sb, torch.arange(b.n, device=DEV)) + assert torch.equal(alloc, before), "allocation differs from the native placement" + rc, rs = b.read(view) + assert torch.equal(_decode(rc, rs), _decode(codes, sb)) + shift = OWNER_SHIFT.get(owner, 0) + last = ( + (shift + int(page[-1])) * block_kv * ENTRY_BYTES + + block_kv * CODE_BYTES + + int(slot[-1]) * SCALE_BYTES + + 3 + ) + print( + f"[engine pool r={ratio} owner={owner}] view {tuple(view.shape)}, allocation {alloc.numel()} B, " + f"last byte written {last}: bit-exact, nothing else changed" + ) + assert last == alloc.numel() - 1 or shift < 2 and ratio == 2 + del before + + +@pytest.mark.parametrize("ratio", [1, 2]) +def test_every_row_count_prefix_bitwise(ratio): + """Every row count on the target's call form, over the engine-size pool (owner 14 at ratio 2, owner + 20 at ratio 1). One 16384-row step: 128 requests x 128 entries over pages drawn from the whole + pool (far end included), request 1 writing the last 128 compressed entries of a 1,048,590-token + request (entries up to Lm // r - 1, on its 8193rd page), with every 97th row reserved (-1). For + every n in 1..16384 the call writes rows 0..n-1 and nothing else: the whole allocation equals + the snapshot with exactly the first n rows of the 16384-row placement (native torch). The sweep + alternates two key sets, so a row the call failed to write would still hold the other set's + bytes, and a row >= n it wrote would differ from the background. The 16384-row result also + decodes to the quantizer's values.""" + owner = 14 if ratio == 2 else 20 + blocks = ENGINE_BLOCKS[-1] + block_kv = 128 // ratio + alloc = _engine_alloc(ratio, blocks) + view = _owner_view(alloc, ratio, owner) + per_req = N_MAX // 128 + cached, new = _far_step(block_kv, 128, per_req, seed=ratio) + cached[1] = LM_1M // ratio - per_req + b = _engine_batch(ratio, blocks, cached, new, seed=ratio) + assert b.n == N_MAX + long_rows = b.pos[per_req : 2 * per_req] + assert int(long_rows[-1]) == LM_1M // ratio - 1 and int(long_rows[-1]) // block_kv == 8192 + d_map, s_map = b.slot_mappings() + reserved = torch.arange(N_MAX, device=DEV) % 97 == 96 + d_map = torch.where(reserved, -1, d_map).contiguous() + s_map = torch.where(reserved, -1, s_map).contiguous() + keys = [_keys(N_MAX, seed=100 + ratio), _keys(N_MAX, seed=200 + ratio)] + words = [_as_scale_word(sb) for _, sb in keys] + assert (keys[0][0] != keys[1][0]).any(1).float().mean().item() > 0.99 + # flat byte indices of every row's code bytes and scale word, in the whole allocation + page, slot = b.pages_slots() + base = (OWNER_SHIFT.get(owner, 0) + page) * block_kv * ENTRY_BYTES + code_idx = ( + base.unsqueeze(1) + (slot * CODE_BYTES).unsqueeze(1) + torch.arange(CODE_BYTES, device=DEV) + ) + scale_idx = ( + base.unsqueeze(1) + + (block_kv * CODE_BYTES + slot * SCALE_BYTES).unsqueeze(1) + + torch.arange(SCALE_BYTES, device=DEV) + ) + kept = (~reserved).nonzero().flatten() + kept_below = torch.cumsum((~reserved).long(), 0).tolist() + exp = alloc.clone() + exp_flat = exp.view(-1) + bad = [] + for n in range(1, N_MAX + 1): + v = n & 1 + op(keys[v][0], words[v], view, d_map, s_map, n) + rows = kept[: kept_below[n - 1]] + exp_flat[code_idx[rows].reshape(-1)] = keys[v][0][rows].reshape(-1) + exp_flat[scale_idx[rows].reshape(-1)] = keys[v][1][rows].reshape(-1) + if not torch.equal(alloc, exp): + bad.append(n) + exp.copy_(alloc) # keep sweeping from the op's state + print( + f"[rows r={ratio}] every n in 1..{N_MAX} on {tuple(view.shape)}: mismatching n {bad[:10]} " + f"({len(bad)} of {N_MAX})" + ) + assert not bad + rc, rs = b.read(view) + ok = ~reserved + v = N_MAX & 1 + assert torch.equal(_decode(rc[ok], rs[ok]), _decode(keys[v][0][ok], keys[v][1][ok])) + del exp + + +def test_one_kernel_at_every_row_count(): + """An ahead-of-time compiled kernel whose row count is only its grid: the same single kernel + launches at 1 row and at 16384 (no per-shape variant).""" + from torch.autograd import DeviceType + from torch.profiler import ProfilerActivity, profile + + b = _Batch(128, [0] * 8, [2048] * 8, seed=3) + codes, sb = _keys(b.n, seed=4) + d_map, s_map = b.slot_mappings() + pool = b.pool(seed=5) + seen = set() + for n in (1, 2, 97, 1935, 8192, 8193, 1536, 16384): + with profile(activities=[ProfilerActivity.CUDA]) as prof: + op(codes, _as_scale_word(sb), pool, d_map, s_map, n) + torch.cuda.synchronize() + kernels = [e.name for e in prof.events() if e.device_type == DeviceType.CUDA] + print(f"[kernels n={n}] {kernels}") + assert len(kernels) == 1 and "indexerKCacheScatterUnifiedKernel" in kernels[0], kernels + seen.add(kernels[0]) + assert len(seen) == 1 + + +# Graph batches: Gset, every G an engine of the envelope captures. That is every integer 1..32 +# plus the padded default list for cuda_graph_config {enable_padding: true, max_batch_size: 256} +# (1, 2, 4, 8..128 step 8, 192, 256); the padding-off list at 256 (1..32, 64, 128, 256) and the +# gates' earlier sets (1..8 and {1, 2, 4, 8, 16, 24, 32}) are subsets. Both steps are captured at +# every G: next_n 1 without DSpark and next_n 6 (the DSpark verify step) with it. A next_n-1 +# generation row reserves one row at either ratio (at ratio 2 it holds an entry only when the row +# completes a pair), and a DSpark verify request reserves ceil(6 / r) rows: 6 at ratio 1, 3 at +# ratio 2. So n = G (next_n 1) and n = 6G / 3G (next_n 6) over Gset: up to 1536 rows at ratio 1 +# and 768 at ratio 2. +_GSET = tuple(range(1, 33)) + (40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256) +assert len(_GSET) == 46 +_DECODE_G = _GSET +_DSPARK_G = _GSET +_CAPTURE = [(r, "decode", g) for r in (1, 2) for g in _DECODE_G] + [ + (r, "dspark", g) for r in (1, 2) for g in _DSPARK_G +] + + +def _graph_step(ratio: int, kind: str, G: int, blocks: int, seed: int, padding: bool): + """One captured step's rows: valid entries request by request, then reserved rows (-1), n in + total. Pages over the whole engine pool, the last valid entry on the far end's last slot. With + `padding` the last request is a graph padding request (no entries).""" + block_kv = 128 // ratio + per = 1 if kind == "decode" else -(-6 // ratio) + n = G * per + g = torch.Generator().manual_seed(seed) + pos = torch.randint(0, 30 * block_kv, (G,), generator=g) + if kind == "decode": + emits = torch.ones(G, dtype=torch.int64) if ratio == 1 else (pos % 2 == 1).long() + else: + emits = torch.full((G,), per, dtype=torch.int64) + if padding and G > 1: + emits[-1] = 0 + live = [i for i in range(G) if emits[i] > 0] + if not live: # keep at least one entry so the replay writes something + emits[0] = 1 + live = [0] + cached = [int(p) // ratio for p in pos] + # the last live request's last entry sits on its last page's last slot + j = live[-1] + cached[j] = (int(torch.randint(1, 30, (1,), generator=g)) + 1) * block_kv - int(emits[j]) + reqs = [(cached[i], int(emits[i])) for i in live] + b = _engine_batch(ratio, blocks, [c for c, _ in reqs], [m for _, m in reqs], seed=seed) + d_map, s_map = b.slot_mappings() + d = torch.full((n,), -1, dtype=torch.int64, device=DEV) + s = torch.full((n,), -1, dtype=torch.int64, device=DEV) + d[: b.n] = d_map + s[: b.n] = s_map + return b, d, s, n + + +@pytest.mark.parametrize("ratio,kind,G", _CAPTURE, ids=[f"r{r}-{k}-G{g}" for r, k, g in _CAPTURE]) +def test_cuda_graph_capture_sets(ratio, kind, G): + """The call captured at every graph batch G of Gset, for the next_n-1 step and the DSpark + verify step, at both ratios, on the engine-size pool (owner 14 / owner 20). After an eager + warm-up call, one capture; then 3 replays, each on new keys and new offsets (pages anywhere + in the pool, the far end included; reserved rows -1; a padding request on replays 1 and 3) + written into the captured buffers. Each replay leaves the whole allocation bitwise equal to + an eager call on the same inputs from the same snapshot, and equal to the native placement.""" + owner = 14 if ratio == 2 else 20 + blocks = ENGINE_BLOCKS[-1] + alloc = _engine_alloc(ratio, blocks) + view = _owner_view(alloc, ratio, owner) + per = 1 if kind == "decode" else -(-6 // ratio) + n = G * per + codes_s = torch.zeros(n, CODE_BYTES, dtype=torch.uint8, device=DEV) + word_s = torch.zeros(n, 1, dtype=torch.int32, device=DEV) + d_s = torch.full((n,), -1, dtype=torch.int64, device=DEV) + s_s = torch.full((n,), -1, dtype=torch.int64, device=DEV) + op(codes_s, word_s, view, d_s, s_s, n) # warm-up, every row skipped + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(codes_s, word_s, view, d_s, s_s, n) + torch.cuda.synchronize() + for rep in range(3): + seed = 100_000 * ratio + 100 * G + 10 * rep + (kind == "dspark") + b, d, s, n2 = _graph_step(ratio, kind, G, blocks, seed, padding=rep % 2 == 0) + assert n2 == n + codes, sb = _keys(n, seed=seed) + codes_s.copy_(codes) + word_s.copy_(_as_scale_word(sb)) + d_s.copy_(d) + s_s.copy_(s) + snap = alloc.clone() + graph.replay() + torch.cuda.synchronize() + got = alloc.clone() + alloc.copy_(snap) + op(codes_s, word_s, view, d_s, s_s, n) + torch.cuda.synchronize() + assert torch.equal(alloc, got), f"replay {rep} differs from eager" + assert not torch.equal(got, snap), f"replay {rep} wrote nothing" + _place(_owner_view(snap, ratio, owner), b, codes, sb, torch.arange(b.n, device=DEV)) + assert torch.equal(alloc, snap), f"replay {rep} differs from the native placement" + del snap, got + print( + f"[graph r={ratio} {kind} G={G}] n={n}: 3 replays bitwise equal to eager and to the placement" + ) + + +# Beyond the engine's pools: an allocation past 2^32 bytes, rows on the pages that straddle the +# 2^31- and 2^32-byte offsets and on the last page. The kernel forms int64 byte offsets; a caller +# that computed the offsets in 32 bits would wrap them (the control). +MARGIN_PAGES = {64: 1_000_000, 128: 500_000} # 4.352e9 bytes either way + + +@pytest.mark.parametrize("block_kv", [64, 128]) +def test_past_2_to_the_32_bytes_and_wrap_control(block_kv): + page_bytes = block_kv * ENTRY_BYTES + pages = MARGIN_PAGES[block_kv] + _ENGINE_CACHE.clear() + torch.cuda.empty_cache() + g = torch.Generator(device=DEV).manual_seed(block_kv) + pool = torch.randint( + 0, 256, (pages, block_kv, 1, ENTRY_BYTES), dtype=torch.uint8, device=DEV, generator=g + ) + assert pool.numel() > 2**32 + far = [2**31 // page_bytes, 2**32 // page_bytes, pages - 1] + for p, bound in zip(far[:2], (2**31, 2**32)): + assert p * page_bytes < bound < (p + 1) * page_bytes # the page straddles the boundary + b = _Batch(block_kv, [0, 0, 0], [block_kv] * 3, seed=block_kv, num_pages=3) + b.table = torch.tensor(far, dtype=torch.int64, device=DEV).view(3, 1) + b.num_pages = pages + codes, sb = _keys(b.n, seed=block_kv) + d_map, s_map = b.slot_mappings() + snap = pool.clone() + op(codes, _as_scale_word(sb), pool, d_map, s_map, b.n) + torch.cuda.synchronize() + exp = snap.clone() + _place(exp, b, codes, sb, torch.arange(b.n, device=DEV)) + assert torch.equal(pool, exp), "pool differs from the native placement" + del exp + # control: the same offsets wrapped to signed 32 bits (negative ones are skipped, the others + # land in the pool's first 2^31 bytes) + pool.copy_(snap) + + def wrap(t): + return ((t + 2**31) % 2**32) - 2**31 + + op(codes, _as_scale_word(sb), pool, wrap(d_map), wrap(s_map), b.n) + torch.cuda.synchronize() + frac = _row_mismatch(b, pool, codes, sb) + past = (d_map >= 2**31).float().mean().item() + print( + f"[margin block_kv={block_kv}] pool {pool.numel()} B, pages {far}: bit-exact; int32-wrapped " + f"offsets: {frac:.4f} of rows wrong ({past:.4f} of rows start past 2^31)" + ) + assert frac >= past > 0.5 + del pool, snap + + +def test_control_wrong_owner_view(): + """Owner 14's offsets written through owner 2's view (the ratio-2 views overlap, shifted by one + page each): the bytes land two pages early, and owner 14's reader sees none of them.""" + ratio, blocks = 2, ENGINE_BLOCKS[-1] + alloc = _engine_alloc(ratio, blocks) + cached, new = _far_step(64, 12, 20, seed=9) + b = _engine_batch(ratio, blocks, cached, new, seed=9) + codes, sb = _keys(b.n, seed=10) + d_map, s_map = b.slot_mappings() + snap = alloc.clone() + op(codes, _as_scale_word(sb), _owner_view(alloc, ratio, 2), d_map, s_map, b.n) + torch.cuda.synchronize() + frac = _row_mismatch(b, _owner_view(alloc, ratio, 14), codes, sb) + alloc.copy_(snap) + print(f"[control] owner 14's offsets through owner 2's view: {frac:.4f} of rows wrong") + assert frac > 0.9 + del snap + + +# Pool-size invariance. Every V4.1 persistent role of one cache manager shares one block +# count B, so B is capped by the whole GPU (284208 MiB = 298,013,687,808 B) divided by the manager's +# persistent bytes per block. The cheapest manager that holds an index-K pool is the PP2 stage-1 +# manager in native storage (layers 20-39: 356 B/token, 45,568 B per block, measured); an unsplit +# native manager costs 113,920 B per block and a stage-0 one 68,352. So no V4.1 manager on a GB300 +# holds more than floor(298,013,687,808 / 45,568) = 6,539,977 blocks. The largest pool measured (a +# stage-1 manager built alone at half the GPU) is 3,245,959 blocks, and the unsplit ceiling 2,615,990; both +# are run as intermediate sizes. One B_MAX serves both ratios (the ratio-2 owners live only in +# unsplit and stage-0 managers, whose own ceilings are lower), so the ratio-2 allocation at B_MAX, +# 3 x 6,539,977 pages of 64 entries (85,385,939,712 B), is a superset of any real one. +# The same step, at the same page ids, runs on exact-size allocations (the pool ends where the +# allocation ends) from B_min, the smallest pool that holds its pages, up to that bound. +B_MAX = 6_539_977 +INV_BLOCKS = (481, 1493, 151040, 852832, 2_615_990, 3_245_959, B_MAX) +B_OTHER = 852832 # the control's wrong page count +SENT = 0xA5 +FAR_BYTES = (2**31, 2**32, 2**33, 2**34, 2**35, 2**36) # those below the allocation's size + + +def _sentinel_alloc(ratio: int, blocks: int) -> torch.Tensor: + """The ratio's whole INDEXER_COMPRESS allocation for `blocks` blocks, every byte SENT.""" + _ENGINE_CACHE.clear() + torch.cuda.empty_cache() + pages = 3 * blocks if ratio == 2 else blocks + return torch.full((pages, 128 // ratio, 1, ENTRY_BYTES), SENT, dtype=torch.uint8, device=DEV) + + +def _others_are_sentinel(alloc: torch.Tensor, touched: torch.Tensor) -> bool: + """Whether every byte of the allocation outside the pages `touched` still holds SENT.""" + saved = alloc[touched].clone() + alloc[touched] = SENT + word = torch.full((8,), SENT, dtype=torch.uint8, device=DEV).view(torch.int64) + ok = not bool((alloc.view(-1).view(torch.int64) != word).any()) + alloc[touched] = saved + return ok + + +def _offsets(page: torch.Tensor, slot: torch.Tensor, block_kv: int): + base = page * block_kv * ENTRY_BYTES + return (base + slot * CODE_BYTES).contiguous(), ( + base + block_kv * CODE_BYTES + slot * SCALE_BYTES + ).contiguous() + + +def _wrap32(t: torch.Tensor) -> torch.Tensor: + return (((t + 2**31) % 2**32) - 2**31).contiguous() + + +def _scatter_rows(alloc: torch.Tensor, ratio: int, rows, wrap: bool = False, page_fn=None): + """One call per owner, through that owner's view; rows = [(owner, view page ids, slots, codes, scale bytes)].""" + block_kv = 128 // ratio + for owner, page, slot, codes, sb in rows: + d, s = _offsets(page if page_fn is None else page_fn(page), slot, block_kv) + if wrap: + d, s = _wrap32(d), _wrap32(s) + op(codes, _as_scale_word(sb), _owner_view(alloc, ratio, owner), d, s, codes.size(0)) + torch.cuda.synchronize() + + +def _expected_pages(ratio: int, rows): + """Native-torch placement: the sorted allocation pages the rows touch, each starting as SENT, + with every row's code bytes and scale word placed in the reader's layout.""" + block_kv = 128 // ratio + touched = torch.unique(torch.cat([OWNER_SHIFT.get(o, 0) + p for o, p, *_ in rows])) + exp = torch.full((touched.numel(), block_kv * ENTRY_BYTES), SENT, dtype=torch.uint8, device=DEV) + ar_c = torch.arange(CODE_BYTES, device=DEV) + ar_s = torch.arange(SCALE_BYTES, device=DEV) + for owner, page, slot, codes, sb in rows: + i = torch.searchsorted(touched, OWNER_SHIFT.get(owner, 0) + page) + exp[i.unsqueeze(1), (slot * CODE_BYTES).unsqueeze(1) + ar_c] = codes + exp[i.unsqueeze(1), (block_kv * CODE_BYTES + slot * SCALE_BYTES).unsqueeze(1) + ar_s] = sb + return touched, exp + + +def _matches_placement(alloc: torch.Tensor, touched: torch.Tensor, exp: torch.Tensor) -> bool: + got = alloc[touched].view(touched.numel(), -1) + return torch.equal(got, exp) and _others_are_sentinel(alloc, touched) + + +def _invariance_rows(ratio: int): + """The invariance step: 10 requests x 20 new entries (request 0 starting mid-page, pages crossed, + the last request ending on the last slot of block B_min - 1), over B_min blocks, each block used + once, page ids 3 x block at ratio 2. Every owner of the ratio writes it with its own keys. + Returns (B_min, rows).""" + block_kv = 128 // ratio + cached, new = _far_step(block_kv, 10, 20, seed=40 + ratio) + b_min = sum(max(1, -(-(c + m) // block_kv)) for c, m in zip(cached, new)) + b = _engine_batch(ratio, b_min, cached, new, seed=40 + ratio) + page, slot = b.pages_slots() + assert int(page[-1]) == (b_min - 1) * (3 if ratio == 2 else 1) and int(slot[-1]) == block_kv - 1 + owners = (2, 8, 14) if ratio == 2 else (20,) + rows = [] + for owner in owners: + codes, sb = _keys(b.n, seed=500 + owner) + rows.append((owner, page, slot, codes, sb)) + return b_min, rows + + +@pytest.mark.parametrize("ratio", [1, 2]) +def test_pool_size_invariance(ratio): + """The same step at the same page ids through every owner's view of the manager + layout (ratio 2: one 3B-page allocation, owners 2 / 8 / 14 at 0 / 1 / 2 pages in; ratio 1: + [B, 128, 1, 68]) at B = B_min, 481, 1493, 151040, 852832 and B_MAX, each an exact-size SENT-filled + allocation. At every B the pages the step touches are bitwise identical to B_min's and to the + native placement, and every other byte of the allocation is still SENT. Control: at B_MAX the + step's page ids counted from the end of a B_OTHER-block pool (a B-dependent addressing) fails + the same check.""" + b_min, rows = _invariance_rows(ratio) + touched, exp = _expected_pages(ratio, rows) + sizes = (b_min,) + INV_BLOCKS + assert list(sizes) == sorted(sizes) + first = None + for blocks in sizes: + alloc = _sentinel_alloc(ratio, blocks) + _scatter_rows(alloc, ratio, rows) + got = alloc[touched].view(touched.numel(), -1).clone() + first = got if first is None else first + same, placed, rest = ( + torch.equal(got, first), + torch.equal(got, exp), + _others_are_sentinel(alloc, touched), + ) + print( + f"[invariance r={ratio}] B={blocks}: allocation {tuple(alloc.shape)} ({alloc.numel()} B), " + f"{sum(r[3].size(0) for r in rows)} rows on {touched.numel()} pages: equal to B_min {same}, " + f"to the native placement {placed}, rest untouched {rest}" + ) + assert same and placed and rest, f"B={blocks}: result depends on the pool size" + if blocks != sizes[-1]: + del alloc + k = 3 if ratio == 2 else 1 + alloc.fill_(SENT) + _scatter_rows(alloc, ratio, rows, page_fn=lambda p: k * (B_OTHER - 1) - p) + wrong = _matches_placement(alloc, touched, exp) + print( + f"[invariance control r={ratio}] page ids counted from the end of a {B_OTHER}-block pool on " + f"B={B_MAX}: check passes {wrong}" + ) + assert not wrong + del alloc + + +def _far_pages(ratio: int, blocks: int): + """Allocation pages that straddle each FAR_BYTES offset, plus the last page of block blocks - 1 + of every owner (ratio 2: allocation pages 3B-3, 3B-2, 3B-1; ratio 1: B-1).""" + page_bytes = (128 // ratio) * ENTRY_BYTES + alloc_bytes = (3 * blocks if ratio == 2 else blocks) * page_bytes + marks = [x for x in FAR_BYTES if x < alloc_bytes] + pages = [x // page_bytes for x in marks] + for x, p in zip(marks, pages): + assert p * page_bytes < x < (p + 1) * page_bytes + return pages + ( + [3 * blocks - 3, 3 * blocks - 2, 3 * blocks - 1] if ratio == 2 else [blocks - 1] + ) + + +def _far_rows(ratio: int, blocks: int, seed: int): + """Every slot of each far page, written through the owner whose view holds the page (allocation + page p belongs to owner shift p % 3 at ratio 2, view page id p - shift = 3 x block).""" + block_kv = 128 // ratio + pages = _far_pages(ratio, blocks) + rows = [] + for owner in (2, 8, 14) if ratio == 2 else (20,): + shift = OWNER_SHIFT.get(owner, 0) + mine = [p - shift for p in pages if ratio == 1 or p % 3 == shift] + if not mine: + continue + assert all(m % (3 if ratio == 2 else 1) == 0 for m in mine) + page = torch.tensor(mine, dtype=torch.int64, device=DEV).repeat_interleave(block_kv) + slot = torch.arange(block_kv, device=DEV).repeat(len(mine)) + codes, sb = _keys(page.numel(), seed=seed + owner) + rows.append((owner, page, slot, codes, sb)) + return rows + + +@pytest.mark.parametrize("ratio", [1, 2]) +def test_far_end_at_b_max(ratio): + """Far-end addressing on the B_MAX allocation (ratio 2: 19,619,931 pages, 85,385,939,712 B; + ratio 1: 6,539,977 pages, 56,923,959,808 B): full pages of rows on the pages straddling bytes + 2^31 .. 2^35 (and 2^36 at ratio 2) and on the last page of block B_MAX - 1 of every owner (for owners 14 + and 20 the allocation's last byte). The whole allocation equals the SENT background with exactly + the native placement. Control: the same offsets wrapped to signed 32 bits fail the check.""" + rows = _far_rows(ratio, B_MAX, seed=900 + ratio) + touched, exp = _expected_pages(ratio, rows) + alloc = _sentinel_alloc(ratio, B_MAX) + assert alloc.numel() > 2**35 + _scatter_rows(alloc, ratio, rows) + ok = _matches_placement(alloc, touched, exp) + last = rows[-1] + tail = torch.equal(alloc.view(-1)[-SCALE_BYTES:], last[4][-1]) + print( + f"[far end r={ratio}] B={B_MAX}, allocation {alloc.numel()} B, pages {touched.tolist()}: native " + f"placement {ok}, owner {last[0]}'s last scale word is the allocation's last 4 bytes {tail}" + ) + assert ok and tail + alloc.fill_(SENT) + _scatter_rows(alloc, ratio, rows, wrap=True) + past = ( + torch.cat([_offsets(p, s, 128 // ratio)[0] for _, p, s, _, _ in rows]) + .ge(2**31) + .float() + .mean() + .item() + ) + wrong = _matches_placement(alloc, touched, exp) + print( + f"[far end control r={ratio}] offsets wrapped to signed 32 bits ({past:.4f} of rows start past " + f"2^31): check passes {wrong}" + ) + assert past > 0.5 and not wrong + del alloc diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_mla_rope_inplace.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_mla_rope_inplace.py new file mode 100644 index 000000000000..be65357e7a59 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_mla_rope_inplace.py @@ -0,0 +1,935 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the mla_rope_inplace catalog entry. + +The entry owns what is certified: `CELLS` is the list, `reference` and +`compare` are the gate. This file builds the DeepSeek-V4.1 tables, turns a +cell's spec into tensors, and keeps the checks a cell cannot carry: the +tables against the checkpoint's recipe, the forward/inverse round trip, the +wrong-variant controls, the op's rejections, CUDA-graph replay, every row +count 1..32768 for each core call form, the capture union, every envelope +engine's own table with positions at its far end, and the packed ViT form. +Every call writes between two guard rows that must come back unchanged, and +the `nope` channels must come back bitwise. +""" + +import math +import zlib + +import pytest +import torch +from torch.profiler import ProfilerActivity, profile + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.mla_rope_inplace import ( # noqa: E501 + mla_rope_inplace as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "mla_rope_inplace requires a CUDA device" + +DEV = "cuda" +ROPE = 64 +# The tables have the row count of the largest ones the target builds. The target sizes +# both tables to md.max_seq_len + 64 rows, where md.max_seq_len is the manager length Lm: +# the engine's max_seq_len L, plus 14 with DSpark on. The envelope's per-request lengths +# are below; the largest is the checkpoint's max_position_embeddings. Positions are drawn +# over the whole table, up to 16x the YaRN original context (65536), so the interpolated +# low-frequency channels rotate by angles a no-YaRN table cannot reproduce. +ENVELOPE_L = (16384, 32768, 65536, 131072, 262144, 1048576) +DSPARK_EXTRA_ROWS = 14 +TABLE_PAD_ROWS = 64 +MAX_POS = max(ENVELOPE_L) + DSPARK_EXTRA_ROWS + TABLE_PAD_ROWS # 1048654 table rows +YARN = dict(factor=16.0, original=65536, beta_fast=32, beta_slow=1) + + +def _table(base: float, yarn: bool, max_pos: int = MAX_POS, mscale: float = 1.0) -> torch.Tensor: + """[max_pos, 2, ROPE/2] fp32 (cos block, sin block) built with the checkpoint's recipe.""" + angles = torch.outer(torch.arange(max_pos, dtype=torch.float32, device=DEV), _freqs(base, yarn)) + return (torch.stack([angles.cos(), angles.sin()], dim=1) * mscale).contiguous() + + +def _freqs(base: float, yarn: bool) -> torch.Tensor: + """[ROPE/2] fp32 rotation frequencies: base^(-2i/ROPE), YaRN frequency interpolation only.""" + freqs = 1.0 / (base ** (torch.arange(0, ROPE, 2, dtype=torch.float32, device=DEV) / ROPE)) + if yarn: + + def corrected_dim(rotations: float) -> float: + return ( + ROPE * math.log(YARN["original"] / (rotations * 2 * math.pi)) / (2 * math.log(base)) + ) + + low = max(math.floor(corrected_dim(YARN["beta_fast"])), 0) + high = min(math.ceil(corrected_dim(YARN["beta_slow"])), ROPE - 1) + ramp = ( + (torch.arange(ROPE // 2, dtype=torch.float32, device=DEV) - low) / max(high - low, 1e-3) + ).clamp(0, 1) + smooth = 1 - ramp + freqs = freqs / YARN["factor"] * (1 - smooth) + freqs * smooth + return freqs + + +_TABLE_SPECS = { + # ratio-0 layers (0, 1) and the DSpark stages: rope_theta 10000, no YaRN + "ratio0_theta1e4": (10000.0, False), + # compressed layers (ratio 1 and 2): compress_rope_theta 160000 with YaRN, no mscale + "compressed_theta160k_yarn": (160000.0, True), +} +_TABLES: dict = {} + + +def _get_table(name: str) -> torch.Tensor: + if name not in _TABLES: + base, yarn = _TABLE_SPECS[name] + _TABLES[name] = _table(base, yarn) + return _TABLES[name] + + +@pytest.mark.parametrize("table_name", list(_TABLE_SPECS)) +def test_table_is_the_checkpoint_reference_recipe_at_full_size(table_name): + """Every row of the MAX_POS-row table equals the checkpoint reference's + precompute_freqs_cis rebuilt here as it builds it (int64 positions promoted to + fp32 angles, then torch.polar), on CUDA as the reference does, bitwise.""" + base, yarn = _TABLE_SPECS[table_name] + angles = torch.outer(torch.arange(MAX_POS, device=DEV), _freqs(base, yarn)) + assert angles.dtype == torch.float32 + cis = torch.polar(torch.ones_like(angles), angles) + table = _get_table(table_name) + assert table.shape == (MAX_POS, 2, ROPE // 2) + assert torch.equal(table[:, 0], cis.real) and torch.equal(table[:, 1], cis.imag) + + +def _reference( + x: torch.Tensor, + pos: torch.Tensor, + table: torch.Tensor, + nope: int, + inverse: bool, + is_neox: bool, + first: bool = False, +) -> torch.Tensor: + """`op.reference` in this file's argument order, `rope_dim` = ROPE. + + `first=True` is for a control only: it rotates the leading ROPE channels instead, by + moving them to the tail, rotating there, and moving them back. + """ + if first: + moved = torch.cat([x[..., ROPE:], x[..., :ROPE]], dim=-1) + y = op.reference(moved, pos, table, x.size(1), nope, ROPE, inverse, is_neox) + return torch.cat([y[..., -ROPE:], y[..., :-ROPE]], dim=-1) + return op.reference(x, pos, table, x.size(1), nope, ROPE, inverse, is_neox) + + +def _positions(num_tokens: int, gen: torch.Generator) -> torch.Tensor: + if num_tokens == 1: + return torch.tensor([MAX_POS - 1], dtype=torch.int32, device=DEV) + if num_tokens <= 8: # a decode / DSpark-verify run of consecutive positions + start = int(torch.randint(0, MAX_POS - num_tokens, (1,), generator=gen).item()) + return torch.arange(start, start + num_tokens, dtype=torch.int32).to(DEV) + pos = torch.randint(0, MAX_POS, (num_tokens,), generator=gen, dtype=torch.int32) + pos[0], pos[-1] = 0, MAX_POS - 1 + return pos.to(DEV) + + +def _run( + x: torch.Tensor, + pos: torch.Tensor, + table: torch.Tensor, + nope: int, + inverse: bool, + is_neox: bool, +) -> torch.Tensor: + """Call the op on x placed between two guard rows; lock both sides of the write.""" + num_tokens, num_heads, head_dim = x.shape + buf = torch.randn(num_tokens + 2, num_heads, head_dim, dtype=x.dtype, device=DEV) + buf[1:-1].copy_(x) + before = buf.clone() + data = buf[1:-1] + assert data.is_contiguous() + result = op(data, pos, table, num_heads, nope, ROPE, inverse, is_neox) + assert result is data, "the op rotates in place and returns data" + assert torch.equal(buf[0], before[0]) and torch.equal(buf[-1], before[-1]), "out-of-row write" + assert torch.equal(data[..., :nope], x[..., :nope]), "nope channels were modified" + return data.clone() + + +def _cell_inputs(spec, gen: torch.Generator): + """(x, positions, table) for one cell, built the way the target builds them.""" + form = spec["form"] + if form == "table": + # T = 1 at the table's last row, T = 6 a consecutive run, T = 8192 random with 0 and + # the last row included. + table = _get_table(spec["table"]) + assert table.shape[0] == spec["rows"] == MAX_POS + num_tokens = spec["T"] + pos = _positions(num_tokens, gen) + elif form == "group_first": + # Latent j of a ratio-r layer at position j * r (the group's first token), for every + # group whose first position is in the table. + table = _get_table(spec["table"]) + assert table.shape[0] == spec["rows"] == MAX_POS + num_tokens = spec["rows"] // spec["ratio"] + pos = (torch.arange(num_tokens, dtype=torch.int32) * spec["ratio"]).to(DEV) + elif form == "random_angles": + # NeoX half-split pairs over a whole 64-wide head, an arbitrary per-row table. + num_tokens = spec["T"] + angles = (torch.rand(num_tokens, ROPE // 2, generator=gen) * 64.0).to(DEV) + table = torch.stack([angles.cos(), angles.sin()], dim=1).contiguous() + pos = torch.arange(num_tokens, dtype=torch.int32, device=DEV) + else: + assert form == "vit_packed", form + # The vision tower's per-request call: its own [P + 1] table with the identity row + # at P, positions (t, t, P). + table = _vit_table(spec["P"], gen) + pos = _vit_positions(spec["P"]) + num_tokens = 3 * spec["P"] + x = torch.randn(num_tokens, spec["H"], spec["D"], generator=gen).to(DEV, torch.bfloat16) + return x, pos, table + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + gen = torch.Generator().manual_seed(zlib.crc32(cell.why.encode())) + x, pos, table = _cell_inputs(spec, gen) + out = _run(x, pos, table, spec["nope"], inverse=spec["inverse"], is_neox=spec["is_neox"]) + op.compare( + out, + op.reference( + x, pos, table, spec["H"], spec["nope"], ROPE, spec["inverse"], spec["is_neox"] + ), + ) + if spec["form"] == "vit_packed": + assert torch.equal(out[2::3], x[2::3]), "a v row read through the identity row changed" + + +def test_inverse_undoes_forward(): + gen = torch.Generator().manual_seed(5000) + x = torch.randn(8192, 16, 512, generator=gen).to(DEV, torch.bfloat16) + pos = _positions(8192, gen) + table = _get_table("compressed_theta160k_yarn") + y = _run(x, pos, table, 448, inverse=False, is_neox=False) + back = _run(y, pos, table, 448, inverse=True, is_neox=False) + # Two bf16 roundings (after the forward and after the inverse rotation), each at most half + # an ulp = 2^-9 of the pair's magnitude; the rotation preserves the pair norm, so the error + # of either channel is bounded by 2 * sqrt(2) * 2^-9 = 5.5e-3 of the pair norm. Gate at 2^-6. + pair = x.float()[..., 448:].unflatten(-1, (-1, 2)).norm(dim=-1).repeat_interleave(2, dim=-1) + err = (back.float()[..., 448:] - x.float()[..., 448:]).abs() + assert bool((err <= pair * 2**-6).all()), f"max err / pair norm {(err / pair).max().item()}" + twice = _run(y, pos, table, 448, inverse=False, is_neox=False) + assert not bool(((twice.float()[..., 448:] - x.float()[..., 448:]).abs() <= pair * 2**-6).all()) + + +# Deliberately wrong references. Each must be rejected by the same gate that accepts the +# correct reference, which is what shows the gate can see the checkpoint's RoPE mistakes. +_MSCALE = 0.1 * math.log(YARN["factor"]) + 1.0 +CONTROLS = { + # (table the op runs with, op inverse, positions scale, wrong reference builder) + "rotate_half_pairing": ( + "compressed_theta160k_yarn", + False, + 1, + lambda x, p: (x, p, "same", 448, False, True, False), + ), + "first_64_dims": ( + "compressed_theta160k_yarn", + False, + 1, + lambda x, p: (x, p, "same", 448, False, False, True), + ), + "yarn_with_mscale": ( + "compressed_theta160k_yarn", + False, + 1, + lambda x, p: (x, p, "mscale", 448, False, False, False), + ), + "compressed_without_yarn": ( + "compressed_theta160k_yarn", + False, + 1, + lambda x, p: (x, p, "noyarn160k", 448, False, False, False), + ), + "compressed_with_theta1e4": ( + "compressed_theta160k_yarn", + False, + 1, + lambda x, p: (x, p, "ratio0_theta1e4", 448, False, False, False), + ), + "ratio0_with_yarn_theta160k": ( + "ratio0_theta1e4", + False, + 1, + lambda x, p: (x, p, "compressed_theta160k_yarn", 448, False, False, False), + ), + "ratio0_with_yarn_theta1e4": ( + "ratio0_theta1e4", + False, + 1, + lambda x, p: (x, p, "yarn1e4", 448, False, False, False), + ), + "forward_instead_of_inverse": ( + "compressed_theta160k_yarn", + True, + 1, + lambda x, p: (x, p, "same", 448, False, False, False), + ), + "latent_at_group_last_position": ( + "compressed_theta160k_yarn", + False, + 2, + lambda x, p: (x, p + 1, "same", 448, False, False, False), + ), +} + + +def _control_table(name: str, run_table: str) -> torch.Tensor: + if name == "same": + return _get_table(run_table) + if name == "mscale": + return _table(160000.0, True, mscale=_MSCALE) + if name == "noyarn160k": + return _table(160000.0, False) + if name == "yarn1e4": + return _table(10000.0, True) + return _get_table(name) + + +@pytest.mark.parametrize("control", list(CONTROLS)) +def test_wrong_variant_is_rejected(control): + run_table, inverse, pos_scale, wrong = CONTROLS[control] + gen = torch.Generator().manual_seed(6000) + x = torch.randn(4096, 16, 512, generator=gen).to(DEV, torch.bfloat16) + pos = _positions(4096, gen) // pos_scale * pos_scale # ratio-2 group-first positions + table = _get_table(run_table) + out = _run(x, pos, table, 448, inverse=inverse, is_neox=False) + op.compare(out, _reference(x, pos, table, 448, inverse=inverse, is_neox=False)) + wx, wp, wtab, wnope, winv, wneox, wfirst = wrong(x, pos) + bad = _reference(wx, wp, _control_table(wtab, run_table), wnope, winv, wneox, wfirst) + close = torch.isclose(out.float(), bad.float(), rtol=1.6e-2, atol=1e-5) + frac_bad = 1.0 - close.float().mean().item() + max_err = (out.float() - bad.float()).abs().max().item() + print( + f"control {control}: {frac_bad:.4f} of elements outside the gate, max abs err {max_err:.3e}" + ) + assert not bool(close.all()), f"control {control} was not rejected" + + +def _raises(msg, fn): + with pytest.raises(RuntimeError, match=msg): + fn() + + +@pytest.mark.parametrize( + "case", ["strided_tail_view", "int64_positions", "fp32_data", "flat_cos_sin_table"] +) +def test_rejected_inputs_leave_data_untouched(case): + x = torch.randn(4, 16, 512, device=DEV, dtype=torch.bfloat16) + before = x.clone() + pos = torch.arange(4, dtype=torch.int32, device=DEV) + table = _table(10000.0, False, max_pos=16) + if case == "strided_tail_view": + _raises( + "data must be contiguous", + lambda: op(x[..., 448:], pos, table, 16, 0, ROPE, False, False), + ) + elif case == "int64_positions": + _raises( + "position_ids must be int32", + lambda: op(x, pos.long(), table, 16, 448, ROPE, False, False), + ) + elif case == "fp32_data": + xf = x.float() + _raises( + "unsupported dtype, expected bf16 or fp16", + lambda: op(xf, pos, table, 16, 448, ROPE, False, False), + ) + else: + _raises( + r"cos_sin_cache must have shape \[max_positions, 2, rope_dim / 2\]", + lambda: op(x, pos, table.view(16, ROPE), 16, 448, ROPE, False, False), + ) + torch.cuda.synchronize() + assert torch.equal(x, before) + + +def test_cuda_graph_replay_reads_positions_and_data_at_replay(): + table = _get_table("compressed_theta160k_yarn") + data = torch.zeros(8, 16, 512, device=DEV, dtype=torch.bfloat16) + pos = torch.zeros(8, dtype=torch.int32, device=DEV) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(data, pos, table, 16, 448, ROPE, False, False) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(data, pos, table, 16, 448, ROPE, False, False) + gen = torch.Generator().manual_seed(7000) + for _ in range(3): + x = torch.randn(8, 16, 512, generator=gen).to(DEV, torch.bfloat16) + p = _positions(8, gen) + data.copy_(x) + pos.copy_(p) + graph.replay() + torch.cuda.synchronize() + op.compare(data.clone(), _reference(x, p, table, 448, inverse=False, is_neox=False)) + + +# ─── Every call form of the core target at every row count 1..32768, and the capture union ── +# +# DeepSeek-V4.1's core target rotates, per attention layer: q [T, H, 512] and the +# window kv [T, 1, 512] with the layer's table (ratio-0 on layers 0-1, the +# compressed table otherwise), the attention output [T, H, 512] inversely, the +# index q [T, 32, 128] with the compressed table, and on KV-source layers the +# index k [n, 1, 128] and the compressed latent [n, 1, 512] at compressed +# positions. H is the rank's share of the 64 heads, 64 / tp under attention TP: +# 32 (tp 2: the pp2 x tep2 / tp2 stages), 16 (tp 4: tep4, tp4 and their pp2 +# stages) or 8 (tp 8: tep8); it is 64 under attention DP (dep4, dep8 and their +# pp2 stages) and in the DSpark drafter's q and output at every layout (ratio-0 +# table). Every H meets both tables: a pp2 stage 0 holds layers 0-19 (ratio 0 +# and 2), stage 1 layers 20-39 (compressed table). T and n are at most the +# engine's max_num_tokens, 16384 per rank; in layers 21-39 the decoder bounded +# replay's pass adds the finishing requests' 128-row tails to the generation +# rows, up to 128 x 256 = 32768. The launch grid is (T, ceil(H / heads per +# block)): block (t, g) reads position_ids[t] and rotates row t only, and the +# kernel variant depends on H, not T. The tests below check that on the GPU at +# every row count from 1 to 32768, for every form. + +_N_MAX = 32768 +# (H, head_dim, nope_dim, inverse) of every call the core target makes +TARGET_FORMS = { + "q_h8": (8, 512, 448, False), + "q_h16": (16, 512, 448, False), + "q_h32": (32, 512, 448, False), + "q_h64": (64, 512, 448, False), + "window_kv": (1, 512, 448, False), + "attn_out_h8": (8, 512, 448, True), + "attn_out_h16": (16, 512, 448, True), + "attn_out_h32": (32, 512, 448, True), + "attn_out_h64": (64, 512, 448, True), + "index_q": (32, 128, 64, False), + "index_k": (1, 128, 64, False), + "compressed_latent": (1, 512, 448, False), +} +# q, kv and the output use the layer's table; the indexer and the latent always use the compressed one +_LAYER_TABLE_FORMS = ( + "q_h8", "q_h16", "q_h32", "q_h64", "window_kv", "attn_out_h8", "attn_out_h16", "attn_out_h32", "attn_out_h64", +) # fmt: skip +_FORM_TABLES = [ + (f, t) + for f in TARGET_FORMS + for t in _TABLE_SPECS + if f in _LAYER_TABLE_FORMS or t != "ratio0_theta1e4" +] +assert len(_FORM_TABLES) == 21 +# T values in the core's runtime-value dump (boot gates tep4, dep4, tp4 and tep8) +_DUMPED_T = ( + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 12, + 14, + 16, + 22, + 24, + 48, + 128, + 140, + 888, + 890, + 1018, + 1030, + 1935, + 1942, + 2048, +) +# The capture union per rank. G is every CUDA-graph batch the envelope +# captures: every batch 1..32, plus TRT-LLM's padded list for +# cuda_graph_config {enable_padding: true, max_batch_size: 256} +# (_generate_cuda_graph_batch_sizes in llmapi/llm_args.py: 1, 2, 4, 8, 16..128 +# in steps of 8, 192, 256). The padding-off list at 256 (1..32, 64, 128, 256) +# is a subset. Per G, the target's captured calls rotate these row counts: +# G -- a decode step (next_n 1), every form; a ratio-2 layer reserves one +# compressed row per request, so its latent and index k are G rows too; +# also the DSpark drafter's bonus-row window kv (ratio-0 table); +# 3G -- the compressed latent and index k of a ratio-2 layer in the DSpark +# target-verify step (next_n 6): each request reserves ceil(6 / 2) = 3 +# compressed rows (num_total_compressed_tokens[2] = G * 3); +# 5G -- the DSpark drafter's block: q and the output at 64 heads, and the +# window kv, on the ratio-0 table; +# 6G -- the rest of the verify step: q, window kv, index q, the inverse +# output, and the ratio-1 latent and index k (one row per token). +# Every form and table is captured at every T of the union (126 values, up to +# 1536), a superset of the row counts each form reaches. +_BATCHES = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_CAPTURE_T = tuple(sorted({m * b for b in _BATCHES for m in (1, 3, 5, 6)})) +assert len(_BATCHES) == 46 and len(_CAPTURE_T) == 126 and max(_CAPTURE_T) == 1536 + + +def _sweep_positions(num_tokens: int, gen: torch.Generator) -> torch.Tensor: + """Random positions in [1, MAX_POS), with 1 and MAX_POS - 1 included. + + Position 0 is the identity rotation, so a row at position 0 that the kernel + skipped would look correct; the sweep keeps every row's position nonzero. + """ + pos = torch.randint(1, MAX_POS, (num_tokens,), generator=gen, dtype=torch.int32) + pos[0], pos[-1] = MAX_POS - 1, 1 + return pos.to(DEV) + + +@pytest.mark.parametrize("form,table_name", _FORM_TABLES) +def test_every_row_count_is_the_same_per_row_function(form, table_name): + """Every prefix T = 1..32768 returns exactly the first T rows of the 32768-row + result, which is held to the fp32 reference. + + Each call runs in place on the first T rows of a buffer whose row T is a + guard row holding finite data that must come back bit-identical. A row the + kernel skipped would keep its unrotated value, which differs from the + rotated one because no position is 0. + """ + num_heads, head_dim, nope = TARGET_FORMS[form][:3] + inverse = TARGET_FORMS[form][3] + gen = torch.Generator().manual_seed(8000 + num_heads + head_dim) + x = torch.randn(_N_MAX + 1, num_heads, head_dim, generator=gen).to(DEV, torch.bfloat16) + guard_row = x[_N_MAX].clone() + x = x[:_N_MAX] + pos = _sweep_positions(_N_MAX, gen) + table = _get_table(table_name) + full = _run(x, pos, table, nope, inverse=inverse, is_neox=False) + ref = _reference(x, pos, table, nope, inverse=inverse, is_neox=False) + op.compare(full, ref) + exact = float((full[..., nope:] == ref[..., nope:]).float().mean()) + print( + f"\n[{form} {table_name}] T={_N_MAX}: rope channels bit-exact {exact:.6f} vs the fp32 reference" + ) + del ref + fb, gb = full.view(torch.int16), guard_row.view(torch.int16) + buf = torch.empty(_N_MAX + 1, num_heads, head_dim, dtype=torch.bfloat16, device=DEV) + mismatch = torch.zeros(_N_MAX + 1, dtype=torch.int64, device=DEV) + guard = torch.zeros(_N_MAX + 1, dtype=torch.int64, device=DEV) + differs = torch.empty( + _N_MAX, num_heads, head_dim, dtype=torch.bool, device=DEV + ) # one buffer; .any() keeps the reduction free of a widened copy + for t in range(1, _N_MAX + 1): + buf[:t].copy_(x[:t]) + buf[t].copy_(guard_row) + op(buf[:t], pos[:t], table, num_heads, nope, ROPE, inverse, False) + mismatch[t] = torch.ne(buf[:t].view(torch.int16), fb[:t], out=differs[:t]).any() + guard[t] = (buf[t].view(torch.int16) != gb).sum() + bad = mismatch.nonzero().flatten().tolist() + over = guard.nonzero().flatten().tolist() + assert not bad, ( + f"{len(bad)} row counts differ from the {_N_MAX}-row result's rows, first {bad[:10]}" + ) + assert not over, f"{len(over)} row counts wrote the guard row, first {over[:10]}" + + +def test_row_sweep_harness_sees_a_one_row_shift(): + """Control for the sweep: its bitwise per-row check rejects a read one row off (data and position).""" + num_heads, head_dim, nope = 16, 512, 448 + gen = torch.Generator().manual_seed(8100) + x = torch.randn(2049, num_heads, head_dim, generator=gen).to(DEV, torch.bfloat16) + pos = _sweep_positions(2049, gen) + table = _get_table("compressed_theta160k_yarn") + full = _run(x[:2048], pos[:2048], table, nope, inverse=False, is_neox=False) + for t in (1, 2, 7, 1024, 2048): + shifted = _run(x[1 : t + 1], pos[1 : t + 1], table, nope, inverse=False, is_neox=False) + n_bad = int((shifted.view(torch.int16) != full.view(torch.int16)[:t]).sum().item()) + print(f"[control row shift] T={t}: {n_bad} / {t * num_heads * head_dim} elements differ") + assert n_bad > t * num_heads * head_dim // 2 + # the rotation alone: the right data at the neighbouring row's position + for t in (1, 7, 2048): + wrong_pos = _run(x[:t], pos[1 : t + 1], table, nope, inverse=False, is_neox=False) + n_bad = int((wrong_pos.view(torch.int16) != full.view(torch.int16)[:t]).sum().item()) + print( + f"[control position shift] T={t}: {n_bad} / {t * num_heads * ROPE} rope elements differ" + ) + assert n_bad > t * num_heads * ROPE // 2 + + +@pytest.mark.parametrize("form,table_name", _FORM_TABLES) +def test_dumped_values_against_reference(form, table_name): + """Fresh inputs at every dumped T, each call form and table, against the fp32 reference.""" + num_heads, head_dim, nope, inverse = TARGET_FORMS[form] + table = _get_table(table_name) + for num_tokens in _DUMPED_T: + gen = torch.Generator().manual_seed(9000 + num_tokens) + x = torch.randn(num_tokens, num_heads, head_dim, generator=gen).to(DEV, torch.bfloat16) + pos = _positions(num_tokens, gen) + out = _run(x, pos, table, nope, inverse=inverse, is_neox=False) + op.compare(out, _reference(x, pos, table, nope, inverse=inverse, is_neox=False)) + + +def test_one_compiled_kernel_per_form_at_every_row_count(): + """Probe: each call form launches the same single kernel at every T (compiled ahead of time, no JIT).""" + table = _get_table("compressed_theta160k_yarn") + gen = torch.Generator().manual_seed(9500) + pos = _sweep_positions(_N_MAX, gen) + counts = sorted( + set(_DUMPED_T) + | set(_CAPTURE_T) + | {9, 64, 65, 129, 4096, 8191, 8192, 8193, 16383, 16384, 16385, _N_MAX - 1, _N_MAX} + ) + for form, (num_heads, head_dim, nope, inverse) in TARGET_FORMS.items(): + data = torch.randn(_N_MAX, num_heads, head_dim, generator=gen).to(DEV, torch.bfloat16) + names = set() + for t in counts: + with profile(activities=[ProfilerActivity.CUDA]) as prof: + op(data[:t], pos[:t], table, num_heads, nope, ROPE, inverse, False) + torch.cuda.synchronize() + launched = [ + e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA + ] + assert len(launched) == 1, (form, t, launched) + names.add(launched[0]) + print(f"{form}: {sorted(names)}") + assert len(names) == 1 and "mlaRoPEInplaceKernel" in next(iter(names)), (form, names) + + +@pytest.mark.parametrize("num_tokens", _CAPTURE_T) +def test_capture_union_replay_is_bitwise_eager(num_tokens): + """Every call form and table captured in one graph at T, after an eager call at that + shape, replayed on new data and positions: each output bitwise equal to an eager call + on that data. + + The replay on the + capture-time data must differ from the replay on new data for every form, + so a replay that reused stale inputs would fail. + """ + t = num_tokens + gen = torch.Generator().manual_seed(10000 + t) + calls = [(f, tn) for f, tn in _FORM_TABLES] + data = { + c: torch.randn(t, TARGET_FORMS[c[0]][0], TARGET_FORMS[c[0]][1], generator=gen).to( + DEV, torch.bfloat16 + ) + for c in calls + } + pos = _positions(t, gen) + # The engine runs each graph shape eagerly (its warm-up) before it captures it. + for c in calls: + num_heads, _, nope, inverse = TARGET_FORMS[c[0]] + op(data[c], pos, _get_table(c[1]), num_heads, nope, ROPE, inverse, False) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for c in calls: + num_heads, _, nope, inverse = TARGET_FORMS[c[0]] + op(data[c], pos, _get_table(c[1]), num_heads, nope, ROPE, inverse, False) + graph.replay() + torch.cuda.synchronize() + old = {c: v.clone() for c, v in data.items()} + new = {c: torch.randn(v.shape, generator=gen).to(DEV, torch.bfloat16) for c, v in data.items()} + new_pos = _positions(t, gen) + for c in calls: + data[c].copy_(new[c]) + pos.copy_(new_pos) + graph.replay() + torch.cuda.synchronize() + for c in calls: + num_heads, _, nope, inverse = TARGET_FORMS[c[0]] + table = _get_table(c[1]) + eager = _run(new[c], new_pos, table, nope, inverse=inverse, is_neox=False) + assert torch.equal(data[c], eager), (c, t) + assert not torch.equal(data[c], old[c]), (c, t, "the replay did not read the new data") + op.compare( + data[c], _reference(new[c], new_pos, table, nope, inverse=inverse, is_neox=False) + ) + + +# ─── Every engine's table size, with positions at its far end ── +# +# The target builds both tables once per engine with md.max_seq_len + 64 rows, +# md.max_seq_len = Lm = L (DSpark off) or L + 14 (DSpark on), for each L in +# ENVELOPE_L. Positions reach about Lm + 5: the drafts of a verify step past a +# request's last token, and the drafter's block. +_ENGINE_TABLES = [(seq_len, dspark) for seq_len in ENVELOPE_L for dspark in (False, True)] +_E4_ROWS = 4096 + + +def _far_end_positions(lm: int, rows: int) -> torch.Tensor: + """Lm - 6 .. Lm + 5 (a full-length request's last tokens and the drafts past + them), the table's last 64 rows, and 2^k - 1, 2^k below the table size.""" + far = set(range(lm - 6, lm + 6)) | set(range(rows - TABLE_PAD_ROWS, rows)) + far |= {v for k in range(10, 21) for v in (2**k - 1, 2**k)} + return torch.tensor(sorted(p for p in far if 0 <= p < rows), dtype=torch.int32) + + +@pytest.mark.parametrize( + "seq_len,dspark", + _ENGINE_TABLES, + ids=[f"L{s}-dspark_{'on' if d else 'off'}" for s, d in _ENGINE_TABLES], +) +def test_engine_table_far_end_positions(seq_len, dspark): + """One engine's own tables (md.max_seq_len + 64 rows), built on the spot: every form and + table call at positions over the whole table with its far end included, eager against the + fp32 reference, and captured in one graph at 6G for G = 256 that replays bitwise equal + to eager on new data and new positions.""" + lm = seq_len + (DSPARK_EXTRA_ROWS if dspark else 0) + rows = lm + TABLE_PAD_ROWS + tables = {} + for name, (base, yarn) in _TABLE_SPECS.items(): + tables[name] = _table(base, yarn, max_pos=rows) + # a shorter table is the first rows of the longest one, bitwise + assert torch.equal(tables[name], _get_table(name)[:rows]), (name, rows) + far = _far_end_positions(lm, rows) + gen = torch.Generator().manual_seed(13000 + rows) + + def positions(num_tokens: int) -> torch.Tensor: + """Random over [0, rows), with the far-end set at the start and, reversed, at the end.""" + pos = torch.randint(0, rows, (num_tokens,), generator=gen, dtype=torch.int32) + pos[: far.numel()] = far + pos[-far.numel() :] = far.flip(0) + return pos.to(DEV) + + pos = positions(_E4_ROWS) + far_rows = torch.cat( + [torch.arange(far.numel()), torch.arange(_E4_ROWS - far.numel(), _E4_ROWS)] + ).to(DEV) + for form, table_name in _FORM_TABLES: + num_heads, head_dim, nope, inverse = TARGET_FORMS[form] + table = tables[table_name] + x = torch.randn(_E4_ROWS, num_heads, head_dim, generator=gen).to(DEV, torch.bfloat16) + out = _run(x, pos, table, nope, inverse=inverse, is_neox=False) + ref = _reference(x, pos, table, nope, inverse=inverse, is_neox=False) + op.compare(out, ref) + op.compare(out[far_rows], ref[far_rows]) + if form == "q_h16" and table_name == "compressed_theta160k_yarn": + # Control: the same gate rejects the far-end rows rotated at the neighbouring position. + wrong = _reference( + x, (pos - 1).clamp(min=0), table, nope, inverse=inverse, is_neox=False + ) + close = torch.isclose( + out[far_rows].float(), wrong[far_rows].float(), rtol=1.6e-2, atol=1e-5 + ) + frac_bad = 1.0 - close.float().mean().item() + print( + f"\n[L {seq_len}, dspark {dspark}] rows {rows}: far-end rows at position - 1: " + f"{frac_bad:.4f} of elements outside the gate" + ) + assert frac_bad > 0.01, frac_bad + + t = max(_CAPTURE_T) + data = { + c: torch.randn(t, TARGET_FORMS[c[0]][0], TARGET_FORMS[c[0]][1], generator=gen).to( + DEV, torch.bfloat16 + ) + for c in _FORM_TABLES + } + cap_pos = positions(t) + for c in _FORM_TABLES: + num_heads, _, nope, inverse = TARGET_FORMS[c[0]] + op(data[c], cap_pos, tables[c[1]], num_heads, nope, ROPE, inverse, False) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for c in _FORM_TABLES: + num_heads, _, nope, inverse = TARGET_FORMS[c[0]] + op(data[c], cap_pos, tables[c[1]], num_heads, nope, ROPE, inverse, False) + new = {c: torch.randn(v.shape, generator=gen).to(DEV, torch.bfloat16) for c, v in data.items()} + new_pos = positions(t) + assert not torch.equal(new_pos, cap_pos) + for c in _FORM_TABLES: + data[c].copy_(new[c]) + cap_pos.copy_(new_pos) + graph.replay() + torch.cuda.synchronize() + for c in _FORM_TABLES: + num_heads, _, nope, inverse = TARGET_FORMS[c[0]] + eager = _run(new[c], new_pos, tables[c[1]], nope, inverse=inverse, is_neox=False) + assert torch.equal(data[c], eager), (c, rows) + op.compare( + data[c], _reference(new[c], new_pos, tables[c[1]], nope, inverse=inverse, is_neox=False) + ) + + +# The DeepSeek-V4.1 vision tower rotates q and k of its packed qkv projection in ONE call: the +# [P, 3072] bf16 output viewed as [3P, 16, 64] (row 3t = q of patch t, 3t+1 = +# k, 3t+2 = v), nope_dim 0, rope_dim 64, NeoX pairs, forward. The fp32 table is +# [P + 1, 2, 32]: rows 0..P-1 are the per-patch 2D RoPE cos/sin, row P is the +# identity (cos 1, sin 0). Positions are int32 (t, t, P) per patch, so q and k +# of patch t read row t and v reads the identity row and comes back bit-exact. +# P is every patch of one request: at most 17037 over the 900 MMMU val +# documents, so T = 3P reaches 51111. +_VIT_P_MAX = 17037 +_VIT_T_MAX = 3 * _VIT_P_MAX +_VIT_HEADS, _VIT_HEAD_DIM = 16, 64 +# The prototype's whole-tower CUDA graph ran at the golden cases' patch counts; +# 17037 is the MMMU maximum. +_VIT_GRAPH_P = (1517, 4680, 5325, 8772, 17037) + + +def _vit_table(num_patches: int, gen: torch.Generator) -> torch.Tensor: + """[P + 1, 2, 32] fp32: random per-patch angles, then the identity row.""" + angles = torch.rand(num_patches, ROPE // 2, generator=gen) * 64.0 + table = torch.stack([angles.cos(), angles.sin()], dim=1) + identity = torch.stack([torch.ones(ROPE // 2), torch.zeros(ROPE // 2)])[None] + return torch.cat([table, identity]).contiguous().to(DEV) + + +def _vit_positions(num_patches: int) -> torch.Tensor: + t = torch.arange(num_patches, dtype=torch.int32) + return torch.stack([t, t, torch.full_like(t, num_patches)], dim=1).reshape(-1).to(DEV) + + +def _vit_call(data: torch.Tensor, pos: torch.Tensor, table: torch.Tensor) -> None: + op(data, pos, table, _VIT_HEADS, 0, ROPE, False, True) + + +def test_vit_packed_qkv_every_row_count(): + """T = 1..51111 in the packed [3P, 16, 64] form: the full call is held to the fp32 + reference with every v row bitwise unchanged, and every prefix T is bitwise the + first T rows of it; the guard row after T is never written.""" + gen = torch.Generator().manual_seed(12000) + table = _vit_table(_VIT_P_MAX, gen) + pos = _vit_positions(_VIT_P_MAX) + x = torch.randn(_VIT_T_MAX + 1, _VIT_HEADS, _VIT_HEAD_DIM, generator=gen).to( + DEV, torch.bfloat16 + ) + guard_row = x[_VIT_T_MAX].clone() + x = x[:_VIT_T_MAX] + full = _run(x, pos, table, 0, inverse=False, is_neox=True) + ref = _reference(x, pos, table, 0, inverse=False, is_neox=True) + op.compare(full, ref) + assert torch.equal(full[2::3], x[2::3]), "a v row read through the identity row changed" + exact_qk = float( + ( + full.view(-1, 3, _VIT_HEADS, _VIT_HEAD_DIM)[:, :2] + == ref.view(-1, 3, _VIT_HEADS, _VIT_HEAD_DIM)[:, :2] + ) + .float() + .mean() + ) + print( + f"\n[vit packed] T={_VIT_T_MAX}: q/k bit-exact {exact_qk:.6f} vs fp32 reference, v rows bitwise unchanged" + ) + fb, gb = full.view(torch.int16), guard_row.view(torch.int16) + buf = torch.empty(_VIT_T_MAX + 1, _VIT_HEADS, _VIT_HEAD_DIM, dtype=torch.bfloat16, device=DEV) + mismatch = torch.zeros(_VIT_T_MAX + 1, dtype=torch.bool, device=DEV) + guard = torch.zeros(_VIT_T_MAX + 1, dtype=torch.int64, device=DEV) + differs = torch.empty(_VIT_T_MAX, _VIT_HEADS, _VIT_HEAD_DIM, dtype=torch.bool, device=DEV) + for t in range(1, _VIT_T_MAX + 1): + buf[:t].copy_(x[:t]) + buf[t].copy_(guard_row) + _vit_call(buf[:t], pos[:t], table) + mismatch[t] = torch.ne(buf[:t].view(torch.int16), fb[:t], out=differs[:t]).any() + guard[t] = (buf[t].view(torch.int16) != gb).sum() + bad = mismatch.nonzero().flatten().tolist() + over = guard.nonzero().flatten().tolist() + assert not bad, ( + f"{len(bad)} row counts differ from the {_VIT_T_MAX}-row result's rows, first {bad[:10]}" + ) + assert not over, f"{len(over)} row counts wrote the guard row, first {over[:10]}" + print(f"[vit packed] T = 1..{_VIT_T_MAX}: every prefix bitwise equal, guard row untouched") + + +def test_vit_packed_qkv_controls(): + """The same gates reject: v read at the last patch's row instead of the identity row, + GPT-J pairs instead of NeoX, and a read one row off.""" + num_patches = 4680 + gen = torch.Generator().manual_seed(12200) + table = _vit_table(num_patches, gen) + pos = _vit_positions(num_patches) + x = torch.randn(3 * num_patches, _VIT_HEADS, _VIT_HEAD_DIM, generator=gen).to( + DEV, torch.bfloat16 + ) + good = _run(x, pos, table, 0, inverse=False, is_neox=True) + wrong_pos = pos.clone() + wrong_pos[2::3] = num_patches - 1 + wrong_v = _run(x, wrong_pos, table, 0, inverse=False, is_neox=True) + v_changed = float((wrong_v[2::3] != x[2::3]).float().mean()) + ref = _reference(x, pos, table, 0, inverse=False, is_neox=True) + gptj = _run(x, pos, table, 0, inverse=False, is_neox=False) + outside = float( + ((gptj.float() - ref.float()).abs() > 1e-5 + 1.6e-2 * ref.float().abs()).float().mean() + ) + shifted = _run(x[1:], pos[1:], table, 0, inverse=False, is_neox=True) + n_bad = float((shifted != good[:-1]).float().mean()) + print( + f"\n[vit packed controls] v at row P-1: {v_changed:.4f} of v elements changed; " + f"GPT-J pairs: {outside:.4f} outside the bf16 gate; one row off: {n_bad:.4f} differ" + ) + assert v_changed > 0.5 and outside > 0.05 and n_bad > 0.5 + + +def test_vit_packed_qkv_one_kernel_at_every_row_count(): + """Probe: the packed form launches the same single kernel at every T probed up to 51111.""" + gen = torch.Generator().manual_seed(12300) + table = _vit_table(_VIT_P_MAX, gen) + pos = _vit_positions(_VIT_P_MAX) + data = torch.randn(_VIT_T_MAX, _VIT_HEADS, _VIT_HEAD_DIM, generator=gen).to(DEV, torch.bfloat16) + counts = ( + 1, + 2, + 3, + 4, + 63, + 64, + 65, + 4551, + 8191, + 8192, + 8193, + 14040, + 15975, + 16383, + 16384, + 16385, + 26316, + 32768, + 51109, + 51110, + 51111, + ) + names = set() + for t in counts: + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _vit_call(data[:t], pos[:t], table) + torch.cuda.synchronize() + launched = [ + e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA + ] + assert len(launched) == 1, (t, launched) + names.add(launched[0]) + print(f"\n[vit packed] {len(counts)} row counts, kernels: {sorted(names)}") + assert len(names) == 1 and "mlaRoPEInplaceKernel" in next(iter(names)), names + + +def test_vit_packed_qkv_cuda_graph_at_the_prototype_sizes(): + """Captured at T = 3P for the prototype's graph sizes; the replay reads the data, table and + positions present at replay time and is bitwise equal to an eager call on them.""" + gen = torch.Generator().manual_seed(12400) + pool = torch.cuda.graph_pool_handle() + graphs = {} + for p in sorted(_VIT_GRAPH_P, reverse=True): + table = _vit_table(p, gen) + pos = _vit_positions(p) + data = torch.randn(3 * p, _VIT_HEADS, _VIT_HEAD_DIM, generator=gen).to(DEV, torch.bfloat16) + _vit_call(data, pos, table) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + _vit_call(data, pos, table) + graphs[p] = (graph, data, pos, table) + for p in _VIT_GRAPH_P: + graph, data, pos, table = graphs[p] + new = torch.randn(data.shape, generator=gen).to(DEV, torch.bfloat16) + new_table = _vit_table(p, gen) + data.copy_(new) + table.copy_(new_table) + graph.replay() + torch.cuda.synchronize() + eager = _run(new, pos, new_table, 0, inverse=False, is_neox=True) + assert torch.equal(data, eager), p + assert torch.equal(data[2::3], new[2::3]), p + op.compare(data, _reference(new, pos, new_table, 0, inverse=False, is_neox=True)) + print( + f"\n[vit packed] captured at P = {_VIT_GRAPH_P}: replays bitwise equal to eager on new data and tables" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_trtllm_sparse_mla_dsv4.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_trtllm_sparse_mla_dsv4.py new file mode 100644 index 000000000000..eceb08b031ab --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_trtllm_sparse_mla_dsv4.py @@ -0,0 +1,6190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the trtllm_sparse_mla_dsv4 entry. + +The entry owns what is certified: `CELLS` is the small-pool list (decode, DSpark verify, varlen +prefill, the window-only placeholder and the non-causal draft block, at every per-rank head count +H 8, 16, 32 and 64), and `reference` and `compare` are the gate. This file turns a cell's spec into +tensors and keeps every check a cell cannot carry: the engine-size pools, the engine envelope, +the native-storage compact reads, CUDA-graph capture, the guard-page trials, the workspace high-water mark +and the wrong-reading controls. + +The reference (`op.reference`) gathers the pool rows each index row names, takes fp32 scores scaled +by softmax_scale, fp32 max and sum, casts the probabilities to bf16 for the PV product (the precision +of DeepSeek-V4.1's reference sparse_attn kernel), and adds the per-head sink to the denominator only, +in natural-log units. `_reference` below reads it with explicit per-row window counts; its +deliberately wrong readings (`_wrong_reading`) are the controls. + +Pools are sized as a cache manager would size them -- independent of the batch -- and every request's +rows are reached through a randomly permuted page table: SWA pool 96 pages x 128 rows (12288 slots), +compressed pool 160 pages x 64 rows (10240 slots). Those widths were chosen here, plainly larger than +any batch in this file; they are not taken from a runtime-value dump. + +The `test_real_pools_*` cases use the pool geometry the core DeepSeek-V4.1 target's cache manager +allocates: SWA 188160 x 128 rows, compressed 453120 x 64 (ratio 2) and 151040 x 128 (ratio 1), and +attend rows in the last pages of both pools -- past the 2^31 / 2^32-byte and 2^31-element offsets of +a 1024-byte row. + +The `test_kvq_*` cases hold the native-storage target's compact read (its two persistent bf16 read buffers, +with the index renumbered onto them and -1 window columns backed by zero rows) bitwise to the +container read on the same values. + +The engine-pool, engine-envelope (`test_e_*`), workspace, pool-size and guard-page cases run at every +per-rank head count of the target's members, H 8, 16, 32 and 64 (MEMBER_HEADS / ALL_HEADS). +""" + +import json +import math +import os +import random +import re +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.trtllm_sparse_mla_dsv4 import ( # noqa: E501 + trtllm_sparse_mla_dsv4 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "trtllm_sparse_mla_dsv4 requires a CUDA device" + +# The guard-page and pool-size trials run this file as a script in a child process, where the +# repo's `__extra_import_path__` hook is not installed: the child reaches `_validating` through +# PYTHONPATH instead (set only in the child's environment). +_CHILD_ENV = dict( + os.environ, + PYTHONPATH=os.pathsep.join( + p for p in (str(Path(__file__).resolve().parents[4]), os.environ.get("PYTHONPATH")) if p + ), +) + +HEAD_DIM = 512 +WINDOW = 128 +SCALE = HEAD_DIM**-0.5 +LOG2E = 1.4426950408889634 +SWA_PAGES, SWA_PAGE = 96, 128 +COMP_PAGES, COMP_PAGE = 160, 64 +RATIO = 2 +INDEX_TOPK = 512 +DEV = torch.device("cuda") +# Rows of mapped memory every pool the call may load a -1 index from must span from its first row +# (Precondition 8; measured in the guard-page tests below). +POOL_MIN_ROWS = 128 + +# The engine envelope (per rank). E1: the CUDA-graph capture set -- every G 1..32 plus TRT-LLM's +# padded default list for cuda_graph_config {enable_padding: true, max_batch_size: 256} (llm_args.py, +# _generate_cuda_graph_batch_sizes: 1, 2, 4, 8, 16..128 step 8, then 192, 256); the padding-off list +# at 256 (1..32, 64, 128, 256) is a subset. 46 values. The decode (next_n 1) and verify (next_n 6) +# graphs of the target's generation call and the DSpark drafter's graph (N = 5G rows) are captured at +# every one of them. E3: a step holds up to 256 requests, and an eager step any G 1..256. +GSET = sorted(set(range(1, 33)) | set(range(40, 129, 8)) | {192, 256}) +assert len(GSET) == 46 +E_B = 256 +# Per-rank query heads of the target's trunk call, 64 / attention TP: 16 at tep4 / tp4, 64 at dep4 / +# dep8 (attention DP), 8 at tep8 (attention TP 8 over two NVL72 trays), 32 at pp2 x {tep2, tp2}. The +# DSpark drafter's call runs all 64 heads on every member. The envelope cases below run at all four. +MEMBER_HEADS = (8, 32) +ALL_HEADS = (8, 16, 32, 64) + +# The entry's gate (`op.compare`, derivation in its docstring): rtol 1.6e-2 plus atol +# 2^-8 * max|V|. RTOL, `_atol` and `_excess` restate it to report an output as a multiple of the gate; +# `_atol` takes max|V| over whole pools, a bound on the attended rows' max. +RTOL = 1.6e-2 + + +def _atol(*pools: torch.Tensor) -> float: + return 2.0**-8 * max(p.float().abs().max().item() for p in pools) + + +def _workspace() -> torch.Tensor: + return torch.zeros(128 * 1024 * 1024, dtype=torch.uint8, device=DEV) + + +def _pools(gen: torch.Generator): + swa = torch.randn(SWA_PAGES, 1, SWA_PAGE, HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) + comp = torch.randn(COMP_PAGES, 1, COMP_PAGE, HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) + return swa, comp + + +def _sinks(num_heads: int) -> torch.Tensor: + # Spans "negligible" to "dominant" against sum(exp(s - max)) over up to 640 columns. + return torch.linspace(-2.0, 9.0, num_heads, device=DEV, dtype=torch.float32) + + +def _page_tables(gen: torch.Generator): + swa_bt = torch.randperm(SWA_PAGES, generator=gen, device=DEV) + comp_bt = torch.randperm(COMP_PAGES, generator=gen, device=DEV) + return swa_bt, comp_bt + + +def _build_rows( + requests, + gen: torch.Generator, + topk: int, + swa_pages_per_req: int, + comp_pages_per_req: int, + *, + swa_bt: torch.Tensor | None = None, + comp_bt: torch.Tensor | None = None, + swa_page: int = SWA_PAGE, + comp_page: int = COMP_PAGE, + ratio: int = RATIO, +): + """Index rows the way a DeepSeek-V4.1 layer builds them. + + requests: list of (cached_len, q_len). Query j of a request sits at pos = cached_len + j and sees + window positions max(0, pos-127)..pos (all valid, in position order) and min(512, (pos+1)//2) + compressed positions (a random subset, sorted, as the indexer top-k is), compacted at the front; + sparse_topk_lens = 128 + that count. topk=128 builds window-only rows. + Each request owns swa_pages_per_req SWA pages used as a windowed ring (position p lives in page + (p // 128) % swa_pages_per_req) and comp_pages_per_req compressed pages used the same way, all + drawn from a random permutation of the pool's pages -- or, when swa_bt / comp_bt are given, + taken in order from those page lists (request b gets entries [b * per_req, (b + 1) * per_req)). + Returns sparse_indices [N, topk] int32 (pool token rows, -1 padding), sparse_topk_lens [N], + seq_lens [B] (= cached_len + q_len), cum_seq_lens_q [B+1]. + """ + if swa_bt is None: + swa_bt, comp_bt = _page_tables(gen) + assert len(requests) * swa_pages_per_req <= swa_bt.numel() + assert comp_pages_per_req == 0 or len(requests) * comp_pages_per_req <= comp_bt.numel() + rows_idx, lens = [], [] + for b, (cached, q_len) in enumerate(requests): + swa_pages = swa_bt[b * swa_pages_per_req : (b + 1) * swa_pages_per_req] + if comp_pages_per_req: + comp_pages = comp_bt[b * comp_pages_per_req : (b + 1) * comp_pages_per_req] + for j in range(q_len): + pos = cached + j + row = torch.full((topk,), -1, dtype=torch.int32, device=DEV) + win = torch.arange(max(0, pos - WINDOW + 1), pos + 1, device=DEV) + page = swa_pages[(win // swa_page) % swa_pages_per_req] + row[: win.numel()] = (page * swa_page + win % swa_page).int() + k = 0 + if topk > WINDOW: + visible = (pos + 1) // ratio + k = min(topk - WINDOW, visible) + sel = torch.randperm(visible, generator=gen, device=DEV)[:k].sort().values + page = comp_pages[(sel // comp_page) % comp_pages_per_req] + row[WINDOW : WINDOW + k] = (page * comp_page + sel % comp_page).int() + rows_idx.append(row) + lens.append(WINDOW + k) + idx = torch.stack(rows_idx).contiguous() + active = torch.tensor(lens, dtype=torch.int32, device=DEV) + seq_lens = torch.tensor([c + q for c, q in requests], dtype=torch.int32, device=DEV) + cu = torch.tensor([0] + [q for _, q in requests], dtype=torch.int32, device=DEV).cumsum(0).int() + return idx, active, seq_lens, cu + + +def _window_counts(seq_lens: torch.Tensor, cu: torch.Tensor) -> torch.Tensor: + """Per query row: min(128, seq_lens[b] - q_len_b + j + 1).""" + cu = cu.long() + q_lens = cu[1:] - cu[:-1] + req = torch.repeat_interleave(torch.arange(seq_lens.numel(), device=DEV), q_lens) + j = torch.arange(req.numel(), device=DEV) - cu[:-1][req] + return (seq_lens.long()[req] - q_lens[req] + j + 1).clamp(max=WINDOW) + + +def _reference( + q: torch.Tensor, + swa: torch.Tensor, + comp: torch.Tensor, + idx: torch.Tensor, + active: torch.Tensor, + window_count: torch.Tensor, + sinks: torch.Tensor, + *, + sink: str = "denominator", + negative_index: str = "zero_row", + dedup: bool = False, +) -> torch.Tensor: + """q [N,H,512] bf16 -> [N,H,512] fp32, with an explicit window count per query row. + + The correct reading is the entry's `op.reference`, given each row as its own q_len-1 request + whose virtual `seq_lens` is that row's window count -- the device the contract names for a + per-row count (min(128, seq_lens - 1 + 0 + 1) is the count itself). Setting `sink`, + `negative_index` or `dedup` away from the correct reading selects a control (`_wrong_reading`). + """ + if (sink, negative_index, dedup) != ("denominator", "zero_row", False): + return _wrong_reading( + q, + swa, + comp, + idx, + active, + window_count, + sinks, + sink=sink, + negative_index=negative_index, + dedup=dedup, + ) + q3 = q.reshape(-1, q.shape[-2], HEAD_DIM) + n = q3.shape[0] + expected, _ = op.reference( + query=q3, + swa_kv_cache=swa, + compressed_kv_cache=comp, + workspace_buffer=None, + sparse_indices=idx, + sparse_topk_lens=active, + seq_lens=window_count.to(torch.int32), + sinks=sinks, + softmax_scale=SCALE, + cum_seq_lens_q=torch.arange(n + 1, dtype=torch.int32, device=q.device), + max_q_len=1, + ) + return expected + + +def _wrong_reading( + q: torch.Tensor, + swa: torch.Tensor, + comp: torch.Tensor, + idx: torch.Tensor, + active: torch.Tensor, + window_count: torch.Tensor, + sinks: torch.Tensor, + *, + sink: str = "denominator", + negative_index: str = "zero_row", + dedup: bool = False, +) -> torch.Tensor: + """The controls' deliberately wrong readings; q [N,H,512] bf16 -> [N,H,512] fp32. + + Column c of row r is attended iff c < active[r] and (c >= 128 or c < window_count[r]). + negative_index "masked" drops an attended -1 instead of reading it as an all-zero row; sink + "numerator" / "log2e" / "scaled" / "ignored" misplace the sink; dedup attends a repeated row + once. Only reached through `_reference` with one of those set: the correct reading is + `op.reference`. + """ + n, h, _ = q.shape + t = idx.shape[1] + swa_f = swa.reshape(-1, HEAD_DIM) + comp_f = comp.reshape(-1, HEAD_DIM) + col = torch.arange(t, device=DEV) + out = torch.empty(n, h, HEAD_DIM, dtype=torch.float32, device=DEV) + for s0 in range(0, n, 256): + s1 = min(n, s0 + 256) + ib = idx[s0:s1].long() + kv = torch.cat( + [swa_f[ib[:, :WINDOW].clamp_min(0)], comp_f[ib[:, WINDOW:].clamp_min(0)]], 1 + ).float() + attended = (col[None] < active[s0:s1, None]) & ( + (col[None] >= WINDOW) | (col[None] < window_count[s0:s1, None]) + ) + if negative_index == "zero_row": + kv = kv * (ib >= 0)[..., None] + else: + attended = attended & (ib >= 0) + if dedup: + same_pool = swa.data_ptr() == comp.data_ptr() + key = ib if same_pool else torch.where(col[None] < WINDOW, ib, ib + (1 << 40)) + for r in range(ib.shape[0]): + seen = set() + for c in attended[r].nonzero().flatten().tolist(): + if int(key[r, c]) in seen: + attended[r, c] = False + seen.add(int(key[r, c])) + s = torch.einsum("nhd,ntd->nht", q[s0:s1].float(), kv) * SCALE + s = s.masked_fill(~attended[:, None, :], float("-inf")) + m = s.amax(-1) + m = torch.where(torch.isfinite(m), m, torch.zeros_like(m)) + p = torch.exp(s - m[..., None]) + denom = p.sum(-1) + num = torch.einsum("nht,ntd->nhd", p.to(torch.bfloat16).float(), kv) + sk = sinks[None, :].float() + if sink == "denominator": + denom = denom + torch.exp(sk - m) + elif sink == "numerator": + e = torch.exp(sk - m) + denom = denom + e + num = num + e[..., None] + elif sink == "log2e": + denom = denom + torch.exp(sk * LOG2E - m) + elif sink == "scaled": + denom = denom + torch.exp(sk * SCALE - m) + elif sink != "ignored": + raise ValueError(sink) + o = num / denom[..., None] + out[s0:s1] = torch.where(attended.any(-1)[:, None, None], o, torch.zeros_like(o)) + return out + + +def _run(q, swa, comp, idx, active, seq_lens, sinks, cu=None, max_q_len=None, out=None): + return op( + query=q, + swa_kv_cache=swa, + compressed_kv_cache=comp, + workspace_buffer=_workspace(), + sparse_indices=idx, + sparse_topk_lens=active, + seq_lens=seq_lens, + sinks=sinks, + softmax_scale=SCALE, + cum_seq_lens_q=cu, + max_q_len=max_q_len, + out=out, + ) + + +def _check(o, ref, swa, comp): + """The entry's gate, with max|V| taken over both pools.""" + op.compare(o.reshape(ref.shape), (ref, _atol(swa, comp) * 2.0**8)) + + +# --------------------------------------------------------------------------------------------- +# the entry's cells: small-pool decode, DSpark verify, varlen prefill, window-only, draft block +# --------------------------------------------------------------------------------------------- +# Compressed pages each request's ring holds (`_build_rows` wraps positions onto them), per kind. +_CELL_COMP_PAGES = {"decode": 17, "verify": 30, "prefill": 50, "window_only": 0} + + +def _cell_inputs(spec) -> dict: + """The op's keyword arguments for one cell, on the file's small pools. + + Seeds and draws are per kind, so a cell's data does not depend on which other cells run. + """ + kind, h = spec["kind"], spec["H"] + sinks = _sinks(h) + if kind == "draft_block": + # Draft block of 5 queries that all see the same ring rows plus all 5 block rows + # (non-causal). Each draft query is passed as its own q_len-1 request whose seq_lens is + # the number of valid ring rows, so the window count equals exactly that number; the 5 + # block rows sit in the compressed segment (padded to 8 columns with -1 beyond + # sparse_topk_lens = 128 + 5). + gen = torch.Generator(device=DEV).manual_seed(5) + swa, _ = _pools(gen) + block_pool = torch.randn(4, 1, COMP_PAGE, HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) + blk = spec["block"] + rows, seq = [], [] + for req, n_ring in enumerate(spec["ring_rows"]): + ring = torch.randperm(SWA_PAGES * SWA_PAGE, generator=gen, device=DEV)[:n_ring].int() + block_rows = (req * COMP_PAGE + torch.arange(blk, device=DEV)).int() + for _ in range(blk): + r = torch.full((spec["width"],), -1, dtype=torch.int32, device=DEV) + r[:n_ring] = ring + r[WINDOW : WINDOW + blk] = block_rows + rows.append(r) + seq.append(n_ring) + idx = torch.stack(rows).contiguous() + n = idx.shape[0] + q = torch.randn(n, 1, h, HEAD_DIM, generator=gen, device=DEV).mul_(2).to(torch.bfloat16) + return dict( + query=q, + swa_kv_cache=swa, + compressed_kv_cache=block_pool, + workspace_buffer=_workspace(), + sparse_indices=idx, + sparse_topk_lens=torch.full((n,), WINDOW + blk, dtype=torch.int32, device=DEV), + seq_lens=torch.tensor(seq, dtype=torch.int32, device=DEV), + sinks=sinks, + softmax_scale=SCALE, + ) + + requests, topk = [tuple(r) for r in spec["requests"]], spec["topk"] + seed = { + "decode": 1000 + h + topk, + "verify": 11 + h, + "prefill": sum(q_len for _, q_len in requests) + h + topk, + "window_only": 7, + }[kind] + gen = torch.Generator(device=DEV).manual_seed(seed) + swa, comp = _pools(gen) + if kind == "window_only": + # Ratio-0 layers: a dummy one-page compressed pool full of NaN, never read at width 128. + comp = torch.full( + (1, 1, COMP_PAGE, HEAD_DIM), float("nan"), dtype=torch.bfloat16, device=DEV + ) + idx, active, seq_lens, cu = _build_rows( + requests, gen, topk, swa_pages_per_req=2, comp_pages_per_req=_CELL_COMP_PAGES[kind] + ) + b, n = len(requests), idx.shape[0] + args = dict( + swa_kv_cache=swa, + compressed_kv_cache=comp, + workspace_buffer=_workspace(), + sparse_indices=idx, + sparse_topk_lens=active, + seq_lens=seq_lens, + sinks=sinks, + softmax_scale=SCALE, + ) + if kind == "prefill": + q = torch.randn(n, h, HEAD_DIM, generator=gen, device=DEV).mul_(2).to(torch.bfloat16) + args.update(cum_seq_lens_q=cu, max_q_len=max(q_len for _, q_len in requests)) + else: + q_len = requests[0][1] + q = torch.randn(b, q_len, h, HEAD_DIM, generator=gen, device=DEV).mul_(2) + q = q.to(torch.bfloat16) + args["query"] = q + return args + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + args = _cell_inputs(cell.spec) + o = op(**args) + assert o.shape == args["query"].shape, (o.shape, args["query"].shape) + if cell.spec["kind"] == "window_only": + assert torch.isfinite(o).all(), "the NaN placeholder pool was read" + op.compare(o, op.reference(**args)) + + +# --------------------------------------------------------------------------------------------- +# DSpark drafter draft block at the drafter's own surface: 64 heads, index width 136 +# --------------------------------------------------------------------------------------------- +# DeepSeek-V4.1's DSpark drafter runs one attention call per stage over G requests x 5 draft +# queries, each query passed as its own q_len-1 request (dense [G*5, 1, 64, 512]). swa_kv_cache is +# the speculative worker's ring [ring_slots, 3, 128, 512] bf16 viewed as +# [ring_slots * 3, 1, 128, 512]: request g's stage-s rows are page slot_g * 3 + s, compact from row +# 0. compressed_kv_cache is the block buffer holding the G*5 transient block K/V rows, row +# g * 5 + j. Index row [base + 0..127 | g*5 + 0..4, -1 x3] (width 136), sparse_topk_lens 133, +# seq_lens = the ring's valid-row count repeated 5x (window_count = min(128, seq_lens)). +# ring_slots is the worker's allocation, max_bs + 1 (one scratch row for CUDA-graph dummies): 33 +# at the engine's max batch size 32, and 257 (max_bs 256) chosen here so the ring is plainly +# larger than any batch. Requests take random slots, the scratch row among them. Ring rows past a +# request's valid count hold a finite decoy (30.0, a score that would dominate the softmax), as a +# reused slot holds a previous request's rows. +DS_HEADS, DS_BLOCK, DS_STAGES = 64, 5, 3 +DS_WIDTH = WINDOW + 8 # 136: window 128 | block 5 | -1 x3 (the width must be a multiple of 4) +DS_TOPK = WINDOW + DS_BLOCK # 133 +DS_RING_SLOTS = (33, 257) +DS_BATCHES = tuple(range(1, 33)) +# The worker's valid-row count per request: 1 at the first draft (the bonus row alone), short and +# edge counts, and counts past the window (the count is not clamped by the worker; the kernel +# saturates it at 128), up to the E4 extreme 1,048,595 (Lm + 5 at max_seq_len 1,048,576). +DS_VALID = (1, 2, 3, 5, 6, 64, 126, 127, 128, 129, 4000, 16398, 262163, 1048595) +DS_DECOY = 30.0 + + +def _ds_ring(ring_slots: int, gen: torch.Generator) -> torch.Tensor: + return torch.randn(ring_slots, DS_STAGES, WINDOW, HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) + + +def _ds_inputs(G, ring, gen, shift=0, block_form="page", valid=None): + """Write fresh ring rows for G requests into `ring` and build one stage's call. + + Returns (q [G*5, 1, 64, 512], block buffer, idx [G*5, 136], active [G*5], seq_lens [G*5]). + block_form "page" passes the block buffer as [1, 1, G*5, 512] (one page), "row" as + [G*5, 1, 1, 512] (one row per page). + """ + slots = torch.randperm(ring.shape[0], generator=gen, device=DEV)[:G] + stage = (G + shift) % DS_STAGES + if valid is None: + valid = [DS_VALID[(g + shift) % len(DS_VALID)] for g in range(G)] + valid = torch.tensor(valid, dtype=torch.int32, device=DEV) + ring[slots] = torch.randn(G, DS_STAGES, WINDOW, HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) + for g in range(G): + ring[slots[g], stage, min(WINDOW, int(valid[g])) :] = DS_DECOY + rows = G * DS_BLOCK + # The block rows sit at the head of an allocation of at least POOL_MIN_ROWS rows (Precondition 8). + backing = torch.zeros(max(rows, POOL_MIN_ROWS), HEAD_DIM, dtype=torch.bfloat16, device=DEV) + backing[:rows] = torch.randn(rows, HEAD_DIM, generator=gen, device=DEV).to(torch.bfloat16) + block = backing[:rows] + block = ( + block.view(1, 1, rows, HEAD_DIM) + if block_form == "page" + else block.view(rows, 1, 1, HEAD_DIM) + ) + base = (slots.long() * DS_STAGES + stage) * WINDOW + idx = torch.full((rows, DS_WIDTH), -1, dtype=torch.int32, device=DEV) + idx[:, :WINDOW] = ( + base.repeat_interleave(DS_BLOCK)[:, None] + torch.arange(WINDOW, device=DEV) + ).int() + idx[:, WINDOW:DS_TOPK] = ( + torch.arange(rows, device=DEV).view(G, DS_BLOCK).repeat_interleave(DS_BLOCK, 0).int() + ) + active = torch.full((rows,), DS_TOPK, dtype=torch.int32, device=DEV) + seq_lens = valid.repeat_interleave(DS_BLOCK) + q = ( + torch.randn(rows, 1, DS_HEADS, HEAD_DIM, generator=gen, device=DEV) + .mul_(2) + .to(torch.bfloat16) + ) + return q, block, idx, active, seq_lens + + +def _ds_counts(seq_lens: torch.Tensor) -> torch.Tensor: + return seq_lens.long().clamp(max=WINDOW) + + +def _ds_atol(ring_view, block, idx, counts) -> float: + """The file's gate, 2^-8 * max|V|, over the rows the call attends (the decoys are not attended).""" + col = torch.arange(WINDOW, device=DEV) + w = idx[:, :WINDOW][col[None] < counts[:, None]] + m = ring_view.view(-1, HEAD_DIM)[torch.unique(w).long()].float().abs().max() + m = torch.maximum(m, block.float().abs().max()) + return 2.0**-8 * m.item() + + +def _ds_reference(q, ring_view, block, idx, active, counts, sinks, **kw): + return _reference( + q.view(-1, DS_HEADS, HEAD_DIM), ring_view, block, idx, active, counts, sinks, **kw + ) + + +def test_dspark_draft_block_h64_every_batch(): + """Eager draft block at 64 heads, both block-buffer forms: every G 1..32 on the 33-slot ring + (max batch 32) and every G 1..256 on the 257-slot ring (max batch 256; N = 5G up to 1280).""" + gen = torch.Generator(device=DEV).manual_seed(4100) + sinks = _sinks(DS_HEADS) + for ring_slots in DS_RING_SLOTS: + ring = _ds_ring(ring_slots, gen) + view = ring.view(ring_slots * DS_STAGES, 1, WINDOW, HEAD_DIM) + for G in range(1, ring_slots): + form = "page" if G % 2 else "row" + q, block, idx, active, seq_lens = _ds_inputs( + G, ring, gen, shift=ring_slots + G, block_form=form + ) + o = _run(q, view, block, idx, active, seq_lens, sinks) + assert o.shape == (G * DS_BLOCK, 1, DS_HEADS, HEAD_DIM) + counts = _ds_counts(seq_lens) + ref = _ds_reference(q, view, block, idx, active, counts, sinks) + _check_atol(o, ref, _ds_atol(view, block, idx, counts)) + + +@pytest.mark.parametrize("G", GSET) +def test_dspark_draft_block_h64_cuda_graph(G): + """The draft block captured at every drafter batch G of the capture set GSET (46 values, 1..256; + N = 5G = 5..1280 query rows), on the 33-slot ring up to G 32 and the 257-slot ring above. + + As the engine does it: one eager call at this G (its warm-up), then capture with a persistent + workspace and `out`. The first replay equals the eager call bitwise; two more replays on new q, + new ring rows and slots (new index rows), new valid counts and a new block buffer, all written + into the captured buffers, equal a fresh eager call on the same inputs bitwise and match the + reference. + """ + gen = torch.Generator(device=DEV).manual_seed(4200 + G) + sinks = _sinks(DS_HEADS) + ring = _ds_ring(DS_RING_SLOTS[0] if G < DS_RING_SLOTS[0] else DS_RING_SLOTS[1], gen) + view = ring.view(-1, 1, WINDOW, HEAD_DIM) + form = "page" if G % 2 else "row" + q, block, idx, active, seq_lens = _ds_inputs(G, ring, gen, block_form=form) + ws = _workspace() + out = torch.empty_like(q) + + def call(): + return op( + query=q, + swa_kv_cache=view, + compressed_kv_cache=block, + workspace_buffer=ws, + sparse_indices=idx, + sparse_topk_lens=active, + seq_lens=seq_lens, + sinks=sinks, + softmax_scale=SCALE, + out=out, + ) + + call() + eager = out.clone() + counts = _ds_counts(seq_lens) + _check_atol( + eager, + _ds_reference(q, view, block, idx, active, counts, sinks), + _ds_atol(view, block, idx, counts), + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, eager) + for shift in (1, 2): + nq, nblock, nidx, nactive, nseq = _ds_inputs(G, ring, gen, shift=shift, block_form=form) + q.copy_(nq) + block.copy_(nblock) + idx.copy_(nidx) + active.copy_(nactive) + seq_lens.copy_(nseq) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + replayed = out.clone() + call() # a fresh eager call on the same inputs + torch.cuda.synchronize() + assert torch.equal(replayed, out), f"replay != eager at shift {shift}" + ncounts = _ds_counts(nseq) + ref = _ds_reference(nq, view, nblock, nidx, nactive, ncounts, sinks) + _check_atol(out, ref, _ds_atol(view, nblock, nidx, ncounts)) + del graph + + +def test_dspark_draft_block_h64_wrong_variants_fail_the_gate(): + """At the drafter surface (G 32, 64 heads) the gate sees every plausible wrong reading. + + Wrong: a causal block (query j attends block rows 0..j), the whole ring page attended (seq_lens + ignored: the stale decoy rows join), the ring count one short, the block rows dropped + (sparse_topk_lens 128), the sink ignored. + """ + gen = torch.Generator(device=DEV).manual_seed(4300) + sinks = _sinks(DS_HEADS) + ring = _ds_ring(DS_RING_SLOTS[0], gen) + view = ring.view(-1, 1, WINDOW, HEAD_DIM) + G = 32 + q, block, idx, active, seq_lens = _ds_inputs(G, ring, gen) + o = _run(q, view, block, idx, active, seq_lens, sinks) + counts = _ds_counts(seq_lens) + atol = _ds_atol(view, block, idx, counts) + + def excess(ref): + return ((o.float().reshape(ref.shape) - ref).abs() / (atol + RTOL * ref.abs())).max().item() + + correct = excess(_ds_reference(q, view, block, idx, active, counts, sinks)) + causal = (WINDOW + 1 + torch.arange(DS_BLOCK, device=DEV).repeat(G)).int() + wrong = { + "causal block": excess(_ds_reference(q, view, block, idx, causal, counts, sinks)), + "whole ring page": excess( + _ds_reference(q, view, block, idx, active, torch.full_like(counts, WINDOW), sinks) + ), + "ring count one short": excess( + _ds_reference(q, view, block, idx, active, counts - 1, sinks) + ), + "block rows dropped": excess( + _ds_reference(q, view, block, idx, torch.full_like(active, WINDOW), counts, sinks) + ), + "sink ignored": excess( + _ds_reference(q, view, block, idx, active, counts, sinks, sink="ignored") + ), + } + print( + f"draft block h64: gate-normalized error correct {correct:.3f}, " + f"wrong { {k: round(v, 1) for k, v in wrong.items()} }" + ) + assert correct <= 1.0, correct + for name, e in wrong.items(): + assert e > 10.0 * max(correct, 0.1), ( + f"wrong variant {name!r} not discriminated: {e:.2f} vs {correct:.2f}" + ) + + +def test_dspark_draft_block_zero_seq_lens_reads_a_full_window(): + """seq_lens 0 is not an empty window: every one of the 128 window columns is attended. + + Measured on sm_103 (probe, then pinned here): the output matches the window_count-128 reading + and fails the window_count-0 reading. A caller must not pass 0 to mean "no ring rows"; rows + it wants unattended are excluded only by a count >= 1. + """ + gen = torch.Generator(device=DEV).manual_seed(4400) + sinks = _sinks(DS_HEADS) + ring = _ds_ring(DS_RING_SLOTS[0], gen) + view = ring.view(-1, 1, WINDOW, HEAD_DIM) + G = 4 + q, block, idx, active, _ = _ds_inputs(G, ring, gen, valid=[WINDOW] * G) + zero = torch.zeros(G * DS_BLOCK, dtype=torch.int32, device=DEV) + o = _run(q, view, block, idx, active, zero, sinks) + full = torch.full((G * DS_BLOCK,), WINDOW, dtype=torch.int64, device=DEV) + atol = _ds_atol(view, block, idx, full) + _check_atol(o, _ds_reference(q, view, block, idx, active, full, sinks), atol) + with pytest.raises(AssertionError): + _check_atol( + o, _ds_reference(q, view, block, idx, active, torch.zeros_like(full), sinks), atol + ) + + +# --------------------------------------------------------------------------------------------- +# pool extent: a loaded -1 index reads one row in 96..127 of its pool (guard-page tests) +# --------------------------------------------------------------------------------------------- +# The kernel loads index columns in groups of 4. A loaded column holding -1 -- the masked tail of a +# compressed segment whose sparse_topk_lens is not a multiple of 4 (the DSpark draft block: lens +# 133, columns 133..135 are -1), or a -1 anywhere in the 128 window columns (all of them are +# loaded) -- makes the call read one row in 96..127 of that pool, and discard it. On a pool with +# fewer than 128 rows that read lands past the pool's end: harmless while the next bytes happen to +# be mapped, an illegal memory access when the allocation ends a mapped range. +# +# The guard-page tests place a pool with the CUDA VMM API so that exactly M rows are mapped from +# its first row and the VA after them is reserved but unmapped: any access at row >= M faults, on +# every run. A fault poisons the CUDA context, so every guard trial runs in a child process (this +# file run as a script) and the parent asserts on the child's output and exit status. + + +def _vmm_place(nrows: int, mapped_rows: int): + """A [nrows, 512] bf16 tensor whose first row starts `mapped_rows` rows before the end of a + mapping, followed by >= 64 MiB of reserved, unmapped VA. Returns (tensor, release).""" + from cuda.bindings import driver as cu + + def ck(res): + err = res[0] if isinstance(res, tuple) else res + assert int(getattr(err, "value", err)) == 0, f"CUDA driver error {err}" + return res[1] if isinstance(res, tuple) and len(res) >= 2 else None + + need = mapped_rows * HEAD_DIM * 2 + prop = cu.CUmemAllocationProp() + prop.type = cu.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.location.type = cu.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + prop.location.id = torch.cuda.current_device() + gran = ck( + cu.cuMemGetAllocationGranularity( + prop, cu.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM + ) + ) + size = (need + gran - 1) // gran * gran + rsize = size + 34 * gran + va = ck(cu.cuMemAddressReserve(rsize, 0, 0, 0)) + handle = ck(cu.cuMemCreate(size, prop, 0)) + ck(cu.cuMemMap(int(va), size, 0, handle, 0)) + acc = cu.CUmemAccessDesc() + acc.location = prop.location + acc.flags = cu.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + ck(cu.cuMemSetAccess(int(va), size, [acc], 1)) + base = int(va) + size - need + + class _Cai: + def __init__(self): + self.__cuda_array_interface__ = { + "shape": (need,), + "typestr": "|u1", + "data": (base, False), + "version": 3, + } + + raw = torch.as_tensor(_Cai(), device=DEV) + raw.zero_() + t = raw.view(torch.bfloat16).view(mapped_rows, HEAD_DIM)[:nrows] + + def release(): + torch.cuda.synchronize() + ck(cu.cuMemUnmap(int(va), size)) + ck(cu.cuMemRelease(handle)) + ck(cu.cuMemAddressFree(va, rsize)) + + return t, release + + +def _guard_case(case): + """One trial's inputs: (inputs, which pool is under test, that pool's rows [n, 512]).""" + gen = torch.Generator(device=DEV).manual_seed(case["seed"]) + sinks = _sinks(DS_HEADS) + ring = _ds_ring(DS_RING_SLOTS[0], gen) + swa = ring.view(-1, 1, WINDOW, HEAD_DIM) + kind = case["kind"] + if kind == "draft": + # The DSpark draft-block call; the block buffer is the pool under test. + q, block, idx, active, seq_lens = _ds_inputs(case["G"], ring, gen, block_form=case["form"]) + if case.get("tail") == "repeat": # masked columns 133..135 name column 132's row, not -1 + idx[:, DS_TOPK:] = idx[:, DS_TOPK - 1 : DS_TOPK] + inp = dict( + q=q, swa=swa, idx=idx, active=active, seq_lens=seq_lens, sinks=sinks, form=case["form"] + ) + return inp, "comp", block.reshape(-1, HEAD_DIM) + if kind == "engine": + # An engine-size pool (runtime-value dump figures) whose last row ends the mapping; the call attends its + # last rows with a -1 loaded from the same pool. Only the attended rows are nonzero. + rows_n, page = case["rows"], case["page"] + pool = torch.zeros(rows_n, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + b, h, lens = 4, case["heads"], case["lens"] + q = torch.randn(b, 1, h, HEAD_DIM, generator=gen, device=DEV).mul_(2).to(torch.bfloat16) + idx = torch.full((b, 640), -1, dtype=torch.int32, device=DEV) + small = torch.randn(5 * WINDOW, HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) # >= 512 rows + if case["pool"] == "comp": + idx[:, :WINDOW] = torch.arange(WINDOW, device=DEV, dtype=torch.int32) + idx[:, WINDOW:lens] = (rows_n - 1 - torch.arange(lens - WINDOW, device=DEV)).int() + seq_lens = torch.full((b,), WINDOW, dtype=torch.int32, device=DEV) + inp = dict(swa=small.view(-1, 1, WINDOW, HEAD_DIM), form=page) + else: + k = case["window"] # valid window rows in the pool's last page, then -1 + idx[:, :k] = (rows_n - WINDOW + torch.arange(k, device=DEV)).int() + idx[:, WINDOW:lens] = torch.arange(lens - WINDOW, device=DEV, dtype=torch.int32) + seq_lens = torch.full((b,), k, dtype=torch.int32, device=DEV) + inp = dict(comp=small.view(-1, 1, 1, HEAD_DIM), swa_page=page) + used = idx[:, :WINDOW] if case["pool"] == "swa" else idx[:, WINDOW:] + used = torch.unique(used[used >= 0]).long() + pool[used] = torch.randn(used.numel(), HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) + active = torch.full((b,), lens, dtype=torch.int32, device=DEV) + inp.update(q=q, idx=idx, active=active, seq_lens=seq_lens, sinks=_sinks(h)) + return inp, case["pool"], pool + # The worst / dummy / swa trials run at 64 heads unless the trial names another head count. + hq = case.get("heads", DS_HEADS) + sinks = _sinks(hq) + q = torch.randn(5, 1, hq, HEAD_DIM, generator=gen, device=DEV).mul_(2).to(torch.bfloat16) + idx = torch.full((5, DS_WIDTH), -1, dtype=torch.int32, device=DEV) + if kind == "worst": + # Every query's partial compressed group is [v, .., -1]: v mod 64 in 60..63 reads row 127. + v, lens = case["v"], case["lens"] + comp = torch.randn(v + 1, HEAD_DIM, generator=gen, device=DEV).to(torch.bfloat16) + idx[:, :WINDOW] = torch.arange(WINDOW, device=DEV, dtype=torch.int32) + idx[:, WINDOW:lens] = v + active = torch.full((5,), lens, dtype=torch.int32, device=DEV) + seq_lens = torch.full((5,), WINDOW, dtype=torch.int32, device=DEV) + inp = dict(q=q, swa=swa, idx=idx, active=active, seq_lens=seq_lens, sinks=sinks, form="row") + return inp, "comp", comp + if kind == "dummy": + # The core's window-only call: index width 128, lens 128, a [1, 1, rows, 512] placeholder as + # the compressed pool, window rows [0..k-1, -1, ..] with count k (a prefill window's tail). + k = case["window"] + dummy = torch.full((case["rows"], HEAD_DIM), float("nan"), dtype=torch.bfloat16, device=DEV) + idx = torch.full((5, WINDOW), -1, dtype=torch.int32, device=DEV) + idx[:, :k] = torch.arange(k, device=DEV, dtype=torch.int32) + active = torch.full((5,), WINDOW, dtype=torch.int32, device=DEV) + seq_lens = torch.full((5,), k, dtype=torch.int32, device=DEV) + inp = dict( + q=q, swa=swa, idx=idx, active=active, seq_lens=seq_lens, sinks=sinks, form="page" + ) + return inp, "comp", dummy + assert kind == "swa" + # A 5-row SWA pool, window row [0..4, -1 x 123] with count 5; the compressed pool is compliant. + rows = torch.randn(5, HEAD_DIM, generator=gen, device=DEV).to(torch.bfloat16) + comp = torch.randn(POOL_MIN_ROWS, HEAD_DIM, generator=gen, device=DEV).to(torch.bfloat16) + idx[:, :5] = torch.arange(5, device=DEV, dtype=torch.int32) + idx[:, WINDOW:DS_TOPK] = torch.arange(5, device=DEV, dtype=torch.int32) + active = torch.full((5,), DS_TOPK, dtype=torch.int32, device=DEV) + seq_lens = torch.full((5,), 5, dtype=torch.int32, device=DEV) + inp = dict( + q=q, + comp=comp.view(-1, 1, 1, HEAD_DIM), + idx=idx, + active=active, + seq_lens=seq_lens, + sinks=sinks, + ) + return inp, "swa", rows + + +def _guard_call(inp, which, rows): + n = rows.shape[0] + if which == "comp": + page = {"row": 1, "page": n}.get(inp["form"], inp["form"]) + swa, comp = inp["swa"], rows.view(n // page, 1, page, HEAD_DIM) + else: + page = inp.get("swa_page", n) + swa, comp = rows.view(n // page, 1, page, HEAD_DIM), inp["comp"] + return _run(inp["q"], swa, comp, inp["idx"], inp["active"], inp["seq_lens"], inp["sinks"]) + + +def _guard_child(trials): + """Run guard trials in order in this process; print TRY before and OK after each call.""" + for case in trials: + inp, which, rows = _guard_case(case) + n = rows.shape[0] + if n < POOL_MIN_ROWS: + safe = torch.zeros(POOL_MIN_ROWS, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + safe[:n] = rows + rows = safe[:n] + want = _guard_call(inp, which, rows).clone() + torch.cuda.synchronize() + placed, release = _vmm_place(n, case["mapped"]) + placed.copy_(rows) + print(f"TRY {case['name']}", flush=True) + got = _guard_call(inp, which, placed) + torch.cuda.synchronize() + same = torch.equal(got.view(torch.int16), want.view(torch.int16)) + if case["kind"] == "dummy": # the NaN placeholder is never read at all + assert torch.isfinite(got).all(), case["name"] + release() + print(f"OK {case['name']} bitwise={int(same)}", flush=True) + + +def _guard_run(trials): + p = subprocess.run( + [sys.executable, os.path.abspath(__file__), "--guard-child", json.dumps(trials)], + env=_CHILD_ENV, + capture_output=True, + text=True, + timeout=900, + ) + ok = { + m.group(1): m.group(2) == "1" + for m in re.finditer(r"^OK (\S+) bitwise=(\d)$", p.stdout, re.M) + } + tried = re.findall(r"^TRY (\S+)$", p.stdout, re.M) + return p, ok, tried + + +def _draft_trial(G, form, mapped, seed, tail="-1"): + return dict( + name=f"draft-G{G}-{form}-tail{tail}-M{mapped}", + kind="draft", + G=G, + form=form, + tail=tail, + mapped=mapped, + seed=seed, + ) + + +def _worst_trial(v, lens, mapped, heads=None): + if heads is None: + return dict( + name=f"worst-v{v}-lens{lens}-M{mapped}", + kind="worst", + v=v, + lens=lens, + mapped=mapped, + seed=5100 + v + lens, + ) + return dict( + name=f"worst-v{v}-lens{lens}-M{mapped}-H{heads}", + kind="worst", + v=v, + lens=lens, + mapped=mapped, + seed=5100 + v + lens, + heads=heads, + ) + + +def _guard_compliant(): + """Trials that must all complete: every pool spans >= 128 rows from its first row (the block + buffer padded to exactly 128 where 5G < 128, or 5G >= 128 rows of its own), or no -1 is loaded + from it.""" + trials = [ + _draft_trial(G, form, POOL_MIN_ROWS, 5000 + G) + for G in range(1, 26) + for form in ("row", "page") + ] + trials += [ + _draft_trial(G, form, 5 * G, 5200 + G) for G in range(26, 33) for form in ("row", "page") + ] + trials += [_worst_trial(v, lens, POOL_MIN_ROWS) for v in (60, 62, 63) for lens in (129, 131)] + trials += [_draft_trial(G, "row", 5 * G, 5300 + G, tail="repeat") for G in (1, 8, 25)] + trials += [dict(name="swa-5rows-M128", kind="swa", mapped=POOL_MIN_ROWS, seed=5400)] + trials += [ + dict( + name=f"window-only-dummy{rows}-window{k}", + kind="dummy", + rows=rows, + window=k, + mapped=rows, + seed=5450 + k, + ) + for rows in (64, 1) + for k in (1, 5, 128) + ] + # The members' head counts (other cubins): the worst index values at 128 mapped rows, and the + # window-only call's NaN placeholder ending a mapping. + trials += [ + _worst_trial(v, lens, POOL_MIN_ROWS, heads=h) + for h in MEMBER_HEADS + for v in (60, 62, 63) + for lens in (129, 131) + ] + trials += [ + dict( + name=f"window-only-dummy{rows}-window{k}-H{h}", + kind="dummy", + rows=rows, + window=k, + mapped=rows, + seed=5460 + k, + heads=h, + ) + for h in MEMBER_HEADS + for rows in (64, 1) + for k in (1, 5, 128) + ] + engine = [ + ("r2-comp", "comp", REAL_COMP[2], 129, 16, 0), + ("r2-comp", "comp", REAL_COMP[2], 639, 64, 0), + ("r1-comp", "comp", REAL_COMP[1], 131, 16, 0), + ("r1-comp", "comp", REAL_COMP[1], 133, 64, 0), + ("swa-window5", "swa", REAL_SWA, 133, 16, 5), + ("swa-window128", "swa", REAL_SWA, 640, 64, 128), + ] + engine += [ + (tag, pool, geom, lens, h, window) + for h in MEMBER_HEADS + for tag, pool, geom, lens, _, window in engine + ] + for tag, pool, (pages, page), lens, h, window in engine: + trials.append( + dict( + name=f"engine-{tag}-lens{lens}-H{h}", + kind="engine", + pool=pool, + rows=pages * page, + page=page, + lens=lens, + heads=h, + window=window, + mapped=pages * page, + seed=5800 + lens + h, + ) + ) + return trials + + +# Each control runs one compliant call, then the same call with the pool ending a mapping short +# of 128 rows; the second must fault. +GUARD_FAULTS = { + "draft_block_G1_row_unpadded": [ + _draft_trial(1, "row", POOL_MIN_ROWS, 5500), + _draft_trial(1, "row", 5, 5500), + ], + "draft_block_G8_page_unpadded": [ + _draft_trial(8, "page", POOL_MIN_ROWS, 5508), + _draft_trial(8, "page", 40, 5508), + ], + "127_rows_are_not_enough": [_worst_trial(62, 129, POOL_MIN_ROWS), _worst_trial(62, 129, 127)], + "127_rows_are_not_enough_H8": [ + _worst_trial(62, 129, POOL_MIN_ROWS, heads=8), + _worst_trial(62, 129, 127, heads=8), + ], + "127_rows_are_not_enough_H32": [ + _worst_trial(62, 129, POOL_MIN_ROWS, heads=32), + _worst_trial(62, 129, 127, heads=32), + ], + "swa_pool_minus_one_window": [ + dict(name="swa-5rows-M128", kind="swa", mapped=POOL_MIN_ROWS, seed=5400), + dict(name="swa-5rows-M5", kind="swa", mapped=5, seed=5400), + ], +} + + +def test_pool_extent_128_rows_never_fault(): + """With >= 128 mapped rows from each pool's first row nothing faults, and the output is bitwise + the output on a normal allocation: the DSpark block buffer at every G 1..25 padded to exactly + 128 rows and at every G 26..32 unpadded (130..160 rows of its own), both forms; the worst index + values (read row 127) at 3 and 1 masked columns; the -1 tail replaced by column 132's row, + unpadded (no -1 is loaded, so nothing is read past the buffer); a 5-row SWA pool with a -1 + window tail, padded; the window-only call's 64-row (and a 1-row) NaN placeholder ending a + mapping, with -1 window tails (never read); and the three engine-size pools (runtime-value + dump figures), each ending a mapping at its last row while the call attends that row and loads + a -1 from the same pool: nothing is read past a pool's end.""" + trials = _guard_compliant() + p, ok, tried = _guard_run(trials) + names = [c["name"] for c in trials] + assert p.returncode == 0, f"child failed after {tried[-1:]}: {p.stderr[-2000:]}" + assert list(ok) == names, (list(ok), names) + assert all(ok.values()), [k for k, v in ok.items() if not v] + print( + f"guard page: {len(names)} compliant trials completed, each bitwise = the normal allocation" + ) + + +@pytest.mark.parametrize("control", list(GUARD_FAULTS)) +def test_pool_extent_guard_page_controls_fault(control): + """The same calls with the pool ending a mapping short of 128 rows fault: the unpadded DSpark + block buffer (G 1 [5,1,1,512], G 8 [1,1,40,512]), the worst index value with 127 mapped rows + (row 127 is read, so 127 rows are not enough), and a 5-row SWA pool whose window row carries -1. + The harness sees the defect: the compliant call before it completed in the same process.""" + first, second = GUARD_FAULTS[control] + p, ok, tried = _guard_run([first, second]) + out = p.stdout + p.stderr + assert ok.get(first["name"]) is True, out[-2000:] + assert tried == [first["name"], second["name"]] and second["name"] not in ok, out[-2000:] + assert p.returncode != 0 + assert "illegal memory access" in out, out[-2000:] + print(f"{control}: compliant call ok, short pool faulted (rc {p.returncode})") + + +def test_overread_rows_never_change_the_output(): + """The row a loaded -1 reads is discarded: with the pool's rows at the head of a 256-row + allocation, filling every row past them with 0, NaN, Inf or 30.0 (a score that would dominate + the softmax) leaves the output bitwise unchanged, and it matches the reference.""" + cases = [ + _draft_trial(G, form, 0, 5600 + G) for G, form in ((1, "row"), (8, "page"), (25, "row")) + ] + cases += [_worst_trial(62, lens, 0) for lens in (129, 130, 131)] + for case in cases: + inp, which, rows = _guard_case(case) + n = rows.shape[0] + backing = torch.zeros(2 * POOL_MIN_ROWS, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + backing[:n] = rows + outs = [] + for fill in (0.0, float("nan"), float("inf"), 30.0): + backing[n:] = fill + outs.append(_guard_call(inp, which, backing[:n])) + for fill, o in zip(("nan", "inf", "30.0"), outs[1:]): + assert torch.equal(o.view(torch.int16), outs[0].view(torch.int16)), ( + f"{case['name']}: {fill} fill changed the output" + ) + counts = _ds_counts(inp["seq_lens"]) + ref = _ds_reference( + inp["q"], inp["swa"], rows, inp["idx"], inp["active"], counts, inp["sinks"] + ) + _check_atol(outs[0], ref, _ds_atol(inp["swa"], rows, inp["idx"], counts)) + + +def test_repeated_tail_equals_minus_one_tail(): + """Filling the draft block's masked columns 133..135 with column 132's row instead of -1 (which + loads no -1, so reads nothing past the buffer) gives the bitwise-identical output.""" + for G in (1, 8, 25): + a, which, rows = _guard_case(_draft_trial(G, "row", 0, 5700 + G)) + b, _, rows_b = _guard_case(_draft_trial(G, "row", 0, 5700 + G, tail="repeat")) + assert torch.equal(rows, rows_b) and not torch.equal(a["idx"], b["idx"]) + backing = torch.zeros(POOL_MIN_ROWS, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + backing[: rows.shape[0]] = rows + oa = _guard_call(a, which, backing[: rows.shape[0]]) + ob = _guard_call(b, which, backing[: rows.shape[0]]) + assert torch.equal(oa.view(torch.int16), ob.view(torch.int16)), G + + +# --------------------------------------------------------------------------------------------- +# index semantics the contract states +# --------------------------------------------------------------------------------------------- +def test_lengths_mask_the_tails(): + """Window columns >= min(128, seq_lens - q_len + j + 1) and all columns >= sparse_topk_lens are + masked: here they name decoy rows (every element 30.0, a score that would dominate the softmax) + and the output still matches the count-bounded reference.""" + gen = torch.Generator(device=DEV).manual_seed(21) + swa, comp = _pools(gen) + decoy_swa = 5 * SWA_PAGE + 7 + decoy_comp = 9 * COMP_PAGE + 3 + swa.view(-1, HEAD_DIM)[decoy_swa] = 30.0 + comp.view(-1, HEAD_DIM)[decoy_comp] = 30.0 + h, b = 16, 6 + seq = torch.tensor([1, 10, 100, 127, 128, 4000], dtype=torch.int32, device=DEV) + idx = torch.empty(b, 640, dtype=torch.int32, device=DEV) + idx[:, :WINDOW] = decoy_swa + idx[:, WINDOW:] = decoy_comp + active = torch.tensor([128, 200, 129, 640, 300, 131], dtype=torch.int32, device=DEV) + counts = seq.long().clamp(max=WINDOW) + good = torch.randperm(SWA_PAGES * SWA_PAGE, generator=gen, device=DEV) + good = good[good != decoy_swa] + goodc = torch.randperm(COMP_PAGES * COMP_PAGE, generator=gen, device=DEV) + goodc = goodc[goodc != decoy_comp] + for r in range(b): + idx[r, : counts[r]] = good[r * WINDOW : r * WINDOW + int(counts[r])].int() + nc = int(active[r]) - WINDOW + idx[r, WINDOW : WINDOW + nc] = goodc[r * 512 : r * 512 + nc].int() + q = torch.randn(b, 1, h, HEAD_DIM, generator=gen, device=DEV).mul_(2).to(torch.bfloat16) + sinks = _sinks(h) + o = _run(q, swa, comp, idx, active, seq, sinks) + ref = _reference(q.view(b, h, HEAD_DIM), swa, comp, idx, active, counts, sinks) + swa_real = swa.clone() + swa_real.view(-1, HEAD_DIM)[decoy_swa] = 0 + comp_real = comp.clone() + comp_real.view(-1, HEAD_DIM)[decoy_comp] = 0 + _check(o, ref, swa_real, comp_real) + + +def test_masked_window_columns_are_still_loaded(): + """A masked window column is loaded all the same: a NaN row it names poisons the output + (0 * NaN in the PV product). Masked columns must hold -1 or name finite rows.""" + gen = torch.Generator(device=DEV).manual_seed(22) + swa, comp = _pools(gen) + nan_row = 3 * SWA_PAGE + 1 + swa.view(-1, HEAD_DIM)[nan_row] = float("nan") + h = 16 + idx = torch.full((1, 640), -1, dtype=torch.int32, device=DEV) + idx[0, :10] = torch.arange(10, device=DEV).int() + idx[0, 100] = nan_row # beyond window_count = 10 + idx[0, WINDOW : WINDOW + 200] = torch.arange(200, device=DEV).int() + q = torch.randn(1, 1, h, HEAD_DIM, generator=gen, device=DEV).to(torch.bfloat16) + seq = torch.tensor([10], dtype=torch.int32, device=DEV) + active = torch.tensor([WINDOW + 200], dtype=torch.int32, device=DEV) + o = _run(q, swa, comp, idx, active, seq, _sinks(h)) + assert not torch.isfinite(o).all() + idx[0, 100] = -1 + o = _run(q, swa, comp, idx, active, seq, _sinks(h)) + assert torch.isfinite(o).all() + + +def test_minus_one_inside_the_active_range_reads_as_a_zero_row(): + """An attended -1 is NOT a mask: it behaves as an all-zero K/V row (score 0, so exp(0 - max) + joins the denominator). Callers must compact valid entries and bound them with the lengths.""" + gen = torch.Generator(device=DEV).manual_seed(31) + swa, comp = _pools(gen) + h, b = 16, 4 + seq = torch.full((b,), 4000, dtype=torch.int32, device=DEV) + idx = torch.empty(b, 640, dtype=torch.int32, device=DEV) + for r in range(b): + idx[r, :WINDOW] = torch.randperm(SWA_PAGES * SWA_PAGE, generator=gen, device=DEV)[ + :WINDOW + ].int() + idx[r, WINDOW:] = torch.randperm(COMP_PAGES * COMP_PAGE, generator=gen, device=DEV)[ + :512 + ].int() + idx[0, 0:128:2] = -1 # holes inside the window + idx[1, 300:640] = -1 # compressed tail left inside sparse_topk_lens + idx[2, 5] = -1 + idx[3, :] = -1 # nothing but holes -> exact zeros + active = torch.full((b,), 640, dtype=torch.int32, device=DEV) + counts = torch.full((b,), WINDOW, dtype=torch.int64, device=DEV) + q = torch.randn(b, 1, h, HEAD_DIM, generator=gen, device=DEV).mul_(2).to(torch.bfloat16) + sinks = _sinks(h) + o = _run(q, swa, comp, idx, active, seq, sinks) + qf = q.view(b, h, HEAD_DIM) + _check( + o, + _reference(qf, swa, comp, idx, active, counts, sinks, negative_index="zero_row"), + swa, + comp, + ) + assert torch.equal(o[3], torch.zeros_like(o[3])) + masked = _reference(qf, swa, comp, idx, active, counts, sinks, negative_index="masked") + with pytest.raises(AssertionError): + _check(o[:2], masked[:2], swa, comp) + + +def test_overlap_is_not_deduplicated(): + """A row named twice (the window and the compressed segment over one aliased pool, or twice in a + segment) is attended twice; the kernel does not de-duplicate.""" + gen = torch.Generator(device=DEV).manual_seed(41) + pool, _ = _pools(gen) + h, b = 16, 3 + idx = torch.empty(b, 256, dtype=torch.int32, device=DEV) + for r in range(b): + w = torch.randperm(SWA_PAGES * SWA_PAGE, generator=gen, device=DEV)[:WINDOW].int() + c = torch.randperm(SWA_PAGES * SWA_PAGE, generator=gen, device=DEV)[:WINDOW].int() + c[:64] = w[:64] + c[64:96] = c[96:128] + idx[r] = torch.cat([w, c]) + active = torch.full((b,), 256, dtype=torch.int32, device=DEV) + seq = torch.full((b,), 4000, dtype=torch.int32, device=DEV) + counts = torch.full((b,), WINDOW, dtype=torch.int64, device=DEV) + q = torch.randn(b, 1, h, HEAD_DIM, generator=gen, device=DEV).mul_(2).to(torch.bfloat16) + sinks = _sinks(h) + o = _run(q, pool, pool, idx, active, seq, sinks) + qf = q.view(b, h, HEAD_DIM) + _check(o, _reference(qf, pool, pool, idx, active, counts, sinks), pool, pool) + with pytest.raises(AssertionError): + _check(o, _reference(qf, pool, pool, idx, active, counts, sinks, dedup=True), pool, pool) + + +def test_wrong_variants_fail_the_gate(): + """Discrimination of the tolerance: the correct reference passes with margin, and every + deliberately wrong reading of the sink or of the index set fails it by a wide margin.""" + gen = torch.Generator(device=DEV).manual_seed(51) + swa, comp = _pools(gen) + atol = _atol(swa, comp) + + def excess( + o, ref + ): # max over elements of |o - ref| / (atol + rtol * |ref|); <= 1 passes the gate + return ((o.float().reshape(ref.shape) - ref).abs() / (atol + RTOL * ref.abs())).max().item() + + requests = [(p, 1) for p in (130, 700, 3000, 9000)] + [(0, 200)] + idx, active, seq_lens, cu = _build_rows( + requests, gen, 640, swa_pages_per_req=2, comp_pages_per_req=30 + ) + n, h = idx.shape[0], 16 + q = torch.randn(n, h, HEAD_DIM, generator=gen, device=DEV).mul_(2).to(torch.bfloat16) + sinks = _sinks(h) + o = _run(q, swa, comp, idx, active, seq_lens, sinks, cu=cu, max_q_len=200) + counts = _window_counts(seq_lens, cu) + correct = [excess(o, _reference(q, swa, comp, idx, active, counts, sinks))] + wrong = { + name: excess(o, _reference(q, swa, comp, idx, active, counts, sinks, sink=name)) + for name in ("ignored", "numerator", "log2e", "scaled") + } + # de-duplicated overlap: 96 window rows repeated in the compressed segment of one aliased pool + dreq = [(p, 1) for p in (130, 700, 3000, 9000)] + didx, dactive, dseq, dcu = _build_rows( + dreq, gen, 640, swa_pages_per_req=2, comp_pages_per_req=30 + ) + didx[:, WINDOW : WINDOW + 96] = didx[:, :96] + dq = ( + torch.randn(len(dreq), 1, h, HEAD_DIM, generator=gen, device=DEV).mul_(2).to(torch.bfloat16) + ) + do = _run(dq, swa, swa, didx, dactive, dseq, sinks) + dqf = dq.view(len(dreq), h, HEAD_DIM) + dcounts = _window_counts(dseq, dcu) + correct.append(excess(do, _reference(dqf, swa, swa, didx, dactive, dcounts, sinks))) + wrong["dedup"] = excess( + do, _reference(dqf, swa, swa, didx, dactive, dcounts, sinks, dedup=True) + ) + worst = max(correct) + print( + f"gate-normalized error: correct {[round(c, 3) for c in correct]}, wrong " + f"{ {k: round(v, 1) for k, v in wrong.items()} }" + ) + assert worst <= 1.0, f"correct reference outside the gate: {correct}" + for name, e in wrong.items(): + assert e > 10.0 * max(worst, 0.1), ( + f"wrong variant {name!r} not discriminated: {e:.2f} vs {worst:.2f}" + ) + + +# --------------------------------------------------------------------------------------------- +# call mechanics +# --------------------------------------------------------------------------------------------- +def test_out_buffer_and_cuda_graph_replay(): + gen = torch.Generator(device=DEV).manual_seed(61) + swa, comp = _pools(gen) + requests = [(p, 1) for p in (5, 400, 3000, 9000)] + idx, active, seq_lens, cu = _build_rows( + requests, gen, 640, swa_pages_per_req=2, comp_pages_per_req=40 + ) + b, h = len(requests), 16 + q = torch.randn(b, 1, h, HEAD_DIM, generator=gen, device=DEV).to(torch.bfloat16) + sinks = _sinks(h) + ws = _workspace() + out = torch.empty(b, 1, h, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + + def call(): + return op( + query=q, + swa_kv_cache=swa, + compressed_kv_cache=comp, + workspace_buffer=ws, + sparse_indices=idx, + sparse_topk_lens=active, + seq_lens=seq_lens, + sinks=sinks, + softmax_scale=SCALE, + out=out, + ) + + ret = call() + assert ret.data_ptr() == out.data_ptr() + eager = out.clone() + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + call() + torch.cuda.current_stream().wait_stream(s) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, eager) + # new inputs in the same buffers replay correctly + q.copy_(torch.randn(q.shape, generator=gen, device=DEV).to(torch.bfloat16)) + graph.replay() + torch.cuda.synchronize() + ref = _reference( + q.view(b, h, HEAD_DIM), swa, comp, idx, active, _window_counts(seq_lens, cu), sinks + ) + _check(out, ref, swa, comp) + assert math.isfinite(out.float().abs().max().item()) + + +# --------------------------------------------------------------------------------------------- +# The context (varlen prefill) call over the engine's configuration ENVELOPE: up to +# 8192 tokens per step (max_num_tokens), up to 32 requests per step (max_batch_size), cached prefix +# + q_len <= 16384 per request (max_seq_len), H 16 / 64, index width 128 (ratio-0, window only) +# and 640 (window + 512 compressed, ratio 2 or 1). Calls are made exactly as the core target makes +# them in a step: the context call on the first tc rows (q, index rows, lengths and out are views +# of step-sized tensors, seq_lens = kv_lens[:nc], cum_seq_lens_q = cu[:nc + 1]) and, in a mixed +# step, the decode call on the views that start at row tc (seq_lens = kv_lens[nc:B]). +# --------------------------------------------------------------------------------------------- +ENV_T = 8192 +ENV_B = 32 +ENV_SEQ = 16384 +ENV_LAYERS = {"r0": (0, 64), "r2": (2, 64), "r1": (1, 128)} # ratio, compressed page rows +ENV_SWA_PAGES = 2 * ENV_B + 32 # a ring of 2 SWA pages per request: a 128-row window spans <= 2 +ENV_COMP_PAGES = {2: ENV_B * (ENV_SEQ // 2 // 64) + 64, 1: ENV_B * (ENV_SEQ // 128) + 64} + +# Every context call the runtime-value dumps showed: (tc, nc, max_q_len) per head count, from the +# boot gates' and the accuracy-gate previews' logs, parsed line by line with an arity check; each +# cell ran at width 128 and 640 in the engine. Three tuples that cannot be a step (tc > nc * +# max_q_len, or max_q_len > tc) are interleaving garbles and are left out. +DUMPED_CTX_CELLS = { + 16: ( + "1/1/1 2/1/2 48/1/48 493/1/493 494/1/494 503/1/503 529/1/529 549/1/549 555/1/555 556/1/556 " + "562/1/562 569/1/569 571/1/571 573/1/573 578/1/578 579/1/579 581/1/581 586/1/586 590/1/590 " + "593/1/593 594/1/594 595/1/595 598/1/598 599/1/599 601/1/601 603/1/603 606/1/606 609/1/609 " + "612/1/612 614/1/614 616/1/616 617/1/617 619/1/619 625/1/625 626/1/626 627/1/627 629/1/629 " + "631/1/631 632/1/632 634/1/634 635/1/635 637/1/637 640/1/640 641/1/641 642/1/642 643/1/643 " + "644/1/644 646/1/646 650/1/650 651/1/651 652/1/652 653/1/653 654/1/654 656/1/656 660/1/660 " + "661/1/661 662/1/662 663/1/663 665/1/665 666/1/666 667/1/667 668/1/668 670/1/670 674/1/674 " + "675/1/675 676/1/676 679/1/679 680/1/680 681/1/681 682/1/682 683/1/683 684/1/684 687/1/687 " + "688/1/688 689/1/689 690/1/690 691/1/691 692/1/692 693/1/693 694/1/694 696/1/696 697/1/697 " + "698/1/698 699/1/699 700/1/700 701/1/701 702/1/702 703/1/703 704/1/704 706/1/706 708/1/708 " + "710/1/710 711/1/711 712/1/712 713/1/713 714/1/714 718/1/718 719/1/719 720/1/720 721/1/721 " + "722/1/722 723/1/723 724/1/724 725/1/725 726/1/726 727/1/727 728/1/728 729/1/729 730/1/730 " + "732/1/732 733/1/733 734/1/734 735/1/735 736/1/736 737/1/737 738/1/738 739/1/739 740/1/740 " + "741/1/741 742/1/742 743/1/743 745/1/745 746/1/746 747/1/747 748/1/748 749/1/749 751/1/751 " + "752/1/752 753/1/753 755/1/755 756/1/756 757/1/757 758/1/758 759/1/759 760/1/760 761/1/761 " + "762/1/762 763/1/763 764/1/764 765/1/765 766/1/766 767/1/767 768/1/768 769/1/769 770/1/770 " + "771/1/771 772/1/772 773/1/773 774/1/774 775/1/775 776/1/776 777/1/777 778/1/778 779/1/779 " + "780/1/780 781/1/781 782/1/782 783/1/783 784/1/784 785/1/785 786/1/786 787/1/787 788/1/788 " + "790/1/790 791/1/791 792/1/792 794/1/794 795/1/795 796/1/796 797/1/797 798/1/798 799/1/799 " + "800/1/800 801/1/801 802/1/802 803/1/803 804/1/804 805/1/805 806/1/806 807/1/807 808/1/808 " + "809/1/809 810/1/810 811/1/811 812/1/812 813/1/813 814/1/814 815/1/815 816/1/816 817/1/817 " + "818/1/818 819/1/819 820/1/820 822/1/822 823/1/823 824/1/824 825/1/825 826/1/826 827/1/827 " + "828/1/828 829/1/829 830/1/830 831/1/831 832/1/832 833/1/833 834/1/834 835/1/835 836/1/836 " + "837/1/837 838/1/838 839/1/839 840/1/840 841/1/841 843/1/843 844/1/844 845/1/845 846/1/846 " + "847/1/847 848/1/848 849/1/849 850/1/850 851/1/851 852/1/852 853/1/853 854/1/854 855/1/855 " + "856/1/856 857/1/857 858/1/858 859/1/859 860/1/860 861/1/861 862/1/862 863/1/863 864/1/864 " + "866/1/866 867/1/867 868/1/868 869/1/869 870/1/870 871/1/871 872/1/872 873/1/873 874/1/874 " + "875/1/875 876/1/876 877/1/877 878/1/878 879/1/879 880/1/880 881/1/881 882/1/882 883/1/883 " + "884/1/884 885/1/885 886/1/886 887/1/887 888/1/888 889/1/889 890/1/890 891/1/891 892/1/892 " + "893/1/893 894/1/894 895/1/895 896/1/896 897/1/897 898/1/898 899/1/899 900/1/900 901/1/901 " + "902/1/902 903/1/903 904/1/904 906/1/906 907/1/907 908/1/908 909/1/909 910/1/910 911/1/911 " + "912/1/912 913/1/913 914/1/914 915/1/915 916/1/916 917/1/917 918/1/918 919/1/919 920/1/920 " + "921/1/921 922/1/922 923/1/923 924/1/924 925/1/925 926/1/926 927/1/927 928/1/928 929/1/929 " + "930/1/930 931/1/931 932/1/932 933/1/933 935/1/935 936/1/936 937/1/937 938/1/938 939/1/939 " + "940/1/940 941/1/941 942/1/942 943/1/943 944/1/944 945/1/945 946/1/946 947/1/947 949/1/949 " + "950/1/950 951/1/951 952/1/952 953/1/953 954/1/954 955/1/955 956/1/956 957/1/957 958/1/958 " + "959/1/959 960/1/960 961/1/961 962/1/962 963/1/963 964/1/964 965/1/965 966/1/966 967/1/967 " + "968/1/968 969/1/969 970/1/970 972/1/972 973/1/973 974/1/974 976/1/976 977/1/977 979/1/979 " + "980/1/980 982/1/982 983/1/983 984/1/984 986/1/986 987/1/987 988/1/988 989/1/989 990/1/990 " + "991/1/991 992/1/992 993/1/993 994/1/994 996/1/996 997/1/997 998/1/998 999/1/999 1000/1/1000 " + "1002/1/1002 1003/1/1003 1004/1/1004 1005/1/1005 1007/1/1007 1008/1/1008 1011/1/1011 1012/1/1012 " + "1013/1/1013 1014/1/1014 1015/1/1015 1016/1/1016 1018/1/1018 1019/1/1019 1020/1/1020 1021/1/1021 " + "1022/1/1022 1023/1/1023 1025/1/1025 1026/1/1026 1028/1/1028 1031/1/1031 1032/1/1032 1034/1/1034 " + "1036/1/1036 1038/1/1038 1039/1/1039 1040/1/1040 1042/1/1042 1043/1/1043 1045/1/1045 1046/1/1046 " + "1048/1/1048 1050/1/1050 1055/1/1055 1057/1/1057 1058/1/1058 1062/1/1062 1065/1/1065 1066/1/1066 " + "1067/1/1067 1068/1/1068 1070/1/1070 1072/1/1072 1076/1/1076 1079/1/1079 1080/1/1080 1081/1/1081 " + "1082/1/1082 1084/1/1084 1085/1/1085 1086/1/1086 1087/1/1087 1090/1/1090 1091/1/1091 1092/1/1092 " + "1093/1/1093 1094/1/1094 1095/1/1095 1098/1/1098 1099/1/1099 1100/1/1100 1101/1/1101 1104/1/1104 " + "1106/1/1106 1109/1/1109 1111/1/1111 1115/1/1115 1118/1/1118 1119/1/1119 1120/1/1120 1122/1/1122 " + "1125/1/1125 1126/1/1126 1128/1/1128 1129/1/1129 1137/1/1137 1139/1/1139 1140/1/1140 1141/1/1141 " + "1143/1/1143 1145/1/1145 1146/1/1146 1149/1/1149 1150/1/1150 1157/1/1157 1170/1/1170 1171/1/1171 " + "1172/1/1172 1178/1/1178 1182/1/1182 1185/1/1185 1187/1/1187 1198/1/1198 1200/1/1200 1204/1/1204 " + "1211/1/1211 1217/1/1217 1226/1/1226 1236/1/1236 1241/2/629 1261/1/1261 1273/1/1273 1275/2/697 " + "1283/2/780 1298/1/1298 1301/1/1301 1333/2/670 1345/2/746 1355/2/698 1374/2/768 1388/2/757 " + "1391/2/730 1395/2/738 1397/2/811 1407/2/858 1408/2/738 1410/2/855 1419/2/723 1423/2/734 " + "1432/2/938 1433/2/787 1448/2/756 1462/2/774 1464/2/782 1466/2/825 1467/2/877 1473/2/777 " + "1477/2/792 1481/2/777 1482/2/788 1484/2/835 1494/2/895 1496/2/854 1497/2/818 1498/2/779 " + "1501/1/1501 1505/2/882 1510/2/820 1513/2/788 1514/2/759 1514/2/844 1518/2/776 1519/2/793 " + "1520/2/786 1523/2/832 1524/2/825 1525/2/831 1525/2/952 1533/2/780 1537/2/839 1539/2/859 " + "1542/2/799 1543/2/854 1546/2/863 1546/2/906 1547/2/784 1547/2/917 1550/2/815 1551/2/809 " + "1551/2/895 1554/2/799 1554/2/850 1561/2/801 1562/2/828 1566/2/784 1567/2/788 1567/2/811 " + "1567/2/911 1568/2/815 1569/2/888 1570/2/851 1571/2/831 1574/2/811 1574/2/834 1575/2/848 " + "1579/2/796 1583/2/1034 1584/2/799 1585/2/807 1585/2/815 1591/2/804 1593/2/862 1594/2/801 " + "1594/2/875 1594/2/918 1595/2/848 1595/2/870 1601/2/849 1601/2/909 1605/2/824 1606/2/847 " + "1606/2/914 1608/2/873 1610/2/853 1614/2/820 1616/2/941 1620/2/878 1620/2/880 1621/2/967 " + "1624/2/866 1625/2/869 1626/2/920 1627/2/834 1627/2/931 1630/2/889 1636/2/883 1637/2/822 " + "1637/2/888 1639/2/954 1640/2/914 1641/2/868 1642/2/1005 1644/2/869 1646/2/873 1647/2/887 " + "1649/2/871 1654/2/924 1655/2/964 1656/2/845 1661/2/942 1661/2/1064 1662/2/833 1663/2/866 " + "1663/2/870 1664/2/903 1664/2/962 1665/2/871 1666/2/943 1667/2/910 1668/2/1036 1669/2/896 " + "1670/2/900 1672/2/841 1672/2/854 1672/2/888 1675/2/856 1675/2/871 1675/2/926 1676/2/880 " + "1680/2/939 1682/2/930 1685/2/911 1685/2/916 1685/2/919 1686/2/925 1687/2/907 1688/2/851 " + "1691/2/907 1693/2/996 1695/2/876 1695/2/1109 1696/2/1008 1699/2/1019 1700/2/864 1705/2/876 " + "1706/2/959 1707/2/856 1709/2/987 1710/2/934 1710/2/1008 1712/2/877 1713/2/906 1714/2/858 " + "1716/2/877 1719/2/1007 1720/2/917 1720/2/921 1725/2/925 1725/2/973 1726/2/918 1728/2/871 " + "1728/2/895 1731/2/877 1731/2/903 1737/2/930 1738/2/954 1740/2/1003 1741/2/907 1742/2/950 " + "1747/2/966 1747/2/990 1748/2/909 1749/2/942 1750/2/918 1752/2/1100 1753/2/993 1754/2/884 " + "1754/2/915 1755/2/888 1756/2/898 1756/2/951 1757/2/986 1758/2/923 1758/2/1098 1760/2/988 " + "1764/2/952 1769/2/893 1769/2/944 1770/2/960 1771/2/1091 1772/2/992 1772/2/1048 1773/2/1005 " + "1775/2/926 1775/2/958 1775/2/1119 1776/2/953 1781/2/893 1781/2/976 1781/2/1019 1783/2/1137 " + "1785/2/924 1790/2/966 1791/2/930 1797/2/1198 1799/2/910 1801/2/912 1802/2/934 1802/2/980 " + "1803/2/996 1804/2/914 1805/2/951 1808/2/1079 1812/2/988 1814/2/953 1814/2/983 1817/2/910 " + "1817/2/912 1817/2/941 1817/2/1089 1818/2/1043 1819/2/977 1820/2/1003 1821/2/1048 1823/2/912 " + "1824/2/963 1825/2/929 1826/2/1067 1832/2/1143 1838/2/1021 1844/2/1045 1850/2/961 1850/2/1149 " + "1852/2/928 1855/2/975 1856/2/951 1856/2/977 1857/2/998 1859/2/1000 1863/2/945 1863/2/975 " + "1864/2/1074 1868/2/947 1869/2/974 1871/2/1026 1873/2/1125 1873/2/1176 1875/2/1122 1876/2/1201 " + "1878/2/957 1887/2/945 1891/2/1136 1892/2/966 1892/2/992 1895/2/1003 1901/2/1050 1902/2/1087 " + "1910/2/1011 1910/2/1028 1912/2/1068 1913/2/1086 1920/2/1081 1921/2/1201 1924/2/1004 1925/2/1154 " + "1926/2/977 1926/2/1031 1926/2/1099 1931/2/1043 1934/2/1920 1935/2/1058 1939/2/1100 1964/2/1058 " + "1968/2/994 1968/2/1091 1972/2/1076 1976/2/988 1976/2/1008 1977/2/1119 1979/2/1053 1980/2/1015 " + "1980/2/1036 1982/2/1038 1982/2/1053 1987/2/1045 1988/2/1011 1993/2/1024 1997/2/1032 1998/2/999 " + "2004/2/1049 2007/2/1176 2009/2/1130 2013/2/1097 2016/2/1154 2017/2/1038 2017/2/1104 2019/2/1094 " + "2026/2/1118 2026/3/811 2034/2/1040 2045/2/1080 2045/2/1092 2048/1/2048 2051/2/1057 2054/2/1099 " + "2072/2/1170 2073/2/1080 2077/2/1066 2085/3/730 2092/2/1198 2094/2/1167 2117/2/1109 2121/2/1082 " + "2135/2/1501 2140/2/1128 2164/2/1185 2165/2/1273 2179/2/1136 2193/2/1177 2213/3/900 2228/2/1139 " + "2232/3/835 2260/2/1145 2313/3/910 2349/2/1237 2367/3/1018 2369/3/891 2401/3/930 2411/3/925 " + "2413/2/1298 2419/3/846 2433/3/936 2457/3/896 2465/3/923 2466/3/896 2480/3/936 2489/3/883 " + "2498/3/884 2506/3/1023 2541/3/936 2556/3/872 2626/3/1065 2643/3/1095 2654/3/963 2687/3/999 " + "2693/3/945 2713/3/999 2715/3/942 2721/3/1140 2728/3/1045 2749/3/1177 2769/3/1080 2813/3/1170 " + "2819/3/989 2851/3/1081 2866/3/1085 2910/3/1049 2968/3/1261 3206/4/1114 3250/4/911 3289/4/867 " + "3336/3/1237 8102/10/1085 8141/9/1163 8153/10/1318 8192/1/8192 " + ), + 64: ( + "1/1/1 2/1/2 14/1/14 16/1/16 48/1/48 549/1/549 550/1/550 562/1/562 571/1/571 586/1/586 590/1/590 " + "597/1/597 599/1/599 601/1/601 606/1/606 609/1/609 616/1/616 619/1/619 623/1/623 626/1/626 " + "627/1/627 630/1/630 631/1/631 632/1/632 634/1/634 635/1/635 637/1/637 643/1/643 644/1/644 " + "646/1/646 649/1/649 650/1/650 651/1/651 652/1/652 653/1/653 654/1/654 656/1/656 661/1/661 " + "662/1/662 663/1/663 666/1/666 667/1/667 668/1/668 670/1/670 672/1/672 674/1/674 675/1/675 " + "676/1/676 680/1/680 681/1/681 682/1/682 683/1/683 684/1/684 685/1/685 688/1/688 689/1/689 " + "690/1/690 691/1/691 694/1/694 696/1/696 697/1/697 698/1/698 699/1/699 701/1/701 702/1/702 " + "704/1/704 706/1/706 711/1/711 712/1/712 713/1/713 715/1/715 719/1/719 720/1/720 721/1/721 " + "722/1/722 723/1/723 725/1/725 726/1/726 727/1/727 728/1/728 729/1/729 730/1/730 732/1/732 " + "733/1/733 734/1/734 735/1/735 737/1/737 738/1/738 739/1/739 740/1/740 741/1/741 742/1/742 " + "743/1/743 745/1/745 746/1/746 747/1/747 748/1/748 749/1/749 751/1/751 752/1/752 753/1/753 " + "754/1/754 755/1/755 756/1/756 757/1/757 758/1/758 759/1/759 760/1/760 761/1/761 762/1/762 " + "763/1/763 764/1/764 765/1/765 766/1/766 767/1/767 768/1/768 769/1/769 770/1/770 771/1/771 " + "773/1/773 774/1/774 775/1/775 776/1/776 777/1/777 778/1/778 779/1/779 780/1/780 781/1/781 " + "782/1/782 783/1/783 784/1/784 785/1/785 786/1/786 787/1/787 788/1/788 790/1/790 791/1/791 " + "792/1/792 793/1/793 794/1/794 795/1/795 796/1/796 797/1/797 798/1/798 799/1/799 800/1/800 " + "801/1/801 802/1/802 803/1/803 804/1/804 805/1/805 806/1/806 807/1/807 808/1/808 809/1/809 " + "810/1/810 811/1/811 812/1/812 813/1/813 814/1/814 815/1/815 816/1/816 817/1/817 818/1/818 " + "819/1/819 820/1/820 822/1/822 823/1/823 824/1/824 825/1/825 826/1/826 827/1/827 828/1/828 " + "829/1/829 830/1/830 831/1/831 832/1/832 833/1/833 834/1/834 835/1/835 836/1/836 837/1/837 " + "838/1/838 839/1/839 840/1/840 841/1/841 843/1/843 844/1/844 845/1/845 846/1/846 847/1/847 " + "848/1/848 849/1/849 850/1/850 851/1/851 852/1/852 853/1/853 854/1/854 855/1/855 856/1/856 " + "857/1/857 858/1/858 859/1/859 860/1/860 861/1/861 862/1/862 863/1/863 864/1/864 866/1/866 " + "867/1/867 868/1/868 869/1/869 870/1/870 871/1/871 872/1/872 873/1/873 874/1/874 875/1/875 " + "876/1/876 877/1/877 878/1/878 879/1/879 880/1/880 881/1/881 882/1/882 883/1/883 884/1/884 " + "885/1/885 886/1/886 887/1/887 888/1/888 889/1/889 890/1/890 891/1/891 892/1/892 893/1/893 " + "894/1/894 895/1/895 896/1/896 897/1/897 898/1/898 899/1/899 900/1/900 901/1/901 902/1/902 " + "903/1/903 904/1/904 906/1/906 907/1/907 908/1/908 909/1/909 910/1/910 911/1/911 912/1/912 " + "913/1/913 914/1/914 915/1/915 916/1/916 917/1/917 918/1/918 919/1/919 920/1/920 921/1/921 " + "922/1/922 923/1/923 924/1/924 925/1/925 926/1/926 927/1/927 928/1/928 929/1/929 930/1/930 " + "931/1/931 932/1/932 933/1/933 934/1/934 936/1/936 937/1/937 939/1/939 940/1/940 941/1/941 " + "942/1/942 943/1/943 944/1/944 945/1/945 946/1/946 947/1/947 949/1/949 950/1/950 951/1/951 " + "952/1/952 953/1/953 954/1/954 955/1/955 957/1/957 958/1/958 959/1/959 961/1/961 962/1/962 " + "963/1/963 964/1/964 965/1/965 966/1/966 967/1/967 968/1/968 970/1/970 972/1/972 973/1/973 " + "974/1/974 975/1/975 976/1/976 977/1/977 979/1/979 980/1/980 983/1/983 984/1/984 986/1/986 " + "987/1/987 988/1/988 989/1/989 990/1/990 991/1/991 992/1/992 993/1/993 994/1/994 996/1/996 " + "997/1/997 998/1/998 999/1/999 1000/1/1000 1002/1/1002 1003/1/1003 1004/1/1004 1005/1/1005 " + "1007/1/1007 1008/1/1008 1011/1/1011 1012/1/1012 1013/1/1013 1014/1/1014 1015/1/1015 1016/1/1016 " + "1018/1/1018 1019/1/1019 1020/1/1020 1021/1/1021 1022/1/1022 1023/1/1023 1025/1/1025 1026/1/1026 " + "1028/1/1028 1031/1/1031 1032/1/1032 1034/1/1034 1036/1/1036 1038/1/1038 1039/1/1039 1040/1/1040 " + "1043/1/1043 1046/1/1046 1048/1/1048 1050/1/1050 1055/1/1055 1057/1/1057 1058/1/1058 1062/1/1062 " + "1064/1/1064 1065/1/1065 1067/1/1067 1068/1/1068 1070/1/1070 1079/1/1079 1080/1/1080 1081/1/1081 " + "1082/1/1082 1084/1/1084 1085/1/1085 1089/1/1089 1090/1/1090 1092/1/1092 1093/1/1093 1094/1/1094 " + "1095/1/1095 1099/1/1099 1100/1/1100 1101/1/1101 1104/1/1104 1106/1/1106 1111/1/1111 1115/1/1115 " + "1119/1/1119 1120/1/1120 1122/1/1122 1126/1/1126 1130/1/1130 1135/1/1135 1136/1/1136 1137/1/1137 " + "1139/1/1139 1141/1/1141 1142/1/1142 1143/1/1143 1145/1/1145 1149/1/1149 1150/1/1150 1154/1/1154 " + "1157/1/1157 1160/1/1160 1170/1/1170 1171/1/1171 1172/1/1172 1177/1/1177 1178/1/1178 1182/1/1182 " + "1185/1/1185 1187/1/1187 1198/1/1198 1200/1/1200 1201/1/1201 1204/1/1204 1211/1/1211 1217/1/1217 " + "1226/1/1226 1237/1/1237 1261/1/1261 1273/1/1273 1301/1/1301 1313/2/764 1344/2/698 1403/2/740 " + "1424/2/730 1439/2/729 1462/2/770 1464/2/753 1472/2/742 1507/2/839 1508/2/847 1511/2/780 " + "1512/2/784 1528/2/771 1535/2/848 1538/2/775 1562/2/841 1566/2/854 1574/2/797 1574/2/834 " + "1575/2/848 1580/2/854 1584/2/826 1596/2/860 1601/2/909 1605/2/842 1606/2/836 1615/2/858 " + "1615/2/859 1626/2/920 1627/2/834 1629/2/825 1637/2/913 1647/2/887 1656/2/845 1661/2/942 " + "1662/2/943 1663/2/858 1663/2/908 1668/2/849 1678/2/904 1680/2/908 1683/2/924 1684/2/916 " + "1685/2/916 1687/2/923 1689/2/965 1694/2/848 1696/2/1008 1701/2/883 1702/2/1003 1703/2/919 " + "1704/2/879 1705/2/856 1707/2/856 1709/2/857 1709/2/987 1710/2/1129 1712/2/910 1716/2/912 " + "1720/2/917 1723/2/990 1724/2/900 1734/2/929 1734/2/950 1735/2/998 1742/2/893 1742/2/950 " + "1743/2/917 1743/2/964 1744/2/1004 1747/2/891 1747/2/1091 1750/2/921 1750/2/926 1752/2/1087 " + "1757/2/986 1758/2/923 1758/2/1098 1759/2/1023 1764/2/952 1766/2/935 1780/2/929 1782/2/916 " + "1791/2/956 1793/2/915 1794/2/1076 1799/2/1119 1802/2/1146 1808/2/969 1810/2/1067 1812/2/988 " + "1820/2/952 1820/2/984 1820/2/1003 1825/2/929 1825/2/930 1828/2/962 1828/2/1060 1836/2/1122 " + "1837/2/974 1843/2/947 1856/2/951 1856/2/1122 1861/2/955 1864/2/1074 1869/2/941 1873/2/965 " + "1873/2/1109 1873/2/1176 1882/2/1170 1885/2/966 1911/2/969 1913/2/1086 1914/2/1025 1921/2/1045 " + "1921/2/1070 1937/2/982 1939/2/1100 1947/2/1150 1952/2/1008 1953/2/1085 1954/2/1066 1956/2/999 " + "1969/2/1024 1969/2/1031 1972/2/1076 1973/2/1012 1973/2/1118 1975/2/1055 1984/2/1094 2008/2/1080 " + "2013/2/1097 2019/2/1125 2020/2/1319 2048/1/2048 2054/4/688 2077/2/1066 2082/2/1112 2105/2/1128 " + "2107/2/1058 2117/2/1109 2117/2/1140 2135/2/1501 2146/2/1236 2155/3/914 2160/2/1139 2163/2/1091 " + "2333/3/836 2413/2/1298 2511/3/919 2523/3/977 2605/3/938 2629/3/877 2644/3/1000 2698/3/1015 " + "2791/3/977 2851/3/1081 2864/3/1045 3114/5/702 3162/5/701 3278/5/697 3601/4/1057 8071/12/833 " + "8095/10/927 8117/9/1097 8119/10/942 8135/10/950 8136/10/944 8150/8/1320 8156/11/835 8160/11/842 " + "8161/11/823 8178/8/1141 8178/8/1165 8192/1/8192 " + ), +} + +# The same for the parallel members' head counts: every context call (tc, nc, max_q_len) in mn's +# per-rank runtime-value dumps at H 8 (tep8 boot and GSM8K) and H 32 (pp2 x tp2 GSM8K), the union +# parsed tuple by tuple with an arity check (the members' runtime-value dumps). Every parsed cell is a +# possible step (max_q_len <= tc <= nc * max_q_len); each ran at width 128 and 640 in the engine. +MEMBER_DUMPED_CTX_CELLS = { + 8: ( + "1/1/1 2/1/2 16/1/16 48/1/48 128/1/128 256/2/128 384/3/128 512/4/128 549/1/549 581/1/581 590/1/590 " + "597/1/597 599/1/599 609/1/609 616/1/616 619/1/619 621/1/621 622/1/622 623/1/623 626/1/626 " + "630/1/630 634/1/634 635/1/635 637/1/637 640/1/640 640/5/128 643/1/643 644/1/644 646/1/646 " + "649/1/649 650/1/650 653/1/653 656/1/656 661/1/661 663/1/663 665/1/665 666/1/666 667/1/667 " + "668/1/668 670/1/670 674/1/674 675/1/675 679/1/679 680/1/680 681/1/681 683/1/683 684/1/684 " + "685/1/685 687/1/687 688/1/688 689/1/689 690/1/690 691/1/691 692/1/692 694/1/694 696/1/696 " + "697/1/697 698/1/698 699/1/699 701/1/701 702/1/702 704/1/704 706/1/706 707/1/707 710/1/710 " + "711/1/711 712/1/712 714/1/714 718/1/718 719/1/719 721/1/721 722/1/722 723/1/723 724/1/724 " + "725/1/725 726/1/726 727/1/727 728/1/728 729/1/729 730/1/730 731/1/731 733/1/733 734/1/734 " + "735/1/735 736/1/736 737/1/737 738/1/738 740/1/740 741/1/741 742/1/742 743/1/743 745/1/745 " + "747/1/747 749/1/749 752/1/752 753/1/753 755/1/755 756/1/756 757/1/757 758/1/758 759/1/759 " + "760/1/760 762/1/762 763/1/763 764/1/764 765/1/765 766/1/766 767/1/767 768/1/768 768/6/128 " + "769/1/769 770/1/770 771/1/771 773/1/773 774/1/774 775/1/775 776/1/776 777/1/777 778/1/778 " + "779/1/779 780/1/780 781/1/781 782/1/782 783/1/783 784/1/784 785/1/785 786/1/786 787/1/787 " + "788/1/788 789/1/789 790/1/790 791/1/791 792/1/792 793/1/793 795/1/795 796/1/796 797/1/797 " + "799/1/799 800/1/800 801/1/801 802/1/802 803/1/803 806/1/806 807/1/807 808/1/808 809/1/809 " + "810/1/810 811/1/811 812/1/812 814/1/814 815/1/815 816/1/816 817/1/817 818/1/818 819/1/819 " + "820/1/820 822/1/822 823/1/823 824/1/824 825/1/825 826/1/826 827/1/827 828/1/828 829/1/829 " + "830/1/830 831/1/831 832/1/832 833/1/833 834/1/834 835/1/835 836/1/836 837/1/837 838/1/838 " + "839/1/839 841/1/841 842/1/842 844/1/844 845/1/845 846/1/846 847/1/847 848/1/848 849/1/849 " + "850/1/850 851/1/851 852/1/852 853/1/853 854/1/854 855/1/855 856/1/856 857/1/857 858/1/858 " + "859/1/859 860/1/860 861/1/861 862/1/862 863/1/863 864/1/864 866/1/866 867/1/867 868/1/868 " + "869/1/869 870/1/870 871/1/871 872/1/872 873/1/873 874/1/874 875/1/875 876/1/876 877/1/877 " + "878/1/878 879/1/879 880/1/880 881/1/881 882/1/882 883/1/883 884/1/884 887/1/887 888/1/888 " + "889/1/889 890/1/890 892/1/892 893/1/893 894/1/894 895/1/895 896/1/896 897/1/897 899/1/899 " + "900/1/900 901/1/901 903/1/903 904/1/904 905/1/905 906/1/906 907/1/907 908/1/908 909/1/909 " + "910/1/910 911/1/911 912/1/912 913/1/913 914/1/914 915/1/915 916/1/916 917/1/917 918/1/918 " + "919/1/919 920/1/920 921/1/921 922/1/922 923/1/923 924/1/924 925/1/925 926/1/926 927/1/927 " + "928/1/928 929/1/929 930/1/930 931/1/931 933/1/933 934/1/934 936/1/936 938/1/938 939/1/939 " + "941/1/941 942/1/942 943/1/943 944/1/944 945/1/945 946/1/946 947/1/947 949/1/949 950/1/950 " + "951/1/951 952/1/952 953/1/953 954/1/954 956/1/956 957/1/957 958/1/958 959/1/959 960/1/960 " + "961/1/961 963/1/963 964/1/964 965/1/965 966/1/966 967/1/967 968/1/968 969/1/969 970/1/970 " + "975/1/975 977/1/977 979/1/979 980/1/980 981/1/981 982/1/982 983/1/983 984/1/984 986/1/986 " + "987/1/987 988/1/988 989/1/989 991/1/991 992/1/992 993/1/993 994/1/994 996/1/996 998/1/998 " + "999/1/999 1000/1/1000 1003/1/1003 1004/1/1004 1005/1/1005 1007/1/1007 1008/1/1008 1011/1/1011 " + "1012/1/1012 1013/1/1013 1014/1/1014 1015/1/1015 1016/1/1016 1018/1/1018 1019/1/1019 1021/1/1021 " + "1024/1/1024 1025/1/1025 1026/1/1026 1031/1/1031 1036/1/1036 1038/1/1038 1039/1/1039 1040/1/1040 " + "1042/1/1042 1043/1/1043 1046/1/1046 1048/1/1048 1049/1/1049 1050/1/1050 1055/1/1055 1057/1/1057 " + "1058/1/1058 1062/1/1062 1064/1/1064 1066/1/1066 1067/1/1067 1068/1/1068 1070/1/1070 1074/1/1074 " + "1076/1/1076 1081/1/1081 1082/1/1082 1084/1/1084 1085/1/1085 1087/1/1087 1089/1/1089 1091/1/1091 " + "1092/1/1092 1093/1/1093 1094/1/1094 1098/1/1098 1099/1/1099 1100/1/1100 1101/1/1101 1102/1/1102 " + "1104/1/1104 1106/1/1106 1111/1/1111 1112/1/1112 1115/1/1115 1119/1/1119 1120/1/1120 1122/1/1122 " + "1128/1/1128 1130/1/1130 1136/1/1136 1137/1/1137 1139/1/1139 1140/1/1140 1141/1/1141 1143/1/1143 " + "1145/1/1145 1146/1/1146 1149/1/1149 1152/9/128 1170/1/1170 1171/1/1171 1172/1/1172 1178/1/1178 " + "1187/1/1187 1198/1/1198 1200/1/1200 1201/1/1201 1204/1/1204 1211/1/1211 1236/1/1236 1237/1/1237 " + "1241/2/629 1261/1/1261 1273/1/1273 1275/2/697 1285/2/688 1298/1/1298 1313/2/764 1326/2/727 " + "1336/2/676 1350/2/719 1374/2/768 1376/2/757 1380/2/745 1388/2/757 1397/2/811 1412/2/780 " + "1419/2/723 1423/2/734 1423/2/773 1433/2/787 1439/2/729 1442/2/723 1447/2/751 1450/2/796 " + "1452/2/730 1458/2/833 1460/2/780 1464/2/797 1466/2/805 1472/2/742 1473/2/740 1476/2/747 " + "1477/2/742 1477/2/758 1477/2/792 1492/2/876 1495/2/794 1497/2/807 1497/2/836 1499/2/829 " + "1506/2/944 1507/2/784 1507/2/839 1511/2/780 1512/2/854 1513/2/788 1523/2/785 1525/2/847 " + "1526/2/807 1526/2/826 1526/2/850 1527/2/836 1528/2/771 1530/2/898 1531/2/799 1531/2/831 " + "1533/2/776 1535/2/801 1536/2/858 1536/2/874 1538/2/793 1539/2/859 1539/2/877 1540/2/909 " + "1544/2/901 1546/2/863 1546/2/895 1547/2/805 1547/2/834 1547/2/847 1547/2/893 1548/2/791 " + "1548/2/815 1549/2/807 1564/2/898 1564/2/938 1566/2/784 1567/2/788 1568/2/815 1569/2/873 " + "1570/2/828 1571/2/874 1572/2/874 1573/2/846 1574/2/793 1574/2/811 1575/2/789 1577/2/824 " + "1581/2/823 1581/2/832 1581/2/950 1582/2/824 1584/2/825 1585/2/804 1585/2/807 1587/2/797 " + "1589/2/990 1590/2/824 1590/2/889 1592/2/881 1594/2/825 1594/2/875 1594/2/918 1595/2/832 " + "1595/2/874 1596/2/863 1597/2/1019 1599/2/859 1603/2/803 1608/2/862 1608/2/868 1610/2/853 " + "1610/2/926 1613/2/910 1614/2/858 1615/2/893 1616/2/888 1618/2/920 1619/2/943 1620/2/880 " + "1620/2/926 1622/2/858 1623/2/922 1624/2/866 1625/2/869 1626/2/920 1627/2/931 1628/2/841 " + "1628/2/850 1630/2/939 1631/2/867 1632/2/854 1634/2/881 1636/2/925 1637/2/822 1638/2/847 " + "1638/2/879 1640/2/994 1641/2/868 1641/2/919 1642/2/1005 1642/2/1011 1644/2/869 1645/2/830 " + "1645/2/869 1646/2/873 1647/2/887 1647/2/998 1649/2/864 1650/2/923 1652/2/874 1652/2/877 " + "1653/2/900 1655/2/878 1655/2/964 1656/2/919 1659/2/831 1660/2/872 1660/2/875 1660/2/912 " + "1661/2/871 1662/2/833 1664/2/897 1665/2/871 1666/2/848 1667/2/914 1670/2/900 1672/2/841 " + "1675/2/871 1675/2/963 1680/2/953 1682/2/930 1684/2/847 1684/2/924 1685/2/910 1685/2/911 " + "1685/2/919 1686/2/944 1687/2/849 1687/2/923 1688/2/877 1688/2/928 1689/2/883 1689/2/965 " + "1690/2/941 1692/2/895 1692/2/917 1692/2/929 1693/2/852 1696/2/970 1697/2/896 1697/2/921 " + "1698/2/886 1700/2/867 1702/2/856 1703/2/883 1704/2/879 1705/2/928 1705/2/991 1706/2/953 " + "1706/2/959 1709/2/880 1709/2/987 1710/2/1008 1712/2/908 1712/2/910 1713/2/943 1714/2/909 " + "1716/2/877 1716/2/912 1717/2/898 1717/2/965 1719/2/921 1722/2/915 1724/2/900 1724/2/954 " + "1725/2/973 1726/2/876 1726/2/918 1727/2/965 1728/2/866 1730/2/887 1730/2/929 1731/2/877 " + "1731/2/911 1731/2/962 1734/2/929 1734/2/950 1737/2/930 1738/2/954 1739/2/891 1742/2/925 " + "1742/2/961 1744/2/1004 1745/2/937 1745/2/941 1747/2/999 1751/2/945 1752/2/1100 1753/2/993 " + "1754/2/884 1756/2/951 1758/2/923 1758/2/1098 1760/2/911 1762/2/882 1763/2/901 1764/2/884 " + "1765/2/924 1765/2/926 1766/2/935 1769/2/942 1770/2/940 1770/2/960 1771/2/894 1772/2/977 " + "1775/2/926 1775/2/958 1775/2/969 1776/2/953 1778/2/920 1780/2/891 1781/2/976 1782/2/896 " + "1782/2/1081 1784/2/1007 1786/2/910 1786/2/1067 1786/2/1130 1789/2/931 1789/2/994 1790/2/943 " + "1791/2/911 1791/2/930 1791/2/1031 1792/2/939 1792/2/1065 1794/2/994 1797/2/910 1798/2/994 " + "1799/2/904 1800/2/949 1800/2/992 1801/2/907 1801/2/983 1802/2/920 1802/2/980 1803/2/911 " + "1804/2/1008 1805/2/951 1806/2/1040 1806/2/1085 1807/2/911 1808/2/969 1808/2/1079 1809/2/935 " + "1810/2/1067 1812/2/988 1812/2/1020 1813/2/982 1814/2/926 1814/2/947 1814/2/953 1814/2/983 " + "1814/2/1034 1816/2/944 1816/2/1025 1817/2/909 1817/2/910 1818/2/1043 1819/2/912 1819/2/966 " + "1820/2/920 1821/2/1048 1824/2/963 1824/2/1026 1824/2/1140 1825/2/930 1827/2/1058 1828/2/915 " + "1829/2/940 1832/2/997 1833/2/918 1837/2/945 1837/2/1002 1839/2/927 1840/2/989 1841/2/1157 " + "1844/2/973 1845/2/988 1846/2/1099 1847/2/979 1851/2/967 1855/2/939 1858/2/991 1859/2/1000 " + "1860/2/1094 1861/2/932 1861/2/955 1862/2/1177 1863/2/975 1865/2/972 1867/2/1081 1868/2/1172 " + "1869/2/1003 1871/2/1003 1872/2/968 1873/2/1109 1873/2/1125 1874/2/1143 1875/2/958 1877/2/958 " + "1877/2/1070 1879/2/983 1881/2/983 1882/2/1008 1883/2/1019 1889/2/1101 1891/2/1068 1895/2/955 " + "1895/2/1003 1897/2/1094 1901/2/1057 1902/2/1067 1908/2/1062 1909/2/956 1909/2/965 1909/2/982 " + "1910/2/1028 1911/2/1057 1913/2/984 1913/2/1024 1915/2/1045 1918/2/1020 1921/2/1201 1924/2/1004 " + "1926/2/1031 1928/2/974 1934/2/1046 1935/2/1058 1936/2/1920 1939/2/1011 1940/2/1092 1941/2/1093 " + "1943/2/980 1945/2/988 1950/2/984 1950/2/1004 1951/2/1122 1952/2/1008 1956/2/986 1957/2/1050 " + "1960/2/1066 1966/2/1036 1968/2/994 1969/2/1031 1975/2/1055 1976/2/988 1979/2/1053 1980/2/1036 " + "1982/2/1053 1983/2/1149 1986/2/1072 1987/2/1045 1993/2/1014 1995/2/1139 2000/2/1122 2007/2/1176 " + "2008/2/1080 2010/2/1064 2011/2/1021 2013/2/1097 2015/2/1211 2018/2/1137 2026/2/1118 2029/2/1085 " + "2039/2/1021 2042/2/1080 2048/1/2048 2050/2/1115 2059/2/1141 2067/2/1143 2072/2/1170 2072/3/811 " + "2077/2/1066 2078/3/742 2079/2/1109 2093/2/1085 2093/3/721 2100/2/1102 2102/2/1141 2104/2/1204 " + "2118/2/1141 2121/2/1182 2132/2/1100 2139/3/720 2142/2/1200 2156/2/1120 2168/2/1143 2174/2/1182 " + "2191/2/1143 2193/2/1177 2199/3/873 2202/2/1126 2203/3/777 2206/3/787 2210/2/1119 2214/3/794 " + "2230/2/1185 2247/2/1128 2250/3/826 2252/2/1501 2286/3/832 2289/3/881 2306/3/842 2307/2/1217 " + "2316/3/895 2318/3/811 2327/3/859 2332/3/880 2335/3/990 2342/3/877 2346/3/848 2353/3/895 " + "2354/3/944 2355/3/815 2364/3/900 2366/3/858 2367/3/881 2369/3/855 2369/3/891 2370/3/881 " + "2371/3/875 2375/3/874 2379/3/828 2379/3/909 2386/3/926 2388/3/898 2404/3/914 2405/3/858 " + "2407/3/984 2410/3/831 2412/3/913 2421/3/900 2423/3/998 2424/2/1226 2424/3/857 2427/3/913 " + "2428/3/834 2429/3/875 2433/3/848 2441/3/924 2442/3/980 2446/3/852 2447/3/959 2449/3/964 " + "2452/3/844 2455/3/950 2467/3/872 2468/3/952 2470/3/909 2472/3/963 2480/3/1000 2485/3/893 " + "2485/3/1038 2496/3/912 2502/3/910 2514/3/981 2515/3/914 2517/3/914 2521/3/923 2523/3/977 " + "2524/3/856 2524/3/1019 2525/3/885 2535/3/958 2541/3/936 2542/3/925 2546/3/876 2552/3/911 " + "2553/3/856 2560/3/887 2562/3/929 2562/3/945 2563/3/910 2563/3/915 2564/3/930 2569/3/910 " + "2571/3/951 2575/3/952 2577/3/950 2581/3/996 2597/3/933 2600/3/888 2610/3/941 2611/3/903 " + "2615/3/977 2616/3/947 2622/3/1048 2623/3/1031 2625/3/884 2625/3/1154 2626/3/921 2626/3/954 " + "2629/3/933 2629/3/1032 2631/3/927 2639/3/994 2640/3/1015 2643/3/973 2652/3/996 2666/3/908 " + "2672/3/998 2673/3/937 2673/3/973 2677/3/924 2678/3/1008 2681/3/947 2684/3/953 2689/3/1028 " + "2689/3/1095 2691/3/1099 2697/3/1146 2698/3/951 2698/3/980 2701/3/1176 2706/3/1067 2713/3/928 " + "2717/3/999 2719/3/1086 2720/3/1034 2726/3/997 2739/3/1084 2743/3/1019 2744/3/1046 2763/3/941 " + "2764/3/1036 2774/3/999 2777/3/1099 2780/3/1118 2789/3/1126 2792/3/957 2793/3/1109 2800/3/1000 " + "2813/3/957 2814/3/1015 2843/3/1048 2850/3/1139 2858/3/1044 2866/3/1119 2867/3/1048 2869/3/1109 " + "2870/3/1273 2874/3/1082 2886/3/1295 2886/3/1501 2890/3/1021 2910/3/1049 2918/3/1154 2928/3/1067 " + "2942/3/999 2946/3/1099 2956/3/1129 2987/3/1150 3007/3/1080 3027/4/831 3031/3/1295 3047/3/1100 " + "3053/4/850 3068/3/1145 3083/3/1115 3174/4/828 3178/3/1139 3185/4/855 3234/3/1185 3281/4/925 " + "3283/4/1021 3290/4/876 3323/4/1034 3325/3/1298 3336/3/1237 3342/4/1116 3392/4/1032 3401/4/908 " + "3423/4/991 3448/4/1040 3470/4/1129 3480/4/1065 3488/4/938 3503/4/1074 3536/4/1038 3538/4/1070 " + "3549/4/969 3551/4/1081 3555/4/1023 3555/4/1200 3557/4/988 3568/4/1003 3658/4/1301 3672/4/1042 " + "3690/4/1170 3694/4/986 3779/4/1236 3815/4/1070 3822/4/1136 3849/4/1066 3895/5/929 3931/5/891 " + "4001/4/1198 4060/5/848 4077/5/910 4093/4/1150 4137/5/993 4224/5/957 4338/5/1100 4358/4/1217 " + "4419/5/1044 4425/5/1087 4427/5/962 4509/5/1141 4663/5/1226 5222/6/964 8044/10/1320 8121/10/1087 " + "8159/9/1165 8192/1/8192 " + ), + 32: ( + "1/1/1 2/1/2 128/1/128 256/2/128 384/3/128 512/4/128 549/1/549 586/1/586 590/1/590 597/1/597 " + "599/1/599 609/1/609 612/1/612 619/1/619 623/1/623 626/1/626 631/1/631 632/1/632 635/1/635 " + "637/1/637 640/5/128 656/1/656 661/1/661 662/1/662 663/1/663 667/1/667 668/1/668 670/1/670 " + "676/1/676 680/1/680 683/1/683 684/1/684 685/1/685 688/1/688 691/1/691 696/1/696 697/1/697 " + "698/1/698 699/1/699 700/1/700 701/1/701 702/1/702 704/1/704 707/1/707 711/1/711 712/1/712 " + "713/1/713 714/1/714 719/1/719 720/1/720 721/1/721 722/1/722 723/1/723 724/1/724 725/1/725 " + "726/1/726 727/1/727 728/1/728 729/1/729 730/1/730 731/1/731 733/1/733 735/1/735 736/1/736 " + "737/1/737 739/1/739 740/1/740 741/1/741 742/1/742 743/1/743 745/1/745 746/1/746 747/1/747 " + "748/1/748 749/1/749 751/1/751 752/1/752 753/1/753 757/1/757 758/1/758 759/1/759 761/1/761 " + "762/1/762 763/1/763 764/1/764 768/1/768 768/6/128 769/1/769 770/1/770 771/1/771 773/1/773 " + "774/1/774 775/1/775 776/1/776 777/1/777 778/1/778 780/1/780 781/1/781 782/1/782 784/1/784 " + "785/1/785 787/1/787 790/1/790 792/1/792 794/1/794 795/1/795 796/1/796 797/1/797 798/1/798 " + "799/1/799 801/1/801 802/1/802 804/1/804 805/1/805 806/1/806 807/1/807 811/1/811 812/1/812 " + "815/1/815 817/1/817 818/1/818 819/1/819 820/1/820 822/1/822 823/1/823 824/1/824 825/1/825 " + "826/1/826 827/1/827 828/1/828 829/1/829 830/1/830 831/1/831 832/1/832 833/1/833 834/1/834 " + "835/1/835 836/1/836 837/1/837 838/1/838 839/1/839 842/1/842 844/1/844 845/1/845 846/1/846 " + "847/1/847 848/1/848 849/1/849 850/1/850 851/1/851 852/1/852 853/1/853 854/1/854 855/1/855 " + "856/1/856 857/1/857 858/1/858 859/1/859 860/1/860 861/1/861 862/1/862 864/1/864 868/1/868 " + "869/1/869 870/1/870 871/1/871 872/1/872 873/1/873 874/1/874 875/1/875 876/1/876 877/1/877 " + "878/1/878 879/1/879 880/1/880 881/1/881 882/1/882 883/1/883 884/1/884 887/1/887 888/1/888 " + "889/1/889 890/1/890 891/1/891 892/1/892 893/1/893 894/1/894 895/1/895 896/1/896 898/1/898 " + "900/1/900 901/1/901 903/1/903 906/1/906 907/1/907 908/1/908 909/1/909 910/1/910 911/1/911 " + "912/1/912 913/1/913 914/1/914 915/1/915 916/1/916 918/1/918 919/1/919 920/1/920 921/1/921 " + "922/1/922 923/1/923 924/1/924 926/1/926 927/1/927 928/1/928 929/1/929 930/1/930 931/1/931 " + "933/1/933 934/1/934 935/1/935 936/1/936 938/1/938 939/1/939 940/1/940 941/1/941 942/1/942 " + "944/1/944 945/1/945 946/1/946 947/1/947 949/1/949 950/1/950 951/1/951 952/1/952 953/1/953 " + "955/1/955 957/1/957 958/1/958 959/1/959 963/1/963 965/1/965 968/1/968 969/1/969 970/1/970 " + "973/1/973 977/1/977 979/1/979 980/1/980 981/1/981 982/1/982 983/1/983 984/1/984 987/1/987 " + "988/1/988 990/1/990 991/1/991 992/1/992 994/1/994 997/1/997 998/1/998 999/1/999 1000/1/1000 " + "1003/1/1003 1004/1/1004 1005/1/1005 1008/1/1008 1011/1/1011 1015/1/1015 1016/1/1016 1018/1/1018 " + "1019/1/1019 1020/1/1020 1021/1/1021 1024/1/1024 1025/1/1025 1031/1/1031 1034/1/1034 1036/1/1036 " + "1038/1/1038 1039/1/1039 1040/1/1040 1043/1/1043 1045/1/1045 1048/1/1048 1050/1/1050 1055/1/1055 " + "1057/1/1057 1058/1/1058 1064/1/1064 1065/1/1065 1067/1/1067 1068/1/1068 1070/1/1070 1074/1/1074 " + "1076/1/1076 1079/1/1079 1080/1/1080 1081/1/1081 1082/1/1082 1084/1/1084 1085/1/1085 1089/1/1089 " + "1091/1/1091 1093/1/1093 1095/1/1095 1097/1/1097 1099/1/1099 1100/1/1100 1102/1/1102 1106/1/1106 " + "1109/1/1109 1112/1/1112 1120/1/1120 1122/1/1122 1129/1/1129 1130/1/1130 1139/1/1139 1140/1/1140 " + "1141/1/1141 1143/1/1143 1150/1/1150 1152/9/128 1154/1/1154 1170/1/1170 1171/1/1171 1172/1/1172 " + "1177/1/1177 1182/1/1182 1187/1/1187 1198/1/1198 1200/1/1200 1204/1/1204 1211/1/1211 1217/1/1217 " + "1236/1/1236 1237/1/1237 1280/10/128 1295/1/1295 1374/2/706 1374/2/721 1374/2/768 1388/2/757 " + "1421/2/765 1423/2/734 1447/2/751 1448/2/756 1460/2/780 1464/2/763 1472/2/742 1477/2/763 " + "1481/2/900 1483/2/827 1485/2/835 1487/2/798 1491/2/817 1493/2/850 1506/2/944 1508/2/817 " + "1513/2/788 1513/2/861 1515/2/836 1517/2/841 1524/2/825 1526/2/807 1526/2/826 1531/2/799 " + "1535/2/848 1542/2/778 1543/2/854 1543/2/914 1544/2/893 1546/2/863 1549/2/793 1560/2/938 " + "1562/2/841 1564/2/898 1567/2/788 1569/2/888 1570/2/786 1571/2/831 1572/2/874 1577/2/809 " + "1579/2/796 1581/2/950 1582/2/824 1587/2/797 1588/2/910 1590/2/889 1594/2/875 1594/2/967 " + "1595/2/848 1596/2/805 1596/2/863 1597/2/831 1599/2/859 1601/2/848 1603/2/803 1603/2/877 " + "1606/2/836 1606/2/914 1614/2/858 1615/2/837 1620/2/840 1620/2/926 1622/2/820 1627/2/834 " + "1627/2/931 1630/2/875 1631/2/867 1633/2/827 1634/2/826 1636/2/896 1640/2/855 1646/2/944 " + "1647/2/859 1648/2/913 1651/2/917 1651/2/1007 1656/2/993 1661/2/942 1663/2/866 1664/2/897 " + "1666/2/943 1667/2/914 1668/2/1036 1670/2/836 1670/2/900 1671/2/855 1671/2/965 1672/2/841 " + "1680/2/889 1680/2/953 1685/2/911 1688/2/848 1689/2/965 1690/2/986 1696/2/970 1697/2/896 " + "1697/2/921 1698/2/886 1699/2/929 1704/2/879 1704/2/895 1705/2/856 1710/2/1008 1716/2/912 " + "1717/2/898 1719/2/872 1725/2/925 1727/2/876 1727/2/965 1731/2/903 1734/2/924 1734/2/929 " + "1735/2/992 1736/2/920 1738/2/954 1742/2/893 1742/2/961 1743/2/964 1745/2/921 1745/2/937 " + "1749/2/901 1752/2/1087 1757/2/986 1758/2/1098 1760/2/996 1765/2/926 1768/2/900 1772/2/947 " + "1772/2/957 1773/2/993 1774/2/894 1775/2/1119 1776/2/893 1778/2/920 1781/2/976 1782/2/1040 " + "1782/2/1081 1789/2/994 1790/2/943 1791/2/1031 1791/2/1048 1792/2/922 1792/2/939 1792/2/1021 " + "1794/2/1076 1795/2/1022 1799/2/965 1801/2/912 1807/2/1032 1813/2/982 1814/2/926 1816/2/950 " + "1817/2/912 1817/2/999 1819/2/912 1822/2/942 1824/2/963 1827/2/967 1828/2/962 1829/2/940 " + "1830/2/916 1845/2/988 1847/2/955 1855/2/939 1856/2/977 1858/2/991 1863/2/936 1865/2/972 " + "1865/2/983 1869/2/1003 1871/2/1026 1874/2/1143 1875/2/958 1876/2/1201 1881/2/1126 1882/2/1008 " + "1884/2/1104 1889/2/1101 1890/2/1038 1902/2/1067 1910/2/1011 1911/2/1066 1920/2/1081 1921/2/1003 " + "1922/2/1044 1926/2/1099 1928/2/974 1934/2/1046 1938/2/1120 1950/2/984 1951/2/1122 1952/2/1008 " + "1959/2/1034 1966/2/1000 1979/2/1053 1983/2/1149 1984/2/1094 1988/2/1011 1993/2/1014 1997/2/1032 " + "1998/2/999 2004/2/1049 2007/2/1176 2009/2/1170 2013/2/1043 2015/2/1100 2023/2/1200 2037/2/1226 " + "2039/2/1139 2042/2/1080 2051/2/1057 2060/3/729 2062/2/1146 2067/2/1143 2102/2/1141 2155/2/1100 " + "2169/2/1261 2176/2/1090 2179/2/1136 2179/2/1273 2221/3/785 2236/3/846 2260/2/1145 2270/3/854 " + "2284/3/876 2289/3/858 2321/3/810 2324/3/994 2333/3/906 2339/3/801 2369/3/877 2392/3/824 " + "2392/3/919 2394/3/904 2410/3/908 2415/3/848 2449/3/883 2452/3/844 2479/3/887 2482/3/998 " + "2491/3/992 2499/3/908 2499/3/918 2502/3/882 2505/3/958 2517/3/1019 2527/3/877 2530/3/930 " + "2550/3/857 2556/3/966 2560/3/887 2584/3/974 2584/3/1028 2595/3/891 2621/3/1007 2629/3/933 " + "2629/3/964 2632/3/1031 2634/3/1000 2644/3/969 2650/3/968 2664/3/1137 2673/3/1023 2677/3/924 " + "2677/3/980 2687/3/1122 2692/3/994 2698/3/980 2710/3/1067 2757/3/932 2769/3/1013 2780/3/987 " + "2780/3/1118 2803/3/1002 2816/3/1062 2893/3/1092 2912/3/1111 2932/4/975 2943/3/1125 2968/3/1115 " + "3008/3/1099 3008/3/1157 3040/3/1091 3103/4/1018 3110/4/801 3158/3/1198 3209/3/1185 3210/3/1143 " + "3241/4/862 3254/4/1019 3326/4/910 3334/4/881 3334/4/885 3369/4/910 3374/4/1046 3407/3/1298 " + "3496/4/1094 3515/4/996 3538/4/963 3646/4/1178 3652/4/1050 3659/4/989 3666/4/953 4016/5/954 " + "4117/4/1128 4136/5/906 4210/4/1301 5615/6/1501 8040/10/993 8044/10/1320 8069/11/1141 8119/10/1116 " + "8121/10/1087 8159/9/1165 8192/1/8192 " + ), +} + +# Kernel map of the context call on sm_103 (152 SMs; the pinned FlashInfer's trtllm-gen cubins), from a +# profiler probe of 94,776 calls: every max_q_len 1..8192 at one request, every (max_q_len, B) with +# max_q_len * B <= 2048 for B 2..32 (at two totals T each), and 3,000 random partitions per (H, +# width). The kernel(s) of a call are a function of (H, width, B, max_q_len) alone, never of T. +# VARLEN_BANDS[(H, width)][B] = ((lowest max_q_len, highest max_q_len, VARLEN_KERNELS index), ...); +# the last band of every B runs to the envelope's top, 8192 - (B - 1). Names drop the common prefix +# fmhaSm100fKernel_QkvBfloat16OBfloat16H512. +VARLEN_KERNELS = ( + ( + "HVPerCta128PagedKvDenseDynamicTokenSparseP1MultiCtasKvGmemSepVarSeqQ64Kv128StaticKeepsAbForGen", + "fmhaReductionKernel<64, 128, false, __nv_bfloat16, __nv_bfloat16>", + ), + ("HVPerCta128PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ16Kv128StaticSwapsAbForGen",), + ("HVPerCta128PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ8Kv128StaticSwapsAbForGen",), + ("HVPerCta128PagedKvDenseDynamicTokenSparseP1VarSeqQ16Kv128PersistentSwapsAbForGen",), + ("HVPerCta128PagedKvDenseDynamicTokenSparseP1VarSeqQ64Kv128PersistentKeepsAbForGen",), + ("HVPerCta128PagedKvDenseDynamicTokenSparseP1VarSeqQ8Kv128PersistentSwapsAbForGen",), + ( + "HVPerCta256PagedKvDenseDynamicTokenSparseP1MultiCtasKvGmemSepVarSeqQ64Kv128StaticKeepsAbForGen", + "fmhaReductionKernel<64, 256, false, __nv_bfloat16, __nv_bfloat16>", + ), + ("HVPerCta256PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ16Kv128StaticSwapsAbForGen",), + ("HVPerCta256PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ8Kv128StaticSwapsAbForGen",), + ("HVPerCta256PagedKvDenseDynamicTokenSparseP1VarSeqQ16Kv128PersistentSwapsAbForGen",), + ("HVPerCta256PagedKvDenseDynamicTokenSparseP1VarSeqQ64Kv128PersistentKeepsAbForGen",), + ("HVPerCta256PagedKvDenseDynamicTokenSparseP1VarSeqQ8Kv128PersistentSwapsAbForGen",), + ( + "PagedKvDenseDynamicTokenSparseP1MultiCtasKvGmemSepVarSeqQ64Kv128StaticKeepsAbForGen", + "fmhaReductionKernel<64, 512, false, __nv_bfloat16, __nv_bfloat16>", + ), + ("PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ16Kv128StaticSwapsAbForGen",), + ("PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ8Kv128StaticSwapsAbForGen",), + ("PagedKvDenseDynamicTokenSparseP1VarSeqQ16Kv128PersistentSwapsAbForGen",), + ("PagedKvDenseDynamicTokenSparseP1VarSeqQ64Kv128PersistentKeepsAbForGen",), + ("PagedKvDenseDynamicTokenSparseP1VarSeqQ8Kv128PersistentSwapsAbForGen",), +) +VARLEN_BANDS = { + (16, 128): { + 1: ((1, 19, 5), (20, 38, 11), (39, 8192, 17)), + 2: ((1, 9, 5), (10, 19, 11), (20, 8191, 17)), + 3: ((1, 6, 5), (7, 12, 11), (13, 8190, 17)), + 4: ((1, 4, 5), (5, 9, 11), (10, 8189, 17)), + 5: ((1, 3, 5), (4, 7, 11), (8, 8188, 17)), + 6: ((1, 3, 5), (4, 6, 11), (7, 8187, 17)), + 7: ((1, 2, 5), (3, 5, 11), (6, 8186, 17)), + 8: ((1, 2, 5), (3, 4, 11), (5, 8185, 17)), + 9: ((1, 2, 5), (3, 4, 11), (5, 8184, 17)), + 10: ((1, 1, 5), (2, 3, 11), (4, 8183, 17)), + 11: ((1, 1, 5), (2, 3, 11), (4, 8182, 17)), + 12: ((1, 1, 5), (2, 3, 11), (4, 8181, 17)), + 13: ((1, 1, 5), (2, 2, 11), (3, 8180, 17)), + 14: ((1, 1, 5), (2, 2, 11), (3, 8179, 17)), + 15: ((1, 1, 5), (2, 2, 11), (3, 8178, 17)), + 16: ((1, 1, 5), (2, 2, 11), (3, 8177, 17)), + 17: ((1, 1, 5), (2, 2, 11), (3, 8176, 17)), + 18: ((1, 1, 5), (2, 2, 11), (3, 8175, 17)), + 19: ((1, 1, 5), (2, 2, 11), (3, 8174, 17)), + 20: ((1, 1, 3), (2, 3, 9), (4, 8173, 15)), + 21: ((1, 1, 3), (2, 3, 9), (4, 8172, 15)), + 22: ((1, 1, 3), (2, 3, 9), (4, 8171, 15)), + 23: ((1, 1, 3), (2, 3, 9), (4, 8170, 15)), + 24: ((1, 1, 3), (2, 3, 9), (4, 8169, 15)), + 25: ((1, 1, 3), (2, 3, 9), (4, 8168, 15)), + 26: ((1, 1, 3), (2, 2, 9), (3, 8167, 15)), + 27: ((1, 1, 3), (2, 2, 9), (3, 8166, 15)), + 28: ((1, 1, 3), (2, 2, 9), (3, 8165, 15)), + 29: ((1, 1, 3), (2, 2, 9), (3, 8164, 15)), + 30: ((1, 1, 3), (2, 2, 9), (3, 8163, 15)), + 31: ((1, 1, 3), (2, 2, 9), (3, 8162, 15)), + 32: ((1, 1, 3), (2, 2, 9), (3, 8161, 15)), + }, + (16, 640): { + 1: ((1, 6, 2), (7, 12, 8), (13, 38, 14), (39, 8192, 17)), + 2: ((1, 3, 2), (4, 6, 8), (7, 19, 14), (20, 8191, 17)), + 3: ((1, 2, 2), (3, 4, 8), (5, 12, 14), (13, 8190, 17)), + 4: ((1, 3, 1), (4, 6, 7), (7, 19, 13), (20, 8189, 15)), + 5: ((1, 2, 1), (3, 5, 7), (6, 15, 13), (16, 8188, 15)), + 6: ((1, 2, 1), (3, 4, 7), (5, 12, 13), (13, 8187, 15)), + 7: ((1, 1, 1), (2, 3, 7), (4, 10, 13), (11, 8186, 15)), + 8: ((1, 1, 1), (2, 3, 7), (4, 9, 13), (10, 8185, 15)), + 9: ((1, 1, 1), (2, 2, 7), (3, 8, 13), (9, 8184, 15)), + 10: ((1, 1, 1), (2, 2, 7), (3, 7, 13), (8, 8183, 15)), + 11: ((1, 1, 1), (2, 2, 7), (3, 6, 13), (7, 8182, 15)), + 12: ((1, 1, 1), (2, 2, 7), (3, 6, 13), (7, 8181, 15)), + 13: ((1, 1, 7), (2, 5, 13), (6, 8180, 15)), + 14: ((1, 1, 7), (2, 5, 13), (6, 8179, 15)), + 15: ((1, 1, 7), (2, 5, 13), (6, 8178, 15)), + 16: ((1, 1, 7), (2, 4, 13), (5, 8177, 15)), + 17: ((1, 1, 7), (2, 4, 13), (5, 8176, 15)), + 18: ((1, 1, 7), (2, 4, 13), (5, 8175, 15)), + 19: ((1, 1, 7), (2, 4, 13), (5, 8174, 15)), + 20: ((1, 1, 7), (2, 3, 13), (4, 8173, 15)), + 21: ((1, 1, 7), (2, 3, 13), (4, 8172, 15)), + 22: ((1, 1, 7), (2, 3, 13), (4, 8171, 15)), + 23: ((1, 1, 7), (2, 3, 13), (4, 8170, 15)), + 24: ((1, 1, 7), (2, 3, 13), (4, 8169, 15)), + 25: ((1, 1, 7), (2, 3, 13), (4, 8168, 15)), + 26: ((1, 2, 13), (3, 8167, 15)), + 27: ((1, 2, 13), (3, 8166, 15)), + 28: ((1, 2, 13), (3, 8165, 15)), + 29: ((1, 2, 13), (3, 8164, 15)), + 30: ((1, 2, 13), (3, 8163, 15)), + 31: ((1, 2, 13), (3, 8162, 15)), + 32: ((1, 2, 13), (3, 8161, 15)), + }, + (64, 128): { + 1: ((1, 38, 4), (39, 76, 10), (77, 8192, 16)), + 2: ((1, 19, 4), (20, 38, 10), (39, 8191, 16)), + 3: ((1, 12, 4), (13, 25, 10), (26, 8190, 16)), + 4: ((1, 9, 4), (10, 19, 10), (20, 8189, 16)), + 5: ((1, 7, 4), (8, 15, 10), (16, 8188, 16)), + 6: ((1, 6, 4), (7, 12, 10), (13, 8187, 16)), + 7: ((1, 5, 4), (6, 10, 10), (11, 8186, 16)), + 8: ((1, 4, 4), (5, 9, 10), (10, 8185, 16)), + 9: ((1, 4, 4), (5, 8, 10), (9, 8184, 16)), + 10: ((1, 3, 4), (4, 7, 10), (8, 8183, 16)), + 11: ((1, 3, 4), (4, 6, 10), (7, 8182, 16)), + 12: ((1, 3, 4), (4, 6, 10), (7, 8181, 16)), + 13: ((1, 2, 4), (3, 5, 10), (6, 8180, 16)), + 14: ((1, 2, 4), (3, 5, 10), (6, 8179, 16)), + 15: ((1, 2, 4), (3, 5, 10), (6, 8178, 16)), + 16: ((1, 2, 4), (3, 4, 10), (5, 8177, 16)), + 17: ((1, 2, 4), (3, 4, 10), (5, 8176, 16)), + 18: ((1, 2, 4), (3, 4, 10), (5, 8175, 16)), + 19: ((1, 2, 4), (3, 4, 10), (5, 8174, 16)), + 20: ((1, 1, 4), (2, 3, 10), (4, 8173, 16)), + 21: ((1, 1, 4), (2, 3, 10), (4, 8172, 16)), + 22: ((1, 1, 4), (2, 3, 10), (4, 8171, 16)), + 23: ((1, 1, 4), (2, 3, 10), (4, 8170, 16)), + 24: ((1, 1, 4), (2, 3, 10), (4, 8169, 16)), + 25: ((1, 1, 4), (2, 3, 10), (4, 8168, 16)), + 26: ((1, 1, 4), (2, 2, 10), (3, 8167, 16)), + 27: ((1, 1, 4), (2, 2, 10), (3, 8166, 16)), + 28: ((1, 1, 4), (2, 2, 10), (3, 8165, 16)), + 29: ((1, 1, 4), (2, 2, 10), (3, 8164, 16)), + 30: ((1, 1, 4), (2, 2, 10), (3, 8163, 16)), + 31: ((1, 1, 4), (2, 2, 10), (3, 8162, 16)), + 32: ((1, 1, 4), (2, 2, 10), (3, 8161, 16)), + }, + (64, 640): { + 1: ((1, 12, 0), (13, 25, 6), (26, 76, 12), (77, 8192, 16)), + 2: ((1, 6, 0), (7, 12, 6), (13, 38, 12), (39, 8191, 16)), + 3: ((1, 4, 0), (5, 8, 6), (9, 25, 12), (26, 8190, 16)), + 4: ((1, 3, 0), (4, 6, 6), (7, 19, 12), (20, 8189, 16)), + 5: ((1, 2, 0), (3, 5, 6), (6, 15, 12), (16, 8188, 16)), + 6: ((1, 2, 0), (3, 4, 6), (5, 12, 12), (13, 8187, 16)), + 7: ((1, 1, 0), (2, 3, 6), (4, 10, 12), (11, 8186, 16)), + 8: ((1, 1, 0), (2, 3, 6), (4, 9, 12), (10, 8185, 16)), + 9: ((1, 1, 0), (2, 2, 6), (3, 8, 12), (9, 8184, 16)), + 10: ((1, 1, 0), (2, 2, 6), (3, 7, 12), (8, 8183, 16)), + 11: ((1, 1, 0), (2, 2, 6), (3, 6, 12), (7, 8182, 16)), + 12: ((1, 1, 0), (2, 2, 6), (3, 6, 12), (7, 8181, 16)), + 13: ((1, 1, 6), (2, 5, 12), (6, 8180, 16)), + 14: ((1, 1, 6), (2, 5, 12), (6, 8179, 16)), + 15: ((1, 1, 6), (2, 5, 12), (6, 8178, 16)), + 16: ((1, 1, 6), (2, 4, 12), (5, 8177, 16)), + 17: ((1, 1, 6), (2, 4, 12), (5, 8176, 16)), + 18: ((1, 1, 6), (2, 4, 12), (5, 8175, 16)), + 19: ((1, 1, 6), (2, 4, 12), (5, 8174, 16)), + 20: ((1, 1, 6), (2, 3, 12), (4, 8173, 16)), + 21: ((1, 1, 6), (2, 3, 12), (4, 8172, 16)), + 22: ((1, 1, 6), (2, 3, 12), (4, 8171, 16)), + 23: ((1, 1, 6), (2, 3, 12), (4, 8170, 16)), + 24: ((1, 1, 6), (2, 3, 12), (4, 8169, 16)), + 25: ((1, 1, 6), (2, 3, 12), (4, 8168, 16)), + 26: ((1, 2, 12), (3, 8167, 16)), + 27: ((1, 2, 12), (3, 8166, 16)), + 28: ((1, 2, 12), (3, 8165, 16)), + 29: ((1, 2, 12), (3, 8164, 16)), + 30: ((1, 2, 12), (3, 8163, 16)), + 31: ((1, 2, 12), (3, 8162, 16)), + 32: ((1, 2, 12), (3, 8161, 16)), + }, +} + + +def _variant_of(names) -> tuple: + """The trtllm-gen kernels of one call, as recorded in VARLEN_KERNELS.""" + out = [] + for n in names: + if n.startswith("fmhaSm100fKernel_QkvBfloat16OBfloat16H512"): + out.append(n[len("fmhaSm100fKernel_QkvBfloat16OBfloat16H512") :]) + elif "fmhaReductionKernel<" in n: + out.append(re.search(r"fmhaReductionKernel<[^>]*>", n).group(0)) + return tuple(out) + + +@pytest.fixture(scope="module") +def env_pools(): + """Random pools large enough for 32 requests of 16384 positions without any aliasing of a + request's compressed rows; every request reaches its rows through a random page table.""" + gen = torch.Generator(device=DEV).manual_seed(4242) + swa = torch.randn(ENV_SWA_PAGES, 1, SWA_PAGE, HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) + pools = {"swa": swa, "comp": {}, "atol": {}} + for layer, (ratio, page) in ENV_LAYERS.items(): + if ratio: + comp = torch.randn( + ENV_COMP_PAGES[ratio], 1, page, HEAD_DIM, generator=gen, device=DEV + ).to(torch.bfloat16) + pools["atol"][layer] = _atol(swa, comp) + else: + # The core target's ratio-0 placeholder shape; NaN, so that any read of it would show. + comp = torch.full((1, 1, 64, HEAD_DIM), float("nan"), dtype=torch.bfloat16, device=DEV) + pools["atol"][layer] = _atol(swa) + pools["comp"][layer] = comp + yield pools + del pools + torch.cuda.empty_cache() + + +def _env_tables(requests, gen, ratio, page, swa_pages, comp_pages, cap_pages=None): + """Random page tables: a 2-page SWA ring per request, and enough compressed pages per request + for every compressed position it can see (distinct rows). With cap_pages, a request that can see + more compressed positions than cap_pages pages hold gets cap_pages pages used as a ring + (position e on page (e // page) % cap_pages, so far positions alias nearer rows).""" + b = len(requests) + swa_bt = torch.randperm(swa_pages, generator=gen, device=DEV)[: 2 * b].view(b, 2) + comp_bts = None + if ratio: + need = [max(1, -(-((c + q) // ratio) // page)) for c, q in requests] + if cap_pages is not None: + need = [min(n, cap_pages) for n in need] + perm = torch.randperm(comp_pages, generator=gen, device=DEV) + assert sum(need) <= comp_pages + starts = [0] + for n in need: + starts.append(starts[-1] + n) + comp_bts = [perm[starts[i] : starts[i + 1]] for i in range(b)] + return swa_bt, comp_bts + + +def _env_rows(requests, gen, width, ratio, comp_page, swa_bt, comp_bts, swa_page=SWA_PAGE): + """Vectorized _build_rows (same layout rules): query j of request (cached, q_len) sits at + pos = cached + j, sees window positions max(0, pos-127)..pos in order (-1 after them), and, + for width 640, min(512, (pos+1)//ratio) compressed positions (all of them while <= 512, else a + random sorted subset), compacted at the front. Returns idx [N, width] int32, active [N] int32, + seq_lens [B] int32, cu [B+1] int32 and the window counts [N] (min(128, pos + 1)). + A request's compressed position e lives on page (e // comp_page) % (its page count): distinct + rows while its pages hold every position it can see, a ring over them otherwise (_env_tables' + cap_pages). Past 65536 visible positions the subset is drawn one position per stratum + [s * v // 512, (s + 1) * v // 512) (distinct and sorted), not by a [rows, v] score matrix.""" + ar = torch.arange(WINDOW, device=DEV) + idx_l, act_l, cnt_l = [], [], [] + for b, (cached, q_len) in enumerate(requests): + pos = torch.arange(cached, cached + q_len, device=DEV) + wpos = (pos - (WINDOW - 1)).clamp_min(0)[:, None] + ar[None] + ring = swa_bt[b] + widx = ring[(wpos // swa_page) % ring.numel()] * swa_page + wpos % swa_page + widx = torch.where(wpos <= pos[:, None], widx, -1).int() + cnt_l.append((pos + 1).clamp(max=WINDOW)) + if width == WINDOW: + idx_l.append(widx) + act_l.append(torch.full((q_len,), WINDOW, dtype=torch.int32, device=DEV)) + continue + k_max = width - WINDOW + vis = (pos + 1) // ratio + vmax = (cached + q_len) // ratio + pages = comp_bts[b] + sels, valids = [], [] + for s0 in range(0, q_len, 1024): + v = vis[s0 : s0 + 1024] + if vmax <= k_max: + sel = torch.arange(k_max, device=DEV)[None].expand(v.numel(), k_max) + valid = sel < v[:, None] + elif vmax > 65536: + strata = torch.arange(k_max + 1, device=DEV) + lo = (strata[None, :-1] * v[:, None]) // k_max + hi = (strata[None, 1:] * v[:, None]) // k_max + u = torch.rand(v.numel(), k_max, generator=gen, device=DEV) + strat = lo + (u * (hi - lo)).long() + sel = torch.where(v[:, None] > k_max, strat, torch.arange(k_max, device=DEV)[None]) + valid = sel < v[:, None] + else: + score = torch.rand(v.numel(), vmax, generator=gen, device=DEV) + score.masked_fill_(torch.arange(vmax, device=DEV)[None] >= v[:, None], -1.0) + val, sel = score.topk(k_max, dim=1) + sel = torch.where(val >= 0, sel, vmax).sort(dim=1).values + valid = sel < vmax + sels.append(sel) + valids.append(valid) + sel, valid = torch.cat(sels), torch.cat(valids) + sel = torch.where(valid, sel, 0) + cpage = pages[(sel // comp_page) % pages.numel()] + cidx = torch.where(valid, cpage * comp_page + sel % comp_page, -1).int() + idx_l.append(torch.cat([widx, cidx], 1)) + act_l.append((WINDOW + valid.sum(1)).int()) + idx = torch.cat(idx_l).contiguous() + active = torch.cat(act_l) + seq_lens = torch.tensor([c + q for c, q in requests], dtype=torch.int32, device=DEV) + cu = torch.tensor([0] + [q for _, q in requests], dtype=torch.int32, device=DEV).cumsum(0).int() + return idx, active, seq_lens, cu, torch.cat(cnt_l) + + +class _Step: + """One engine step: nc context requests (their rows first), then G generation requests of + next_n rows each (next_n 1: decode; 6: the DSpark target verify). + + The per-row and per-request tensors are prefixes of buffers sized for the envelope, as the + runtime's metadata buffers are; the two calls get the target's views of them.""" + + def __init__( + self, + ctx, + dec_positions, + num_heads, + layer, + pools, + gen, + swa_bt=None, + comp_bts=None, + next_n=1, + cap_pages=None, + ): + ratio, page = ENV_LAYERS[layer] + self.layer = layer + self.width = WINDOW + (INDEX_TOPK if ratio else 0) + self.next_n = next_n + self.requests = list(ctx) + [(p, next_n) for p in dec_positions] + self.nc, self.tc = len(ctx), sum(q for _, q in ctx) + self.B, self.G = len(self.requests), len(dec_positions) + self.max_ctx_q = max((q for _, q in ctx), default=0) + self.swa, self.comp = pools["swa"], pools["comp"][layer] + if swa_bt is None: + swa_bt, comp_bts = _env_tables( + self.requests, gen, ratio, page, self.swa.shape[0], self.comp.shape[0], cap_pages + ) + idx, active, seq_lens, cu, self.counts = _env_rows( + self.requests, gen, self.width, ratio, page, swa_bt, comp_bts + ) + self.T = idx.shape[0] + self.idx = idx + self.active = torch.full((max(self.T, ENV_T),), -7, dtype=torch.int32, device=DEV)[: self.T] + self.active.copy_(active) + self.kv_lens = torch.full((max(self.B, ENV_B) + 1,), -7, dtype=torch.int32, device=DEV) + self.kv_lens[: self.B] = seq_lens + self.cu = torch.full((max(self.B, ENV_B) + 2,), -7, dtype=torch.int32, device=DEV) + self.cu[: self.B + 1] = cu + self.q3 = ( + torch.randn(self.T, num_heads, HEAD_DIM, generator=gen, device=DEV) + .mul_(2) + .to(torch.bfloat16) + ) + self.sinks = _sinks(num_heads) + + def run(self, ws, out=None, swa=None, comp=None, idx=None): + """The target's two calls; swa / comp / idx replace the container pools and index (a native + read form of the same step).""" + tc, gen_rows, h, nn = self.tc, self.T - self.tc, self.q3.shape[1], self.next_n + swa = self.swa if swa is None else swa + comp = self.comp if comp is None else comp + idx = self.idx if idx is None else idx + o = ( + torch.empty(self.T, h, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + if out is None + else out + ) + if tc: + op( + torch.split(self.q3, [tc, gen_rows], 0)[0], + swa, + comp, + ws, + torch.split(idx, [tc, gen_rows], 0)[0], + torch.split(self.active, [tc, gen_rows], 0)[0], + self.kv_lens[: self.nc], + self.sinks, + SCALE, + 1.0, + cum_seq_lens_q=self.cu[: self.nc + 1], + max_q_len=self.max_ctx_q, + out=torch.split(o, [tc, gen_rows], 0)[0], + ) + if gen_rows: + g = self.G + op( + torch.split(self.q3, [tc, gen_rows], 0)[1].reshape(g, nn, h, HEAD_DIM), + swa, + comp, + ws, + torch.split(idx, [tc, gen_rows], 0)[1], + torch.split(self.active, [tc, gen_rows], 0)[1], + self.kv_lens[self.nc : self.B], + self.sinks, + SCALE, + 1.0, + out=torch.split(o, [tc, gen_rows], 0)[1].reshape(g, nn, h, HEAD_DIM), + ) + return o + + def reference(self, idx=None): + return _reference( + self.q3, + self.swa, + self.comp, + self.idx if idx is None else idx, + self.active, + self.counts, + self.sinks, + ) + + +def _excess(o, ref, atol) -> float: + """max |o - ref| / (atol + RTOL * |ref|): <= 1 is inside the file's gate.""" + return ((o.float().reshape(ref.shape) - ref).abs_() / (atol + RTOL * ref.abs())).amax().item() + + +def _fresh_or_cached(rng, q_len): + return 0 if rng.random() < 0.4 else rng.randint(0, ENV_SEQ - q_len) + + +@pytest.mark.parametrize("cached", [0, 8192]) +@pytest.mark.parametrize("layer", ["r0", "r2"]) +@pytest.mark.parametrize("num_heads", [16, 64]) +def test_prefill_every_max_q_len_one_request(env_pools, num_heads, layer, cached): + """One context request of every q_len 1..8192 (max_q_len = T = q_len) against the fp32 + reference. Query row j's result does not depend on q_len (its window count is + min(128, cached + j + 1) and its index row is its own), so each call is held to the first + q_len rows of ONE reference computed for the 8192-row request.""" + gen = torch.Generator(device=DEV).manual_seed(31000 + num_heads + cached + ENV_LAYERS[layer][0]) + st = _Step([(cached, ENV_T)], [], num_heads, layer, env_pools, gen) + ref = st.reference() + den = env_pools["atol"][layer] + RTOL * ref.abs() + ws = _workspace() + out = torch.empty_like(st.q3) + worst = torch.zeros(ENV_T, device=DEV) + seq = torch.zeros(1, dtype=torch.int32, device=DEV) + cu = torch.zeros(2, dtype=torch.int32, device=DEV) + diff = torch.empty_like(ref) # one fp32 scratch buffer, so the sweep allocates nothing per call + for n in range(1, ENV_T + 1): + seq.fill_(cached + n) + cu[1].fill_(n) + op( + st.q3[:n], + st.swa, + st.comp, + ws, + st.idx[:n], + st.active[:n], + seq, + st.sinks, + SCALE, + 1.0, + cum_seq_lens_q=cu, + max_q_len=n, + out=out[:n], + ) + d = diff[:n] + d.copy_(out[:n]).sub_(ref[:n]).abs_().div_(den[:n]) + torch.amax(d, out=worst[n - 1]) + torch.cuda.synchronize() + bad = (worst > 1.0).nonzero().flatten() + print(f"every-q_len H{num_heads} {layer} cached {cached}: max excess {worst.max().item():.3f}") + assert bad.numel() == 0, ( + f"q_len {(bad + 1).tolist()[:20]} exceed the gate (max {worst.max().item():.2f})" + ) + assert torch.isfinite(out).all() + + +def _band_cases(num_heads, width, rng): + """Every kernel band of VARLEN_BANDS at every request count 1..32: its lowest, one interior and + its highest max_q_len. The other B-1 requests get random q_len in [1, max_q_len] within the + 8192-token step budget.""" + cases = [] + for b, bands in VARLEN_BANDS[(num_heads, width)].items(): + for lo, hi, vid in bands: + for mq in sorted({lo, (lo + hi) // 2, hi}): + budget = ENV_T - mq + others = [rng.randint(1, mq) for _ in range(b - 1)] + while sum(others) > budget: + i = max(range(b - 1), key=lambda k: others[k]) + others[i] = max(1, others[i] - max(1, (sum(others) - budget))) + q_lens = [mq] + others + rng.shuffle(q_lens) + cases.append(([(_fresh_or_cached(rng, q), q) for q in q_lens], vid)) + return cases + + +@pytest.mark.parametrize("width", [128, 640]) +@pytest.mark.parametrize("num_heads", [16, 64]) +def test_prefill_kernel_variant_bands(env_pools, num_heads, width): + """Kernel bands of the context call. The kernel is picked on the host from (H, width, number of + requests, max_q_len) (probe: 94,776 profiled calls, never two variants for one such cell, total + T irrelevant); VARLEN_BANDS is that map. Every band is run at its bottom, an interior point and + its top at every request count 1..32, with the kernels read back from the profiler, and every + call is held to the fp32 reference.""" + rng = random.Random(num_heads * 1000 + width) + gen = torch.Generator(device=DEV).manual_seed(num_heads * 1000 + width) + ws = _workspace() + cases = _band_cases(num_heads, width, rng) + seen, worst, wrong = [], 0.0, [] + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for i, (ctx, vid) in enumerate(cases): + layer = "r0" if width == WINDOW else ("r2", "r1")[i % 2] + st = _Step(ctx, [], num_heads, layer, env_pools, gen) + torch.cuda.synchronize() + torch.cuda._sleep(10) + o = st.run(ws) + torch.cuda._sleep(10) + e = _excess(o, st.reference(), env_pools["atol"][layer]) + worst = max(worst, e) + if e > 1.0: + wrong.append((ctx, e)) + torch.cuda.synchronize() + ev = sorted( + (e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA), + key=lambda e: e.time_range.start, + ) + open_call, cur = False, [] + for e in ev: + if "spin_kernel" in e.name or "sleep" in e.name: + if open_call: + seen.append(_variant_of(cur)) + open_call, cur = not open_call, [] + elif open_call: + cur.append(e.name) + assert len(seen) == len(cases), (len(seen), len(cases)) + mism = [(c, VARLEN_KERNELS[v], s) for (c, v), s in zip(cases, seen) if VARLEN_KERNELS[v] != s] + used = {v for _, v in cases} + print( + f"bands H{num_heads} W{width}: {len(cases)} calls, {len(used)} variants, max excess {worst:.3f}" + ) + assert not mism, mism[:3] + assert not wrong, wrong[:3] + + +def _partition_case(seed): + """A random step: k in 1..32 context requests over T in 1..8192 tokens (log-uniform or uniform), + random q_len split and order, fresh or cached prefix (cached + q_len <= 16384); for half the + seeds, G random decode rows after the context rows (k + G <= 32, T + G <= 8192).""" + rng = random.Random(900000 + seed) + num_heads = (16, 64)[seed % 2] + layer = ("r0", "r2", "r1")[(seed // 2) % 3] + k = rng.randint(1, ENV_B) + if rng.random() < 0.5: + t = int(round(math.exp(rng.uniform(0.0, math.log(ENV_T))))) + else: + t = rng.randint(1, ENV_T) + t = max(t, k) + cuts = sorted(rng.sample(range(1, t), k - 1)) if k > 1 else [] + q_lens = [b - a for a, b in zip([0] + cuts, cuts + [t])] + ctx = [(_fresh_or_cached(rng, q), q) for q in q_lens] + dec = [] + if seed % 4 >= 2 and k < ENV_B and t < ENV_T: + g = rng.randint(1, min(ENV_B - k, ENV_T - t)) + dec = [rng.randint(0, ENV_SEQ - 1) for _ in range(g)] + return num_heads, layer, ctx, dec + + +PARTITION_SEEDS = list(range(240)) + + +@pytest.mark.parametrize("seed", PARTITION_SEEDS) +def test_prefill_random_partition(env_pools, seed): + """240 seeded random partitions of the context envelope, half of them mixed steps whose + decode call reads the views starting at row tc, all against the fp32 reference.""" + num_heads, layer, ctx, dec = _partition_case(seed) + gen = torch.Generator(device=DEV).manual_seed(500000 + seed) + st = _Step(ctx, dec, num_heads, layer, env_pools, gen) + o = st.run(_workspace()) + e = _excess(o, st.reference(), env_pools["atol"][layer]) + print( + f"partition {seed}: H{num_heads} {layer} T {st.T} tc {st.tc} nc {st.nc} G {st.G} " + f"max_q_len {st.max_ctx_q} excess {e:.3f}" + ) + assert e <= 1.0, (seed, e) + + +def test_prefill_random_partitions_reach_the_envelope(): + """The 240 partitions span the envelope: k 1..32, T up to 8192, cached prefixes past 8192, and + mixed steps whose decode views start at odd rows / odd requests.""" + cases = [_partition_case(s) for s in PARTITION_SEEDS] + ks = {len(c[2]) for c in cases} + ts = [sum(q for _, q in c[2]) for c in cases] + mixed = [c for c in cases if c[3]] + odd_tc = sum(sum(q for _, q in c[2]) % 2 for c in mixed) + odd_nc = sum(len(c[2]) % 2 for c in mixed) + assert min(ks) == 1 and max(ks) >= 30, sorted(ks) + assert min(ts) <= 4 and max(ts) >= 8000, (min(ts), max(ts)) + assert max(c + q for case in cases for c, q in case[2]) >= 16000 + assert len(mixed) >= 100 and odd_tc >= 20 and odd_nc >= 20, (len(mixed), odd_tc, odd_nc) + assert {(c[0], c[1]) for c in cases} == {(h, layer) for h in (16, 64) for layer in ENV_LAYERS} + + +def test_prefill_partition_wrong_readings_fail_the_gate(env_pools): + """The partition harness can see a step read wrongly. On the first six mixed partitions that + have a fresh context request among several, the kernel's output sits inside the gate while the + reference built on each wrong reading lands far outside it: every context query given its + request's LAST window count (queries not causal within a request), each request's window count + taken from its neighbour's seq_lens (requests paired with the wrong kv_lens entry), and the decode + rows' q read from row 0 of the step instead of row tc (the row-offset view lost).""" + picked = [] + for seed in PARTITION_SEEDS: + num_heads, layer, ctx, dec = _partition_case(seed) + if dec and len(ctx) >= 2 and any(c == 0 and q >= 2 for c, q in ctx): + picked.append(seed) + if len(picked) == 6: + break + assert len(picked) == 6 + for seed in picked: + num_heads, layer, ctx, dec = _partition_case(seed) + gen = torch.Generator(device=DEV).manual_seed(500000 + seed) + st = _Step(ctx, dec, num_heads, layer, env_pools, gen) + o = st.run(_workspace()) + atol = env_pools["atol"][layer] + correct = _excess(o, st.reference(), atol) + q_lens = [q for _, q in st.requests] + seq = st.kv_lens[: st.B].long() + last = torch.repeat_interleave(seq.clamp(max=WINDOW), torch.tensor(q_lens, device=DEV)) + rolled = seq.roll(1) + pos_rolled = torch.cat( + [torch.arange(q, device=DEV) + int(s) - q for s, q in zip(rolled.tolist(), q_lens)] + ) + wrong = { + "non_causal_window": _reference( + st.q3, st.swa, st.comp, st.idx, st.active, last, st.sinks + ), + "neighbour_seq_lens": _reference( + st.q3, + st.swa, + st.comp, + st.idx, + st.active, + (pos_rolled + 1).clamp(1, WINDOW), + st.sinks, + ), + } + q_lost = st.q3.clone() + q_lost[st.tc :] = st.q3[: st.G] + wrong["decode_view_from_row_0"] = _reference( + q_lost, st.swa, st.comp, st.idx, st.active, st.counts, st.sinks + ) + excess = {k: _excess(o, r, atol) for k, r in wrong.items()} + print( + f"partition controls seed {seed} (nc {st.nc}, G {st.G}): correct {correct:.3f}, wrong " + f"{ {k: round(v, 1) for k, v in excess.items()} }" + ) + assert correct <= 1.0, correct + for name, e in excess.items(): + assert e > 10.0 * max(correct, 0.1), ( + f"seed {seed}: {name} not discriminated ({e:.2f} vs {correct:.2f})" + ) + + +def _dumped_cells(num_heads): + cells = {**DUMPED_CTX_CELLS, **MEMBER_DUMPED_CTX_CELLS}[num_heads] + return [tuple(int(x) for x in tok.split("/")) for tok in "".join(cells).split()] + + +def _dumped_partition(tc, nc, mq, rng): + """A step with this (tc, nc, max_q_len): one request of max_q_len, the other nc - 1 splitting + the rest with every q_len in [1, max_q_len], random order and cached prefixes.""" + parts, extra = [1] * (nc - 1), tc - mq - (nc - 1) + while extra > 0: + i = rng.randrange(nc - 1) + add = min(extra, mq - parts[i], rng.randint(1, extra)) + parts[i] += add + extra -= add + q_lens = [mq] + parts + rng.shuffle(q_lens) + return [(_fresh_or_cached(rng, q), q) for q in q_lens] + + +@pytest.mark.parametrize("width", [128, 640]) +@pytest.mark.parametrize("num_heads", ALL_HEADS) +def test_prefill_every_dumped_cell(env_pools, num_heads, width): + """Every context call (tc, H, nc, max_q_len) the boot gates and the accuracy-gate previews + dumped (DUMPED_CTX_CELLS; H 8 / 32 from the parallel members' dumps, MEMBER_DUMPED_CTX_CELLS), + at width 128 (ratio 0) and 640 (ratio 2 and ratio 1 on alternate cells), as a step with that + many requests, against the fp32 reference.""" + rng = random.Random(77000 + num_heads + width) + gen = torch.Generator(device=DEV).manual_seed(77000 + num_heads + width) + ws = _workspace() + cells = _dumped_cells(num_heads) + worst, bad = 0.0, [] + for i, (tc, nc, mq) in enumerate(cells): + layer = "r0" if width == WINDOW else ("r2", "r1")[i % 2] + ctx = _dumped_partition(tc, nc, mq, rng) + assert (sum(q for _, q in ctx), len(ctx), max(q for _, q in ctx)) == (tc, nc, mq) + st = _Step(ctx, [], num_heads, layer, env_pools, gen) + e = _excess(st.run(ws), st.reference(), env_pools["atol"][layer]) + worst = max(worst, e) + if e > 1.0: + bad.append(((tc, nc, mq), layer, e)) + print(f"dumped cells H{num_heads} W{width}: {len(cells)} calls, max excess {worst:.3f}") + assert not bad, bad[:5] + + +# --------------------------------------------------------------------------------------------- +# Pool-size invariance. The kernel receives each pool's base pointer and row +# stride, and the compressed pool's row count (mNumPagesInMemPool, which sizes a TMA descriptor); +# these tests run the same logical rows at the same page indices at B = the smallest pool holding +# them, the engine's size and the largest pool the GPU can allocate, in a child process that owns +# the whole device. +# --------------------------------------------------------------------------------------------- +# Page rows and the engine's pool size (runtime-value dump, as REAL_SWA / REAL_COMP below) per role. +POOL_ROLES = {"swa": (SWA_PAGE, 188160), "r2": (64, 453120), "r1": (128, 151040)} +POOL_RESERVE = 6 << 30 # bytes left free next to the largest pool, for q, out, workspace, reference +# The largest pool is sized from the free memory of the run; it must reach at least this many rows +# (262.5 GiB of 1024-byte rows, what the contract certifies as its upper bound). +POOL_MIN_MAX_ROWS = 2150400 * 128 +POOL_BOUNDARIES = tuple(1 << k for k in range(21, 31)) + + +POOL_CTX = [(0, 300), (4999, 240), (131, 1), (16000, 64)] +POOL_DEC = [16383, 77, 4000] + + +def _pool_tables(role, pages, gen, other): + """Page tables of POOL_CTX + POOL_DEC whose target-role rows live only in `pages`, in order.""" + page = POOL_ROLES[role][0] + reqs = POOL_CTX + [(p, 1) for p in POOL_DEC] + if role == "swa": + _, comp_bts = _env_tables(reqs, gen, 2, 64, 64, other.shape[0]) + return pages[: 2 * len(reqs)].view(len(reqs), 2), comp_bts + ratio = ENV_LAYERS[role][0] + swa_bt, _ = _env_tables(reqs, gen, 0, page, other.shape[0], 1) + need = [max(1, -(-((c + q) // ratio) // page)) for c, q in reqs] + assert sum(need) == pages.numel() + starts = [0] + for n in need: + starts.append(starts[-1] + n) + return swa_bt, [pages[starts[i] : starts[i + 1]] for i in range(len(reqs))] + + +def _pool_role(role, results): + """Pool-size invariance for one pool role, in a process that owns the device.""" + page, engine_pages = POOL_ROLES[role] + layer = "r2" if role == "swa" else role + gen = torch.Generator(device=DEV).manual_seed({"swa": 1, "r2": 2, "r1": 3}[role]) + other_page = 64 if role == "swa" else SWA_PAGE + other = torch.randn( + 4096 if role == "swa" else 256, 1, other_page, HEAD_DIM, generator=gen, device=DEV + ) + other = other.to(torch.bfloat16) + reqs = POOL_CTX + [(p, 1) for p in POOL_DEC] + ratio = ENV_LAYERS[layer][0] + n_pages = ( + 2 * len(reqs) + if role == "swa" + else sum(max(1, -(-((c + q) // ratio) // page)) for c, q in reqs) + ) + # Distinct logical pages scattered over [0, 4 * n_pages); the highest one ends the smallest pool. + logical = torch.cat( + [ + torch.randperm(4 * n_pages - 1, generator=gen, device=DEV)[: n_pages - 1], + torch.tensor([4 * n_pages - 1], device=DEV), + ] + ) + logical = logical[torch.randperm(n_pages, generator=gen, device=DEV)] + b_small = 4 * n_pages + data = torch.randn(n_pages, 1, page, HEAD_DIM, generator=gen, device=DEV).to(torch.bfloat16) + atol = 2.0**-8 * max(data.float().abs().max().item(), other.float().abs().max().item()) + tables = _pool_tables(role, logical, gen, other) + qs = { + h: torch.randn( + sum(q for _, q in POOL_CTX) + len(POOL_DEC), h, HEAD_DIM, generator=gen, device=DEV + ).to(torch.bfloat16) + for h in (16, 64) + } + # The members' head counts draw their q from a generator of their own (the H 16 / 64 q stay as + # they were); the far-end placement below renames the pages every head count's step attends. + gen_m = torch.Generator(device=DEV).manual_seed(7 + {"swa": 1, "r2": 2, "r1": 3}[role]) + qs.update( + { + h: torch.randn( + sum(q for _, q in POOL_CTX) + len(POOL_DEC), + h, + HEAD_DIM, + generator=gen_m, + device=DEV, + ).to(torch.bfloat16) + for h in MEMBER_HEADS + } + ) + ws = _workspace() + + def step(pool, h, tables): + pools = { + "swa": pool if role == "swa" else other, + "comp": {layer: other if role == "swa" else pool}, + } + st = _Step( + POOL_CTX, POOL_DEC, h, layer, pools, torch.Generator(device=DEV).manual_seed(h), *tables + ) + st.q3.copy_(qs[h]) + return st + + def run_all(pool, tables=tables): + return {h: step(pool, h, tables).run(ws).clone() for h in ALL_HEADS} + + def make(b): + pool = torch.zeros(b, 1, page, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + pool[logical] = data + return pool + + torch.cuda.empty_cache() + free, total = torch.cuda.mem_get_info() + b_max = (free - POOL_RESERVE) // (page * HEAD_DIM * 2) + b_mid = engine_pages + assert b_small < b_mid < b_max and b_max * page >= POOL_MIN_MAX_ROWS, (b_small, b_mid, b_max) + entry = { + "B_small": b_small, + "B_mid": b_mid, + "B_max": b_max, + "page_rows": page, + "device_free_bytes": free, + "device_total_bytes": total, + } + pool = make(b_small) + ref_outs = run_all(pool) + # Rows of the role pool each head count's step attends (for the far-end placement below). + attended = None + for h in ALL_HEADS: + st = step(pool, h, tables) + entry.setdefault("ref_excess", 0.0) + entry["ref_excess"] = max(entry["ref_excess"], _excess(ref_outs[h], st.reference(), atol)) + col = torch.arange(st.idx.shape[1], device=DEV) + att = (col[None] < st.active[:, None].long()) & ( + (col[None] >= WINDOW) | (col[None] < st.counts[:, None]) + ) + sl = slice(0, WINDOW) if role == "swa" else slice(WINDOW, None) + rows = st.idx[:, sl][att[:, sl] & (st.idx[:, sl] >= 0)] + pages_h = set((rows // page).unique().tolist()) + attended = pages_h if attended is None else attended & pages_h + same = {"small": True} + del pool + torch.cuda.empty_cache() + pool = make(b_mid) + same["mid"] = all(torch.equal(o, ref_outs[h]) for h, o in run_all(pool).items()) + del pool + torch.cuda.empty_cache() + pool = make(b_max) + same["max"] = all(torch.equal(o, ref_outs[h]) for h, o in run_all(pool).items()) + for name, b in (("max_view_small", b_small), ("max_view_mid", b_mid)): + same[name] = all(torch.equal(o, ref_outs[h]) for h, o in run_all(pool[:b]).items()) + # The same logical step with its pages renamed onto the far end of the largest pool: the last + # page, and pages straddling rows 2^21 .. 2^30, go to logical pages the step attends. + total_rows = b_max * page + must = [b_max - 1] + [r // page for r in POOL_BOUNDARIES if r < total_rows] + must += [r // page - 1 for r in POOL_BOUNDARIES if r < total_rows] + att_logical = [p for p in logical.tolist() if p in attended] + rest_logical = [p for p in logical.tolist() if p not in attended] + lut = torch.full((b_small,), -1, dtype=torch.int64, device=DEV) + order = att_logical + rest_logical + far_pool = list(dict.fromkeys(must[: len(att_logical)])) + filler = ( + b_max - 1 - 16 * n_pages + torch.randperm(16 * n_pages, generator=gen, device=DEV) + ).tolist() + far_pool += [p for p in filler if p not in set(far_pool)] + far = torch.tensor(far_pool[:n_pages], dtype=torch.int64, device=DEV) + lut[torch.tensor(order, device=DEV)] = far + pool[lut[logical]] = data + far_tables = ( + (lut[tables[0]], tables[1]) if role == "swa" else (tables[0], [lut[c] for c in tables[1]]) + ) + far_named = set(far.tolist()) + alias = sorted({(p * page % WRAP_ROWS) // page for p in far_named} - far_named) + a = torch.tensor(alias, dtype=torch.int64, device=DEV) + pool[a] = torch.randn(a.numel(), 1, page, HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) + far_ex, wrap_ex, max_row = 0.0, float("inf"), -1 + far_same = True + for h in ALL_HEADS: + st = step(pool, h, far_tables) + o = st.run(ws) + far_same &= torch.equal(o, ref_outs[h]) + far_ex = max(far_ex, _excess(o, st.reference(), atol)) + sl = slice(0, WINDOW) if role == "swa" else slice(WINDOW, None) + col = torch.arange(st.idx.shape[1], device=DEV) + att = (col[None] < st.active[:, None].long()) & ( + (col[None] >= WINDOW) | (col[None] < st.counts[:, None]) + ) + max_row = max(max_row, int(torch.where(att[:, sl], st.idx[:, sl], -1).max())) + wrapped = st.idx.clone() + seg = wrapped[:, sl] + wrapped[:, sl] = torch.where(seg >= 0, seg % WRAP_ROWS, seg) + wrap_ex = min(wrap_ex, _excess(o, st.reference(wrapped), atol)) + same["max_far_pages_renamed"] = bool(far_same) + entry.update( + bitwise=same, + far_excess=far_ex, + wrap_excess=wrap_ex, + max_attended_row=max_row, + last_row=total_rows - 1, + must=must, + must_attended=sorted(set(must) & set(lut[torch.tensor(att_logical, device=DEV)].tolist())), + pool_shape=list(pool.shape), + ) + results[role] = entry + del pool + torch.cuda.empty_cache() + + +def _pool_child(): + results = {} + for role in POOL_ROLES: + _pool_role(role, results) + print("POOL_CHILD_RESULT " + json.dumps(results), flush=True) + + +@pytest.fixture(scope="module") +def pool_invariance(): + torch.cuda.empty_cache() + proc = subprocess.run( + [sys.executable, os.path.abspath(__file__), "--pool-child"], + env=_CHILD_ENV, + capture_output=True, + text=True, + timeout=1800, + ) + line = [ln for ln in proc.stdout.splitlines() if ln.startswith("POOL_CHILD_RESULT ")] + assert proc.returncode == 0 and line, ( + proc.returncode, + proc.stdout[-3000:], + proc.stderr[-3000:], + ) + return json.loads(line[-1][len("POOL_CHILD_RESULT ") :]) + + +@pytest.mark.parametrize("role", list(POOL_ROLES)) +def test_pool_size_invariance(pool_invariance, role): + """A mixed step (context + decode, H 8, 16, 32 and 64) over the same logical rows at the same page + indices is bitwise identical at the smallest pool that holds them, the engine's pool size and + the largest pool the GPU can allocate (separate allocations, and views of the largest one); at + that largest pool, rows in its last page and in pages straddling rows 2^21..2^30 match the + reference, and reading them at (row mod 2^22) instead fails the gate by a wide margin.""" + r = pool_invariance[role] + print(f"pool invariance {role}: {r}") + assert all(r["bitwise"].values()), r["bitwise"] + assert r["ref_excess"] <= 1.0 and r["far_excess"] <= 1.0, r + assert r["max_attended_row"] >= r["last_row"] - r["page_rows"] + 1, r + assert r["B_max"] - 1 in r["must_attended"], r + assert r["wrap_excess"] > 10.0, r + + +# --------------------------------------------------------------------------------------------- +# the engine's real pool sizes and decode batch sizes (runtime-value dump) +# --------------------------------------------------------------------------------------------- +# Pool geometry the core DeepSeek-V4.1 target's cache manager allocated per rank in its boot and +# eval runs: one SWA pool shared by all 40 layers ([188160, 128, 512], 24.7 GB bf16) and one +# compressed pool per KV-source owner ([453120, 64, 512] at ratio 2, 29.7 GB; [151040, 128, 512] at +# ratio 1, 19.8 GB). The target hands them over as [pages, 1, page, 512] views. A 1024-byte row puts +# row 2^21 at byte 2^31, row 2^22 at byte 2^32 and at element 2^31, so a 32-bit offset anywhere in +# the address path would wrap from those rows on. Every case below attends rows in the LAST page of +# each pool it uses, plus pages straddling rows 2^21 / 2^22 / 2^23 / 2^24, plus the first page. +REAL_SWA = (188160, 128) +REAL_COMP = {2: (453120, 64), 1: (151040, 128)} +WRAP_ROWS = 1 << 22 # rows per 2^32 bytes (and per 2^31 bf16 elements) at 1024 B per row +ROW_BOUNDARIES = (1 << 21, 1 << 22, 1 << 23, 1 << 24) +# Decode batch sizes the engine ran: eager mixed steps saw G in 1..10, 16..19, 24, 27..32. The decode +# CUDA graphs are captured per rank at every batch 1..8 (the boot gate, padding off) and at 1, 2, 4, +# 8, 16, 24, 32 (the accuracy gate, padding on): the capture set. The native-storage compact-read +# capture tests below still use this set (their buffers are sized for max_batch_size 32). +GRAPH_BATCHES = [1, 2, 3, 4, 5, 6, 7, 8, 16, 24, 32] +assert set(GRAPH_BATCHES) <= set(GSET) +# DSpark target verify: 1 + 5 draft tokens per request, dense [G, 6, H, 512] with next_n 6. +VERIFY_Q = 6 +# Decode positions of requests 1.. (request 0 is always a long one, so that its first page -- the +# pool's last page -- holds window rows, the SWA pool's last row among them, and compressed rows). +REAL_POSITIONS = [0, 3, 60, 126, 127, 128, 500, 1500, 5000, 9000, 65000] +# Envelope E4: the kv lengths the engine can hand this call at max_seq_len L in {16384, 32768, 65536, +# 131072, 262144, 1048576}: L itself (DSpark off), the cache manager's length Lm = L + 14 and Lm + 5 +# (DSpark on: positions reach about Lm + 5). Decode / verify requests cycle through REAL_POSITIONS +# interleaved with these, placed so that the request's seq_lens IS the listed kv length. +LONG_KV = sorted( + {L + d for L in (16384, 32768, 65536, 131072, 262144, 1048576) for d in (0, 14, 19)} +) + + +def _cached_positions(q_len: int): + """The LONG_KV lengths (as cached prefixes for a q_len request) interleaved with REAL_POSITIONS, + a long one first.""" + long = [kv - q_len for kv in LONG_KV] + out = [] + for i in range(max(len(REAL_POSITIONS), len(long))): + out += long[i : i + 1] + REAL_POSITIONS[i : i + 1] + return out + + +def _paged_rows(requests, gen, width, ratio, swa_bt, comp_bt, swa_page, comp_page): + """_build_rows, vectorized over every query row: request b's pages are the rows swa_bt[b] and + comp_bt[b], used as rings (position p on page (p // page) % pages per request). Query j of + request (cached, q_len) sits at pos = cached + j and sees window positions max(0, pos-127)..pos + in order (-1 after them) and min(512, (pos + 1) // ratio) compressed positions: all of them + while there are at most 512, else one random position per stratum [s * v // 512, + (s + 1) * v // 512) (distinct and sorted, as the indexer top-k is), -1 after them. + Returns (idx [N, width] int32, active [N] int32, seq_lens [B] int32, cu [B + 1] int32).""" + q_lens = torch.tensor([q for _, q in requests], dtype=torch.int64, device=DEV) + cached = torch.tensor([c for c, _ in requests], dtype=torch.int64, device=DEV) + cu = torch.cat([torch.zeros(1, dtype=torch.int64, device=DEV), q_lens.cumsum(0)]) + n = int(cu[-1]) + req = torch.repeat_interleave(torch.arange(len(requests), device=DEV), q_lens) + pos = cached[req] + torch.arange(n, device=DEV) - cu[req] + wpos = (pos - (WINDOW - 1)).clamp_min(0)[:, None] + torch.arange(WINDOW, device=DEV)[None] + ring = swa_bt[req] + widx = ring.gather(1, (wpos // swa_page) % ring.shape[1]) * swa_page + wpos % swa_page + widx = torch.where(wpos <= pos[:, None], widx, -1) + if width == WINDOW: + idx, active = widx, torch.full((n,), WINDOW, dtype=torch.int64, device=DEV) + else: + k = width - WINDOW + vis = (pos + 1) // ratio + strata = torch.arange(k + 1, device=DEV) + lo = (strata[None, :-1] * vis[:, None]) // k + hi = (strata[None, 1:] * vis[:, None]) // k + u = torch.rand(n, k, generator=gen, device=DEV) + sel = torch.where(vis[:, None] > k, lo + (u * (hi - lo)).long(), strata[None, :-1]) + valid = sel < vis[:, None] + pages = comp_bt[req] + cpage = pages.gather(1, (sel // comp_page) % pages.shape[1]) + cidx = torch.where(valid, cpage * comp_page + sel % comp_page, -1) + idx, active = torch.cat([widx, cidx], 1), WINDOW + valid.sum(1) + seq_lens = cached + q_lens + return idx.int().contiguous(), active.int(), seq_lens.int(), cu.int() + + +class _RealPools: + """Zero-initialized pools of the engine's size; each case fills only the pages it names.""" + + def __init__(self, layer: str): + self.layer = layer + self.ratio = {"r2": 2, "r1": 1, "r0": 0}[layer] + self.swa_pages, self.swa_page = REAL_SWA + self.swa = torch.zeros( + self.swa_pages, 1, self.swa_page, HEAD_DIM, dtype=torch.bfloat16, device=DEV + ) + if self.ratio: + self.comp_pages, self.comp_page = REAL_COMP[self.ratio] + self.comp = torch.zeros( + self.comp_pages, 1, self.comp_page, HEAD_DIM, dtype=torch.bfloat16, device=DEV + ) + self.topk = WINDOW + INDEX_TOPK + else: + # Window-only layers pass a one-page placeholder; NaN, so that a read of it would show. + self.comp_pages, self.comp_page = 1, 64 + self.comp = torch.full( + (1, 1, 64, HEAD_DIM), float("nan"), dtype=torch.bfloat16, device=DEV + ) + self.topk = WINDOW + + @staticmethod + def _pick(gen: torch.Generator, total: int, page: int, n: int) -> torch.Tensor: + """n distinct pages: the pool's last page first, then pages straddling the row boundaries, the + first page, and random pages from the pool's last stretch -- shuffled after the first.""" + must = [total - 1] + for r in ROW_BOUNDARIES: + if r < total * page: + must += [r // page - 1, r // page] + must.append(0) + far = (total - 1 - 8 * n + torch.randperm(8 * n, generator=gen, device=DEV)).tolist() + pages = list(dict.fromkeys(must + far))[:n] + rest = torch.tensor(pages[1:], dtype=torch.int64, device=DEV) + rest = rest[torch.randperm(rest.numel(), generator=gen, device=DEV)] + return torch.cat([torch.tensor([pages[0]], dtype=torch.int64, device=DEV), rest]) + + @staticmethod + def _fill(gen: torch.Generator, pool: torch.Tensor, pages: torch.Tensor) -> None: + """Random rows in the named pages, and independent random rows in the pages a 32-bit wrap of + their row offsets would land on (row mod 2^22), so a wrapped read is a wrong value, not 0.""" + page = pool.shape[2] + pool[pages] = torch.randn(pages.numel(), 1, page, HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) + named = set(pages.tolist()) + alias = sorted({(p * page % WRAP_ROWS) // page for p in named} - named) + if alias: + a = torch.tensor(alias, dtype=torch.int64, device=DEV) + pool[a] = torch.randn(a.numel(), 1, page, HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) + + def rows(self, requests, gen: torch.Generator, swa_per_req: int = 2, comp_per_req: int = 4): + """Index rows for `requests`: request b's pages are entries [b * per_req, (b + 1) * per_req) of + the picked page lists (request 0 gets the pool's last page first).""" + b = len(requests) + swa_bt = self._pick(gen, self.swa_pages, self.swa_page, swa_per_req * b) + self._fill(gen, self.swa, swa_bt) + comp_bt = None + if self.ratio: + comp_bt = self._pick(gen, self.comp_pages, self.comp_page, comp_per_req * b) + self._fill(gen, self.comp, comp_bt) + return _paged_rows( + requests, + gen, + self.topk, + max(self.ratio, 1), + swa_bt.view(b, swa_per_req), + comp_bt.view(b, comp_per_req) if comp_bt is not None else None, + self.swa_page, + self.comp_page, + ) + + def atol(self, idx: torch.Tensor) -> float: + """The file's gate, 2^-8 * max|V|, over the rows this call names (the rest of the pool is 0).""" + w = idx[:, :WINDOW] + m = self.swa.view(-1, HEAD_DIM)[torch.unique(w[w >= 0]).long()].float().abs().max() + c = idx[:, WINDOW:] + if (c >= 0).any(): + m = torch.maximum( + m, self.comp.view(-1, HEAD_DIM)[torch.unique(c[c >= 0]).long()].float().abs().max() + ) + return 2.0**-8 * m.item() + + def assert_far_end_attended(self, idx, active, counts) -> None: + """The call really attends rows in the last page of each pool, past row 2^22.""" + col = torch.arange(idx.shape[1], device=DEV) + att = (col[None] < active[:, None].long()) & ( + (col[None] >= WINDOW) | (col[None] < counts[:, None]) + ) + w = torch.where(att[:, :WINDOW], idx[:, :WINDOW], -1) + swa_last = (self.swa_pages - 1) * self.swa_page + assert int(w.max()) >= swa_last >= WRAP_ROWS, (int(w.max()), swa_last) + if self.ratio: + c = torch.where(att[:, WINDOW:], idx[:, WINDOW:], -1) + comp_last = (self.comp_pages - 1) * self.comp_page + assert int(c.max()) >= comp_last >= WRAP_ROWS, (int(c.max()), comp_last) + + def release(self) -> None: + del self.swa, self.comp + torch.cuda.empty_cache() + + +@pytest.fixture(scope="module", params=["r2", "r1", "r0"]) +def real_pools(request): + pools = _RealPools(request.param) + yield pools + pools.release() + + +def _decode_requests(batch: int, shift: int = 0): + # Position 20095 + 256k: offset 127 of ring page 0, i.e. the SWA pool's very last row. + first = (20095 + 256 * shift, 1) + pos = _cached_positions(1) + return [first] + [(pos[(i + shift) % len(pos)], 1) for i in range(batch - 1)] + + +def _check_atol(o, ref, atol): + """The entry's gate at an explicit atol (2^-8 * max|V| over the rows the caller names).""" + op.compare(o.reshape(ref.shape), (ref, atol * 2.0**8)) + + +@pytest.mark.parametrize("num_heads", ALL_HEADS) +def test_real_pools_decode_every_batch_size(real_pools, num_heads): + """Eager decode at every batch size 1..256 over the engine's pools, far-end rows attended, kv + lengths up to the E4 extreme (LONG_KV) among the requests.""" + p = real_pools + gen = torch.Generator(device=DEV).manual_seed(7000 + num_heads + p.ratio) + sinks = _sinks(num_heads) + for batch in range(1, E_B + 1): + requests = _decode_requests(batch, shift=batch) + idx, active, seq_lens, cu = p.rows(requests, gen) + counts = _window_counts(seq_lens, cu) + p.assert_far_end_attended(idx, active, counts) + q = ( + torch.randn(batch, 1, num_heads, HEAD_DIM, generator=gen, device=DEV) + .mul_(2) + .to(torch.bfloat16) + ) + o = _run(q, p.swa, p.comp, idx, active, seq_lens, sinks) + assert o.shape == (batch, 1, num_heads, HEAD_DIM) + assert torch.isfinite(o).all(), f"B={batch}" + ref = _reference( + q.view(batch, num_heads, HEAD_DIM), p.swa, p.comp, idx, active, counts, sinks + ) + _check_atol(o, ref, p.atol(idx)) + + +@pytest.mark.parametrize("num_heads", ALL_HEADS) +def test_real_pools_prefill_varlen(real_pools, num_heads): + """Varlen prefill (fresh 2048-token prompt + chunked continuation + 1-token row) over the + engine's pools, far-end rows attended.""" + p = real_pools + gen = torch.Generator(device=DEV).manual_seed(8000 + num_heads + p.ratio) + requests = [(0, 2048), (6000, 300), (90, 1)] + idx, active, seq_lens, cu = p.rows(requests, gen, swa_per_req=2, comp_per_req=8) + counts = _window_counts(seq_lens, cu) + p.assert_far_end_attended(idx, active, counts) + n = idx.shape[0] + q = torch.randn(n, num_heads, HEAD_DIM, generator=gen, device=DEV).mul_(2).to(torch.bfloat16) + sinks = _sinks(num_heads) + o = _run(q, p.swa, p.comp, idx, active, seq_lens, sinks, cu=cu, max_q_len=2048) + ref = _reference(q, p.swa, p.comp, idx, active, counts, sinks) + _check_atol(o, ref, p.atol(idx)) + + +@pytest.mark.parametrize("batch", GSET) +@pytest.mark.parametrize("num_heads", ALL_HEADS) +def test_real_pools_decode_cuda_graph(real_pools, num_heads, batch): + """Decode captured into a CUDA graph at every batch of the envelope's capture set GSET (46 + values, 1..256); replay equals eager bitwise, then two replays on new q / index rows / lengths + (new far-end pages) equal a fresh eager call on the same inputs bitwise and match the + reference.""" + p = real_pools + gen = torch.Generator(device=DEV).manual_seed(9000 + 100 * batch + num_heads + p.ratio) + sinks = _sinks(num_heads) + + def make(shift): + requests = _decode_requests(batch, shift) + idx, active, seq_lens, cu = p.rows(requests, gen) + p.assert_far_end_attended(idx, active, _window_counts(seq_lens, cu)) + q = ( + torch.randn(batch, 1, num_heads, HEAD_DIM, generator=gen, device=DEV) + .mul_(2) + .to(torch.bfloat16) + ) + return q, idx, active, seq_lens, cu + + q, idx, active, seq_lens, cu = make(0) + ws = _workspace() + out = torch.empty(batch, 1, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + + def call(): + return op( + query=q, + swa_kv_cache=p.swa, + compressed_kv_cache=p.comp, + workspace_buffer=ws, + sparse_indices=idx, + sparse_topk_lens=active, + seq_lens=seq_lens, + sinks=sinks, + softmax_scale=SCALE, + out=out, + ) + + call() + eager = out.clone() + ref = _reference( + q.view(batch, num_heads, HEAD_DIM), + p.swa, + p.comp, + idx, + active, + _window_counts(seq_lens, cu), + sinks, + ) + _check_atol(eager, ref, p.atol(idx)) + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + call() + torch.cuda.current_stream().wait_stream(s) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, eager) + for shift in (1, 2): + nq, nidx, nactive, nseq, ncu = make(shift) + q.copy_(nq) + idx.copy_(nidx) + active.copy_(nactive) + seq_lens.copy_(nseq) + out.zero_() + graph.replay() + torch.cuda.synchronize() + replayed = out.clone() + call() # a fresh eager call on the same inputs + torch.cuda.synchronize() + assert torch.equal(replayed, out), f"replay != eager at shift {shift}" + ref = _reference( + nq.view(batch, num_heads, HEAD_DIM), + p.swa, + p.comp, + nidx, + nactive, + _window_counts(nseq, ncu), + sinks, + ) + _check_atol(out, ref, p.atol(nidx)) + del graph + + +def _verify_requests(batch: int, shift: int = 0): + # Request 0's last verify query sits at 20095 + 256k: offset 127 of ring page 0, the SWA pool's + # very last row. + first = (20090 + 256 * shift, VERIFY_Q) + pos = _cached_positions(VERIFY_Q) + return [first] + [(pos[(i + shift) % len(pos)], VERIFY_Q) for i in range(batch - 1)] + + +@pytest.mark.parametrize("num_heads", ALL_HEADS) +def test_real_pools_verify_every_batch_size(real_pools, num_heads): + """Eager DSpark target verify (q_len 6, dense) at every batch size 1..256 over the engine's pools, + kv lengths up to the E4 extreme (LONG_KV) among the requests.""" + p = real_pools + gen = torch.Generator(device=DEV).manual_seed(7500 + num_heads + p.ratio) + sinks = _sinks(num_heads) + for batch in range(1, E_B + 1): + idx, active, seq_lens, cu = p.rows(_verify_requests(batch, shift=batch), gen) + counts = _window_counts(seq_lens, cu) + p.assert_far_end_attended(idx, active, counts) + q = ( + torch.randn(batch, VERIFY_Q, num_heads, HEAD_DIM, generator=gen, device=DEV) + .mul_(2) + .to(torch.bfloat16) + ) + o = _run(q, p.swa, p.comp, idx, active, seq_lens, sinks) + assert o.shape == (batch, VERIFY_Q, num_heads, HEAD_DIM) + ref = _reference(q.view(-1, num_heads, HEAD_DIM), p.swa, p.comp, idx, active, counts, sinks) + _check_atol(o, ref, p.atol(idx)) + + +@pytest.mark.parametrize("batch", GSET) +@pytest.mark.parametrize("num_heads", ALL_HEADS) +def test_real_pools_verify_q_len_6_cuda_graph(real_pools, num_heads, batch): + """DSpark verify (q_len 6) captured at every batch of GSET (1..256) on the engine's pools. The + first replay equals the eager call bitwise; three further replays on new q, index rows, per-row + lengths and seq_lens (new far-end pages) equal a fresh eager call on the same inputs bitwise and + match the reference. Two wrong readings of the 6-query block must fail the gate.""" + p = real_pools + gen = torch.Generator(device=DEV).manual_seed(9500 + 100 * batch + num_heads + p.ratio) + sinks = _sinks(num_heads) + rows = VERIFY_Q * batch + + def make(shift): + idx, active, seq_lens, cu = p.rows(_verify_requests(batch, shift), gen) + counts = _window_counts(seq_lens, cu) + p.assert_far_end_attended(idx, active, counts) + q = ( + torch.randn(batch, VERIFY_Q, num_heads, HEAD_DIM, generator=gen, device=DEV) + .mul_(2) + .to(torch.bfloat16) + ) + return q, idx, active, seq_lens, counts + + q, idx, active0, seq0, counts = make(0) + # Row / request prefixes of step buffers sized for the largest graph batch, as the target passes them. + lens_buf = torch.full((VERIFY_Q * GSET[-1],), WINDOW, dtype=torch.int32, device=DEV) + kv_lens_buf = torch.full((GSET[-1],), VERIFY_Q, dtype=torch.int32, device=DEV) + active, seq_lens = lens_buf[:rows], kv_lens_buf[:batch] + active.copy_(active0) + seq_lens.copy_(seq0) + ws = _workspace() + + def call(out): + return op( + query=q, + swa_kv_cache=p.swa, + compressed_kv_cache=p.comp, + workspace_buffer=ws, + sparse_indices=idx, + sparse_topk_lens=active, + seq_lens=seq_lens, + sinks=sinks, + softmax_scale=SCALE, + out=out, + ) + + shape = (batch, VERIFY_Q, num_heads, HEAD_DIM) + eager = call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) # the engine's eager warm-up + ref = _reference(q.view(rows, num_heads, HEAD_DIM), p.swa, p.comp, idx, active, counts, sinks) + _check_atol(eager, ref, p.atol(idx)) + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + torch.cuda.current_stream().wait_stream(s) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + # Allocated inside the captured forward, as the target allocates its [T, H, 512] output. + o = torch.empty(rows, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + out = call(o.view(shape)) + assert out.data_ptr() == o.data_ptr() + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, eager) + for shift in (1, 2, 3): + nq, nidx, nactive, nseq, ncounts = make(shift) + q.copy_(nq) + idx.copy_(nidx) + active.copy_(nactive) + seq_lens.copy_(nseq) + out.fill_(float("nan")) + graph.replay() + fresh = call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + torch.cuda.synchronize() + assert torch.equal(out, fresh), f"replay != eager at shift {shift}" + qf = nq.view(rows, num_heads, HEAD_DIM) + ref = _reference(qf, p.swa, p.comp, nidx, nactive, ncounts, sinks) + atol = p.atol(nidx) + _check_atol(out, ref, atol) + + def excess(r): + return ((out.float().reshape(r.shape) - r).abs() / (atol + RTOL * r.abs())).max().item() + + correct = excess(ref) + if shift != 1 or batch < 2: + print( + f"verify-replay {p.layer} H{num_heads} G{batch} shift {shift}: correct {correct:.3f}" + ) + continue + # q read as [6, G] instead of [G, 6]: index row b*6+j paired with query row j*G+b. + swapped = ( + qf.reshape(VERIFY_Q, batch, num_heads, HEAD_DIM) + .transpose(0, 1) + .reshape(rows, num_heads, HEAD_DIM) + ) + wrong = { + "q_block_transposed": excess( + _reference(swapped, p.swa, p.comp, nidx, nactive, ncounts, sinks) + ) + } + # Non-causal block: every query of a request gets the window count of its last query. + last = torch.repeat_interleave(nseq.long().clamp(max=WINDOW), VERIFY_Q) + if not torch.equal(last, ncounts): + wrong["block_not_causal"] = excess( + _reference(qf, p.swa, p.comp, nidx, nactive, last, sinks) + ) + print( + f"verify-controls {p.layer} H{num_heads} G{batch}: correct {correct:.3f}, wrong " + f"{ {k: round(v, 1) for k, v in wrong.items()} }" + ) + assert correct <= 1.0, correct + for name, e in wrong.items(): + assert e > 10.0 * max(correct, 0.1), ( + f"{name} not discriminated: {e:.2f} vs {correct:.2f}" + ) + del graph + + +def test_real_pools_wrap_control(real_pools): + """The harness can see a 32-bit address wrap: reading every attended row at (row mod 2^22) -- the + row a wrapped 32-bit byte offset (or 31-bit element offset) of a 1024-byte row lands on, filled + here with independent data -- fails the same gate by a wide margin, while the true rows pass.""" + p = real_pools + gen = torch.Generator(device=DEV).manual_seed(9900 + p.ratio) + batch, h = 32, 16 + idx, active, seq_lens, cu = p.rows(_decode_requests(batch), gen) + counts = _window_counts(seq_lens, cu) + p.assert_far_end_attended(idx, active, counts) + q = torch.randn(batch, 1, h, HEAD_DIM, generator=gen, device=DEV).mul_(2).to(torch.bfloat16) + sinks = _sinks(h) + o = _run(q, p.swa, p.comp, idx, active, seq_lens, sinks) + qf = q.view(batch, h, HEAD_DIM) + atol = p.atol(idx) + + def excess(ref): + return ((o.float().reshape(ref.shape) - ref).abs() / (atol + RTOL * ref.abs())).max().item() + + correct = excess(_reference(qf, p.swa, p.comp, idx, active, counts, sinks)) + wrapped = torch.where(idx >= 0, idx % WRAP_ROWS, idx) + assert (wrapped != idx).any() + wrong = excess(_reference(qf, p.swa, p.comp, wrapped, active, counts, sinks)) + print( + f"{p.layer}: pools swa {tuple(p.swa.shape)} comp {tuple(p.comp.shape)} " + f"({torch.cuda.memory_allocated() / 2**30:.1f} GiB allocated), max attended row " + f"swa {int(idx[:, :WINDOW].max())} comp {int(idx[:, WINDOW:].max()) if idx.shape[1] > WINDOW else -1}; " + f"gate-normalized error correct {correct:.3f}, wrapped-address reading {wrong:.1f}" + ) + assert correct <= 1.0, correct + assert wrong > 10.0 * max(correct, 0.1), (wrong, correct) + + +# --------------------------------------------------------------------------------------------- +# The native-storage compact read +# --------------------------------------------------------------------------------------------- +# Under native KV storage the target stores the KV quantized and, every step, dequantizes +# exactly the rows its attention reads into two persistent bf16 buffers, allocated once per cache +# manager (zero-filled, outside capture) and never cleared; it hands THOSE buffers to this call: +# +# - window: swa_c [max_num_tokens * 128, 512], passed as [max_num_tokens, 1, 128, 512]. Query t's +# window column c is dequantized to row t * 128 + c, a -1 column as an all-zero (+0.0) row, and +# the index table's window half becomes that numbering: arange(T * 128).view(T, 128), a prefix +# view [:T] of one persistent [max_num_tokens, 128] arange buffer. The window half then holds no +# -1; where the container path has a -1 the compact path names a zero row. +# - compressed: comp_c [max_batch_size * 132 * 128, 512]. Request b's entry e (every e < kv_len // r, +# b the request's position in the step) at row b * L_r + e, L_r = 132 * 128 / r, passed per ratio +# as [max_batch_size * 132, 1, 128 / r, 512] over its first max_batch_size * 132 * 128 / r rows. +# The compressed half of the index is b * L_r + e for a valid top-k entry and -1 after them (the +# index op maps the top-k through a synthetic page table; certified in that entry). +# - window-only (ratio-0) layers pass the renumbered window half itself ([T, 128], the arange view) +# and a zeros [1, 1, 64, 512] placeholder as the compressed pool; ratio 1 / 2 layers pass the +# concatenation [T, 640]. +# +# Sizes: max_num_tokens 2048 or 8192 and max_batch_size 8 or 32 in every envelope configuration +# (the metadata's swa_local_indices / compress_block_tables, measured in the dequantizer entries' +# native-storage tests): swa_c 262,144 / 1,048,576 rows, comp_c 135,168 / 540,672 rows. +# +# The container path (bf16 KV storage) is the call this file certifies above: the same values read +# from the cache manager's pools through the block tables, with -1 in the index. Here the container +# pools are env_pools' random pools, the compact buffers are filled from them by native torch gathers +# (what the two dequantizers write; the dequantizers themselves are certified bitwise to the bf16 +# container in their own entries), and every compact read is held BITWISE to the container read on +# the same values, and the container read to the fp32 reference. The buffers are pre-filled with +# finite random rows (a previous step's content), so a row a step fails to rewrite would show. +KVQ_COLS = 132 +KVQ_L = {2: KVQ_COLS * 64, 1: KVQ_COLS * 128} +KVQ_CONFIGS = ((8192, 32), (8192, 8), (2048, 32), (2048, 8)) # (max_num_tokens, max_batch_size) +KVQ_TAIL = 16397 # DSpark: generation positions reach max_seq_len 16398 - 1; contexts stay < 16384 + + +class _KvqBuffers: + """The target's two read buffers for one (max_num_tokens, max_batch_size), with the views and the + window numbering it hands the call.""" + + def __init__(self, max_t, seqs, gen): + self.max_t, self.seqs = max_t, seqs + self.swa_c = torch.randn( + max_t * WINDOW, HEAD_DIM, generator=gen, device=DEV, dtype=torch.bfloat16 + ) + self.comp_c = torch.randn( + seqs * KVQ_COLS * 128, HEAD_DIM, generator=gen, device=DEV, dtype=torch.bfloat16 + ) + self.swa_c4 = self.swa_c.view(max_t, 1, WINDOW, HEAD_DIM) + self.comp_c4 = { + r: self.comp_c[: seqs * KVQ_COLS * (128 // r)].view( + seqs * KVQ_COLS, 1, 128 // r, HEAD_DIM + ) + for r in (1, 2) + } + self.win_idx = torch.arange(max_t * WINDOW, dtype=torch.int32, device=DEV).view( + max_t, WINDOW + ) + self.dummy = torch.zeros(1, 1, 64, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + for v in self.comp_c4.values(): + assert ( + v.data_ptr() == self.comp_c.data_ptr() and v.shape[0] * v.shape[2] >= POOL_MIN_ROWS + ) + + +@pytest.fixture(scope="module") +def kvq_buffers(): + gen = torch.Generator(device=DEV).manual_seed(1_400_000) + bufs = {cfg: _KvqBuffers(*cfg, gen) for cfg in KVQ_CONFIGS} + yield bufs + del bufs + torch.cuda.empty_cache() + + +def _kvq_rows(requests, floors, gen, width, ratio, comp_page, swa_bt, comp_bts): + """_env_rows with the local compressed positions kept (a random subset in score order when more + than 512 are visible, the newest visible entry always among them), and an optional window floor per request + (the metadata's swa_floor, an encoder-replay step's replay start: the window of position p is + max(floor, p - 127) .. p, then -1, while the call's window count stays min(128, p + 1), so the + -1s past the floored window sit INSIDE the count and are attended). Returns idx [N, width] + (container rows, -1), active [N], seq_lens [B], cu [B+1], counts [N], req [N] (the request's + position in the step) and loc [N, 512] (the compressed positions e, -1) or None.""" + ar = torch.arange(WINDOW, device=DEV) + idx_l, act_l, cnt_l, req_l, loc_l = [], [], [], [], [] + for b, ((cached, q_len), floor) in enumerate(zip(requests, floors)): + pos = torch.arange(cached, cached + q_len, device=DEV) + start = (pos - (WINDOW - 1)).clamp_min(floor).clamp_min(0) + wpos = start[:, None] + ar[None] + ring = swa_bt[b] + widx = ring[(wpos // SWA_PAGE) % ring.numel()] * SWA_PAGE + wpos % SWA_PAGE + widx = torch.where(wpos <= pos[:, None], widx, -1).int() + cnt_l.append((pos + 1).clamp(max=WINDOW)) + req_l.append(torch.full((q_len,), b, dtype=torch.int64, device=DEV)) + if width == WINDOW: + idx_l.append(widx) + act_l.append(torch.full((q_len,), WINDOW, dtype=torch.int32, device=DEV)) + continue + k_max = width - WINDOW + vis = (pos + 1) // ratio + vmax = (cached + q_len) // ratio + pages = comp_bts[b] + assert pages.numel() * comp_page >= vmax + sels, valids = [], [] + for s0 in range(0, q_len, 1024): + v = vis[s0 : s0 + 1024] + if vmax <= k_max: + sel = torch.arange(k_max, device=DEV)[None].expand(v.numel(), k_max) + valid = sel < v[:, None] + else: + score = torch.rand(v.numel(), vmax, generator=gen, device=DEV) + col = torch.arange(vmax, device=DEV)[None] + score.masked_fill_( + col == (v - 1)[:, None], 2.0 + ) # the newest visible entry is always picked + score.masked_fill_(col >= v[:, None], -1.0) + val, sel = score.topk(k_max, dim=1) + valid = val >= 0 # score order: valid entries first, then -1 + sels.append(sel) + valids.append(valid) + sel, valid = torch.cat(sels), torch.cat(valids) + sel0 = torch.where(valid, sel, 0) + cidx = torch.where(valid, pages[sel0 // comp_page] * comp_page + sel0 % comp_page, -1).int() + idx_l.append(torch.cat([widx, cidx], 1)) + loc_l.append(torch.where(valid, sel, -1).int()) + act_l.append((WINDOW + valid.sum(1)).int()) + idx = torch.cat(idx_l).contiguous() + seq_lens = torch.tensor([c + q for c, q in requests], dtype=torch.int32, device=DEV) + cu = torch.tensor([0] + [q for _, q in requests], dtype=torch.int32, device=DEV).cumsum(0).int() + loc = torch.cat(loc_l) if loc_l else None + return idx, torch.cat(act_l), seq_lens, cu, torch.cat(cnt_l), torch.cat(req_l), loc + + +class _KvqStep: + """One engine step (nc context requests, then G generation requests of next_n rows each) in both + read forms: the container form (env_pools' pools, idx with -1) and the native-storage compact form + (bufs' views, the renumbered index cidx). Per-row and per-request tensors are prefixes of + envelope-sized buffers, as the runtime's metadata buffers are.""" + + def __init__(self, ctx, gen_pos, next_n, num_heads, layer, pools, bufs, gen, floors=None): + ratio, page = ENV_LAYERS[layer] + self.ratio, self.layer, self.bufs = ratio, layer, bufs + self.width = WINDOW + (INDEX_TOPK if ratio else 0) + self.requests = list(ctx) + [(p, next_n) for p in gen_pos] + self.floors = list(floors if floors is not None else [0] * len(ctx)) + [0] * len(gen_pos) + self.nc, self.tc = len(ctx), sum(q for _, q in ctx) + self.B, self.G, self.next_n = len(self.requests), len(gen_pos), next_n + assert self.B <= bufs.seqs + self.max_ctx_q = max((q for _, q in ctx), default=0) + self.swa, self.comp = pools["swa"], pools["comp"][layer] + self.swa_bt, self.comp_bts = _env_tables( + self.requests, gen, ratio, page, self.swa.shape[0], self.comp.shape[0] + ) + idx, active, seq_lens, cu, self.counts, self.req, self.loc = _kvq_rows( + self.requests, self.floors, gen, self.width, ratio, page, self.swa_bt, self.comp_bts + ) + self.T = idx.shape[0] + assert self.T <= bufs.max_t + self.idx = idx + self.active = torch.full((max(self.T, ENV_T),), -7, dtype=torch.int32, device=DEV)[: self.T] + self.active.copy_(active) + self.kv_lens = torch.full((max(self.B, ENV_B) + 1,), -7, dtype=torch.int32, device=DEV) + self.kv_lens[: self.B] = seq_lens + self.cu = torch.full((max(self.B, ENV_B) + 2,), -7, dtype=torch.int32, device=DEV) + self.cu[: self.B + 1] = cu + self.q3 = ( + torch.randn(self.T, num_heads, HEAD_DIM, generator=gen, device=DEV) + .mul_(2) + .to(torch.bfloat16) + ) + self.sinks = _sinks(num_heads) + win = bufs.win_idx[: self.T] + if ratio: + comp = torch.where(self.loc >= 0, self.req[:, None] * KVQ_L[ratio] + self.loc, -1).int() + self.cidx = torch.cat([win, comp], 1) + assert int(comp.max()) < bufs.comp_c4[ratio].shape[0] * bufs.comp_c4[ratio].shape[2] + else: + self.cidx = win + + def fill(self, minus_one=None, skip_request=None): + """What the two dequantizers write this step (native torch gathers from the container): + row t * 128 + c of swa_c for every query t < T and window column c -- a -1 column as +0.0, + or as the constant `minus_one` for a control -- and row b * L_r + e of comp_c for every + entry e < kv_len // r of every request b (except `skip_request`, for a control). Nothing + else in the buffers changes.""" + bufs, T = self.bufs, self.T + w = self.idx[:, :WINDOW].long() + rows = self.swa.view(-1, HEAD_DIM)[w.clamp_min(0)] + fillv = torch.zeros((), dtype=torch.bfloat16, device=DEV) + if minus_one is not None: + fillv = torch.tensor(minus_one, dtype=torch.bfloat16, device=DEV) + bufs.swa_c[: T * WINDOW] = torch.where((w < 0)[..., None], fillv, rows).view( + T * WINDOW, HEAD_DIM + ) + if self.ratio: + r, bk, L = self.ratio, SWA_PAGE // self.ratio, KVQ_L[self.ratio] + dst, src = [], [] + for b, (c, q) in enumerate(self.requests): + n = (c + q) // r + if n and b != skip_request: + e = torch.arange(n, device=DEV) + src.append(self.comp_bts[b][e // bk] * bk + e % bk) + dst.append(b * L + e) + if dst: + bufs.comp_c[torch.cat(dst)] = self.comp.view(-1, HEAD_DIM)[torch.cat(src)] + + def run(self, ws, compact, idx=None, out=None): + """The target's two calls: the context call on the first tc rows, then the generation call on + the views from row tc as [G, next_n, H, 512] with seq_lens = kv_lens[nc:B].""" + tc, gr, h = self.tc, self.T - self.tc, self.q3.shape[1] + if compact: + swa = self.bufs.swa_c4 + comp = self.bufs.comp_c4[self.ratio] if self.ratio else self.bufs.dummy + idx = self.cidx if idx is None else idx + else: + swa, comp = self.swa, self.comp + idx = self.idx if idx is None else idx + o = ( + torch.empty(self.T, h, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + if out is None + else out + ) + if tc: + op( + torch.split(self.q3, [tc, gr], 0)[0], + swa, + comp, + ws, + torch.split(idx, [tc, gr], 0)[0], + torch.split(self.active, [tc, gr], 0)[0], + self.kv_lens[: self.nc], + self.sinks, + SCALE, + 1.0, + cum_seq_lens_q=self.cu[: self.nc + 1], + max_q_len=self.max_ctx_q, + out=torch.split(o, [tc, gr], 0)[0], + ) + if gr: + g, nn = self.G, self.next_n + op( + torch.split(self.q3, [tc, gr], 0)[1].reshape(g, nn, h, HEAD_DIM), + swa, + comp, + ws, + torch.split(idx, [tc, gr], 0)[1], + torch.split(self.active, [tc, gr], 0)[1], + self.kv_lens[self.nc : self.B], + self.sinks, + SCALE, + 1.0, + out=torch.split(o, [tc, gr], 0)[1].reshape(g, nn, h, HEAD_DIM), + ) + return o + + def reference(self, **kw): + return _reference( + self.q3, self.swa, self.comp, self.idx, self.active, self.counts, self.sinks, **kw + ) + + def attended_minus_one(self) -> int: + """Window columns inside the window count whose container index is -1 (floored windows).""" + col = torch.arange(WINDOW, device=DEV) + return int(((col[None] < self.counts[:, None]) & (self.idx[:, :WINDOW] < 0)).sum()) + + +def _bits(o): + return o.view(torch.int16) + + +def _kvq_case(seed): + """A random step of the engine envelope for one configuration: seeds 0..199 at (max_num_tokens, + max_batch_size) = (8192, 32), 200..259 at the three smaller ones in turn. k in 1..max_batch_size + requests, prefill-only (30 %), generation-only (20 %) or mixed; generation at next_n 1 or 6 + (DSpark verify), a quarter of generation requests in the DSpark tail up to position 16397; + contexts split T at random (log-uniform or uniform), fresh or cached (cached + q_len <= 16384); + a quarter of cached contexts carry a window floor.""" + rng = random.Random(1_500_000 + seed) + num_heads = (16, 64)[seed % 2] + layer = ("r0", "r2", "r1")[(seed // 2) % 3] + cfg = KVQ_CONFIGS[0] if seed < 200 else KVQ_CONFIGS[1 + (seed - 200) % 3] + max_t, seqs = cfg + k = rng.randint(1, seqs) + next_n = rng.choice((1, 6)) + kind = rng.random() + n_gen = 0 if kind < 0.3 else (k if kind < 0.5 else rng.randint(1, k)) + n_ctx = k - n_gen + ctx = [] + if n_ctx: + budget = max_t - n_gen * next_n + if rng.random() < 0.5: + t = int(round(math.exp(rng.uniform(0.0, math.log(budget))))) + else: + t = rng.randint(n_ctx, budget) + t = min(max(t, n_ctx), budget) + cuts = sorted(rng.sample(range(1, t), n_ctx - 1)) if n_ctx > 1 else [] + q_lens = [b - a for a, b in zip([0] + cuts, cuts + [t])] + ctx = [(_fresh_or_cached(rng, q), q) for q in q_lens] + top = KVQ_TAIL + 1 - next_n + gen = [ + rng.randint(ENV_SEQ - next_n, top) if rng.random() < 0.25 else rng.randint(0, top) + for _ in range(n_gen) + ] + floors = [ + rng.randint(max(1, c - WINDOW + 1), c) if c and rng.random() < 0.25 else 0 for c, _ in ctx + ] + return num_heads, layer, cfg, ctx, gen, next_n, floors + + +KVQ_SEEDS = list(range(260)) + + +@pytest.mark.parametrize("seed", KVQ_SEEDS) +def test_kvq_compact_read_random_partition(env_pools, kvq_buffers, seed): + """Random steps for the compact read: 260 seeded random steps (200 on the (8192, 32) buffers, 20 on + each smaller configuration's), each run in both forms: the compact read is bitwise the container + read, and the container read is inside the gate against the fp32 reference.""" + num_heads, layer, cfg, ctx, gen_pos, next_n, floors = _kvq_case(seed) + g = torch.Generator(device=DEV).manual_seed(1_600_000 + seed) + st = _KvqStep(ctx, gen_pos, next_n, num_heads, layer, env_pools, kvq_buffers[cfg], g, floors) + st.fill() + ws = _workspace() + o_pool = st.run(ws, compact=False) + o_comp = st.run(ws, compact=True) + e = _excess(o_pool, st.reference(), env_pools["atol"][layer]) + same = torch.equal(_bits(o_comp), _bits(o_pool)) + print( + f"kvq partition {seed}: cfg {cfg} H{num_heads} {layer} T {st.T} tc {st.tc} nc {st.nc} G {st.G}x{next_n} " + f"attended -1 {st.attended_minus_one()} excess {e:.3f} compact==container {same}" + ) + assert same, (seed, int((_bits(o_comp) != _bits(o_pool)).sum())) + assert e <= 1.0, (seed, e) + + +def test_kvq_compact_read_random_partitions_reach_the_envelope(): + """The 260 steps span the envelope: every (H, layer), every configuration, k from 1 to >= 30 and + T to >= 8000 on (8192, 32), prefill-only, decode-only and verify-only, mixed at next_n 1 and 6, + generation in the DSpark tail, and floored windows.""" + cases = [_kvq_case(s) for s in KVQ_SEEDS] + big = [c for c in cases if c[2] == KVQ_CONFIGS[0]] + ks = [len(c[3]) + len(c[4]) for c in big] + ts = [sum(q for _, q in c[3]) + len(c[4]) * c[5] for c in big] + kinds = {(bool(c[3]), bool(c[4]), c[5] if c[4] else 0) for c in big} + assert len(big) == 200 and min(ks) == 1 and max(ks) >= 30, (len(big), min(ks), max(ks)) + assert min(ts) <= 4 and max(ts) >= 8000, (min(ts), max(ts)) + assert { + (True, False, 0), + (False, True, 1), + (False, True, 6), + (True, True, 1), + (True, True, 6), + } <= kinds, kinds + assert sum(1 for c in cases for p in c[4] if p + c[5] - 1 >= ENV_SEQ) >= 50 + assert sum(1 for c in cases for f in c[6] if f) >= 50 + assert {c[2] for c in cases} == set(KVQ_CONFIGS) + assert all( + sum(q for _, q in c[3]) + len(c[4]) * c[5] <= c[2][0] and len(c[3]) + len(c[4]) <= c[2][1] + for c in cases + ) + assert {(c[0], c[1]) for c in cases} == {(h, layer) for h in (16, 64) for layer in ENV_LAYERS} + + +# --------------------------------------------------------------------------------------------- +# Floored windows passed with floor-adjusted seq_lens (seq_lens = kv_len - floor) +# --------------------------------------------------------------------------------------------- +def _floored_steps(): + """The random steps of _kvq_case that hold a cached context, re-floored: EVERY cached context + (cached c > 0) gets a window floor, half of them at c itself (an encoder-replay chunk: its window + starts at its own first row) and half at a random point of [max(1, c - 127), c - 1] (a decoder + replay tail, whose first row can lie above its floor), fresh contexts none. The first 40 such + seeds on the (8192, 32) buffers and the first 8 on the three smaller configurations.""" + out = [] + for lo, hi, n in ((0, 200, 40), (200, len(KVQ_SEEDS), 8)): + k = 0 + for seed in KVQ_SEEDS[lo:hi]: + num_heads, layer, cfg, ctx, gen, next_n, _ = _kvq_case(seed) + if not any(c for c, _ in ctx): + continue + rng = random.Random(1_800_000 + seed) + floors = [ + 0 + if not c + else ( + c + if rng.random() < 0.5 + else rng.randint(max(1, c - WINDOW + 1), c - 1) + if c > 1 + else c + ) + for c, _ in ctx + ] + out.append((seed, (num_heads, layer, cfg, ctx, gen, next_n, floors))) + k += 1 + if k == n: + break + return out + + +KVQ_FLOOR_CASES = _floored_steps() + + +def _floor_adjusted(st): + """Set the step's context seq_lens to kv_len - floor and return the per-row window counts that + gives, min(128, p - floor + 1), with each row's position p and floor.""" + pos_l, flo_l = [], [] + for b, ((c, q), f) in enumerate(zip(st.requests, st.floors)): + pos_l.append(torch.arange(c, c + q, device=DEV)) + flo_l.append(torch.full((q,), f, dtype=torch.int64, device=DEV)) + if b < st.nc: + st.kv_lens[b] = c + q - f + pos, flo = torch.cat(pos_l), torch.cat(flo_l) + return (pos - flo + 1).clamp(max=WINDOW), pos, flo + + +@pytest.mark.parametrize("seed", [s for s, _ in KVQ_FLOOR_CASES]) +def test_kvq_floored_windows_with_floor_adjusted_seq_lens(env_pools, kvq_buffers, seed): + """A context request whose window is floored (the window of position p is max(floor, p - 127) + .. p, then -1) passed with seq_lens = kv_len - floor. To the kernel that is a context with + kv_len - floor - q_len cached rows (none for an encoder-replay chunk): its window count + min(128, p - floor + 1) is exactly the floored window, so no -1 is attended and the result is + the MASKED reading of the floored window (the zero-row reading of seq_lens = kv_len is the + other one). In both read forms against the fp32 reference, compact == container bitwise; and + every row whose window reaches no position below its floor -- fresh contexts, the generation + rows, a floored request's rows at p - 127 >= floor -- is bitwise the same step called with + seq_lens = kv_len.""" + _floored_step(env_pools, kvq_buffers, seed) + + +@pytest.mark.parametrize("seed", [s for s, _ in KVQ_FLOOR_CASES]) +def test_kvq_floored_windows_at_the_members_head_counts(env_pools, kvq_buffers, seed): + """The same 48 floored steps with the members' head counts in place of the case's (H 16 -> 8, + H 64 -> 32), the same checks.""" + _floored_step(env_pools, kvq_buffers, seed, {16: 8, 64: 32}) + + +def _floored_step(env_pools, kvq_buffers, seed, heads_map=None): + num_heads, layer, cfg, ctx, gen_pos, next_n, floors = dict(KVQ_FLOOR_CASES)[seed] + if heads_map is not None: + num_heads = heads_map[num_heads] + g = torch.Generator(device=DEV).manual_seed(1_900_000 + seed) + st = _KvqStep(ctx, gen_pos, next_n, num_heads, layer, env_pools, kvq_buffers[cfg], g, floors) + st.fill() + ws = _workspace() + o_kv = st.run(ws, compact=True) + old_counts = st.counts + counts, pos, flo = _floor_adjusted(st) + col = torch.arange(WINDOW, device=DEV) + attended_neg = int(((col[None] < counts[:, None]) & (st.idx[:, :WINDOW] < 0)).sum()) + valid = (st.idx[:, :WINDOW] >= 0).sum(1) + assert attended_neg == 0 and torch.equal(counts, valid) and int(counts.min()) >= 1 + assert all(int(st.kv_lens[b]) >= q for b, (_, q) in enumerate(st.requests[: st.nc])) + o_pool = st.run(ws, compact=False) + o_comp = st.run(ws, compact=True) + ref = _reference(st.q3, st.swa, st.comp, st.idx, st.active, counts, st.sinks) + masked = _reference( + st.q3, st.swa, st.comp, st.idx, st.active, old_counts, st.sinks, negative_index="masked" + ) + assert torch.equal(ref, masked), "the floor-adjusted count is not the masked floored window" + e = _excess(o_pool, ref, env_pools["atol"][layer]) + same = torch.equal(_bits(o_comp), _bits(o_pool)) + below = (pos - (WINDOW - 1) < flo) & (flo > 0) # rows whose window the floor truncates + keep = ~below + unchanged = torch.equal(_bits(o_comp)[keep], _bits(o_kv)[keep]) + print( + f"floored {seed}: cfg {cfg} H{num_heads} {layer} T {st.T} tc {st.tc} nc {st.nc} G {st.G}x{next_n} floors " + f"{sum(1 for f in floors if f)} (at the first row {sum(1 for (c, _), f in zip(ctx, floors) if f and f == c)}) " + f"truncated rows {int(below.sum())} excess {e:.3f} compact==container {same} other rows unchanged {unchanged}" + ) + assert same, (seed, int((_bits(o_comp) != _bits(o_pool)).sum())) + assert e <= 1.0, (seed, e) + assert unchanged, (seed, int((_bits(o_comp)[keep] != _bits(o_kv)[keep]).sum())) + + +def test_kvq_floored_windows_reach_the_forms(): + """The 48 steps cover both floor forms (at the first row, and below it), H 16 / 64 x every layer, + every buffer configuration, mixed steps at next_n 1 and 6, and a floored request whose first row + sits >= 127 past its floor or at it with >= 127 cached rows below (the first row then counts 127 + columns fewer).""" + cases = [c for _, c in KVQ_FLOOR_CASES] + assert len(cases) == 48 + at_first = sum(1 for c in cases for (cc, _), f in zip(c[3], c[6]) if f and f == cc) + below = sum(1 for c in cases for (cc, _), f in zip(c[3], c[6]) if f and f < cc) + assert at_first >= 20 and below >= 20, (at_first, below) + assert {(c[0], c[1]) for c in cases} == {(h, layer) for h in (16, 64) for layer in ENV_LAYERS} + assert {c[5] for c in cases if c[4]} == {1, 6} + assert {c[2] for c in cases} == set(KVQ_CONFIGS) + assert ( + sum( + 1 for c in cases for (cc, _), f in zip(c[3], c[6]) if f and f == cc and cc >= WINDOW - 1 + ) + >= 10 + ) + + +def test_kvq_floored_windows_kv_len_reading_fails_the_gate(env_pools, kvq_buffers): + """The harness sees the difference the floor-adjusted seq_lens makes. On two encoder-replay chunks + (floor = the first row, 127 and 3970, 129 rows each) beside a fresh context and generation rows, + at H 16 / 64 and every layer: seq_lens = kv_len - floor is inside the gate against the masked + reference, while seq_lens = kv_len (a chunk's first rows attend up to 127 below-floor -1 columns + as zero rows) lands outside it on those rows; both calls agree bitwise on every other row. The + chunk at 127 has few compressed columns (64 / 128 visible at ratio 2 / 1), so the zero rows' + weight shows at every layer; at 3970 the 512 compressed columns dilute it (about the gate at + ratio 2 / 1, printed).""" + ws = _workspace() + for num_heads in (16, 64): + for layer in ENV_LAYERS: + g = torch.Generator(device=DEV).manual_seed(2_100_000 + num_heads + 7 * len(layer)) + ctx, floors = [(127, 129), (3970, 129), (0, 300)], [127, 3970, 0] + st = _KvqStep( + ctx, + [5000, 700], + 6, + num_heads, + layer, + env_pools, + kvq_buffers[KVQ_CONFIGS[0]], + g, + floors, + ) + st.fill() + atol = env_pools["atol"][layer] + o_kv = st.run(ws, compact=True) + counts, pos, flo = _floor_adjusted(st) + o_fix = st.run(ws, compact=True) + ref = _reference(st.q3, st.swa, st.comp, st.idx, st.active, counts, st.sinks) + hit = (pos - (WINDOW - 1) < flo) & (flo > 0) + wrong = { + f: _excess(o_kv[hit & (flo == f)], ref[hit & (flo == f)], atol) for f in (127, 3970) + } + right = _excess(o_fix, ref, atol) + print( + f"floored control H{num_heads} {layer}: seq_lens = kv_len - floor {right:.3f} of the gate, " + f"seq_lens = kv_len {wrong[127]:.1f} (floor 127) / {wrong[3970]:.1f} (floor 3970) on the " + f"{int(hit.sum())} truncated rows" + ) + assert right <= 1.0 and wrong[127] > 10.0 * max(right, 0.1), ( + num_heads, + layer, + right, + wrong, + ) + assert torch.equal(_bits(o_kv)[~hit], _bits(o_fix)[~hit]) + + +@pytest.mark.parametrize("width", [128, 640]) +@pytest.mark.parametrize("num_heads", [16, 64]) +def test_kvq_compact_read_kernel_variant_bands(env_pools, kvq_buffers, num_heads, width): + """Kernel bands for the compact read. The context call's kernel is picked from (H, width, number of + requests, max_q_len), and the compact read changes none of those. Every band of VARLEN_BANDS at + its bottom, one interior point and its top at every request count 1..32 (_band_cases) runs as a + context step on the (8192, 32) compact buffers: the profiler reads back exactly the band's + kernels for every compact call, and every compact call is bitwise the container read on the + same values.""" + rng = random.Random(num_heads * 1000 + width + 17) + gen = torch.Generator(device=DEV).manual_seed(num_heads * 1000 + width + 17) + ws = _workspace() + bufs = kvq_buffers[KVQ_CONFIGS[0]] + cases = _band_cases(num_heads, width, rng) + seen, differ = [], [] + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for i, (ctx, vid) in enumerate(cases): + layer = "r0" if width == WINDOW else ("r2", "r1")[i % 2] + st = _KvqStep(ctx, [], 1, num_heads, layer, env_pools, bufs, gen) + st.fill() + o_pool = st.run(ws, compact=False) + torch.cuda.synchronize() + torch.cuda._sleep(10) + o = st.run(ws, compact=True) + torch.cuda._sleep(10) + if not torch.equal(_bits(o), _bits(o_pool)): + differ.append((ctx, vid)) + torch.cuda.synchronize() + ev = sorted( + (e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA), + key=lambda e: e.time_range.start, + ) + open_call, cur = False, [] + for e in ev: + if "spin_kernel" in e.name or "sleep" in e.name: + if open_call: + seen.append(_variant_of(cur)) + open_call, cur = not open_call, [] + elif open_call: + cur.append(e.name) + assert len(seen) == len(cases), (len(seen), len(cases)) + mism = [(c, VARLEN_KERNELS[v], s) for (c, v), s in zip(cases, seen) if VARLEN_KERNELS[v] != s] + used = {v for _, v in cases} + print( + f"kvq bands H{num_heads} W{width}: {len(cases)} compact calls, {len(used)} variants, " + f"kernel mismatches {len(mism)}, not bitwise the container read {len(differ)}" + ) + assert not mism, mism[:3] + assert not differ, differ[:3] + + +@pytest.mark.parametrize("width", [128, 640]) +@pytest.mark.parametrize("num_heads", [16, 64]) +def test_kvq_compact_read_every_dumped_cell(env_pools, kvq_buffers, num_heads, width): + """The dumped partitions, for the compact read: every context call (tc, H, nc, max_q_len) + of DUMPED_CTX_CELLS as a step of nc requests (_dumped_partition) on the (8192, 32) compact + buffers, at width 128 (ratio 0) and 640 (ratio 2 and ratio 1 on alternate cells): each compact + read is bitwise the container read on the same values, which is inside the gate.""" + rng = random.Random(78000 + num_heads + width) + gen = torch.Generator(device=DEV).manual_seed(78000 + num_heads + width) + ws = _workspace() + bufs = kvq_buffers[KVQ_CONFIGS[0]] + cells = _dumped_cells(num_heads) + worst, bad, differ = 0.0, [], [] + for i, (tc, nc, mq) in enumerate(cells): + layer = "r0" if width == WINDOW else ("r2", "r1")[i % 2] + ctx = _dumped_partition(tc, nc, mq, rng) + st = _KvqStep(ctx, [], 1, num_heads, layer, env_pools, bufs, gen) + st.fill() + o_pool = st.run(ws, compact=False) + o_comp = st.run(ws, compact=True) + if not torch.equal(_bits(o_comp), _bits(o_pool)): + differ.append((tc, nc, mq, layer)) + e = _excess(o_pool, st.reference(), env_pools["atol"][layer]) + worst = max(worst, e) + if e > 1.0: + bad.append(((tc, nc, mq), layer, e)) + print( + f"kvq dumped cells H{num_heads} W{width}: {len(cells)} steps, not bitwise the container read " + f"{len(differ)}, max excess {worst:.3f}" + ) + assert not differ, differ[:5] + assert not bad, bad[:5] + + +def _kvq_far_steps(max_t, seqs): + """Steps that reach both buffers' far ends: T = max_num_tokens with max_batch_size requests, the + last a DSpark verify ending at 16397 (the window half names swa_c's last row, the compressed half + row (seqs - 1) * L_r + 16398 // r - 1), without and with window floors on every cached context; + and a verify-only step of max_batch_size requests at the tail.""" + n_ctx = seqs - 1 + rest = max_t - 6 - 3 * (n_ctx - 1) + ctx = [(ENV_SEQ - rest, rest)] + [(131 + 97 * i, 3) for i in range(n_ctx - 1)] + tail = [KVQ_TAIL - 5] + floors = [c - 40 for c, _ in ctx] + verify = [KVQ_TAIL - 5 - (i % 3) for i in range(seqs)] + return [(ctx, tail, 6, None), (ctx, tail, 6, floors), ([], verify, 6, None)] + + +@pytest.mark.parametrize("cfg", KVQ_CONFIGS, ids=[f"{t}x{s}" for t, s in KVQ_CONFIGS]) +def test_kvq_compact_read_far_end_of_every_buffer_size(env_pools, kvq_buffers, cfg): + """At each buffer size the target allocates, the steps of _kvq_far_steps at H 16 and 64 and + every layer kind: the window half names the window buffer's last row, the compressed half the + largest row the compressed buffer can hold for this configuration (asserted); the compact read is + bitwise the container read (the floored steps with attended -1s among them), inside the gate.""" + bufs = kvq_buffers[cfg] + max_t, seqs = cfg + ws = _workspace() + worst, n = 0.0, 0 + for i, (ctx, gen_pos, next_n, floors) in enumerate(_kvq_far_steps(max_t, seqs)): + for num_heads in (16, 64): + for layer in ENV_LAYERS: + g = torch.Generator(device=DEV).manual_seed( + 1_700_000 + 100 * i + num_heads + ENV_LAYERS[layer][0] + ) + st = _KvqStep(ctx, gen_pos, next_n, num_heads, layer, env_pools, bufs, g, floors) + if i < 2: + assert ( + st.T == max_t + and st.B == seqs + and int(st.cidx[:, :WINDOW].max()) == max_t * WINDOW - 1 + ) + if st.ratio: + far = (seqs - 1) * KVQ_L[st.ratio] + (KVQ_TAIL + 1) // st.ratio - 1 + assert int(st.cidx[:, WINDOW:].max()) == far + if floors is not None: + assert st.attended_minus_one() > 0 + st.fill() + o_pool = st.run(ws, compact=False) + o_comp = st.run(ws, compact=True) + assert torch.equal(_bits(o_comp), _bits(o_pool)), (cfg, i, num_heads, layer) + e = _excess(o_pool, st.reference(), env_pools["atol"][layer]) + assert e <= 1.0, (cfg, i, num_heads, layer, e) + worst, n = max(worst, e), n + 1 + print(f"kvq far end {cfg}: {n} steps, compact == container bitwise, max excess {worst:.3f}") + + +@pytest.mark.parametrize("layer", list(ENV_LAYERS)) +@pytest.mark.parametrize("num_heads", [16, 64]) +def test_kvq_compact_read_cuda_graph_capture_set(env_pools, kvq_buffers, num_heads, layer): + """The capture set on the compact buffers: the target's generation call (dense [G, next_n, H, 512], + seq_lens = kv_lens[0:G], index / lengths / seq_lens / q prefixes of step buffers, the output + allocated inside the capture) at G in 1..8, 16, 24, 32 x next_n 1 and 6 on the (8192, 32) + buffers and G 1..8 x next_n 1 and 6 on the (2048, 8) buffers, each captured after one eager + call at that shape. The first replay equals the eager call bitwise; three more on new steps (new + compact rows, q, compressed half, lengths, seq_lens; request 0 in the DSpark tail) each equal a + fresh eager compact call AND the container read on the same values bitwise, and match the + reference.""" + ratio = ENV_LAYERS[layer][0] + ws = _workspace() + sinks = _sinks(num_heads) + n_graphs, worst = 0, 0.0 + for cfg, batches in ((KVQ_CONFIGS[0], GRAPH_BATCHES), (KVQ_CONFIGS[3], GRAPH_BATCHES[:8])): + bufs = kvq_buffers[cfg] + max_t, seqs = cfg + width = WINDOW + (INDEX_TOPK if ratio else 0) + idx_buf = torch.full((max_t, width), -1, dtype=torch.int32, device=DEV) + lens_buf = torch.full((max_t,), WINDOW, dtype=torch.int32, device=DEV) + kvl_buf = torch.full((seqs + 1,), 1, dtype=torch.int32, device=DEV) + q_buf = torch.zeros(max_t, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + comp = bufs.comp_c4[ratio] if ratio else bufs.dummy + for next_n in (1, 6): + for G in batches: + T = G * next_n + gen = torch.Generator(device=DEV).manual_seed( + 1_800_000 + 1000 * G + 100 * next_n + num_heads + ratio + ) + rng = random.Random(1_900_000 + 1000 * G + 100 * next_n + num_heads + ratio) + idx = bufs.win_idx[:T] if not ratio else idx_buf[:T] + + def make(step): + pos = [rng.randint(0, KVQ_TAIL + 1 - next_n) for _ in range(G)] + pos[0] = (KVQ_TAIL if step % 2 == 0 else ENV_SEQ - 1) + 1 - next_n + st = _KvqStep([], pos, next_n, num_heads, layer, env_pools, bufs, gen) + st.fill() + q_buf[:T].copy_(st.q3) + if ratio: + idx_buf[:T].copy_(st.cidx) + lens_buf[:T].copy_(st.active) + kvl_buf[:G].copy_(st.kv_lens[:G]) + return st + + def call(out): + return op( + q_buf[:T].view(G, next_n, num_heads, HEAD_DIM), + bufs.swa_c4, + comp, + ws, + idx, + lens_buf[:T], + kvl_buf[:G], + sinks, + SCALE, + 1.0, + out=out, + ) + + shape = (G, next_n, num_heads, HEAD_DIM) + st = make(0) + eager = call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + assert torch.equal( + _bits(eager.view(T, num_heads, HEAD_DIM)), _bits(st.run(ws, compact=False)) + ) + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + torch.cuda.current_stream().wait_stream(s) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + o = torch.empty(T, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + out = call(o.view(shape)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(_bits(out), _bits(eager)), (cfg, G, next_n) + prev = eager.clone() + for it in range(1, 4): + st = make(it) + out.fill_(float("nan")) + graph.replay() + fresh = call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + pool = st.run(ws, compact=False) + torch.cuda.synchronize() + assert torch.equal(_bits(out), _bits(fresh)), ( + cfg, + G, + next_n, + it, + "replay != eager", + ) + assert torch.equal(_bits(out.view(T, num_heads, HEAD_DIM)), _bits(pool)), ( + cfg, + G, + next_n, + it, + ) + assert not torch.equal(_bits(out), _bits(prev)), (cfg, G, next_n, it) + e = _excess(pool, st.reference(), env_pools["atol"][layer]) + assert e <= 1.0, (cfg, G, next_n, it, e) + worst = max(worst, e) + prev = out.clone() + del graph + n_graphs += 1 + print( + f"kvq capture H{num_heads} {layer}: {n_graphs} graphs x 4 replays, replay = eager = container, " + f"max excess {worst:.3f}" + ) + + +def _kvq_control_steps(): + """(H, layer, cfg, ctx, gen, next_n, floors): mixed steps whose cached contexts carry window + floors (attended -1s), a fresh short context (masked -1 window tail), and DSpark-tail + generation rows.""" + ctx = [(300, 200), (5000, 64), (0, 50), (16000, 300), (129, 1)] + floors = [250, 4990, 0, 15950, 60] + return [ + (16, "r2", KVQ_CONFIGS[0], ctx, [16392, 77, 4000], 6, floors), + (64, "r1", KVQ_CONFIGS[0], ctx, [16397, 128, 9000], 1, floors), + (16, "r0", KVQ_CONFIGS[3], ctx, [16392, 3], 6, floors), + (64, "r2", KVQ_CONFIGS[2], ctx, [16397, 500, 12000, 7], 1, floors), + ] + + +def test_kvq_compact_read_wrong_backing_and_numbering_fail_the_gate(env_pools, kvq_buffers): + """The harness sees each way the compact read can go wrong. On mixed steps with floored windows + (attended -1s) the correct compact read is bitwise the container read and inside the gate, while + the reference lands far outside it for: + - '-1 not zero-backed': the rows standing for -1 window columns hold 3.0 instead of +0.0; + - the window half renumbered one row late (t * 128 + c + 1); + - the compressed half renumbered one entry late (b * L_r + e + 1), or at the other ratio's row + stride (b * L_{3-r} + e); + - one request's compressed rows not rewritten this step (stale). + On the same steps without floors (only masked -1s), the 3.0 backing leaves the output bitwise + unchanged (masked columns carry no weight), while NaN backing poisons it (masked columns are + loaded): the backing must be finite everywhere and zero where a -1 is attended.""" + ws = _workspace() + for num_heads, layer, cfg, ctx, gen_pos, next_n, floors in _kvq_control_steps(): + bufs = kvq_buffers[cfg] + ratio = ENV_LAYERS[layer][0] + atol = env_pools["atol"][layer] + for floored in (True, False): + g = torch.Generator(device=DEV).manual_seed(2_000_000 + num_heads + 10 * ratio) + st = _KvqStep( + ctx, + gen_pos, + next_n, + num_heads, + layer, + env_pools, + bufs, + g, + floors if floored else None, + ) + ref = st.reference() + st.fill() + o_pool = st.run(ws, compact=False) + o_comp = st.run(ws, compact=True) + assert torch.equal(_bits(o_comp), _bits(o_pool)) + correct = _excess(o_comp, ref, atol) + assert correct <= 1.0, correct + neg = st.attended_minus_one() + st.fill(minus_one=3.0) + backed = st.run(ws, compact=True) + st.fill(minus_one=float("nan")) + nan_backed = st.run(ws, compact=True) + st.fill() + if not floored: + assert neg == 0 + assert torch.equal(_bits(backed), _bits(o_pool)), ( + "a finite masked row changed the output" + ) + assert not torch.isfinite(nan_backed).all(), ( + "a NaN masked row did not poison the output" + ) + print( + f"kvq controls H{num_heads} {layer} {cfg} unfloored: 3.0 backing bitwise unchanged, NaN backing " + f"poisons {int((~torch.isfinite(nan_backed)).any(-1).any(-1).sum())} of {st.T} rows" + ) + continue + assert neg > 0 + wrong = {"minus_one_not_zero_backed": _excess(backed, ref, atol)} + w1 = st.cidx.clone() + w1[:, :WINDOW] += 1 + wrong["window_renumbered_one_row_late"] = _excess( + st.run(ws, compact=True, idx=w1), ref, atol + ) + if ratio: + c = st.cidx[:, WINDOW:] + w2 = st.cidx.clone() + w2[:, WINDOW:] = torch.where(c >= 0, c + 1, c) + wrong["compressed_renumbered_one_entry_late"] = _excess( + st.run(ws, compact=True, idx=w2), ref, atol + ) + w3 = st.cidx.clone() + w3[:, WINDOW:] = torch.where( + st.loc >= 0, st.req[:, None] * KVQ_L[3 - ratio] + st.loc, -1 + ).int() + assert int(w3.max()) < bufs.comp_c4[ratio].shape[0] * bufs.comp_c4[ratio].shape[2] + wrong["other_ratio_row_stride"] = _excess( + st.run(ws, compact=True, idx=w3), ref, atol + ) + L = KVQ_L[ratio] + bufs.comp_c[L : 2 * L] = torch.randn( + L, HEAD_DIM, generator=g, device=DEV, dtype=torch.bfloat16 + ) + st.fill( + skip_request=1 + ) # request 1's rows keep the content above (an earlier step's) + wrong["request_1_rows_stale"] = _excess(st.run(ws, compact=True), ref, atol) + st.fill() + print( + f"kvq controls H{num_heads} {layer} {cfg}: attended -1 {neg}, correct {correct:.3f}, wrong " + f"{ {k: round(v, 1) for k, v in wrong.items()} }" + ) + for name, e in wrong.items(): + assert e > 10.0 * max(correct, 0.1), ( + f"{name} not discriminated: {e:.2f} vs {correct:.2f}" + ) + + +def test_kvq_compact_buffers_ending_a_mapping_never_fault(): + """Precondition 8 at the compact buffers' exact sizes. Each view the target hands over is placed + so that its last row ends a mapping (reserved, unmapped VA after it): the window buffer + [2048 | 8192, 1, 128, 512] with its last 128 rows attended (the last query of a full step; the + compact window half holds no -1); the compressed views [8 | 32 x 132, 1, 64 | 128, 512] with + their last rows attended and a -1 loaded from them (lengths 129 / 131 / 133 / 639). Nothing + faults, and every output is bitwise the output on a normal allocation. The smallest of these is + 67,584 rows: a loaded -1's over-read (a row in 96..127 from the first row) lands inside it.""" + trials = [] + for max_t, h in ((2048, 16), (8192, 64)): + rows = max_t * WINDOW + trials.append( + dict( + name=f"kvq-swa-{max_t}-H{h}", + kind="engine", + pool="swa", + rows=rows, + page=WINDOW, + lens=WINDOW + INDEX_TOPK, + heads=h, + window=WINDOW, + mapped=rows, + seed=2_100_000 + max_t, + ) + ) + for (seqs, r), lens, h in ( + ((8, 2), 129, 16), + ((32, 2), 131, 64), + ((8, 1), 133, 64), + ((32, 1), 639, 16), + ): + page = 128 // r + rows = seqs * KVQ_COLS * page + assert rows >= 67_584 + trials.append( + dict( + name=f"kvq-comp-r{r}-seqs{seqs}-lens{lens}-H{h}", + kind="engine", + pool="comp", + rows=rows, + page=page, + lens=lens, + heads=h, + window=0, + mapped=rows, + seed=2_100_000 + lens, + ) + ) + p, ok, tried = _guard_run(trials) + names = [c["name"] for c in trials] + assert p.returncode == 0, f"child failed after {tried[-1:]}: {p.stderr[-2000:]}" + assert list(ok) == names, (list(ok), names) + assert all(ok.values()), [k for k, v in ok.items() if not v] + print( + f"kvq guard page: {len(names)} compact-buffer trials completed, each bitwise = the normal allocation" + ) + + +# --------------------------------------------------------------------------------------------- +# The native-storage compact read over max_batch_size +# --------------------------------------------------------------------------------------------- +# The compressed read buffer is sized by the engine's max_batch_size s, and the window buffer by its +# max_num_tokens only. Measured in the dequantizer entries' tests (every max_num_tokens 2048 / 8192 x +# s 1..32 x DSpark off / on, max_seq_len 16384 / 16398): the metadata's compress_block_tables[r] are +# [s, 132] and swa_local_indices [max_num_tokens, 128], and the CUDA-graph metadata copy keeps both. +# So the target hands this call +# window [max_num_tokens, 1, 128, 512] (262,144 / 1,048,576 rows, whatever s is) +# compressed [s * 132, 1, 128 / r, 512] (16,896 * s / r rows; 8,448 / 16,896 at s = 1) +# with entry e of request b at row b * L_r + e, L_r = 132 * 128 / r. The largest compressed row a step +# can name is (s - 1) * L_r + 16398 // r - 1. The engine captures every decode batch G in 1..s +# (padding off; padding on captures a subset plus s). The tests below run the compact read at every +# s in 1..32, bitwise against the container read on the same values. The witness engine's +# configuration, max_num_tokens 2048 with max_batch_size 1, runs as its own configuration. +KVQ_MBS = tuple(range(1, 33)) +KVQ_MBS_LISTED = (1, 2, 3, 4, 8, 16, 32) +KVQ_WITNESS = (2048, 1) +# (max_num_tokens, max_batch_size) -> random steps: 200 at every s (8192 tokens) and at the witness. +KVQ_MBS_STEPS = {(ENV_T, s): 200 for s in KVQ_MBS} +KVQ_MBS_STEPS[KVQ_WITNESS] = 200 +KVQ_MBS_CAPTURE = [(ENV_T, s) for s in KVQ_MBS] + [KVQ_WITNESS] +KVQ_MBS_SMALLEST_ROWS = KVQ_COLS * 64 # the ratio-2 view at s = 1 + + +class _KvqAxisBuffers: + """_KvqBuffers for one (max_num_tokens, s) whose window buffer is shared by every s at that + max_num_tokens (the window buffer does not depend on s). The compressed buffer is exactly + s * 132 * 128 rows, its ratio views exactly [s * 132, 1, 128 / r, 512].""" + + def __init__(self, swa_c, seqs, gen): + self.max_t, self.seqs = swa_c.shape[0] // WINDOW, seqs + self.swa_c = swa_c + self.comp_c = torch.randn( + seqs * KVQ_COLS * 128, HEAD_DIM, generator=gen, device=DEV, dtype=torch.bfloat16 + ) + self.swa_c4 = self.swa_c.view(self.max_t, 1, WINDOW, HEAD_DIM) + self.comp_c4 = { + r: self.comp_c[: seqs * KVQ_COLS * (128 // r)].view( + seqs * KVQ_COLS, 1, 128 // r, HEAD_DIM + ) + for r in (1, 2) + } + self.win_idx = torch.arange(self.max_t * WINDOW, dtype=torch.int32, device=DEV).view( + self.max_t, WINDOW + ) + self.dummy = torch.zeros(1, 1, 64, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + for r, v in self.comp_c4.items(): + assert ( + v.data_ptr() == self.comp_c.data_ptr() + and v.shape[0] * v.shape[2] == 16_896 * seqs // r + ) + + +@pytest.fixture(scope="module") +def kvq_axis_buffers(): + """Lazily built _KvqAxisBuffers per (max_num_tokens, s): about 1.3 GB of window buffers and 8.7 GB + of compressed buffers once every s has been used.""" + gen = torch.Generator(device=DEV).manual_seed(3_400_000) + swa = {} + bufs = {} + + def get(cfg): + if cfg not in bufs: + max_t, seqs = cfg + if max_t not in swa: + swa[max_t] = torch.randn( + max_t * WINDOW, HEAD_DIM, generator=gen, device=DEV, dtype=torch.bfloat16 + ) + bufs[cfg] = _KvqAxisBuffers(swa[max_t], seqs, gen) + return bufs[cfg] + + yield get + bufs.clear() + swa.clear() + torch.cuda.empty_cache() + + +def _kvq_axis_case(cfg, i): + """Step i of configuration cfg = (max_num_tokens, s): _kvq_case's distribution with k <= s. Step 0 + is a full batch (k = s; mixed when s >= 2), step 1 a single request, step 2 a verify-only full batch + in the DSpark tail; the rest draw k from 1..s. H and the layer cycle with i.""" + max_t, seqs = cfg + rng = random.Random(3_500_000 + 100_000 * (max_t // 2048) + 1000 * seqs + i) + num_heads = (16, 64)[i % 2] + layer = ("r0", "r2", "r1")[(i // 2) % 3] + next_n = rng.choice((1, 6)) + if i == 0: + k = seqs + n_gen = rng.randint(1, k - 1) if k >= 2 else 0 + elif i == 1: + k = 1 + n_gen = rng.randint(0, 1) + elif i == 2: + k, n_gen, next_n = seqs, seqs, 6 + else: + k = rng.randint(1, seqs) + kind = rng.random() + n_gen = 0 if kind < 0.3 else (k if kind < 0.5 else rng.randint(1, k)) + n_ctx = k - n_gen + ctx = [] + if n_ctx: + budget = max_t - n_gen * next_n + if i == 0 or rng.random() < 0.5: + t = budget if i == 0 else int(round(math.exp(rng.uniform(0.0, math.log(budget))))) + else: + t = rng.randint(n_ctx, budget) + t = min(max(t, n_ctx), budget) + cuts = sorted(rng.sample(range(1, t), n_ctx - 1)) if n_ctx > 1 else [] + q_lens = [b - a for a, b in zip([0] + cuts, cuts + [t])] + ctx = [(_fresh_or_cached(rng, q), q) for q in q_lens] + top = KVQ_TAIL + 1 - next_n + if i == 2: # every request ends in 16384..16397 + gen = [rng.randint(ENV_SEQ - next_n + 1, top) for _ in range(n_gen)] + else: + gen = [ + rng.randint(ENV_SEQ - next_n, top) if rng.random() < 0.25 else rng.randint(0, top) + for _ in range(n_gen) + ] + floors = [ + rng.randint(max(1, c - WINDOW + 1), c) if c and rng.random() < 0.25 else 0 for c, _ in ctx + ] + return num_heads, layer, ctx, gen, next_n, floors + + +def _kvq_axis_ids(cfgs): + return [f"{t}x{s}" for t, s in cfgs] + + +@pytest.mark.parametrize("cfg", list(KVQ_MBS_STEPS), ids=_kvq_axis_ids(KVQ_MBS_STEPS)) +def test_kvq_mbs_compact_read_random_partitions(env_pools, kvq_axis_buffers, cfg): + """Random steps on the compact buffers of every max_batch_size s in 1..32 (max_num_tokens 8192) and of + the witness configuration (2048, 1): 200 seeded steps at each (_kvq_axis_case: k <= s requests, + T <= max_num_tokens, prefill, decode, verify and mixed, DSpark tail, floored windows, H 16 / 64 x + every layer kind). Every compact + read is bitwise the container read on the same values, and the container read is inside the + gate.""" + bufs = kvq_axis_buffers(cfg) + ws = _workspace() + n = KVQ_MBS_STEPS[cfg] + worst, rows, attended, ks, kinds = 0.0, 0, 0, set(), set() + for i in range(n): + num_heads, layer, ctx, gen_pos, next_n, floors = _kvq_axis_case(cfg, i) + g = torch.Generator(device=DEV).manual_seed( + 3_600_000 + 1000 * cfg[1] + i + (cfg[0] == 2048) * 500_000 + ) + st = _KvqStep(ctx, gen_pos, next_n, num_heads, layer, env_pools, bufs, g, floors) + assert st.B <= cfg[1] and st.T <= cfg[0] + st.fill() + o_pool = st.run(ws, compact=False) + o_comp = st.run(ws, compact=True) + same = torch.equal(_bits(o_comp), _bits(o_pool)) + assert same, (cfg, i, int((_bits(o_comp) != _bits(o_pool)).sum())) + e = _excess(o_pool, st.reference(), env_pools["atol"][layer]) + assert e <= 1.0, (cfg, i, e) + worst, rows, attended = max(worst, e), rows + st.T, attended + st.attended_minus_one() + ks.add(st.B) + kinds.add((bool(ctx), bool(gen_pos), next_n if gen_pos else 0)) + print( + f"kvq mbs partitions {cfg}: {n} steps, {rows} query rows, k {min(ks)}..{max(ks)}, kinds {sorted(kinds)}, " + f"attended -1 {attended}, compact == container bitwise, max excess {worst:.3f}" + ) + assert min(ks) == 1 and max(ks) == cfg[1] + if cfg[1] >= 2: + assert any(c and g for c, g, _ in kinds), kinds + + +def test_kvq_mbs_random_partitions_reach_the_envelope(): + """The axis steps span what each configuration permits: every (H, layer) and a full batch of + k = s requests at every s, a mixed step at every s >= 2, T = max_num_tokens everywhere, generation + in the DSpark tail everywhere, and floored windows.""" + for cfg, n in KVQ_MBS_STEPS.items(): + cases = [_kvq_axis_case(cfg, i) for i in range(n)] + ts = [sum(q for _, q in c[2]) + len(c[3]) * c[4] for c in cases] + assert all(t <= cfg[0] and len(c[2]) + len(c[3]) <= cfg[1] for t, c in zip(ts, cases)), cfg + assert {(c[0], c[1]) for c in cases} == { + (h, layer) for h in (16, 64) for layer in ENV_LAYERS + }, cfg + assert max(ts) == cfg[0], (cfg, max(ts)) + assert max(len(c[2]) + len(c[3]) for c in cases) == cfg[1], cfg + assert any(p + c[4] - 1 >= ENV_SEQ for c in cases for p in c[3]), cfg + if cfg[1] >= 2: + assert any(c[2] and c[3] for c in cases), cfg + big = [_kvq_axis_case(cfg, i) for cfg, n in KVQ_MBS_STEPS.items() if n == 200 for i in range(n)] + assert sum(1 for c in big for f in c[5] if f) >= 50 + + +def _kvq_axis_far_steps(max_t, seqs): + """_kvq_far_steps for every s. At s >= 2 a full step (T = max_num_tokens, s requests, the last a + DSpark verify ending at 16397) names the window buffer's last row and compressed row + (s - 1) * L_r + 16398 // r - 1; with floors, and a verify-only full batch in the tail. At s = 1 the + window buffer's last row is named by one context of max_num_tokens tokens (also with a floor), and + compressed row 16398 // r - 1 by one verify ending at 16397.""" + if seqs >= 2: + return _kvq_far_steps(max_t, seqs) + ctx = [(ENV_SEQ - max_t, max_t)] + return [(ctx, [], 1, None), (ctx, [], 1, [ENV_SEQ - max_t - 40]), ([], [KVQ_TAIL - 5], 6, None)] + + +KVQ_MBS_FAR = [(t, s) for t in (2048, ENV_T) for s in KVQ_MBS] + + +@pytest.mark.parametrize("cfg", KVQ_MBS_FAR, ids=_kvq_axis_ids(KVQ_MBS_FAR)) +def test_kvq_mbs_compact_read_far_end(env_pools, kvq_axis_buffers, cfg): + """The far end of every buffer size the target allocates over max_batch_size 1..32 at both + max_num_tokens: the steps of _kvq_axis_far_steps at H 16 and 64 and every layer kind name the + window buffer's last row and the largest compressed row (s - 1) * L_r + 16398 // r - 1 of this s + (asserted over the steps); the compact read is bitwise the container read (the floored steps with + attended -1s among them), inside the gate.""" + bufs = kvq_axis_buffers(cfg) + max_t, seqs = cfg + ws = _workspace() + worst, n = 0.0, 0 + top_win, top_comp, full = -1, {1: -1, 2: -1}, False + for i, (ctx, gen_pos, next_n, floors) in enumerate(_kvq_axis_far_steps(max_t, seqs)): + for num_heads in (16, 64): + for layer in ENV_LAYERS: + g = torch.Generator(device=DEV).manual_seed( + 3_700_000 + 100 * i + num_heads + ENV_LAYERS[layer][0] + ) + st = _KvqStep(ctx, gen_pos, next_n, num_heads, layer, env_pools, bufs, g, floors) + full |= st.T == max_t and st.B == seqs + top_win = max(top_win, int(st.cidx[:, :WINDOW].max())) + if st.ratio: + top_comp[st.ratio] = max(top_comp[st.ratio], int(st.cidx[:, WINDOW:].max())) + if floors is not None: + assert st.attended_minus_one() > 0 + st.fill() + o_pool = st.run(ws, compact=False) + o_comp = st.run(ws, compact=True) + assert torch.equal(_bits(o_comp), _bits(o_pool)), (cfg, i, num_heads, layer) + e = _excess(o_pool, st.reference(), env_pools["atol"][layer]) + assert e <= 1.0, (cfg, i, num_heads, layer, e) + worst, n = max(worst, e), n + 1 + far = {r: (seqs - 1) * KVQ_L[r] + (KVQ_TAIL + 1) // r - 1 for r in (1, 2)} + print( + f"kvq mbs far end {cfg}: {n} steps, window row {top_win}, compressed rows {top_comp} (far {far}), " + f"compact == container bitwise, max excess {worst:.3f}" + ) + assert full and top_win == max_t * WINDOW - 1 and top_comp == far + + +def _kvq_capture(env_pools, bufs, num_heads, layer, batches, seed): + """test_kvq_compact_read_cuda_graph_capture_set's procedure on one configuration's buffers: + per next_n 1 / 6 and G in `batches`, one eager call, a capture, the first replay == the eager call, + then three replays on new steps each == a fresh eager compact call == the container read, inside + the gate. Returns (graphs, worst excess).""" + ratio = ENV_LAYERS[layer][0] + ws = _workspace() + sinks = _sinks(num_heads) + max_t, seqs = bufs.max_t, bufs.seqs + width = WINDOW + (INDEX_TOPK if ratio else 0) + idx_buf = torch.full((max_t, width), -1, dtype=torch.int32, device=DEV) + lens_buf = torch.full((max_t,), WINDOW, dtype=torch.int32, device=DEV) + kvl_buf = torch.full((seqs + 1,), 1, dtype=torch.int32, device=DEV) + q_buf = torch.zeros(max_t, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + comp = bufs.comp_c4[ratio] if ratio else bufs.dummy + n_graphs, worst = 0, 0.0 + for next_n in (1, 6): + for G in batches: + T = G * next_n + gen = torch.Generator(device=DEV).manual_seed( + seed + 1000 * G + 100 * next_n + num_heads + ratio + ) + rng = random.Random(seed + 7 + 1000 * G + 100 * next_n + num_heads + ratio) + idx = bufs.win_idx[:T] if not ratio else idx_buf[:T] + + def make(step): + pos = [rng.randint(0, KVQ_TAIL + 1 - next_n) for _ in range(G)] + pos[0] = (KVQ_TAIL if step % 2 == 0 else ENV_SEQ - 1) + 1 - next_n + st = _KvqStep([], pos, next_n, num_heads, layer, env_pools, bufs, gen) + st.fill() + q_buf[:T].copy_(st.q3) + if ratio: + idx_buf[:T].copy_(st.cidx) + lens_buf[:T].copy_(st.active) + kvl_buf[:G].copy_(st.kv_lens[:G]) + return st + + def call(out): + return op( + q_buf[:T].view(G, next_n, num_heads, HEAD_DIM), + bufs.swa_c4, + comp, + ws, + idx, + lens_buf[:T], + kvl_buf[:G], + sinks, + SCALE, + 1.0, + out=out, + ) + + shape = (G, next_n, num_heads, HEAD_DIM) + st = make(0) + eager = call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + assert torch.equal( + _bits(eager.view(T, num_heads, HEAD_DIM)), _bits(st.run(ws, compact=False)) + ) + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + torch.cuda.current_stream().wait_stream(s) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + o = torch.empty(T, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + out = call(o.view(shape)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(_bits(out), _bits(eager)), (max_t, seqs, G, next_n) + prev = eager.clone() + for it in range(1, 4): + st = make(it) + out.fill_(float("nan")) + graph.replay() + fresh = call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + pool = st.run(ws, compact=False) + torch.cuda.synchronize() + assert torch.equal(_bits(out), _bits(fresh)), ( + max_t, + seqs, + G, + next_n, + it, + "replay != eager", + ) + assert torch.equal(_bits(out.view(T, num_heads, HEAD_DIM)), _bits(pool)), ( + max_t, + seqs, + G, + next_n, + it, + ) + assert not torch.equal(_bits(out), _bits(prev)), (max_t, seqs, G, next_n, it) + e = _excess(pool, st.reference(), env_pools["atol"][layer]) + assert e <= 1.0, (max_t, seqs, G, next_n, it, e) + worst = max(worst, e) + prev = out.clone() + del graph + n_graphs += 1 + return n_graphs, worst + + +@pytest.mark.parametrize("layer", list(ENV_LAYERS)) +@pytest.mark.parametrize("num_heads", [16, 64]) +@pytest.mark.parametrize("cfg", KVQ_MBS_CAPTURE, ids=_kvq_axis_ids(KVQ_MBS_CAPTURE)) +def test_kvq_mbs_compact_read_cuda_graph_every_batch( + env_pools, kvq_axis_buffers, cfg, num_heads, layer +): + """Capture on the compact buffers of every max_batch_size s in 1..32 (max_num_tokens 8192) and of + the witness configuration (2048, 1): the generation call at EVERY G in 1..s (the engine's + capture set at s with padding off; padding on captures a subset of it) x next_n 1 and 6, at H 16 / 64 + and every layer kind. The first replay equals the eager call; three more replays on new steps + (request 0 in the DSpark tail) each equal a fresh eager compact call and the container read, bitwise.""" + bufs = kvq_axis_buffers(cfg) + n_graphs, worst = _kvq_capture( + env_pools, bufs, num_heads, layer, range(1, cfg[1] + 1), 3_800_000 + cfg[1] + ) + print( + f"kvq mbs capture {cfg} H{num_heads} {layer}: {n_graphs} graphs x 4 replays, replay = eager = container, " + f"max excess {worst:.3f}" + ) + assert n_graphs == 2 * cfg[1] + + +KVQ_MBS_BAND_SEQS = (1, 2, 3, 4, 8, 16) + + +@pytest.mark.parametrize("width", [128, 640]) +@pytest.mark.parametrize("num_heads", [16, 64]) +def test_kvq_mbs_compact_read_kernel_variant_bands(env_pools, kvq_axis_buffers, num_heads, width): + """Kernel bands on the small compact buffers: every band of VARLEN_BANDS with B <= s requests (bottom, + interior, top; _band_cases) as a context step on the (8192, s) buffers for s in 1, 2, 3, 4, 8, 16 -- + at s = 1 every B = 1 band on the smallest buffers. The profiler reads back exactly the band's + kernels for every compact call, and every compact call is bitwise the container read.""" + ws = _workspace() + total, used = 0, set() + for seqs in KVQ_MBS_BAND_SEQS: + rng = random.Random(num_heads * 1000 + width + 37 + seqs) + gen = torch.Generator(device=DEV).manual_seed(num_heads * 1000 + width + 37 + seqs) + bufs = kvq_axis_buffers((ENV_T, seqs)) + cases = [c for c in _band_cases(num_heads, width, rng) if len(c[0]) <= seqs] + seen, differ = [], [] + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for i, (ctx, vid) in enumerate(cases): + layer = "r0" if width == WINDOW else ("r2", "r1")[i % 2] + st = _KvqStep(ctx, [], 1, num_heads, layer, env_pools, bufs, gen) + st.fill() + o_pool = st.run(ws, compact=False) + torch.cuda.synchronize() + torch.cuda._sleep(10) + o = st.run(ws, compact=True) + torch.cuda._sleep(10) + if not torch.equal(_bits(o), _bits(o_pool)): + differ.append((ctx, vid)) + torch.cuda.synchronize() + ev = sorted( + (e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA), + key=lambda e: e.time_range.start, + ) + open_call, cur = False, [] + for e in ev: + if "spin_kernel" in e.name or "sleep" in e.name: + if open_call: + seen.append(_variant_of(cur)) + open_call, cur = not open_call, [] + elif open_call: + cur.append(e.name) + assert len(seen) == len(cases), (seqs, len(seen), len(cases)) + mism = [ + (c, VARLEN_KERNELS[v], s) for (c, v), s in zip(cases, seen) if VARLEN_KERNELS[v] != s + ] + assert not mism, (seqs, mism[:3]) + assert not differ, (seqs, differ[:3]) + assert {len(ctx) for ctx, _ in cases} == set(range(1, seqs + 1)) + total += len(cases) + used |= {v for _, v in cases} + print( + f"kvq mbs bands H{num_heads} W{width} s={seqs}: {len(cases)} compact calls (B 1..{seqs}), " + f"{len({v for _, v in cases})} variants, profiler map = VARLEN_BANDS, compact == container bitwise" + ) + print(f"kvq mbs bands H{num_heads} W{width}: {total} compact calls, {len(used)} variants") + + +def _kvq_axis_control_steps(): + """(H, layer, cfg, ctx, gen, next_n, floors) on the smallest buffers. s = 1: a floored cached + context (attended -1s), a DSpark-tail verify, a fresh context (masked -1 window tails), on both + s = 1 configurations. s = 2: a floored context plus a tail verify (request 1 has compressed rows), + and a fresh context plus a decode row.""" + out = [] + for cfg in ((ENV_T, 1), KVQ_WITNESS): + out += [ + (16, "r2", cfg, [(300, 200)], [], 1, [250]), + (64, "r1", cfg, [(5000, 64)], [], 1, [4990]), + (16, "r0", cfg, [(16000, 300)], [], 1, [15950]), + (64, "r2", cfg, [], [KVQ_TAIL - 5], 6, None), + (16, "r1", cfg, [(0, 50)], [], 1, None), + (64, "r0", cfg, [(0, 50)], [], 1, None), + ] + cfg = (ENV_T, 2) + out += [ + (16, "r2", cfg, [(300, 200)], [KVQ_TAIL - 5], 6, [250]), + (64, "r1", cfg, [(5000, 64)], [KVQ_TAIL], 1, [4990]), + (16, "r0", cfg, [(16000, 300)], [128], 1, [15950]), + (64, "r2", cfg, [(0, 50)], [77], 1, None), + ] + return out + + +def test_kvq_mbs_compact_read_controls_on_the_smallest_buffers(env_pools, kvq_axis_buffers): + """The harness sees a wrong compact read on the smallest buffers (s = 1 at both max_num_tokens, and + s = 2). Each control applies only where it can differ, and the test says which did not apply: + - '-1 not zero-backed' (the rows standing for -1 window columns hold 3.0): needs attended -1s + (floored windows); + - the window half renumbered one row late: every step; + - the compressed half one entry late: ratio 1 / 2 layers; + - the compressed half at the other ratio's row stride (b * L_{3-r} + e) and one request's rows + left stale: need a request at position b >= 1 with compressed rows, so NOT applicable at s = 1 + (there b = 0 and the other stride names the same row); + and on the steps whose only -1s are masked (fresh contexts), 3.0 backing leaves the output bitwise + unchanged and NaN backing poisons it. The correct compact read is bitwise the container read and + inside the gate; every applicable wrong reading lands > 10x the gate.""" + ws = _workspace() + applied = {} + for num_heads, layer, cfg, ctx, gen_pos, next_n, floors in _kvq_axis_control_steps(): + bufs = kvq_axis_buffers(cfg) + ratio = ENV_LAYERS[layer][0] + atol = env_pools["atol"][layer] + g = torch.Generator(device=DEV).manual_seed( + 3_900_000 + num_heads + 10 * ratio + cfg[1] + len(gen_pos) + ) + st = _KvqStep(ctx, gen_pos, next_n, num_heads, layer, env_pools, bufs, g, floors) + ref = st.reference() + st.fill() + o_pool = st.run(ws, compact=False) + o_comp = st.run(ws, compact=True) + assert torch.equal(_bits(o_comp), _bits(o_pool)) + correct = _excess(o_comp, ref, atol) + assert correct <= 1.0, correct + neg = st.attended_minus_one() + masked = int((st.idx[:, :WINDOW] < 0).sum()) - neg + wrong, na = {}, [] + st.fill(minus_one=3.0) + backed = st.run(ws, compact=True) + st.fill(minus_one=float("nan")) + nan_backed = st.run(ws, compact=True) + st.fill() + if neg: + wrong["minus_one_not_zero_backed"] = _excess(backed, ref, atol) + else: + na.append("minus_one_not_zero_backed") + assert torch.equal(_bits(backed), _bits(o_pool)), ( + "a finite masked row changed the output" + ) + if masked: + assert not torch.isfinite(nan_backed).all(), ( + "a NaN masked row did not poison the output" + ) + w1 = st.cidx.clone() + w1[:, :WINDOW] += 1 + wrong["window_renumbered_one_row_late"] = _excess( + st.run(ws, compact=True, idx=w1), ref, atol + ) + if ratio: + c = st.cidx[:, WINDOW:] + w2 = st.cidx.clone() + w2[:, WINDOW:] = torch.where(c >= 0, c + 1, c) + wrong["compressed_renumbered_one_entry_late"] = _excess( + st.run(ws, compact=True, idx=w2), ref, atol + ) + if bool(((st.loc >= 0) & (st.req[:, None] >= 1)).any()): + w3 = st.cidx.clone() + w3[:, WINDOW:] = torch.where( + st.loc >= 0, st.req[:, None] * KVQ_L[3 - ratio] + st.loc, -1 + ).int() + if int(w3.max()) < bufs.comp_c4[ratio].shape[0] * bufs.comp_c4[ratio].shape[2]: + wrong["other_ratio_row_stride"] = _excess( + st.run(ws, compact=True, idx=w3), ref, atol + ) + else: # ratio 2 at small s: the ratio-1 stride points past the view (not a callable index) + na.append("other_ratio_row_stride(past the view)") + L = KVQ_L[ratio] + bufs.comp_c[L : 2 * L] = torch.randn( + L, HEAD_DIM, generator=g, device=DEV, dtype=torch.bfloat16 + ) + st.fill(skip_request=1) + wrong["request_1_rows_stale"] = _excess(st.run(ws, compact=True), ref, atol) + st.fill() + else: + na += ["other_ratio_row_stride", "request_1_rows_stale"] + else: + na += [ + "compressed_renumbered_one_entry_late", + "other_ratio_row_stride", + "request_1_rows_stale", + ] + print( + f"kvq mbs controls {cfg} H{num_heads} {layer} B {st.B}: attended -1 {neg}, masked -1 {masked}, " + f"correct {correct:.3f}, wrong { {k: round(v, 1) for k, v in wrong.items()} }, not applicable {na}" + ) + for name, e in wrong.items(): + assert e > 10.0 * max(correct, 0.1), ( + f"{name} not discriminated: {e:.2f} vs {correct:.2f}" + ) + applied.setdefault((cfg[1], name), 0) + applied[(cfg[1], name)] += 1 + print(f"kvq mbs controls applied (s, control): steps {applied}") + for s in (1, 2): + for name in ( + "minus_one_not_zero_backed", + "window_renumbered_one_row_late", + "compressed_renumbered_one_entry_late", + ): + assert applied.get((s, name), 0) > 0, (s, name) + for name in ("other_ratio_row_stride", "request_1_rows_stale"): + assert applied.get((2, name), 0) > 0 and (1, name) not in applied, name + + +def test_kvq_mbs_smallest_compact_buffers_ending_a_mapping_never_fault(): + """Precondition 8 at the smallest compressed buffers the axis allocates: the views at s = 1, 2, 3 + ([132 s, 1, 64, 512], 8,448 s rows at ratio 2; [132 s, 1, 128, 512], 16,896 s rows at ratio 1), each + placed so that its last row ends a mapping, its last rows attended and a -1 loaded from it (lengths + 129 / 131 / 133 / 639, H 16 and 64). Nothing faults, and every output is bitwise the output on a + normal allocation. The smallest view, 8,448 rows, is 66 times the 128 rows a loaded -1's over-read + (a row in 96..127 from the first row) needs.""" + trials = [] + for seqs in (1, 2, 3): + for r in (2, 1): + page = 128 // r + rows = seqs * KVQ_COLS * page + assert ( + rows == 16_896 * seqs // r and rows >= KVQ_MBS_SMALLEST_ROWS >= 66 * POOL_MIN_ROWS + ) + for j, lens in enumerate((129, 131, 133, 639)): + h = (16, 64)[(j + seqs + r) % 2] + trials.append( + dict( + name=f"kvq-mbs-comp-r{r}-s{seqs}-lens{lens}-H{h}", + kind="engine", + pool="comp", + rows=rows, + page=page, + lens=lens, + heads=h, + window=0, + mapped=rows, + seed=2_200_000 + 100 * seqs + 10 * r + j, + ) + ) + p, ok, tried = _guard_run(trials) + names = [c["name"] for c in trials] + assert p.returncode == 0, f"child failed after {tried[-1:]}: {p.stderr[-2000:]}" + assert list(ok) == names, (list(ok), names) + assert all(ok.values()), [k for k, v in ok.items() if not v] + print( + f"kvq mbs guard page: {len(names)} trials on the s = 1..3 compressed views (smallest " + f"{KVQ_MBS_SMALLEST_ROWS} rows), each bitwise = the normal allocation" + ) + + +# --------------------------------------------------------------------------------------------- +# The context call over the engine envelope (per rank): up to 16384 tokens per step +# (max_num_tokens, E2), up to 256 requests (max_batch_size, E3), and a cached prefix + q_len up to the +# DSpark cache manager's length Lm = L + 14 at max_seq_len L up to 1,048,576 (E4); in a mixed step the +# generation call follows at next_n 1 (DSpark off) or 6 (DSpark on, kv lengths up to Lm + 5). Calls are +# made as the core target makes them (_Step). These tests use their own pools (env_pools_e): a 2-page +# SWA ring per request for 256 requests, and 130 compressed pages per request. A _Step request gets +# pages for every compressed position it can see up to 128 pages -- every position of a 16384-token +# request at ratio 2 (64-row pages) and ratio 1 (128-row pages); a request that can see more uses +# its 128 pages as a ring, so far positions alias nearer rows and an index row can name one row +# twice; the call and the reference both attend it twice. +# --------------------------------------------------------------------------------------------- +E_T = 16384 +E_L = (16384, 32768, 65536, 131072, 262144, 1048576) +E_LM = E_L[-1] + 14 # the DSpark manager length at max_seq_len 1,048,576 +E_CAP_PAGES = 128 +E_SWA_PAGES = 2 * E_B + 32 +# 130 pages per request: the native-storage per-request steps below give every request distinct pages for +# all its compressed positions up to 16398 (129 pages at either ratio). +E_COMP_PAGES = E_B * 130 + 64 +# A 16384-token chunk that ends at position Lm - 1 of the longest DSpark request. +E_LONG_CACHED = E_LM - E_T + + +@pytest.fixture(scope="module") +def env_pools_e(): + """Random pools for the engine envelope (see the section comment).""" + gen = torch.Generator(device=DEV).manual_seed(4343) + swa = torch.randn(E_SWA_PAGES, 1, SWA_PAGE, HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) + pools = {"swa": swa, "comp": {}, "atol": {}} + for layer, (ratio, page) in ENV_LAYERS.items(): + if ratio: + comp = torch.empty(E_COMP_PAGES, 1, page, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + for s0 in range(0, E_COMP_PAGES, 4096): + s1 = min(E_COMP_PAGES, s0 + 4096) + comp[s0:s1] = torch.randn(s1 - s0, 1, page, HEAD_DIM, generator=gen, device=DEV).to( + torch.bfloat16 + ) + pools["atol"][layer] = _atol(swa, comp) + else: + # The core target's ratio-0 placeholder shape; NaN, so that any read of it would show. + comp = torch.full((1, 1, 64, HEAD_DIM), float("nan"), dtype=torch.bfloat16, device=DEV) + pools["atol"][layer] = _atol(swa) + pools["comp"][layer] = comp + yield pools + del pools + torch.cuda.empty_cache() + + +def _e_step(ctx, dec, num_heads, layer, pools, gen, next_n=1): + return _Step(ctx, dec, num_heads, layer, pools, gen, next_n=next_n, cap_pages=E_CAP_PAGES) + + +@pytest.mark.parametrize("cached", [0, E_LONG_CACHED]) +@pytest.mark.parametrize("layer", ["r0", "r2"]) +@pytest.mark.parametrize("num_heads", ALL_HEADS) +def test_e_prefill_every_max_q_len_to_16384(env_pools_e, num_heads, layer, cached): + """One context request of every q_len 1..16384 (max_q_len = T = q_len), fresh and as the last + 16384 tokens of a request of the DSpark manager length 1,048,590 (positions up to 1,048,589), + against the fp32 reference. As in test_prefill_every_max_q_len_one_request, row j's result does + not depend on q_len, so each call is held to the first q_len rows of ONE reference computed for + the 16384-row request.""" + gen = torch.Generator(device=DEV).manual_seed( + 32000 + num_heads + (cached > 0) + ENV_LAYERS[layer][0] + ) + st = _e_step([(cached, E_T)], [], num_heads, layer, env_pools_e, gen) + ref = st.reference() + den = env_pools_e["atol"][layer] + RTOL * ref.abs() + ws = _workspace() + out = torch.empty_like(st.q3) + worst = torch.zeros(E_T, device=DEV) + seq = torch.zeros(1, dtype=torch.int32, device=DEV) + cu = torch.zeros(2, dtype=torch.int32, device=DEV) + diff = torch.empty_like(ref) + for n in range(1, E_T + 1): + seq.fill_(cached + n) + cu[1].fill_(n) + op( + st.q3[:n], + st.swa, + st.comp, + ws, + st.idx[:n], + st.active[:n], + seq, + st.sinks, + SCALE, + 1.0, + cum_seq_lens_q=cu, + max_q_len=n, + out=out[:n], + ) + d = diff[:n] + d.copy_(out[:n]).sub_(ref[:n]).abs_().div_(den[:n]) + torch.amax(d, out=worst[n - 1]) + torch.cuda.synchronize() + bad = (worst > 1.0).nonzero().flatten() + print( + f"every-q_len-16384 H{num_heads} {layer} cached {cached}: max excess {worst.max().item():.3f}" + ) + assert bad.numel() == 0, ( + f"q_len {(bad + 1).tolist()[:20]} exceed the gate (max {worst.max().item():.2f})" + ) + assert torch.isfinite(out).all() + + +def _e_partition_case(seed): + """A random step of the engine envelope: k in 1..256 context requests over T in 1..16384 + tokens (k and T log-uniform or uniform), random q_len split and order; a max_seq_len L drawn from + E_L with DSpark on (Lm = L + 14) or off (Lm = L); each context request fresh (40 %), ending at the + last position Lm - 1 (15 %) or at a random cached prefix (cached + q_len <= Lm). For half the seeds + a mixed step: G generation requests after the context rows (k + G <= 256, T + G * next_n <= + 16384) at next_n 6 with kv lengths up to Lm + 5 (DSpark on) or next_n 1 up to L (off). + Seeds below 1000 run H 16 / 64, seeds from 1000 on the members' H 8 / 32.""" + rng = random.Random(1900000 + seed) + num_heads = ((16, 64) if seed < 1000 else MEMBER_HEADS)[seed % 2] + layer = ("r0", "r2", "r1")[(seed // 2) % 3] + if rng.random() < 0.5: + k = rng.randint(1, E_B) + else: + k = int(round(math.exp(rng.uniform(0.0, math.log(E_B))))) + if rng.random() < 0.5: + t = int(round(math.exp(rng.uniform(0.0, math.log(E_T))))) + else: + t = rng.randint(1, E_T) + t = max(t, k) + cuts = sorted(rng.sample(range(1, t), k - 1)) if k > 1 else [] + q_lens = [b - a for a, b in zip([0] + cuts, cuts + [t])] + L = E_L[rng.randrange(len(E_L))] + dspark = rng.random() < 0.5 + lm = L + 14 if dspark else L + ctx = [] + for q in q_lens: + u = rng.random() + ctx.append((0 if u < 0.4 else (lm - q if u < 0.55 else rng.randint(0, lm - q)), q)) + dec, next_n = [], (6 if dspark else 1) + if seed % 4 >= 2 and k < E_B and E_T - t >= next_n: + g = rng.randint(1, min(E_B - k, (E_T - t) // next_n)) + top = lm - 1 if dspark else L - 1 # verify: cached + 6 <= Lm + 5; decode: pos + 1 <= L + dec = [top if rng.random() < 0.2 else rng.randint(0, top) for _ in range(g)] + return num_heads, layer, ctx, dec, next_n, L, dspark + + +E_PARTITION_SEEDS = list(range(240)) +E_MEMBER_PARTITION_SEEDS = list(range(1000, 1240)) +E_PARTITION_SETS = { + "h16_h64": (E_PARTITION_SEEDS, (16, 64)), + "h8_h32": (E_MEMBER_PARTITION_SEEDS, MEMBER_HEADS), +} + + +def _seed_ids(seeds): + """Test ids: the seed, prefixed 'h8h32-' for the members' head counts (seeds from 1000 on).""" + return [str(s) if s < 1000 else f"h8h32-{s}" for s in seeds] + + +@pytest.mark.parametrize( + "seed", + E_PARTITION_SEEDS + E_MEMBER_PARTITION_SEEDS, + ids=_seed_ids(E_PARTITION_SEEDS + E_MEMBER_PARTITION_SEEDS), +) +def test_e_prefill_random_partition(env_pools_e, pq_buffers, seed): + """At the engine envelope: 240 seeded random partitions of T in [1, 16384] + into k in [1, 256] context requests at H 16 / 64, and 240 more at H 8 / 32, half of them mixed + steps whose generation call (next_n 1 or 6) reads the views starting at row tc, all against the + fp32 reference. Each step is also read in the target's native per-query form (compact window, + compressed rows t * 512 + k at max_num_tokens 16384; see _PerQueryBuffers), bitwise equal to the + container read.""" + num_heads, layer, ctx, dec, next_n, L, dspark = _e_partition_case(seed) + gen = torch.Generator(device=DEV).manual_seed(1500000 + seed) + st = _e_step(ctx, dec, num_heads, layer, env_pools_e, gen, next_n) + ws = _workspace() + o = st.run(ws) + e = _excess(o, st.reference(), env_pools_e["atol"][layer]) + o_pq = _pq_read(st, pq_buffers, ws) + print( + f"e-partition {seed}: H{num_heads} {layer} L {L} dspark {int(dspark)} T {st.T} tc {st.tc} nc {st.nc} " + f"G {st.G} next_n {next_n} max_q_len {st.max_ctx_q} max kv {int(st.kv_lens[: st.B].max())} " + f"excess {e:.3f}, per-query read bitwise {torch.equal(_bits(o_pq), _bits(o))}" + ) + assert torch.isfinite(o).all(), seed + assert e <= 1.0, (seed, e) + assert torch.equal(_bits(o_pq), _bits(o)), seed + + +@pytest.mark.parametrize("seeds", list(E_PARTITION_SETS)) +def test_e_prefill_random_partitions_reach_the_envelope(seeds): + """Each set of 240 partitions spans the engine envelope: k from 1 to >= 250 (and every + decade), T up to >= 16000, total requests up to 256, context requests ending at Lm - 1 for every L + with DSpark on and off, verify rows at kv length Lm + 5 of the longest L, mixed steps at next_n 1 + and 6 whose views start at odd rows / odd requests, and every (H, layer) of its head counts.""" + seed_list, heads = E_PARTITION_SETS[seeds] + cases = [_e_partition_case(s) for s in seed_list] + ks = [len(c[2]) for c in cases] + ts = [sum(q for _, q in c[2]) for c in cases] + mixed = [c for c in cases if c[3]] + assert min(ks) == 1 and max(ks) >= 250, (min(ks), max(ks)) + assert all(any(lo <= k < 10 * lo for k in ks) for lo in (1, 10, 100)), sorted(ks) + assert min(ts) <= 4 and max(ts) >= 16000, (min(ts), max(ts)) + assert max(len(c[2]) + len(c[3]) for c in cases) == E_B + ends = {(c[5], c[6]) for c in cases for cached, q in c[2] if cached + q == c[5] + 14 * c[6]} + assert ends == {(L, d) for L in E_L for d in (False, True)}, sorted(ends) + assert any(p + 6 == E_LM + 5 for c in cases if c[6] for p in c[3]), "no verify row at Lm + 5" + assert {c[4] for c in mixed} == {1, 6} + odd_tc = sum(sum(q for _, q in c[2]) % 2 for c in mixed) + odd_nc = sum(len(c[2]) % 2 for c in mixed) + assert len(mixed) >= 100 and odd_tc >= 20 and odd_nc >= 20, (len(mixed), odd_tc, odd_nc) + assert {(c[0], c[1]) for c in cases} == {(h, layer) for h in heads for layer in ENV_LAYERS} + + +@pytest.mark.parametrize("seeds", list(E_PARTITION_SETS)) +def test_e_prefill_partition_wrong_readings_fail_the_gate(env_pools_e, seeds): + """The partition harness sees a wrong reading at the new envelope too. On the first four mixed + partitions of each seed set with >= 64 requests and a fresh context request of >= 2 tokens, the + reference built on each wrong reading -- every context query given its request's LAST window + count, each request's window count from its neighbour's seq_lens, the generation rows' q read + from row 0 instead of row tc -- lands far outside the gate while the call sits inside it.""" + picked = [] + for seed in E_PARTITION_SETS[seeds][0]: + num_heads, layer, ctx, dec, next_n, L, dspark = _e_partition_case(seed) + if dec and len(ctx) >= 64 and any(c == 0 and q >= 2 for c, q in ctx): + picked.append(seed) + if len(picked) == 4: + break + assert len(picked) == 4, picked + for seed in picked: + num_heads, layer, ctx, dec, next_n, L, dspark = _e_partition_case(seed) + gen = torch.Generator(device=DEV).manual_seed(1500000 + seed) + st = _e_step(ctx, dec, num_heads, layer, env_pools_e, gen, next_n) + o = st.run(_workspace()) + atol = env_pools_e["atol"][layer] + correct = _excess(o, st.reference(), atol) + q_lens = [q for _, q in st.requests] + seq = st.kv_lens[: st.B].long() + last = torch.repeat_interleave(seq.clamp(max=WINDOW), torch.tensor(q_lens, device=DEV)) + rolled = seq.roll(1) + pos_rolled = torch.cat( + [torch.arange(q, device=DEV) + int(s) - q for s, q in zip(rolled.tolist(), q_lens)] + ) + wrong = { + "non_causal_window": _reference( + st.q3, st.swa, st.comp, st.idx, st.active, last, st.sinks + ), + "neighbour_seq_lens": _reference( + st.q3, + st.swa, + st.comp, + st.idx, + st.active, + (pos_rolled + 1).clamp(1, WINDOW), + st.sinks, + ), + } + q_lost = st.q3.clone() + q_lost[st.tc :] = st.q3[: st.T - st.tc] + wrong["gen_view_from_row_0"] = _reference( + q_lost, st.swa, st.comp, st.idx, st.active, st.counts, st.sinks + ) + excess = {k: _excess(o, r, atol) for k, r in wrong.items()} + print( + f"e-partition controls seed {seed} (H{num_heads} {layer} nc {st.nc}, G {st.G}, next_n {next_n}): " + f"correct {correct:.3f}, wrong { {k: round(v, 1) for k, v in excess.items()} }" + ) + assert correct <= 1.0, correct + for name, e in excess.items(): + assert e > 10.0 * max(correct, 0.1), ( + f"seed {seed}: {name} not discriminated ({e:.2f} vs {correct:.2f})" + ) + + +# Kernel map of the context call over the engine envelope (profiler probe of 174,080 calls on +# sm_103, 152 SMs: every (H, width, B) with B 1..256 at every max_q_len 1..160 and at 256, 512, 1024, +# 2048, 4096, 8191, 8192, 8193, 12288 and 16384 - (B - 1), the other B - 1 requests at q_len 1). For B <= 32 it is +# VARLEN_BANDS cell by cell, with the top band running on to 16384 - (B - 1). E_BAND_MAP[(H, width)] = +# 'B0-B1:lo-hi/v,...;...' over runs of B with one band structure; hi '*' is the envelope's top 16384 - (B - 1); v +# indexes E_VARLEN_KERNELS. The H 8 / 32 entries come from the same probe run at those head counts (another +# 174,080 calls); the Q32 kernels (indices 18..23) appear only at H 32. +E_VARLEN_KERNELS = ( + ("HVPerCta128PagedKvDenseDynamicTokenSparseP1VarSeqQ8Kv128PersistentSwapsAbForGen",), + ("HVPerCta256PagedKvDenseDynamicTokenSparseP1VarSeqQ8Kv128PersistentSwapsAbForGen",), + ("PagedKvDenseDynamicTokenSparseP1VarSeqQ8Kv128PersistentSwapsAbForGen",), + ("HVPerCta128PagedKvDenseDynamicTokenSparseP1VarSeqQ16Kv128PersistentSwapsAbForGen",), + ("HVPerCta256PagedKvDenseDynamicTokenSparseP1VarSeqQ16Kv128PersistentSwapsAbForGen",), + ("PagedKvDenseDynamicTokenSparseP1VarSeqQ16Kv128PersistentSwapsAbForGen",), + ("HVPerCta128PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ8Kv128StaticSwapsAbForGen",), + ("HVPerCta256PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ8Kv128StaticSwapsAbForGen",), + ("PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ8Kv128StaticSwapsAbForGen",), + ("HVPerCta128PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ16Kv128StaticSwapsAbForGen",), + ("HVPerCta256PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ16Kv128StaticSwapsAbForGen",), + ("PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ16Kv128StaticSwapsAbForGen",), + ("HVPerCta128PagedKvDenseDynamicTokenSparseP1VarSeqQ64Kv128PersistentKeepsAbForGen",), + ("HVPerCta256PagedKvDenseDynamicTokenSparseP1VarSeqQ64Kv128PersistentKeepsAbForGen",), + ("PagedKvDenseDynamicTokenSparseP1VarSeqQ64Kv128PersistentKeepsAbForGen",), + ( + "HVPerCta128PagedKvDenseDynamicTokenSparseP1MultiCtasKvGmemSepVarSeqQ64Kv128StaticKeepsAbForGen", + "fmhaReductionKernel<64, 128, false, __nv_bfloat16, __nv_bfloat16>", + ), + ( + "HVPerCta256PagedKvDenseDynamicTokenSparseP1MultiCtasKvGmemSepVarSeqQ64Kv128StaticKeepsAbForGen", + "fmhaReductionKernel<64, 256, false, __nv_bfloat16, __nv_bfloat16>", + ), + ( + "PagedKvDenseDynamicTokenSparseP1MultiCtasKvGmemSepVarSeqQ64Kv128StaticKeepsAbForGen", + "fmhaReductionKernel<64, 512, false, __nv_bfloat16, __nv_bfloat16>", + ), + ("HVPerCta128PagedKvDenseDynamicTokenSparseP1VarSeqQ32Kv128PersistentSwapsAbForGen",), + ("HVPerCta256PagedKvDenseDynamicTokenSparseP1VarSeqQ32Kv128PersistentSwapsAbForGen",), + ("PagedKvDenseDynamicTokenSparseP1VarSeqQ32Kv128PersistentSwapsAbForGen",), + ("HVPerCta128PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ32Kv128StaticSwapsAbForGen",), + ("HVPerCta256PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ32Kv128StaticSwapsAbForGen",), + ("PagedKvDenseDynamicTokenSparseP1MultiCtasKvVarSeqQ32Kv128StaticSwapsAbForGen",), +) +E_BAND_MAP = { + (16, 128): ( + "1-1:1-19/0,20-38/1,39-*/2;" + "2-2:1-9/0,10-19/1,20-*/2;" + "3-3:1-6/0,7-12/1,13-*/2;" + "4-4:1-4/0,5-9/1,10-*/2;" + "5-5:1-3/0,4-7/1,8-*/2;" + "6-6:1-3/0,4-6/1,7-*/2;" + "7-7:1-2/0,3-5/1,6-*/2;" + "8-9:1-2/0,3-4/1,5-*/2;" + "10-12:1-1/0,2-3/1,4-*/2;" + "13-19:1-1/0,2-2/1,3-*/2;" + "20-25:1-1/3,2-3/4,4-*/5;" + "26-38:1-1/3,2-2/4,3-*/5;" + "39-76:1-1/4,2-*/5;" + "77-256:1-*/5" + ), + (16, 640): ( + "1-1:1-6/6,7-12/7,13-38/8,39-*/2;" + "2-2:1-3/6,4-6/7,7-19/8,20-*/2;" + "3-3:1-2/6,3-4/7,5-12/8,13-*/2;" + "4-4:1-3/9,4-6/10,7-19/11,20-*/5;" + "5-5:1-2/9,3-5/10,6-15/11,16-*/5;" + "6-6:1-2/9,3-4/10,5-12/11,13-*/5;" + "7-7:1-1/9,2-3/10,4-10/11,11-*/5;" + "8-8:1-1/9,2-3/10,4-9/11,10-*/5;" + "9-9:1-1/9,2-2/10,3-8/11,9-*/5;" + "10-10:1-1/9,2-2/10,3-7/11,8-*/5;" + "11-12:1-1/9,2-2/10,3-6/11,7-*/5;" + "13-15:1-1/10,2-5/11,6-*/5;" + "16-19:1-1/10,2-4/11,5-*/5;" + "20-25:1-1/10,2-3/11,4-*/5;" + "26-38:1-2/11,3-*/5;" + "39-76:1-1/11,2-*/5;" + "77-256:1-*/5" + ), + (64, 128): ( + "1-1:1-38/12,39-76/13,77-*/14;" + "2-2:1-19/12,20-38/13,39-*/14;" + "3-3:1-12/12,13-25/13,26-*/14;" + "4-4:1-9/12,10-19/13,20-*/14;" + "5-5:1-7/12,8-15/13,16-*/14;" + "6-6:1-6/12,7-12/13,13-*/14;" + "7-7:1-5/12,6-10/13,11-*/14;" + "8-8:1-4/12,5-9/13,10-*/14;" + "9-9:1-4/12,5-8/13,9-*/14;" + "10-10:1-3/12,4-7/13,8-*/14;" + "11-12:1-3/12,4-6/13,7-*/14;" + "13-15:1-2/12,3-5/13,6-*/14;" + "16-19:1-2/12,3-4/13,5-*/14;" + "20-25:1-1/12,2-3/13,4-*/14;" + "26-38:1-1/12,2-2/13,3-*/14;" + "39-76:1-1/13,2-*/14;" + "77-256:1-*/14" + ), + (64, 640): ( + "1-1:1-12/15,13-25/16,26-76/17,77-*/14;" + "2-2:1-6/15,7-12/16,13-38/17,39-*/14;" + "3-3:1-4/15,5-8/16,9-25/17,26-*/14;" + "4-4:1-3/15,4-6/16,7-19/17,20-*/14;" + "5-5:1-2/15,3-5/16,6-15/17,16-*/14;" + "6-6:1-2/15,3-4/16,5-12/17,13-*/14;" + "7-7:1-1/15,2-3/16,4-10/17,11-*/14;" + "8-8:1-1/15,2-3/16,4-9/17,10-*/14;" + "9-9:1-1/15,2-2/16,3-8/17,9-*/14;" + "10-10:1-1/15,2-2/16,3-7/17,8-*/14;" + "11-12:1-1/15,2-2/16,3-6/17,7-*/14;" + "13-15:1-1/16,2-5/17,6-*/14;" + "16-19:1-1/16,2-4/17,5-*/14;" + "20-25:1-1/16,2-3/17,4-*/14;" + "26-38:1-2/17,3-*/14;" + "39-76:1-1/17,2-*/14;" + "77-256:1-*/14" + ), + (8, 128): ( + "1-1:1-38/0,39-76/1,77-*/2;" + "2-2:1-19/0,20-38/1,39-*/2;" + "3-3:1-12/0,13-25/1,26-*/2;" + "4-4:1-9/0,10-19/1,20-*/2;" + "5-5:1-7/0,8-15/1,16-*/2;" + "6-6:1-6/0,7-12/1,13-*/2;" + "7-7:1-5/0,6-10/1,11-*/2;" + "8-8:1-4/0,5-9/1,10-*/2;" + "9-9:1-4/0,5-8/1,9-*/2;" + "10-10:1-3/0,4-7/1,8-*/2;" + "11-12:1-3/0,4-6/1,7-*/2;" + "13-15:1-2/0,3-5/1,6-*/2;" + "16-19:1-2/0,3-4/1,5-*/2;" + "20-25:1-1/0,2-3/1,4-*/2;" + "26-38:1-1/0,2-2/1,3-*/2;" + "39-76:1-1/1,2-*/2;" + "77-256:1-*/2" + ), + (8, 640): ( + "1-1:1-12/6,13-25/7,26-76/8,77-*/2;" + "2-2:1-6/6,7-12/7,13-38/8,39-*/2;" + "3-3:1-4/6,5-8/7,9-25/8,26-*/2;" + "4-4:1-3/6,4-6/7,7-19/8,20-*/2;" + "5-5:1-2/6,3-5/7,6-15/8,16-*/2;" + "6-6:1-2/6,3-4/7,5-12/8,13-*/2;" + "7-7:1-1/6,2-3/7,4-10/8,11-*/2;" + "8-8:1-1/6,2-3/7,4-9/8,10-*/2;" + "9-9:1-1/6,2-2/7,3-8/8,9-*/2;" + "10-10:1-1/6,2-2/7,3-7/8,8-*/2;" + "11-12:1-1/6,2-2/7,3-6/8,7-*/2;" + "13-15:1-1/7,2-5/8,6-*/2;" + "16-19:1-1/7,2-4/8,5-*/2;" + "20-25:1-1/7,2-3/8,4-*/2;" + "26-38:1-2/8,3-*/2;" + "39-76:1-1/8,2-*/2;" + "77-256:1-*/2" + ), + (32, 128): ( + "1-1:1-19/3,20-38/4,39-*/5;" + "2-2:1-9/3,10-19/4,20-*/5;" + "3-3:1-6/3,7-12/4,13-*/5;" + "4-4:1-4/3,5-9/4,10-*/5;" + "5-5:1-3/3,4-7/4,8-*/5;" + "6-6:1-3/3,4-6/4,7-*/5;" + "7-7:1-2/3,3-5/4,6-*/5;" + "8-9:1-2/3,3-4/4,5-*/5;" + "10-12:1-1/3,2-3/4,4-*/5;" + "13-19:1-1/3,2-2/4,3-*/5;" + "20-25:1-1/18,2-3/19,4-*/20;" + "26-38:1-1/18,2-2/19,3-*/20;" + "39-76:1-1/19,2-*/20;" + "77-256:1-*/20" + ), + (32, 640): ( + "1-1:1-6/9,7-12/10,13-38/11,39-*/5;" + "2-2:1-3/9,4-6/10,7-19/11,20-*/5;" + "3-3:1-2/9,3-4/10,5-12/11,13-*/5;" + "4-4:1-3/21,4-6/22,7-19/23,20-*/20;" + "5-5:1-2/21,3-5/22,6-15/23,16-*/20;" + "6-6:1-2/21,3-4/22,5-12/23,13-*/20;" + "7-7:1-1/21,2-3/22,4-10/23,11-*/20;" + "8-8:1-1/21,2-3/22,4-9/23,10-*/20;" + "9-9:1-1/21,2-2/22,3-8/23,9-*/20;" + "10-10:1-1/21,2-2/22,3-7/23,8-*/20;" + "11-12:1-1/21,2-2/22,3-6/23,7-*/20;" + "13-15:1-1/22,2-5/23,6-*/20;" + "16-19:1-1/22,2-4/23,5-*/20;" + "20-25:1-1/22,2-3/23,4-*/20;" + "26-38:1-2/23,3-*/20;" + "39-76:1-1/23,2-*/20;" + "77-256:1-*/20" + ), +} + + +def _e_bands(spec: str) -> dict: + """Parse E_BAND_MAP: per (H, width), 'B0-B1:lo-hi/v,...;...' with the last band's hi '*' = + 16384 - (B - 1), the envelope's top for B requests.""" + out = {} + for key, text in spec.items(): + per_b = {} + for seg in text.split(";"): + brange, bands = seg.split(":") + b0, b1 = (int(x) for x in brange.split("-")) + for b in range(b0, b1 + 1): + row = [] + for band in bands.split(","): + lohi, vid = band.split("/") + lo, hi = lohi.split("-") + row.append((int(lo), E_T - (b - 1) if hi == "*" else int(hi), int(vid))) + per_b[b] = tuple(row) + assert sorted(per_b) == list(range(1, E_B + 1)), key + for b, row in per_b.items(): + assert row[0][0] == 1 and row[-1][1] == E_T - (b - 1), (key, b, row) + assert all(row[i][1] + 1 == row[i + 1][0] for i in range(len(row) - 1)), (key, b, row) + out[key] = per_b + return out + + +E_VARLEN_BANDS = _e_bands(E_BAND_MAP) + + +def _e_band_template(num_heads, layer, cached, pools, gen): + """The max_q_len request of every band case: E_T rows at `cached`, with ONE fp32 reference; a + case of max_q_len m uses its first m rows (a row's result does not depend on q_len).""" + st = _e_step([(cached, E_T)], [], num_heads, layer, pools, gen) + return st, st.reference() + + +def _e_band_case(tpl, mq, b, rng, gen, pools, layer): + """One call of b requests: the template's first mq rows as one request, and b - 1 short + requests (q_len 1..min(mq, 8), fresh or cached up to E_LM; a 2-page SWA ring and a 4-page + compressed ring each, as _paged_rows builds them) in random order, within the 16384-row step. + Returns (q, idx, active, seq_lens, cu, the big request's row slice, the short rows' reference).""" + st, _ = tpl + others = [rng.randint(1, min(mq, 8)) for _ in range(b - 1)] + while sum(others) > E_T - mq: + i = max(range(b - 1), key=lambda k: others[k]) + others[i] -= 1 + small = [(0 if rng.random() < 0.4 else rng.randint(0, E_LM - q), q) for q in others] + at = rng.randrange(b) + q_parts, i_parts, a_parts, seq, qls = [], [], [], [], [] + ref_small = None + if small: + ratio, page = ENV_LAYERS[layer] + ns = len(small) + swa_bt = torch.randperm(st.swa.shape[0], generator=gen, device=DEV)[: 2 * ns].view(ns, 2) + comp_bt = ( + torch.randperm(st.comp.shape[0], generator=gen, device=DEV)[: 4 * ns].view(ns, 4) + if ratio + else None + ) + s_idx, s_act, s_seq, s_cu = _paged_rows( + small, gen, st.width, max(ratio, 1), swa_bt, comp_bt, SWA_PAGE, page + ) + s_cnt = _window_counts(s_seq, s_cu) + s_q = ( + torch.randn(s_idx.shape[0], st.q3.shape[1], HEAD_DIM, generator=gen, device=DEV) + .mul_(2) + .to(torch.bfloat16) + ) + ref_small = _reference(s_q, st.swa, st.comp, s_idx, s_act, s_cnt, st.sinks) + cut = int(s_cu[at]) + for part in ("before", "big", "after"): + if part == "big": + q_parts.append(st.q3[:mq]) + i_parts.append(st.idx[:mq]) + a_parts.append(st.active[:mq]) + seq.append(int(st.kv_lens[0]) - E_T + mq) + qls.append(mq) + elif small: + sl = slice(0, cut) if part == "before" else slice(cut, None) + rs = slice(0, at) if part == "before" else slice(at, None) + q_parts.append(s_q[sl]) + i_parts.append(s_idx[sl]) + a_parts.append(s_act[sl]) + seq += s_seq[rs].tolist() + qls += [q for _, q in small[rs]] + big = slice(cut if small else 0, (cut if small else 0) + mq) + cu = torch.tensor([0] + qls, dtype=torch.int32, device=DEV).cumsum(0).int() + return ( + torch.cat(q_parts).contiguous(), + torch.cat(i_parts).contiguous(), + torch.cat(a_parts).contiguous(), + torch.tensor(seq, dtype=torch.int32, device=DEV), + cu, + big, + ref_small, + ) + + +@pytest.mark.parametrize("width", [128, 640]) +@pytest.mark.parametrize("num_heads", ALL_HEADS) +def test_e_prefill_kernel_variant_bands(env_pools_e, num_heads, width): + """Kernel bands of the context call over the engine envelope: every kernel band of + E_VARLEN_BANDS at every request count B 1..256, at its lowest, one interior and its highest + max_q_len (the top band ending at 16384 - (B - 1)), with the kernels read back from the profiler + and every row held to the fp32 reference. The max_q_len request is a prefix of one 16384-row + template per (layer, cached prefix) -- fresh, and the last 16384 positions of a 1,048,590-token + request -- so its reference is computed once; the other B - 1 requests are 1..8 tokens (the + kernel is a function of (H, width, B, max_q_len), not of the other requests' lengths).""" + rng = random.Random(num_heads * 1000 + width + 7) + gen = torch.Generator(device=DEV).manual_seed(num_heads * 1000 + width + 7) + ws = _workspace() + layers = ("r0", "r0") if width == WINDOW else ("r2", "r1") + tpls = [ + (layers[0], _e_band_template(num_heads, layers[0], 0, env_pools_e, gen)), + (layers[1], _e_band_template(num_heads, layers[1], E_LONG_CACHED, env_pools_e, gen)), + ] + cases = [] + for b, bands in E_VARLEN_BANDS[(num_heads, width)].items(): + for lo, hi, vid in bands: + for mq in sorted({lo, (lo + hi) // 2, hi}): + cases.append((b, mq, vid)) + seen, worst, wrong = [], 0.0, [] + out = torch.empty(E_T, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for i, (b, mq, vid) in enumerate(cases): + layer, tpl = tpls[i % 2] + st, ref_big = tpl + q, idx, act, seq, cu, big, ref_small = _e_band_case( + tpl, mq, b, rng, gen, env_pools_e, layer + ) + n = q.shape[0] + torch.cuda.synchronize() + torch.cuda._sleep(10) + op( + q, + st.swa, + st.comp, + ws, + idx, + act, + seq, + st.sinks, + SCALE, + 1.0, + cum_seq_lens_q=cu, + max_q_len=mq, + out=out[:n], + ) + torch.cuda._sleep(10) + atol = env_pools_e["atol"][layer] + e = _excess(out[big], ref_big[:mq], atol) + if ref_small is not None: + o_small = torch.cat([out[: big.start], out[big.stop : n]]) + e = max(e, _excess(o_small, ref_small, atol)) + worst = max(worst, e) + if e > 1.0: + wrong.append((b, mq, layer, e)) + torch.cuda.synchronize() + ev = sorted( + (e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA), + key=lambda e: e.time_range.start, + ) + open_call, cur = False, [] + for e in ev: + if "spin_kernel" in e.name or "sleep" in e.name: + if open_call: + seen.append(_variant_of(cur)) + open_call, cur = not open_call, [] + elif open_call: + cur.append(e.name) + assert len(seen) == len(cases), (len(seen), len(cases)) + mism = [ + (c, E_VARLEN_KERNELS[c[2]], s) for c, s in zip(cases, seen) if E_VARLEN_KERNELS[c[2]] != s + ] + used = {c[2] for c in cases} + print( + f"e-bands H{num_heads} W{width}: {len(cases)} calls, B 1..{E_B}, {len(used)} variants, " + f"max excess {worst:.3f}" + ) + assert not mism, mism[:3] + assert not wrong, wrong[:3] + + +# --------------------------------------------------------------------------------------------- +# The native-storage read at the engine envelope. Under native KV storage the target hands +# this call bf16 read buffers: the compact window buffer [max_num_tokens * 128, 512] (query t's +# window column c at row t * 128 + c, a -1 column as a zero row), and a compressed buffer in one of +# two forms, whichever holds fewer rows: +# - per request (b * L_r + e, the form of the test_kvq_* cases above): seqs x cols x 128 rows; +# - per query: query t's top-k column k at row t * 512 + k of a [max_num_tokens * 512, 512] buffer, +# viewed per ratio as [max_num_tokens * 512 / (128 / r), 1, 128 / r, 512]. The index table's +# compressed half is then the constant arange(max_num_tokens * 512).view(max_num_tokens, 512)[:T]: +# only the columns below sparse_topk_lens are written each step, and every column past it names a +# finite row (zero at allocation, a previous step's values after). +# At max_num_tokens 16384 the per-query buffer is 8,388,608 rows (65536 pages of 128 at ratio 1, +# 131072 of 64 at ratio 2) and the window buffer 2,097,152. Each read below is held BITWISE to the +# container read of the same step (int16 view: signed zeros count); the container read is held to +# the fp32 reference. +# --------------------------------------------------------------------------------------------- +E_PQ_ROWS = E_T * INDEX_TOPK + + +class _PerQueryBuffers: + """The target's native read buffers at max_num_tokens 16384 in the per-query form, prefilled with + finite random rows (a previous step's content), with the constant renumbered index table.""" + + def __init__(self, gen): + self.swa_c = torch.randn( + E_T * WINDOW, HEAD_DIM, generator=gen, device=DEV, dtype=torch.bfloat16 + ) + self.comp_q = torch.empty(E_PQ_ROWS, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + for s0 in range(0, E_PQ_ROWS, 1 << 20): + self.comp_q[s0 : s0 + (1 << 20)] = torch.randn( + 1 << 20, HEAD_DIM, generator=gen, device=DEV, dtype=torch.bfloat16 + ) + self.swa_c4 = self.swa_c.view(E_T, 1, WINDOW, HEAD_DIM) + self.comp_q4 = { + r: self.comp_q.view(E_PQ_ROWS // (128 // r), 1, 128 // r, HEAD_DIM) for r in (1, 2) + } + win = torch.arange(E_T * WINDOW, dtype=torch.int32, device=DEV).view(E_T, WINDOW) + topk = torch.arange(E_PQ_ROWS, dtype=torch.int32, device=DEV).view(E_T, INDEX_TOPK) + self.idx = {0: win, 1: torch.cat([win, topk], 1).contiguous()} + self.idx[2] = self.idx[1] + self.dummy = torch.zeros(1, 1, 64, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + + def fill(self, swa, comp, idx, active): + """What the target's window and top-k dequantizers write for a step whose container index is + `idx`: row t * 128 + c of swa_c for every query t and window column c (a -1 column as +0.0), + and row t * 512 + k of comp_q for every compressed column k < active[t] - 128 (the index op + packs a row's valid entries first; asserted). Rows of the masked columns keep what they held.""" + T = idx.shape[0] + w = idx[:, :WINDOW].long() + rows = swa.view(-1, HEAD_DIM)[w.clamp_min(0)] + rows.masked_fill_((w < 0)[..., None], 0.0) + self.swa_c[: T * WINDOW] = rows.view(T * WINDOW, HEAD_DIM) + del rows + if idx.shape[1] > WINDOW: + c = idx[:, WINDOW:].long() + valid = torch.arange(INDEX_TOPK, device=DEV)[None] < (active.long() - WINDOW)[:, None] + assert bool((c[valid] >= 0).all()) and bool((c[~valid] < 0).all()) + dst = torch.arange(T * INDEX_TOPK, device=DEV).view(T, INDEX_TOPK)[valid] + self.comp_q[dst] = comp.view(-1, HEAD_DIM)[c[valid]] + + def views(self, ratio, T): + """(swa pool, compressed pool, index) the target passes in the per-query form.""" + return self.swa_c4, (self.comp_q4[ratio] if ratio else self.dummy), self.idx[ratio][:T] + + +@pytest.fixture(scope="module") +def pq_buffers(): + bufs = _PerQueryBuffers(torch.Generator(device=DEV).manual_seed(2_300_000)) + yield bufs + del bufs + torch.cuda.empty_cache() + + +def _pq_read(st, bufs, ws, out=None): + """The step `st` (a _Step on env_pools_e) read in the per-query form.""" + ratio = ENV_LAYERS[st.layer][0] + bufs.fill(st.swa, st.comp, st.idx, st.active) + swa, comp, idx = bufs.views(ratio, st.T) + return st.run(ws, out=out, swa=swa, comp=comp, idx=idx) + + +def _attended_max(idx, active, counts, cols): + """The largest index an attended column in `cols` names.""" + col = torch.arange(idx.shape[1], device=DEV) + att = (col[None] < active[:, None].long()) & ( + (col[None] >= WINDOW) | (col[None] < counts[:, None]) + ) + return int(torch.where(att[:, cols], idx[:, cols], -1).max()) + + +E_PQ_FAR = { + # one context request of 16384 tokens ending at position 1,048,589 (the DSpark manager length) + "ctx": ([(E_LONG_CACHED, E_T)], [], 1), + # a mixed step of 16384 rows: two contexts, then 256 - 2 verify requests up to kv length Lm + 5 + "mixed": ( + [(E_LONG_CACHED, E_T - 6 * 254 - 700), (0, 700)], + [E_LM - 1 - (i % 7) for i in range(254)], + 6, + ), +} + + +@pytest.mark.parametrize("case", list(E_PQ_FAR)) +@pytest.mark.parametrize("layer", list(ENV_LAYERS)) +@pytest.mark.parametrize("num_heads", ALL_HEADS) +def test_e_native_per_query_far_end(env_pools_e, pq_buffers, num_heads, layer, case): + """16384-row steps in the per-query form: the last query's attended columns name the window + buffer's last row (2,097,151) and, at ratio 1 / 2, the per-query buffer's last row (8,388,607). + The per-query read is bitwise the container read, which is inside the gate.""" + ctx, dec, next_n = E_PQ_FAR[case] + gen = torch.Generator(device=DEV).manual_seed( + 2_310_000 + num_heads + 10 * ENV_LAYERS[layer][0] + len(dec) + ) + st = _e_step(ctx, dec, num_heads, layer, env_pools_e, gen, next_n) + assert st.T == E_T, st.T + ws = _workspace() + o_pool = st.run(ws) + o_pq = _pq_read(st, pq_buffers, ws) + ratio = ENV_LAYERS[layer][0] + _, _, idx = pq_buffers.views(ratio, st.T) + assert _attended_max(idx, st.active, st.counts, slice(0, WINDOW)) == E_T * WINDOW - 1 + if ratio: + assert _attended_max(idx, st.active, st.counts, slice(WINDOW, None)) == E_PQ_ROWS - 1 + assert torch.equal(_bits(o_pq), _bits(o_pool)), (num_heads, layer, case) + e = _excess(o_pool, st.reference(), env_pools_e["atol"][layer]) + print( + f"per-query far end H{num_heads} {layer} {case}: T {st.T} B {st.B}, bitwise = container, excess {e:.3f}" + ) + assert e <= 1.0, e + + +@pytest.mark.parametrize("layer", ["r2", "r1"]) +@pytest.mark.parametrize("num_heads", ALL_HEADS) +def test_e_native_per_query_generation_cuda_graph(env_pools_e, pq_buffers, num_heads, layer): + """The generation call in the per-query form captured at every G of GSET x next_n 1 and 6: q, + lengths and seq_lens are prefixes of step buffers, the index is the constant renumbered table's + prefix, the output is allocated inside the capture, and one eager call comes first. The first + replay equals the eager call bitwise; three more on new steps (new buffer rows, q, lengths and + seq_lens; kv lengths up to Lm + 5 = 1,048,595) each equal a fresh eager per-query call AND the + container read bitwise, and the container read matches the reference.""" + ratio = ENV_LAYERS[layer][0] + ws = _workspace() + sinks = _sinks(num_heads) + lens_buf = torch.full((E_T,), WINDOW, dtype=torch.int32, device=DEV) + kvl_buf = torch.full((E_B + 1,), 1, dtype=torch.int32, device=DEV) + q_buf = torch.zeros(VERIFY_Q * E_B, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + n_graphs, worst = 0, 0.0 + for next_n in (1, VERIFY_Q): + long = _cached_positions(next_n) + for G in GSET: + T = G * next_n + gen = torch.Generator(device=DEV).manual_seed( + 2_320_000 + 1000 * G + 100 * next_n + num_heads + ratio + ) + rng = random.Random(2_330_000 + 1000 * G + 100 * next_n + num_heads + ratio) + swa, comp, idx = pq_buffers.views(ratio, T) + + def make(step): + pos = [ + long[(i + step) % len(long)] if rng.random() < 0.5 else rng.randint(0, E_LM - 1) + for i in range(G) + ] + st = _e_step([], pos, num_heads, layer, env_pools_e, gen, next_n) + pq_buffers.fill(st.swa, st.comp, st.idx, st.active) + q_buf[:T].copy_(st.q3) + lens_buf[:T].copy_(st.active) + kvl_buf[:G].copy_(st.kv_lens[:G]) + return st + + def call(out): + return op( + q_buf[:T].view(G, next_n, num_heads, HEAD_DIM), + swa, + comp, + ws, + idx, + lens_buf[:T], + kvl_buf[:G], + sinks, + SCALE, + 1.0, + out=out, + ) + + shape = (G, next_n, num_heads, HEAD_DIM) + st = make(0) + eager = call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + assert torch.equal(_bits(eager.view(T, num_heads, HEAD_DIM)), _bits(st.run(ws))) + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + torch.cuda.current_stream().wait_stream(s) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + o = torch.empty(T, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + out = call(o.view(shape)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(_bits(out), _bits(eager)), (G, next_n) + prev = eager.clone() + for it in range(1, 4): + st = make(it) + out.fill_(float("nan")) + graph.replay() + fresh = call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + pool = st.run(ws) + torch.cuda.synchronize() + assert torch.equal(_bits(out), _bits(fresh)), (G, next_n, it, "replay != eager") + assert torch.equal(_bits(out.view(T, num_heads, HEAD_DIM)), _bits(pool)), ( + G, + next_n, + it, + ) + assert not torch.equal(_bits(out), _bits(prev)), (G, next_n, it) + e = _excess(pool, st.reference(), env_pools_e["atol"][layer]) + assert e <= 1.0, (G, next_n, it, e) + worst = max(worst, e) + prev = out.clone() + del graph + n_graphs += 1 + print( + f"per-query capture H{num_heads} {layer}: {n_graphs} graphs x 4 replays, replay = eager = container, " + f"max excess {worst:.3f}" + ) + + +# The per-request form at the envelope corner where it is the smaller one: max_num_tokens 16384, +# 256 sequences, the DSpark page-table width 132 (max_seq_len 16384): 4,325,376 rows, request b's +# entry e at row b * L_r + e. Steps are the test_kvq_* steps (_KvqStep) on env_pools_e. +E_KVQ_CFG = (E_T, E_B) + + +@pytest.fixture(scope="module") +def e_kvq_buffers(): + bufs = _KvqBuffers(*E_KVQ_CFG, torch.Generator(device=DEV).manual_seed(2_400_000)) + yield bufs + del bufs + torch.cuda.empty_cache() + + +def _e_kvq_case(seed): + """_kvq_case at (16384, 256): k in 1..256 requests, prefill-only / generation-only / mixed, + next_n 1 or 6, a quarter of generation requests in the DSpark tail up to 16397, contexts fresh or + cached (cached + q_len <= 16384), a quarter of cached contexts with a window floor. Seeds below + 1000 run H 16 / 64, seeds from 1000 on the members' H 8 / 32.""" + rng = random.Random(2_500_000 + seed) + num_heads = ((16, 64) if seed < 1000 else MEMBER_HEADS)[seed % 2] + layer = ("r0", "r2", "r1")[(seed // 2) % 3] + max_t, seqs = E_KVQ_CFG + k = ( + rng.randint(1, seqs) + if rng.random() < 0.5 + else int(round(math.exp(rng.uniform(0.0, math.log(seqs))))) + ) + next_n = rng.choice((1, 6)) + kind = rng.random() + n_gen = 0 if kind < 0.3 else (k if kind < 0.5 else rng.randint(1, k)) + n_ctx = k - n_gen + ctx = [] + if n_ctx: + budget = max_t - n_gen * next_n + if rng.random() < 0.5: + t = int(round(math.exp(rng.uniform(0.0, math.log(budget))))) + else: + t = rng.randint(n_ctx, budget) + t = min(max(t, n_ctx), budget) + cuts = sorted(rng.sample(range(1, t), n_ctx - 1)) if n_ctx > 1 else [] + q_lens = [b - a for a, b in zip([0] + cuts, cuts + [t])] + ctx = [(_fresh_or_cached(rng, q), q) for q in q_lens] + top = KVQ_TAIL + 1 - next_n + gen = [ + rng.randint(ENV_SEQ - next_n, top) if rng.random() < 0.25 else rng.randint(0, top) + for _ in range(n_gen) + ] + floors = [ + rng.randint(max(1, c - WINDOW + 1), c) if c and rng.random() < 0.25 else 0 for c, _ in ctx + ] + return num_heads, layer, ctx, gen, next_n, floors + + +E_KVQ_SEEDS = list(range(60)) +# The first 60-seed block from 1000 on (in steps of 60) whose steps reach k >= 250 and T >= 16000. +E_MEMBER_KVQ_SEEDS = list(range(1240, 1300)) +E_KVQ_SETS = {"h16_h64": (E_KVQ_SEEDS, (16, 64)), "h8_h32": (E_MEMBER_KVQ_SEEDS, MEMBER_HEADS)} + + +@pytest.mark.parametrize( + "seed", E_KVQ_SEEDS + E_MEMBER_KVQ_SEEDS, ids=_seed_ids(E_KVQ_SEEDS + E_MEMBER_KVQ_SEEDS) +) +def test_e_native_per_request_read(env_pools_e, e_kvq_buffers, seed): + """60 random steps of up to 256 requests and 16384 rows on the per-request buffers at (16384, + 256) at H 16 / 64, and 60 more at H 8 / 32: the compact read is bitwise the container read, + which is inside the gate.""" + num_heads, layer, ctx, gen_pos, next_n, floors = _e_kvq_case(seed) + g = torch.Generator(device=DEV).manual_seed(2_600_000 + seed) + st = _KvqStep(ctx, gen_pos, next_n, num_heads, layer, env_pools_e, e_kvq_buffers, g, floors) + ws = _workspace() + st.fill() + o_pool = st.run(ws, compact=False) + o_comp = st.run(ws, compact=True) + assert torch.equal(_bits(o_comp), _bits(o_pool)), seed + e = _excess(o_pool, st.reference(), env_pools_e["atol"][layer]) + print( + f"per-request {seed}: H{num_heads} {layer} T {st.T} tc {st.tc} nc {st.nc} G {st.G} next_n {next_n} " + f"attended -1 {st.attended_minus_one()} bitwise = container, excess {e:.3f}" + ) + assert e <= 1.0, (seed, e) + + +@pytest.mark.parametrize("seeds", list(E_KVQ_SETS)) +def test_e_native_per_request_reach(seeds): + """Each set of 60 steps reaches k 256 and >= 16000 rows, both next_n, floored contexts and every + (H, layer) of its head counts.""" + seed_list, heads = E_KVQ_SETS[seeds] + cases = [_e_kvq_case(s) for s in seed_list] + ks = [len(c[2]) + len(c[3]) for c in cases] + ts = [sum(q for _, q in c[2]) + len(c[3]) * c[4] for c in cases] + assert max(ks) >= 250 and max(ts) >= 16000, (max(ks), max(ts)) + assert {c[4] for c in cases if c[3]} == {1, 6} and any(any(c[5]) for c in cases) + assert {(c[0], c[1]) for c in cases} == {(h, layer) for h in heads for layer in ENV_LAYERS} + + +def _e_kvq_far_steps(): + """_kvq_far_steps at (16384, 256): T = 16384 with 256 requests -- one long context, 254 3-token + contexts spread over positions 131..16383 (every request stays inside max_seq_len 16384, so its + entries stay inside its L_r rows), and a DSpark verify ending at 16397 -- without and with window + floors on every cached context; and a verify-only step of 256 requests at the tail.""" + max_t, seqs = E_KVQ_CFG + n_ctx = seqs - 1 + rest = max_t - 6 - 3 * (n_ctx - 1) + ctx = [(ENV_SEQ - rest, rest)] + [ + (131 + (97 * i) % (ENV_SEQ - 134), 3) for i in range(n_ctx - 1) + ] + assert all(c + q <= ENV_SEQ for c, q in ctx) + floors = [c - 40 for c, _ in ctx] + verify = [KVQ_TAIL - 5 - (i % 3) for i in range(seqs)] + return [(ctx, [KVQ_TAIL - 5], 6, None), (ctx, [KVQ_TAIL - 5], 6, floors), ([], verify, 6, None)] + + +def test_e_native_per_request_far_end(env_pools_e, e_kvq_buffers): + """_e_kvq_far_steps, H 8 / 16 / 32 / 64 x every layer: the window half names the window buffer's + last row, the compressed half row 255 * L_r + 16398 // r - 1 (asserted); the compact read is + bitwise the container read (with floored windows among them), inside the gate.""" + bufs = e_kvq_buffers + max_t, seqs = E_KVQ_CFG + ws = _workspace() + worst, n = 0.0, 0 + for i, (ctx, gen_pos, next_n, floors) in enumerate(_e_kvq_far_steps()): + for num_heads in ALL_HEADS: + for layer in ENV_LAYERS: + g = torch.Generator(device=DEV).manual_seed( + 2_700_000 + 100 * i + num_heads + ENV_LAYERS[layer][0] + ) + st = _KvqStep(ctx, gen_pos, next_n, num_heads, layer, env_pools_e, bufs, g, floors) + if i < 2: + assert ( + st.T == max_t + and st.B == seqs + and int(st.cidx[:, :WINDOW].max()) == max_t * WINDOW - 1 + ) + if st.ratio: + far = (seqs - 1) * KVQ_L[st.ratio] + (KVQ_TAIL + 1) // st.ratio - 1 + assert int(st.cidx[:, WINDOW:].max()) == far + st.fill() + o_pool = st.run(ws, compact=False) + o_comp = st.run(ws, compact=True) + assert torch.equal(_bits(o_comp), _bits(o_pool)), (i, num_heads, layer) + e = _excess(o_pool, st.reference(), env_pools_e["atol"][layer]) + assert e <= 1.0, (i, num_heads, layer, e) + worst, n = max(worst, e), n + 1 + print( + f"per-request far end {E_KVQ_CFG}: {n} steps, compact == container bitwise, max excess {worst:.3f}" + ) + + +@pytest.mark.parametrize("layer", list(ENV_LAYERS)) +@pytest.mark.parametrize("num_heads", ALL_HEADS) +def test_e_native_per_request_generation_cuda_graph(env_pools_e, e_kvq_buffers, num_heads, layer): + """The generation call on the per-request buffers at (16384, 256) captured at every G of GSET x + next_n 1 and 6, as test_kvq_compact_read_cuda_graph_capture_set does it on the (8192, 32) + buffers: replay = fresh eager compact call = container read, bitwise, inside the gate.""" + ratio = ENV_LAYERS[layer][0] + bufs = e_kvq_buffers + max_t, seqs = E_KVQ_CFG + ws = _workspace() + sinks = _sinks(num_heads) + width = WINDOW + (INDEX_TOPK if ratio else 0) + idx_buf = torch.full((max_t, width), -1, dtype=torch.int32, device=DEV) + lens_buf = torch.full((max_t,), WINDOW, dtype=torch.int32, device=DEV) + kvl_buf = torch.full((seqs + 1,), 1, dtype=torch.int32, device=DEV) + q_buf = torch.zeros(VERIFY_Q * seqs, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + comp = bufs.comp_c4[ratio] if ratio else bufs.dummy + n_graphs, worst = 0, 0.0 + for next_n in (1, VERIFY_Q): + for G in GSET: + T = G * next_n + gen = torch.Generator(device=DEV).manual_seed( + 2_800_000 + 1000 * G + 100 * next_n + num_heads + ratio + ) + rng = random.Random(2_900_000 + 1000 * G + 100 * next_n + num_heads + ratio) + idx = bufs.win_idx[:T] if not ratio else idx_buf[:T] + + def make(step): + pos = [rng.randint(0, KVQ_TAIL + 1 - next_n) for _ in range(G)] + pos[0] = (KVQ_TAIL if step % 2 == 0 else ENV_SEQ - 1) + 1 - next_n + st = _KvqStep([], pos, next_n, num_heads, layer, env_pools_e, bufs, gen) + st.fill() + q_buf[:T].copy_(st.q3) + if ratio: + idx_buf[:T].copy_(st.cidx) + lens_buf[:T].copy_(st.active) + kvl_buf[:G].copy_(st.kv_lens[:G]) + return st + + def call(out): + return op( + q_buf[:T].view(G, next_n, num_heads, HEAD_DIM), + bufs.swa_c4, + comp, + ws, + idx, + lens_buf[:T], + kvl_buf[:G], + sinks, + SCALE, + 1.0, + out=out, + ) + + shape = (G, next_n, num_heads, HEAD_DIM) + st = make(0) + eager = call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + assert torch.equal( + _bits(eager.view(T, num_heads, HEAD_DIM)), _bits(st.run(ws, compact=False)) + ) + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + torch.cuda.current_stream().wait_stream(s) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + o = torch.empty(T, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEV) + out = call(o.view(shape)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(_bits(out), _bits(eager)), (G, next_n) + prev = eager.clone() + for it in range(1, 4): + st = make(it) + out.fill_(float("nan")) + graph.replay() + fresh = call(torch.empty(shape, dtype=torch.bfloat16, device=DEV)) + pool = st.run(ws, compact=False) + torch.cuda.synchronize() + assert torch.equal(_bits(out), _bits(fresh)), (G, next_n, it, "replay != eager") + assert torch.equal(_bits(out.view(T, num_heads, HEAD_DIM)), _bits(pool)), ( + G, + next_n, + it, + ) + assert not torch.equal(_bits(out), _bits(prev)), (G, next_n, it) + e = _excess(pool, st.reference(), env_pools_e["atol"][layer]) + assert e <= 1.0, (G, next_n, it, e) + worst = max(worst, e) + prev = out.clone() + del graph + n_graphs += 1 + print( + f"per-request capture (16384, 256) H{num_heads} {layer}: {n_graphs} graphs x 4 replays, replay = eager = " + f"container, max excess {worst:.3f}" + ) + + +# --------------------------------------------------------------------------------------------- +# Envelope E2': decoder bounded replay (on by default at tep4 / tp4). Layers 21-39 run a decoder pass +# whose context call holds the finishing requests' 128-row tails -- up to 256 requests of at most 128 +# rows each, so up to 32768 rows, beyond max_num_tokens -- and whose generation call follows on the +# generation rows. To the kernel a tail is a context chunk of <= 128 rows at a cached prefix (with a +# floor-adjusted seq_lens it is one with fewer cached rows). +# --------------------------------------------------------------------------------------------- +E_TD = WINDOW * E_B + + +def _e_replay_case(seed): + """Seeds 0..5: the full decoder pass, 256 tails of 128 rows (32768 rows), H 16 / 64 x ratio 0 / 2 / + 1. Seeds 6..35: k in 1..256 tails of 1..128 rows (a third of them exactly 128), each ending at a + random kv length up to Lm (a quarter at Lm) for a random L with DSpark on or off; for half of them + a mixed pass with G generation requests after the tails (k + G <= 256) at next_n 6 (DSpark on) or 1. + Seeds 1000..1035 are the same 36 shapes' kinds at the members' H 8 / 32 (new random draws).""" + rng = random.Random(3_000_000 + seed) + num_heads = ((16, 64) if seed < 1000 else MEMBER_HEADS)[seed % 2] + seed %= 1000 + layer = ("r0", "r2", "r1")[(seed // 2) % 3] + L = E_L[rng.randrange(len(E_L))] + dspark = rng.random() < 0.5 + lm = L + 14 if dspark else L + next_n = 6 if dspark else 1 + if seed < 6: + q_lens = [WINDOW] * E_B + else: + k = ( + rng.randint(1, E_B) + if rng.random() < 0.5 + else int(round(math.exp(rng.uniform(0.0, math.log(E_B))))) + ) + q_lens = [WINDOW if rng.random() < 1 / 3 else rng.randint(1, WINDOW) for _ in range(k)] + ctx = [] + for q in q_lens: + kv = lm if rng.random() < 0.25 else rng.randint(q, lm) + ctx.append((kv - q, q)) + dec = [] + if seed >= 6 and seed % 2 and len(ctx) < E_B: + g = rng.randint(1, E_B - len(ctx)) + top = lm - 1 if dspark else L - 1 + dec = [top if rng.random() < 0.2 else rng.randint(0, top) for _ in range(g)] + return num_heads, layer, ctx, dec, next_n, L, dspark + + +E_REPLAY_SEEDS = list(range(36)) +E_MEMBER_REPLAY_SEEDS = list(range(1000, 1036)) +E_REPLAY_SETS = { + "h16_h64": (E_REPLAY_SEEDS, (16, 64)), + "h8_h32": (E_MEMBER_REPLAY_SEEDS, MEMBER_HEADS), +} + + +@pytest.mark.parametrize( + "seed", + E_REPLAY_SEEDS + E_MEMBER_REPLAY_SEEDS, + ids=_seed_ids(E_REPLAY_SEEDS + E_MEMBER_REPLAY_SEEDS), +) +def test_e_decoder_replay_rows_to_32768(env_pools_e, seed): + """E2': the decoder pass's context call at up to 32768 rows in up to 256 requests of <= 128 rows + (and its generation call after it), against the fp32 reference. For a context-only pass the + kernels are read back from the profiler and equal E_VARLEN_BANDS' cell for (H, width, B, + max_q_len).""" + num_heads, layer, ctx, dec, next_n, L, dspark = _e_replay_case(seed) + gen = torch.Generator(device=DEV).manual_seed(3_100_000 + seed) + st = _e_step(ctx, dec, num_heads, layer, env_pools_e, gen, next_n) + ws = _workspace() + seen = None + if dec: + o = st.run(ws) + else: + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + torch.cuda.synchronize() + torch.cuda._sleep(10) + o = st.run(ws) + torch.cuda._sleep(10) + torch.cuda.synchronize() + ev = sorted( + (e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA), + key=lambda e: e.time_range.start, + ) + names, inside = [], False + for e in ev: + if "spin_kernel" in e.name or "sleep" in e.name: + inside = not inside + elif inside: + names.append(e.name) + seen = _variant_of(names) + width = st.width + band = [ + v for lo, hi, v in E_VARLEN_BANDS[(num_heads, width)][st.nc] if lo <= st.max_ctx_q <= hi + ] + assert seen == E_VARLEN_KERNELS[band[0]], (seed, seen, band) + e = _excess(o, st.reference(), env_pools_e["atol"][layer]) + print( + f"decoder-pass {seed}: H{num_heads} {layer} L {L} dspark {int(dspark)} T {st.T} tc {st.tc} nc {st.nc} " + f"max_q_len {st.max_ctx_q} G {st.G} next_n {next_n} max kv {int(st.kv_lens[: st.B].max())} " + f"kernels {seen} excess {e:.3f}" + ) + assert torch.isfinite(o).all(), seed + assert e <= 1.0, (seed, e) + + +@pytest.mark.parametrize("seeds", list(E_REPLAY_SETS)) +def test_e_decoder_replay_cases_reach_the_shape(seeds): + """Each set of 36 passes reaches 32768 rows in 256 tails of 128 at every (H, layer) of its head + counts, mixed passes at next_n 1 and 6, and tails ending at Lm for the longest L.""" + seed_list, heads = E_REPLAY_SETS[seeds] + cases = [_e_replay_case(s) for s in seed_list] + ts = [sum(q for _, q in c[2]) for c in cases] + assert max(ts) == E_TD and all(max(q for _, q in c[2]) <= WINDOW for c in cases) + full = {(c[0], c[1]) for c, t in zip(cases, ts) if t == E_TD} + assert full == {(h, layer) for h in heads for layer in ENV_LAYERS}, full + assert {c[4] for c in cases if c[3]} == {1, 6} + assert {(c[0], c[1]) for c in cases} == {(h, layer) for h in heads for layer in ENV_LAYERS} + assert any(cached + q == E_LM for c in cases for cached, q in c[2]) + + +# --------------------------------------------------------------------------------------------- +# The workspace over the engine envelope. The target passes one persistent 128 MiB workspace +# (_ATTN_WS_BYTES, zeroed once) to every call, the drafter's included. FlashInfer hands the kernel +# the whole buffer as the multi-CTA split-KV scratch -- per-CTA max/sum statistics ([SMs x 64] +# float2) followed by per-CTA partial O ([CTAs x 64 x headDimPerCta] bf16) -- and checks no size +# (the multi-CTA semaphores live in a separate buffer FlashInfer allocates per call). The kernel +# splits the KV only while the unsplit grid is smaller than the SM count, and never into more CTAs +# than there are SMs, so the scratch does not grow with G, next_n, H, the query rows, the number of +# requests or the kv length: at most SMs x 64 x (8 + 512 x 2) bytes, 10,039,296 on 152 SMs. +# --------------------------------------------------------------------------------------------- +WS_FILLS = (0xA5, 0x5A) + + +def _ws_probe_calls(): + """(label, kind, H, width, B, rows-per-request or q_lens) for every call shape the test runs.""" + calls = [] + for g in range(1, E_B + 1): + for nn in (1, VERIFY_Q): + for h in ALL_HEADS: + for w in (128, 640): + calls.append(("dense", h, w, g, nn)) + calls.append(("draft", DS_HEADS, DS_WIDTH, g, 1)) + for h in ALL_HEADS: + for w in (128, 640): + for mq in list(range(1, 161)) + [512, 2048, 8192, 8193, 12288, E_T]: + calls.append(("varlen", h, w, 1, [mq])) + for b in range(2, E_B + 1): + for mq in (1, 2, 3, 4, 6, 8, 12, 13, 19, 20, 25, 26, 38, 39, 76, 77): + if b * mq <= E_T: + calls.append(("varlen", h, w, b, [mq] * b)) + if b in GSET: + calls.append(("varlen", h, w, b, [E_T - (b - 1)] + [1] * (b - 1))) + calls.append(("varlen", h, w, E_B, [WINDOW] * E_B)) # E2': the 32768-row decoder pass + return calls + + +def test_workspace_high_water_mark_over_the_envelope(): + """The 128 MiB workspace covers the engine envelope, and its content does not matter. + + Every call below runs three times, on the workspace filled with 0xA5, with 0x5A and with zeros; + the high-water mark is the end of the last 8-byte word that differs from the fill. Shapes: dense + decode (next_n 1) and verify (next_n 6) at every G 1..256, H 8 / 16 / 32 / 64, width 128 / 640 + (the target's generation call); the drafter's [5G, 1, 64, 512] at width 136 for every G 1..256; the + context call (H 8 / 16 / 32 / 64) at one request of every max_q_len 1..160 and up to 16384, at + every B 2..256 with B requests of each max_q_len at the kernel's split-KV edges, B in GSET at the + 16384-row top, and the 32768-row decoder pass (256 requests of 128 rows). + kv lengths are 1,048,595 (decode, verify, context) and 129 (drafter).""" + sm = torch.cuda.get_device_properties(DEV).multi_processor_count + bound = sm * 64 * (8 + HEAD_DIM * 2) + gen = torch.Generator(device=DEV).manual_seed(5150) + swa = torch.randn(600, 1, SWA_PAGE, HEAD_DIM, generator=gen, device=DEV).to(torch.bfloat16) + comp = torch.randn(600, 1, 64, HEAD_DIM, generator=gen, device=DEV).to(torch.bfloat16) + rows_max = E_T * 2 + q_buf = torch.randn(rows_max, 64, HEAD_DIM, generator=gen, device=DEV).to(torch.bfloat16) + idx_buf = { + w: torch.randint(0, 600 * 64, (rows_max, w), generator=gen, device=DEV, dtype=torch.int32) + for w in (WINDOW, DS_WIDTH, WINDOW + INDEX_TOPK) + } + ws = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=DEV) + zws = torch.zeros_like(ws) + words = ws.view(torch.int64) + word_idx = torch.arange(words.numel(), device=DEV) + fill_word = {f: int.from_bytes(bytes([f]) * 8, "little", signed=True) for f in WS_FILLS} + worst, touched, bad = (0, None), 0, [] + extremes, hw_by_heads = {}, {} + for kind, h, w, b, per in _ws_probe_calls(): + if kind == "varlen": + n, mq = sum(per), max(per) + q = q_buf[:n, :h] + cu = torch.tensor([0] + per, dtype=torch.int32, device=DEV).cumsum(0).int() + seq = torch.tensor([1048595 - E_T + x for x in per], dtype=torch.int32, device=DEV) + else: + n, mq, cu = b * (DS_BLOCK if kind == "draft" else per), None, None + rows = n if kind == "draft" else b + q = q_buf[:n, :h].reshape(rows, n // rows, h, HEAD_DIM) + seq = torch.full( + (rows,), 129 if kind == "draft" else 1048595, dtype=torch.int32, device=DEV + ) + act = torch.full((n,), DS_TOPK if kind == "draft" else w, dtype=torch.int32, device=DEV) + idx = idx_buf[w][:n] + sinks = _sinks(h) + outs, hw = [], 0 + for f in WS_FILLS + (None,): + if f is None: + buf = zws + else: + ws.fill_(f) + buf = ws + outs.append( + op( + q, + swa, + comp, + buf, + idx, + act, + seq, + sinks, + SCALE, + 1.0, + cum_seq_lens_q=cu, + max_q_len=mq, + ) + ) + if f is not None: + diff = words != fill_word[f] + if bool(diff.any()): + hw = max(hw, (int(torch.where(diff, word_idx, -1).max()) + 1) * 8) + label = (kind, h, w, b, mq if kind == "varlen" else per) + if not ( + torch.equal(outs[0], outs[1]) + and torch.equal(outs[0], outs[2]) + and bool(torch.isfinite(outs[0]).all()) + ): + bad.append(label) + touched += hw > 0 + if kind != "draft": + hw_by_heads[h] = max(hw_by_heads.get(h, 0), hw) + if hw > worst[0]: + worst = (hw, label) + if (kind, b) in (("dense", E_B), ("draft", E_B)) or (kind == "varlen" and n >= E_T): + extremes[label] = hw + print( + f"workspace: {len(_ws_probe_calls())} call shapes, {touched} touched the workspace, high-water mark " + f"{worst[0]} bytes ({worst[0] / 2**20:.3f} MiB) at {worst[1]}; bound {bound} on {sm} SMs; " + f"envelope extremes {sorted(set(extremes.values()))}; high-water mark by H (trunk calls) {hw_by_heads}" + ) + assert not bad, bad[:5] + assert touched > 0 and worst[0] > 0 # the harness sees the scratch writes it measures + assert worst[0] <= bound <= 128 * 1024 * 1024 // 8, (worst, bound) + assert sorted(hw_by_heads) == list(ALL_HEADS) and all(v > 0 for v in hw_by_heads.values()), ( + hw_by_heads + ) + assert set(extremes.values()) == {0}, {k: v for k, v in extremes.items() if v} + + +if __name__ == "__main__": + if len(sys.argv) == 3 and sys.argv[1] == "--guard-child": + _guard_child(json.loads(sys.argv[2])) + elif len(sys.argv) == 2 and sys.argv[1] == "--pool-child": + _pool_child() diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_trtllm_sparse_mla_dsv4_fp8.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_trtllm_sparse_mla_dsv4_fp8.py new file mode 100644 index 000000000000..f6eda65f9cf4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/attention/test_modeling_v2_trtllm_sparse_mla_dsv4_fp8.py @@ -0,0 +1,468 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the trtllm_sparse_mla_dsv4_fp8 entry. + +The entry owns what is certified: `CELLS` (decode and DSpark verify, varlen prefill, the scale pairs and +the per-query compressed buffer) is the list, and `reference` and `compare` are the gate. This file turns +a cell's spec into tensors and keeps the checks a cell cannot carry: the device-tensor `bmm1_scale`, +CUDA-graph capture and replay, and the wrong-reading controls. + +The call computes with E4M3 operands, so the reference (`op.reference`) is taken on the SAME E4M3 +values: the codes are exact in fp32, gathered per index row, fp32 scores scaled by bmm1_scale (= +softmax_scale * q_scale * kv_scale), fp32 max and sum, fp32 probabilities, the per-head sink added to the +denominator only (natural-log units), and the output scaled by bmm2_scale (= kv_scale). What the kernel +adds on top is its own arithmetic: it quantizes the probabilities to E4M3 (scaled by 448) for the PV MMA +and accumulates in its own order. + +Gate (per query row and head): ||o - ref||_2 <= REL * ||ref||_2 + ABS * sqrt(512) * max|V|, with +REL = 2^-4 -- the largest relative rounding of one E4M3 probability, which is the whole error of a row +whose softmax sits on one key -- and ABS = 2^-7 for rows whose output nearly cancels. The discrimination +is measured in test_wrong_variants_fail_the_gate: the correct reference and each deliberately wrong +reading are reported as a multiple of the gate (sm_103: correct <= 0.45, wrong 12-58x). The wrong +readings are the op's own reference on altered arguments (no sink, sink x log2(e), bmm1 without +q_scale, kv_scale twice in bmm2, lengths cut to the window, every window count 128). + +The pools are the native-storage target's compact read buffers, now E4M3: the window buffer +[max_num_tokens * 128, 512] viewed [max_num_tokens, 1, 128, 512] (query t's column c is row t*128+c) +and the compressed buffer [max_batch_size * 132 * 128 / r, 512] viewed [max_batch_size * 132, 1, +128 / r, 512] (request b's entry e is row b * L_r + e), at max_num_tokens 8192 and max_batch_size 32. +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.attention.trtllm_sparse_mla_dsv4_fp8 import ( # noqa: E501 + trtllm_sparse_mla_dsv4_fp8 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "trtllm_sparse_mla_dsv4_fp8 requires a CUDA device" + +DEV = torch.device("cuda") +D = 512 +WIN = 128 +SCALE = D**-0.5 +E4 = torch.float8_e4m3fn +MAX_T = 8192 +MBS = 32 +COLS = 132 +REL = 2.0**-4 +ABS = 2.0**-7 +WS = None + + +def _ws(): + global WS + if WS is None: + WS = torch.zeros(128 * 1024 * 1024, dtype=torch.uint8, device=DEV) + return WS + + +def _gen(seed): + return torch.Generator(device=DEV).manual_seed(seed) + + +def _sinks(H): + return torch.linspace(-2.0, 9.0, H, device=DEV, dtype=torch.float32) + + +class Buffers: + """Persistent E4M3 compact buffers (allocated once, like the target's).""" + + def __init__(self): + self.win = torch.zeros(MAX_T * WIN, D, dtype=E4, device=DEV) + self.comp = torch.zeros(MBS * COLS * WIN, D, dtype=E4, device=DEV) + self.dummy = torch.zeros(1, 1, 64, D, dtype=E4, device=DEV) + + def win4(self): + return self.win.view(MAX_T, 1, WIN, D) + + def comp4(self, r): + return self.comp[: MBS * COLS * WIN // r].view(MBS * COLS, 1, WIN // r, D) + + +BUF = None + + +def _buf(): + global BUF + if BUF is None: + BUF = Buffers() + return BUF + + +def _fill_e4m3(dst, gen, lo=-8, hi=8): + """Random E4M3 rows with per-row magnitudes over 2^lo..2^hi code units (the transcoded KV spans + about 2^-9..2^8 at the target's scale).""" + n = dst.shape[0] + mag = torch.randint(lo, hi + 1, (n, 1), device=DEV, generator=gen).float() + v = torch.randn(n, D, device=DEV, generator=gen) * torch.exp2(mag) * 8 + dst.copy_(v.clamp(-448, 448).to(E4)) + + +def _step(reqs, H, ratio, gen, n_comp_valid=None): + """Build one step: reqs = [(cached, q_len)]; returns inputs for the call and for the reference. + Window columns of query t name rows t*128 + c; rows past the window count hold finite decoys.""" + b = _buf() + T = sum(q for _, q in reqs) + width = WIN + (512 if ratio else 0) + q = torch.randn(T, H, D, device=DEV, generator=gen) * 4 + q8 = q.clamp(-448, 448).to(E4) + _fill_e4m3(b.win[: T * WIN], gen) + idx = torch.full((T, width), -1, dtype=torch.int32, device=DEV) + idx[:, :WIN] = torch.arange(T * WIN, dtype=torch.int32, device=DEV).view(T, WIN) + lens = torch.full((T,), WIN, dtype=torch.int32, device=DEV) + wc = torch.empty(T, dtype=torch.int64, device=DEV) + L_r = COLS * WIN // max(ratio, 1) + t = 0 + for bi, (cached, ql) in enumerate(reqs): + pos = cached + torch.arange(ql, device=DEV) + wc[t : t + ql] = (pos + 1).clamp_max(WIN) + if ratio: + n_c = (pos + 1) // ratio # compressed entries visible to each query + cap = 512 if n_comp_valid is None else n_comp_valid + k = n_c.clamp_max(cap) + n_max = int(n_c.max()) + for c0 in range(0, ql, 1024): + c1 = min(ql, c0 + 1024) + if n_max == 0: + break + # a random subset of size k of 0..n_c-1 per row, sorted, -1 after it + sc = torch.rand(c1 - c0, n_max, device=DEV, generator=gen) + sc = torch.where(torch.arange(n_max, device=DEV)[None] < n_c[c0:c1, None], sc, -1.0) + kk = min(n_max, cap) + top = sc.topk(kk, dim=1).indices.sort(dim=1).values + # rows with n_c < kk: keep only their first k valid entries (all of them, in order) + keep = torch.arange(kk, device=DEV)[None] < k[c0:c1, None] + top = torch.where(keep, top, torch.full_like(top, 1 << 30)).sort(dim=1).values + sel = torch.where(top < (1 << 30), bi * L_r + top, -1).int() + idx[t + c0 : t + c1, WIN : WIN + kk] = sel + lens[t : t + ql] = (WIN + k).int() + t += ql + if ratio: + rows = b.comp4(ratio).view(-1, D) + _fill_e4m3(rows[: len(reqs) * L_r], gen) + seq = torch.tensor([c + q for c, q in reqs], dtype=torch.int32, device=DEV) + cu = torch.tensor( + [0] + list(torch.tensor([q for _, q in reqs]).cumsum(0)), dtype=torch.int32, device=DEV + ) + return dict( + q8=q8, + idx=idx, + lens=lens, + wc=wc, + seq=seq, + cu=cu, + T=T, + width=width, + ratio=ratio, + max_q=max(q for _, q in reqs), + reqs=reqs, + ) + + +def _pools(st): + b = _buf() + comp = b.comp4(st["ratio"]) if st["ratio"] else b.dummy + return b.win4(), comp + + +def _reference( + st, + sq, + skv, + sinks, + win_count=None, + drop_comp=False, + sink_mode="ok", + bmm2=None, + pools=None, +): + """`op.reference` on a step's arguments; returns [T, H, D] fp32. + + The keywords are the controls' wrong readings, each expressed through the op's own arguments: + `win_count` gives every row its own q_len-1 request with that virtual seq_lens, `drop_comp` cuts + every length to the window, `sink_mode` "none" / "log2e" passes no sink / sinks x log2(e), and + `bmm2` replaces bmm2_scale. bmm1_scale is softmax_scale * sq * skv. + """ + swa, comp = _pools(st) if pools is None else pools + T = st["T"] + if win_count is None: + seq, cu = st["seq"], st["cu"] + else: + seq = win_count.to(torch.int32) + cu = torch.arange(T + 1, dtype=torch.int32, device=DEV) + lens = torch.full_like(st["lens"], WIN) if drop_comp else st["lens"] + sk = {"ok": sinks, "none": None, "log2e": sinks * 1.4426950408889634}[sink_mode] + expected, _ = op.reference( + st["q8"], + swa, + comp, + None, + st["idx"], + lens, + seq, + sk, + SCALE * sq * skv, + skv if bmm2 is None else bmm2, + cum_seq_lens_q=cu, + max_q_len=st["max_q"], + ) + return expected + + +def _gate_ratio(o, ref, vmax): + err = (o.float() - ref).norm(dim=-1) + tol = REL * ref.norm(dim=-1) + ABS * math.sqrt(D) * vmax + return (err / tol).max().item() + + +def _vmax(st, skv, pools=None): + if pools is not None: + return max(p.float().abs().max().item() for p in pools) * skv + swa, comp = _pools(st) + T = st["T"] + m = swa.reshape(-1, D)[: T * WIN].float().abs().max().item() + if st["ratio"]: + m = max( + m, + comp.reshape(-1, D)[: len(st["reqs"]) * COLS * WIN // st["ratio"]] + .float() + .abs() + .max() + .item(), + ) + return m * skv + + +def _call_decode(st, H, bmm1, skv, out=None): + swa, comp = _pools(st) + G = len(st["reqs"]) + nn = st["reqs"][0][1] + return op( + st["q8"].view(G, nn, H, D), + swa, + comp, + _ws(), + st["idx"], + st["lens"], + st["seq"], + _sinks(H), + bmm1, + skv, + out=out, + ) + + +def _call_varlen(st, H, bmm1, skv): + swa, comp = _pools(st) + return op( + st["q8"], + swa, + comp, + _ws(), + st["idx"], + st["lens"], + st["seq"], + _sinks(H), + bmm1, + skv, + cum_seq_lens_q=st["cu"], + max_q_len=st["max_q"], + ) + + +SQ, SKV = 2.0**-3, 2.0**-6 + + +def _bmm1(tensor=False): + v = SCALE * SQ * SKV + return torch.tensor([v], dtype=torch.float32, device=DEV) if tensor else v + + +@pytest.mark.parametrize("H", [16, 64]) +def test_tensor_bmm1_equals_float_bmm1_bitwise(H): + gen = _gen(7 + H) + st = _step([(100, 6), (5000, 6), (9, 6)], H, 2, gen) + a = _call_decode(st, H, _bmm1(False), SKV) + b = _call_decode(st, H, _bmm1(True), SKV) + assert torch.equal(a.view(torch.int16), b.view(torch.int16)) + + +@pytest.mark.parametrize("next_n", [1, 6]) +@pytest.mark.parametrize("H", [16, 64]) +def test_cuda_graph_capture_replay(H, next_n): + """Decode capture at the target's capture batches with a device bmm1 scale; each replay equals a + fresh eager call bitwise and passes the reference.""" + for G in (1, 2, 4, 8, 16, 24, 32): + gen = _gen(G + 50 * next_n + H) + st = _step( + [(int(c), next_n) for c in torch.randint(0, 16000, (G,), generator=gen, device=DEV)], + H, + 2, + gen, + ) + bmm1 = _bmm1(True) + swa, comp = _pools(st) + q = st["q8"].view(G, next_n, H, D).clone() + idx, lens, seq = st["idx"].clone(), st["lens"].clone(), st["seq"].clone() + sinks = _sinks(H) + + def call(): + return op(q, swa, comp, _ws(), idx, lens, seq, sinks, bmm1, SKV) + + call() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = call() + for rep in range(3): + gen2 = _gen(1000 + rep + G) + st2 = _step( + [ + (int(c), next_n) + for c in torch.randint(0, 16000, (G,), generator=gen2, device=DEV) + ], + H, + 2, + gen2, + ) + q.copy_(st2["q8"].view(G, next_n, H, D)) + idx.copy_(st2["idx"]) + lens.copy_(st2["lens"]) + seq.copy_(st2["seq"]) + bmm1.fill_(SCALE * SQ * SKV * (1.0 if rep != 1 else 0.5)) + graph.replay() + eager = call() + assert torch.equal(out.view(torch.int16), eager.view(torch.int16)) + sq = SQ * (1.0 if rep != 1 else 0.5) + ref = _reference(st2, sq, SKV, sinks) + op.compare(out.view(st2["T"], H, D), (ref, _vmax(st2, SKV))) + + +def test_wrong_variants_fail_the_gate(): + """Discrimination: the correct reference vs deliberately wrong readings, as multiples of the gate.""" + H = 16 + gen = _gen(99) + st = _step([(3000, 1), (50, 1), (16000, 1), (0, 300), (9000, 64)], H, 2, gen) + o = _call_varlen(st, H, _bmm1(), SKV) + vmax = _vmax(st, SKV) + sinks = _sinks(H) + correct = _reference(st, SQ, SKV, sinks) + good = _gate_ratio(o, correct, vmax) + wrong = { + "sink ignored": _reference(st, SQ, SKV, sinks, sink_mode="none"), + "sink x log2e": _reference(st, SQ, SKV, sinks, sink_mode="log2e"), + "q scale ignored (bmm1 without q_scale)": _reference(st, 1.0, SKV, sinks), + "kv scale twice in bmm2": _reference(st, SQ, SKV, sinks, bmm2=SKV * SKV), + "compressed segment dropped": _reference(st, SQ, SKV, sinks, drop_comp=True), + "window count 128 for every row": _reference( + st, SQ, SKV, sinks, win_count=torch.full_like(st["wc"], WIN) + ), + } + print(f"correct reference: {good:.3f} of the gate") + for name, ref in wrong.items(): + g = _gate_ratio(o, ref, vmax) + print(f" {name}: {g:.1f} x the gate") + assert g > 5.0, name + op.compare(o, (correct, vmax)) + + +PQ = {} + + +def _per_query_buffer(max_t, ratio): + """The target's per-query compressed read buffer (used when it is smaller than the per-request + one): [max_num_tokens * 512, 512] E4M3, viewed as [max_t * 512 / (128 / r), 1, 128 / r, 512].""" + key = (max_t, ratio) + if key not in PQ: + PQ.clear() + PQ[key] = torch.zeros(max_t * 512, D, dtype=E4, device=DEV) + return PQ[key], PQ[key].view(max_t * 512 // (WIN // ratio), 1, WIN // ratio, D) + + +def _per_query_step(max_t, reqs, H, ratio, gen): + """The per-query form: query t's top-k column k at row t * 512 + k, every compressed column + renumbered (a row's valid entries first, then rows past sparse_topk_lens holding stale finite + decoys that would dominate the softmax if attended), long contexts (cached up to 2^20). + + Returns the step (window rows as in the other cells, compressed half built here) and the + compressed view. + """ + st = _step(reqs, H, 0, gen) # window rows as in the other tests; compressed half built here + T = st["T"] + buf, comp4 = _per_query_buffer(max_t, ratio) + buf[: T * 512].copy_(torch.full((T * 512, D), 300.0, device=DEV).to(E4)) # stale decoys + count = torch.empty(T, dtype=torch.int64, device=DEV) + t = 0 + for c, ql in reqs: + pos = c + torch.arange(ql, device=DEV) + count[t : t + ql] = ((pos + 1) // ratio).clamp_max(512) + t += ql + rows = torch.arange(T * 512, device=DEV).view(T, 512) + valid = torch.arange(512, device=DEV)[None] < count[:, None] + mag = torch.randint(-8, 9, (T, 512, 1), device=DEV, generator=gen).float() + vals = ( + (torch.randn(T, 512, D, device=DEV, generator=gen) * torch.exp2(mag) * 8) + .clamp(-448, 448) + .to(E4) + ) + buf[: T * 512].view(T, 512, D)[valid] = vals[valid] + idx = torch.cat([st["idx"][:, :WIN], rows.to(torch.int32)], 1).contiguous() + lens = (WIN + count).to(torch.int32) + st.update(idx=idx, lens=lens, width=WIN + 512, ratio=ratio) + return st, comp4 + + +def _cell_call(spec): + """The step, the op's positional arguments and its keyword arguments for one cell. + + Seeds are per kind and per configuration, so a cell's data does not depend on which other cells + run. + """ + kind, H, sq, skv = spec["kind"], spec["H"], spec["q_scale"], spec["kv_scale"] + if kind == "decode": + G, nn, ratio = spec["G"], spec["next_n"], spec["ratio"] + gen = _gen(G * 1000 + H * 10 + ratio * 3 + nn) + cached = torch.randint(0, 16398 - nn, (G,), generator=gen, device=DEV).tolist() + cached[0] = 16398 - nn # the DSpark tail of max_seq_len 16398 + if G > 1: + cached[1] = 3 + st = _step([(c, nn) for c in cached], H, ratio, gen) + swa, comp = _pools(st) + elif kind == "prefill": + reqs = [tuple(r) for r in spec["requests"]] + st = _step(reqs, H, spec["ratio"], _gen(len(reqs) * 100 + H + spec["ratio"])) + swa, comp = _pools(st) + elif kind == "scales": + reqs = [tuple(r) for r in spec["requests"]] + st = _step(reqs, H, spec["ratio"], _gen(int(1 / sq) + int(1 / skv))) + swa, comp = _pools(st) + else: + reqs = [tuple(r) for r in spec["requests"]] + gen = _gen(spec["max_t"] + len(reqs) + H + spec["ratio"]) + st, comp = _per_query_step(spec["max_t"], reqs, H, spec["ratio"], gen) + swa = _buf().win4() + dense = kind in ("decode", "scales") or (kind == "per_query" and spec["decode"]) + if dense: + q = st["q8"].view(len(st["reqs"]), st["reqs"][0][1], H, D) + kwargs = {} + else: + q = st["q8"] + kwargs = dict(cum_seq_lens_q=st["cu"], max_q_len=st["max_q"]) + args = (q, swa, comp, _ws(), st["idx"], st["lens"], st["seq"], _sinks(H), SCALE * sq * skv, skv) + return st, args, kwargs + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + st, args, kwargs = _cell_call(cell.spec) + o = op(*args, **kwargs) + assert o.shape == args[0].shape, (o.shape, args[0].shape) + if cell.spec["kind"] == "per_query": + assert int(st["idx"][:, WIN:].max()) == st["T"] * 512 - 1, "the buffer's last row is named" + op.compare(o, op.reference(*args, **kwargs)) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/_allreduce_op_matrix.py new file mode 100644 index 000000000000..33a40a3b1d9f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/_allreduce_op_matrix.py @@ -0,0 +1,1403 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the DeepSeek-V4.1 target's allreduce entry. + +A collective cannot be exercised in one process, so this script is its own +launcher: run it plainly and, for each suite, it spawns one worker process per +device named in CUDA_VISIBLE_DEVICES, rendezvoused through torch.distributed's +env://, under a deadline it enforces by killing every worker's session. A +broken collective hangs rather than raising, so the deadline is what keeps a +wedged kernel from taking the calling run down with it. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _allreduce_op_matrix.py [--group-size G] + +The ranks run with MPI disabled, so the entry takes the ProcessGroup path the +product takes then (`allreduce_pg`, bound once through `process_group`): the +group's communicator is the Mapping's `tp_group_pg`, and the IPC workspace is +the one `get_allreduce_workspace` builds from that Mapping. + +The surface is the plain summing all-reduce (`op = 0`) over the payloads a +DeepSeek-V4.1 tensor-parallel forward hands it: fp32 `[T, 5120]` partial sums +(row-parallel output projection, routed-expert output) and bf16 rows +(`[T, 5120]`, and the `[T, 24, 256]` engram rows whose shards are zero off +their owning rank), on every strategy the op implements for a plain sum, +eagerly and replayed from CUDA graphs. Every cell of the entry's `CELLS` for +this rank's layout is driven here. + +The group is a parameter, as the target's collective group is: the world is +cut into contiguous blocks of `--group-size` ranks (default: the whole world), +which is how a pipeline-parallel engine lays out its stages' TP groups +(`Mapping.tp_group`), and every block runs the same checks at once on its own +group. The layouts this covers: + +- world 4, group 4 (tep4 / tp4 on one node) and world 2, group 2 + (pp2 x tep2 / tp2 stage groups, run alone), from this launcher; +- world 4 in groups of 2 (the pp2 x tep2 / tp2 layout itself: groups [0, 1] + and [2, 3], each with the workspace a pp2 mapping builds), from this + launcher with `--group-size 2`; +- world 8 over two NVL72 trays (tep8), and world 8 in groups of 4 (the + pp2 x tep4 / tp4 layout, one group per tray), from an external two-node + `torchrun` that starts the rank body directly (the command is in the + entry's `note`): `--rank-worker [--group-size G]`. + +A group whose ranks do not share one host gets no IPC workspace (the target's +`ar_one_node`), so there the fp32 sums go to NCCL (strategy 0) exactly as the +target sends them, the workspace strategies are not run, and the two mixed +suites (which are about the workspace) do not apply. + +It runs as five jobs, each a fresh multi-rank process set with its own IPC +workspace: the fp32 suite (every strategy on fp32, with bf16 rows only on the +two workspace-free strategies, as a DeepSeek-V4.1 forward interleaves them), +the bf16 suite (every strategy on bf16), the core suite (the DeepSeek-V4.1 core +target's two call forms -- fp32 [T, 5120] on AUTO (NCCL across nodes) at every +row count 1..32768, bf16 [T, 6144] on NCCL at every row count 1..16384 -- and +decode graphs at every per-rank row count G and 6G of the capture batch set), +and two mixed suites, which replay the defect that forces that split -- the +ONESHOT kernel sizes the reset of its previous call's buffer in that call's +element count, so a workspace whose ONESHOT path sees two element sizes reads +stale data -- and assert that this harness sees it (see the entry's `note`); +the second does it with the core target's own layer sequence and the engram +rows on the wrong strategy. + +Not a pytest module, despite the `check_*` bodies. They are one fixed +sequence inside a single multi-rank job: each reads module-global rank state +that only `_run_one_rank` binds. Collected as tests they would run at world +size 1 against unbound globals, so neither this file's name nor its function +names match pytest's collection patterns. + +The collected entry point is `test_modeling_v2_allreduce.py` beside this file: +it starts this job and turns its exit code into an assertion. Both halves are +started by file path: the launcher must carry no CUDA or torch.distributed +state of its own, and the ranks reach the entry by absolute import. +""" + +import hashlib +import os +import random +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "allreduce requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +_GROUP_FLAG = "--group-size" +# Each suite's own deadline. The core suite sweeps every row count up to 32768 +# and captures 84 decode graphs, so it gets the largest share. +SUITE_DEADLINE_S: Dict[str, int] = { + "fp32": 600, + "bf16": 600, + "core": 2400, + "mixed": 600, + "mixed_core": 600, +} +DEADLINE_S = sum(SUITE_DEADLINE_S.values()) +"""Budget for the whole launcher: the sum of the suites' deadlines.""" + +SUITES: Tuple[str, ...] = tuple(SUITE_DEADLINE_S) + +HIDDEN = 5120 # DeepSeek-V4.1-Flash hidden size +ENGRAM_ROW = (24, 256) # one engram value row per token + +# Every strategy the op implements for a plain sum, by the integer it takes. +ALL_STRATEGIES: Dict[int, str] = { + 0: "NCCL", + 8: "NCCL_SYMMETRIC", + 3: "AUTO", + 4: "ONESHOT", + 5: "TWOSHOT", + 1: "MIN_LATENCY", +} +# The ones that read the IPC workspace; the other two run with workspace=None. +WORKSPACE_STRATEGIES = (3, 4, 5, 1) +# The strategies this group runs: all six on one node, the two workspace-free +# ones when the group spans nodes (no IPC workspace exists there). Rebound in +# the rank body. +STRATEGIES: Dict[int, str] = dict(ALL_STRATEGIES) + +# The fp32 [T, 5120], bf16 [T, 5120] and bf16 [T, 24, 256] row counts every +# strategy runs are the entry's cells (`allreduce.CELLS`), driven by +# `check_certified_cells_*`. The row-sharded engram check runs the 3-D form's +# row counts. +ENGRAM_ROWS: Tuple[int, ...] = (1, 8, 32, 4096) + +# One decode graph of a TP4 DeepSeek-V4.1 forward holds two fp32 all-reduces +# per decoder layer (output projection and routed experts) over 40 layers. +SITES_PER_DECODE_GRAPH = 80 +ENGRAM_SITE_EVERY = 10 # plus one bf16 engram row sum per this many fp32 sites +GRAPH_ROWS: Tuple[int, ...] = (1, 2, 3, 4, 8, 16, 32, 64, 128) + +# The DeepSeek-V4.1 core target's two call forms under tensor parallelism +# (`_allreduce32` and `_engram`): fp32 [T, 5120] partial sums on AUTO with the +# workspace when the group is on one node (NCCL when it spans nodes: the +# target's `ar_one_node`), and the engram rows as a 2-D bf16 [T, 24 * 256] view +# on NCCL. +ENGRAM_HIDDEN = 6144 +AUTO, NCCL = 3, 0 +CORE_FP32 = AUTO # the fp32 sums' strategy for this group; rebound in the rank body +# Rows per rank. The fp32 sums run on every layer, and layers 21-39 of a +# tensor-parallel context step can run a decoder pass over more rows than +# max_num_tokens: the finishing requests' 128-row tails plus the generation +# rows, up to 128 x 256 = 32768. The engram sums run on layers 1 and 14 only, +# ahead of that pass, so their rows stop at max_num_tokens (16384). +N_MAX = 32768 +ENGRAM_N_MAX = 16384 +SWEEP_SLACK = 509 # extra table rows (prime): every call reads its rows at its own offset +# What the fp32 [T, 5120] sums resolve to, per group size, on sm_103 (the op's +# lookup table, `selectStrategyLookUpTable`, is indexed by the group size and +# buckets T by floor(log2(T)) and hidden 5120 into its 4096 column; the 64 MiB +# workspace check sends [3277, 5120] fp32 and everything above to NCCL). Read +# from the op's own debug log (`AllReduceOp runtime strategy for rank 0:`). +# A group that spans nodes calls NCCL directly. +AUTO_FP32_BANDS_BY_GROUP: Dict[int, Tuple[Tuple[int, int, str], ...]] = { + 2: ((1, 3276, "ONESHOT"), (3277, N_MAX, "NCCL")), + 4: ((1, 1023, "ONESHOT"), (1024, 3276, "TWOSHOT"), (3277, N_MAX, "NCCL")), +} +NCCL_BANDS: Tuple[Tuple[int, int, str], ...] = ((1, N_MAX, "NCCL"),) +AUTO_FP32_BANDS: Tuple[Tuple[int, int, str], ...] = AUTO_FP32_BANDS_BY_GROUP[ + 4 +] # rebound in the rank body +# One decode forward: 40 layers, an attention and a MoE fp32 sum each, plus the +# engram rows' bf16 sum ahead of the attention sum on layers 1 and 14. +CORE_LAYERS = 40 +ENGRAM_LAYER_IDS = (1, 14) +# The capture batch set, per rank: every G in 1..32 plus the padded default list +# of cuda_graph_config {enable_padding: true, max_batch_size: 256} (1, 2, 4, 8, +# 16..128 step 8, 192, 256) -- 46 values. A captured step is generation-only +# with next_n 1 (DSpark off) or 6 (DSpark on), so the target's sums see T = G +# and T = 6G (82 values, up to 1536). The batches G 1..8, 16, 24, 32 with 5G and +# 6G of those stay in as well: they add 35 and 160. +GSET: Tuple[int, ...] = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_GATE_BATCHES: Tuple[int, ...] = tuple(range(1, 9)) + (16, 24, 32) +CORE_CAPTURE_ROWS: Tuple[int, ...] = tuple( + sorted( + set(GSET) + | {6 * g for g in GSET} + | set(_GATE_BATCHES) + | {5 * b for b in _GATE_BATCHES} + | {6 * b for b in _GATE_BATCHES} + ) +) +# Eager prefill-sized sums replayed between the decode graphs: row counts the +# runtime dumps showed, the decoder pass's maximum, and a row count on each side +# of AUTO's edges. The engram form runs only where its rows reach. +PREFILL_BETWEEN_REPLAYS: Tuple[int, ...] = (1018, 1935, 2048, 3277, 10444, 16384, 32768) + +# Bound inside the rank body, never in the launcher: the launcher must carry no +# CUDA or torch.distributed state, and importing the entry pulls in +# tensorrt_llm. RANK and WORLD are this rank's position in its group and the +# group's size; HOST_GROUP is a gloo group over the group (for host-side checks +# only); GROUP is the group's global ranks, which is what the op takes; LAYOUT +# names the cells' layout this job runs. +allreduce: Any = None +validating: Any = None +HOST_GROUP: Any = None +RANK = 0 +WORLD = 1 +GROUP: List[int] = [0] +WORKSPACE: Optional[torch.Tensor] = None +ONE_NODE = True # every rank of the group on one host: the IPC workspace exists +LAYOUT = "" + + +def _barrier() -> None: + """Host-side barrier over this rank's group.""" + dist.barrier(group=HOST_GROUP) + + +def _host_allgather(obj: Any) -> List[Any]: + """One host-side value from every rank of this rank's group, in group order.""" + gathered: List[Any] = [None] * WORLD + dist.all_gather_object(gathered, obj, group=HOST_GROUP) + return gathered + + +def _payload(rank: int, shape: Sequence[int], seed: int, dtype: torch.dtype) -> torch.Tensor: + """One rank's contribution, distinct per (rank, seed), whose sums are exact. + + Every rank regenerates every other rank's payload from the same seeds, + which is what keeps the reference arithmetic rather than a second + collective. fp32 values are multiples of 1/256 below 4096 in magnitude, so + any sum of four of them, in any order, is exact in fp32 (at most 24 + significant bits); bf16 values are multiples of 1/8 of at most 4, so every + partial sum of four is a multiple of 1/8 of at most 16 and exact in bf16's + 8 significant bits. Every assertion on these payloads can be bitwise. + """ + gen = torch.Generator(device="cuda").manual_seed(seed * 977 + rank + 1) + if dtype == torch.float32: + k = torch.randint( + -(2**20), 2**20 + 1, tuple(shape), generator=gen, device="cuda", dtype=torch.int64 + ) + return (k.double() / 256.0).float() + k = torch.randint(-32, 33, tuple(shape), generator=gen, device="cuda", dtype=torch.int32) + return (k.float() / 8.0).to(dtype) + + +def _sum_ref(shape: Sequence[int], seed: int, dtype: torch.dtype) -> torch.Tensor: + """The arithmetic sum of every rank's `_payload`, exact in `dtype` here. + + Delegates to the entry: what the sum is supposed to be is the entry's + claim, built on this rank alone from the seeds every rank shares. + """ + return allreduce.reference([_payload(r, shape, seed, dtype) for r in range(WORLD)], GROUP)[0] + + +def _random_parts(shape: Sequence[int], seed: int, dtype: torch.dtype) -> List[torch.Tensor]: + """Unrestricted normal payloads for every rank, the shape real partials have.""" + parts = [] + for r in range(WORLD): + gen = torch.Generator(device="cuda").manual_seed(seed * 131 + r) + parts.append( + torch.randn(tuple(shape), generator=gen, device="cuda", dtype=torch.float32).to(dtype) + ) + return parts + + +def _call(x: torch.Tensor, strategy: int, trigger_completion_at_end: bool = True) -> torch.Tensor: + ws = WORKSPACE if strategy in WORKSPACE_STRATEGIES else None + outs = allreduce( + x, None, None, None, None, ws, GROUP, strategy, 0, 1e-6, trigger_completion_at_end + ) + assert isinstance(outs, list) and len(outs) == 1, ( + f"op=0 must return one tensor, got {len(outs)}" + ) + return outs[0] + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + """Gate of exactly zero, against a reference that is exact here. + + The gate itself is the entry's `compare`; `where` only labels the failure. + """ + try: + allreduce.compare([out], [ref]) + except AssertionError as exc: + raise AssertionError(f"{where}: {exc}") from exc + + +def _digest(t: torch.Tensor) -> str: + return hashlib.sha1(t.contiguous().view(torch.uint8).cpu().numpy().tobytes()).hexdigest() + + +def _assert_same_bits_on_every_rank(out: torch.Tensor, where: str) -> None: + """Every rank must hold the same bytes. A host-side check of the result, + not a source of the reference.""" + digests = _host_allgather(_digest(out)) + assert len(set(digests)) == 1, f"{where}: ranks disagree: {digests}" + + +def _layout_cells(suite: str) -> List[Any]: + """The entry's cells for this job's layout and `suite`, in declaration order.""" + cells = [c for c in allreduce.CELLS if c.spec["layout"] == LAYOUT and c.spec["suite"] == suite] + assert cells, f"no {suite} cells for layout {LAYOUT}" + return cells + + +def _drive_certified_cells(suite: str) -> None: + """Every cell of `suite` at this layout: every row count on its strategy. + + Per call form, row count outer and strategy inner, so the strategies take + turns on one workspace as a serving engine's mixed traffic does. Inside + `validating`, so the guard is shown to admit every input a cell names. + Each result must be fresh (not aliasing the input), contiguous, on the + input's device, of the input's shape and dtype, the exact sum, and the + input untouched. + """ + cells = _layout_cells(suite) + forms: List[Tuple[torch.dtype, Tuple[int, ...]]] = [] + for c in cells: + form = (c.spec["dtype"], tuple(c.spec["trailing"])) + if form not in forms: + forms.append(form) + for dtype, trailing in forms: + of_form = [ + c for c in cells if (c.spec["dtype"], tuple(c.spec["trailing"])) == (dtype, trailing) + ] + rows_of_form = tuple(of_form[0].spec["rows"]) + assert all(tuple(c.spec["rows"]) == rows_of_form for c in of_form), [c.why for c in of_form] + for rows in rows_of_form: + for cell in of_form: + strategy = cell.spec["strategy"] + name = ALL_STRATEGIES[strategy] + assert strategy in STRATEGIES, f"{cell.why}: {name} does not apply to this group" + shape = (rows, *trailing) + base = 1000 if dtype == torch.float32 else 2000 + len(shape) + seed = base + rows * 16 + strategy + x = _payload(RANK, shape, seed, dtype) + before = x.clone() + with validating(allreduce): + out = _call(x, strategy) + torch.cuda.synchronize() + where = f"{cell.why} rows={rows}" + assert out.shape == x.shape and out.dtype == dtype, (where, out.shape, out.dtype) + assert out.is_contiguous() and out.device == x.device, where + assert out.data_ptr() != x.data_ptr(), f"{where}: result aliases the input" + _assert_bitwise(out, _sum_ref(shape, seed, dtype), where) + assert torch.equal(x, before), f"{where}: input was modified" + _barrier() + + +def check_certified_cells_fp32() -> None: + """fp32 [T, 5120]: every strategy returns the exact sum, fresh, input untouched.""" + _drive_certified_cells("fp32") + + +def check_certified_cells_bf16() -> None: + """bf16 [T, 5120] and the 3-D engram shape [T, 24, 256], on every strategy.""" + _drive_certified_cells("bf16") + + +def _engram_rows_owned_by_one_rank_come_back_bitwise(strategies: Sequence[int]) -> None: + """Row-sharded lookups: each row is nonzero on exactly one rank. + + That is how a vocabulary-sharded table is combined -- the owner holds the + real row and every other rank a zero row -- so the sum is `row + 0 + 0 + 0` + and must be the owner's row bit for bit, in bf16, whatever the strategy + and however unrestricted the values are. + """ + for rows in ENGRAM_ROWS: + shape = (rows, *ENGRAM_ROW) + owner = torch.Generator(device="cpu").manual_seed(rows) + owners = torch.randint(0, WORLD, (rows,), generator=owner).cuda() + full = _random_parts(shape, 3000 + rows, torch.bfloat16)[0] + for strategy in strategies: + name = STRATEGIES[strategy] + mine = torch.where((owners == RANK).view(rows, 1, 1), full, torch.zeros_like(full)) + out = _call(mine, strategy) + torch.cuda.synchronize() + where = f"engram rows={rows} {name}" + assert torch.equal(out.view(torch.int16), full.view(torch.int16)), ( + f"{where}: not the owner's row bits" + ) + _assert_same_bits_on_every_rank(out, where) + _barrier() + + +def check_engram_rows_owned_by_one_rank_on_the_workspace_free_strategies() -> None: + _engram_rows_owned_by_one_rank_come_back_bitwise((0, 8)) + + +def check_engram_rows_owned_by_one_rank_on_every_strategy() -> None: + _engram_rows_owned_by_one_rank_come_back_bitwise(tuple(STRATEGIES)) + + +def check_fp32_on_the_workspace_interleaved_with_bf16_on_nccl() -> None: + """The DeepSeek-V4.1 mix: fp32 sums on a workspace strategy, bf16 rows on NCCL. + + A tensor-parallel layer sums its fp32 output-projection and routed-expert + partials and, on engram layers, a bf16 row; the bf16 sum goes through a + workspace-free strategy, so the workspace's ONESHOT path only ever sees + fp32. Decode- and prefill-sized steps alternate, as a serving engine + alternates them. + """ + rows_cycle = (1, 8, 32, 4096, 3, 128, 2048, 1) + step = 0 + for strategy in WORKSPACE_STRATEGIES: + for bf16_strategy in (0, 8): + for rows in rows_cycle: + step += 1 + seeds = (7000 + step * 3, 7001 + step * 3, 7002 + step * 3) + wo = _call(_payload(RANK, (rows, HIDDEN), seeds[0], torch.float32), strategy) + eng = _call( + _payload(RANK, (rows, *ENGRAM_ROW), seeds[1], torch.bfloat16), bf16_strategy + ) + moe = _call(_payload(RANK, (rows, HIDDEN), seeds[2], torch.float32), strategy) + torch.cuda.synchronize() + where = f"interleave {STRATEGIES[strategy]}+{STRATEGIES[bf16_strategy]} rows={rows}" + _assert_bitwise( + wo, _sum_ref((rows, HIDDEN), seeds[0], torch.float32), f"{where} fp32#1" + ) + _assert_bitwise( + eng, _sum_ref((rows, *ENGRAM_ROW), seeds[1], torch.bfloat16), f"{where} bf16" + ) + _assert_bitwise( + moe, _sum_ref((rows, HIDDEN), seeds[2], torch.float32), f"{where} fp32#2" + ) + _barrier() + + +def check_fp32_partial_sums_are_reduced_in_fp32() -> None: + """Unrestricted fp32 partials: the result is the fp32 sum, on every rank alike. + + Four fp32 addends in some order: the default fp32 tolerances of + `assert_close` cover the at-most-three roundings any summation order + makes. And every rank must hold the same bits -- a rank that disagrees + with its peers would feed a different residual stream into the next + layer on each rank. The reference is the entry's (the correctly rounded + sum); the gate is not the entry's zero-tolerance `compare`, which is for + payloads whose sums are exact. + """ + for rows in (1, 32, 4096): + shape = (rows, HIDDEN) + parts = _random_parts(shape, 4000 + rows, torch.float32) + ref = allreduce.reference(parts, GROUP)[0] + for strategy, name in STRATEGIES.items(): + out = _call(parts[RANK].clone(), strategy) + torch.cuda.synchronize() + where = f"fp32 random rows={rows} {name}" + torch.testing.assert_close(out, ref, msg=lambda m: f"{where}: {m}") + _assert_same_bits_on_every_rank(out, where) + _barrier() + + +def check_trigger_completion_at_end_false_gives_the_same_bits() -> None: + """The flag moves the programmatic-launch trigger, not the arithmetic.""" + for rows in (1, 8, 2048): + for strategy in WORKSPACE_STRATEGIES: + seed = 5000 + rows + strategy + shape = (rows, HIDDEN) + out = _call( + _payload(RANK, shape, seed, torch.float32), + strategy, + trigger_completion_at_end=False, + ) + torch.cuda.synchronize() + _assert_bitwise( + out, + _sum_ref(shape, seed, torch.float32), + f"trigger=False rows={rows} {STRATEGIES[strategy]}", + ) + _barrier() + + +def check_cuda_graph_decode_graphs_replay_in_any_order() -> None: + """One graph per decode row count, 80 fp32 all-reduce sites each, shared pool. + + Every tenth fp32 site is followed by a bf16 engram-row sum on NCCL, the + composition a DeepSeek-V4.1 decode graph has. + + Per strategy: every graph is captured into one memory pool and kept alive, + then replayed shuffled, twice, each time with a fresh payload copied into + the captured inputs and an eager call of an unrelated size in between (the + workspace strategies keep per-call state in device memory, and an eager + call moves it between replays). Every site's output must be the exact sum + of what its input held at replay time. + """ + rng = random.Random(1234) # identical on every rank, no collective needed + distinct = 8 # inputs per graph, cycled over the sites + for strategy, name in STRATEGIES.items(): + # The group's communicator and, for the workspace strategies, the + # workspace state must exist before a capture; warm on a side stream. + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for rows in GRAPH_ROWS: + _call(_payload(RANK, (rows, HIDDEN), 1, torch.float32), strategy) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + _barrier() + + pool = None + graphs: Dict[int, torch.cuda.CUDAGraph] = {} + inputs: Dict[int, List[torch.Tensor]] = {} + outputs: Dict[int, List[torch.Tensor]] = {} + engram_inputs: Dict[int, List[torch.Tensor]] = {} + engram_outputs: Dict[int, List[torch.Tensor]] = {} + for rows in GRAPH_ROWS: + inputs[rows] = [ + _payload(RANK, (rows, HIDDEN), 7, torch.float32) for _ in range(distinct) + ] + engram_inputs[rows] = [ + _payload(RANK, (rows, *ENGRAM_ROW), 8, torch.bfloat16) for _ in range(2) + ] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + outputs[rows], engram_outputs[rows] = [], [] + for i in range(SITES_PER_DECODE_GRAPH): + outputs[rows].append(_call(inputs[rows][i % distinct], strategy)) + if i % ENGRAM_SITE_EVERY == ENGRAM_SITE_EVERY - 1: + # a bf16 engram row, on NCCL: see the interleave check + engram_outputs[rows].append( + _call(engram_inputs[rows][len(engram_outputs[rows]) % 2], 0) + ) + pool = graph.pool() if pool is None else pool + graphs[rows] = graph + torch.cuda.synchronize() + _barrier() + + for round_ in range(2): + order = list(GRAPH_ROWS) + rng.shuffle(order) + for rows in order: + seeds = [60000 + round_ * 1000 + rows * 10 + i for i in range(distinct)] + for i, seed in enumerate(seeds): + inputs[rows][i].copy_(_payload(RANK, (rows, HIDDEN), seed, torch.float32)) + engram_seeds = [65000 + round_ * 1000 + rows * 10 + i for i in range(2)] + for i, seed in enumerate(engram_seeds): + engram_inputs[rows][i].copy_( + _payload(RANK, (rows, *ENGRAM_ROW), seed, torch.bfloat16) + ) + graphs[rows].replay() + eager_rows = rng.choice((33, 777, 1500)) + eager_seed = 70000 + round_ * 1000 + rows + eager = _call( + _payload(RANK, (eager_rows, HIDDEN), eager_seed, torch.float32), strategy + ) + torch.cuda.synchronize() + where = f"graph {name} rows={rows} round={round_}" + refs = [_sum_ref((rows, HIDDEN), seed, torch.float32) for seed in seeds] + for site, out in enumerate(outputs[rows]): + _assert_bitwise(out, refs[site % distinct], f"{where} site={site}") + engram_refs = [ + _sum_ref((rows, *ENGRAM_ROW), seed, torch.bfloat16) for seed in engram_seeds + ] + for site, out in enumerate(engram_outputs[rows]): + _assert_bitwise(out, engram_refs[site % 2], f"{where} engram site={site}") + _assert_bitwise( + eager, + _sum_ref((eager_rows, HIDDEN), eager_seed, torch.float32), + f"{where} eager", + ) + _barrier() + del graphs, outputs, inputs, engram_outputs, engram_inputs + torch.cuda.synchronize() + _barrier() + + +def check_the_gate_discriminates_a_wrong_reduce() -> None: + """Deliberately wrong reductions, each of which the gates above must reject. + + A zero-tolerance gate cannot be too loose, but it can be blind; these are + the ways an all-reduce of these payloads plausibly goes wrong, each with + the right shape and dtype. The last one is the real op run on a wrong + input -- the fp32 partials rounded to bf16 before the reduction, which is + what an all-reduce in the activation dtype does to them -- and it must + fail the fp32 gate that the correct call passes. + + Each variant is formed in fp64 and rounded once to fp32, the output's + dtype, which is what the entry's `compare` sees. + """ + rows, seed = 32, 88000 + shape = (rows, HIDDEN) + blocks = [_payload(r, shape, seed, torch.float32).double() for r in range(WORLD)] + good = _call(_payload(RANK, shape, seed, torch.float32), 0) + torch.cuda.synchronize() + exact = sum(blocks) + _assert_bitwise(good, _sum_ref(shape, seed, torch.float32), "baseline") + + wrong = { + "one rank never contributed": sum(blocks[:-1]), + "one rank counted twice": exact + blocks[0], + "mean instead of sum": exact / WORLD, + "partials rounded to bf16": sum(b.to(torch.bfloat16).double() for b in blocks), + } + for label, variant in wrong.items(): + variant = variant.float() + fraction = (variant != good).float().mean().item() + assert fraction > 0.9, f"{label}: only {fraction:.4f} of elements differ" + try: + _assert_bitwise(good, variant, label) + except AssertionError: + pass + else: + raise AssertionError(f"{label}: the gate accepted a wrong reduction") + + parts = _random_parts(shape, 89000, torch.float32) + ref = allreduce.reference(parts, GROUP)[0] + in_bf16 = _call(parts[RANK].to(torch.bfloat16), 0).float() + torch.cuda.synchronize() + try: + torch.testing.assert_close(in_bf16, ref) + except AssertionError: + pass + else: + raise AssertionError("an all-reduce of bf16-rounded partials passed the fp32 gate") + _barrier() + + +def check_wrapper_guards_a_non_contiguous_input() -> None: + """The wrapper rejects a strided view: the op reads packed memory, and on + the MPI path a strided view was reduced over the wrong elements. + + Only the guard is checked here. On this path strategy 0 sums a clone of + the input, which is contiguous, so the bare op does not show the defect the + guard is for; the guard is what this test protects. + """ + base = _payload(RANK, (8, 2 * HIDDEN), 91000, torch.float32) + view = base[:, ::2] + assert not view.is_contiguous() + try: + with validating(allreduce): + _call(view, 0) + except AssertionError as err: + assert "contiguous" in str(err), err + else: + raise AssertionError("the wrapper accepted a non-contiguous input") + _barrier() + + +def check_wrapper_guards_float8() -> None: + """The wrapper rejects e4m3, which the MPI path accepted and summed as + bytes rather than floats. + + Only the guard is checked here, for the same reason as above: what the + ProcessGroup's own allreduce does with a float8 tensor is not this entry's + contract. + """ + x = torch.full((4, HIDDEN), float(RANK + 1), device="cuda").to(torch.float8_e4m3fn) + try: + with validating(allreduce): + _call(x, 0) + except AssertionError as err: + assert "float8_e4m3fn" in str(err), err + else: + raise AssertionError("the wrapper accepted a float8_e4m3fn input") + _barrier() + + +def check_oneshot_reads_stale_data_after_an_element_size_change() -> None: + """A known defect, made visible: the harness must see it for a clean run to count. + + ONESHOT resets the buffer its previous call used, sizing the reset as that + call's element count times the CURRENT call's element size. So after an + fp32 call, a bf16 call resets only half of the fp32 call's bytes, and the + call that next reuses that buffer takes the unreset half for data a peer + has already written. With one rank arriving late, every other rank sums + that rank's contribution to an earlier call instead of this one. + + The identical sequence with fp32 throughout is the control, and it must be + exact; the mixed one must be wrong on every on-time rank. This is why the + other suites give each element size its own workspace, and why the + contract routes the bf16 rows of a mixed-dtype forward through NCCL. + """ + rows = 128 + late = WORLD - 1 + + def sequence(mid_dtype: torch.dtype, seed: int) -> float: + _call(_payload(RANK, (rows, HIDDEN), seed, torch.float32), 4) + _call(_payload(RANK, (1, HIDDEN), seed + 1, mid_dtype), 4) + _call(_payload(RANK, (1, HIDDEN), seed + 2, mid_dtype), 4) + x = _payload(RANK, (rows, HIDDEN), seed + 3, torch.float32) + torch.cuda.synchronize() + _barrier() + if RANK == late: + torch.cuda._sleep(200_000_000) # hold this rank's push back + out = _call(x, 4) + torch.cuda.synchronize() + ref = _sum_ref((rows, HIDDEN), seed + 3, torch.float32) + return (out != ref).float().mean().item() + + for trial in range(3): + fractions = _host_allgather(sequence(torch.float32, 95000 + trial * 10)) + assert all(f == 0.0 for f in fractions), ( + f"control (fp32 throughout) trial {trial}: {fractions}" + ) + for trial in range(3): + fractions = _host_allgather(sequence(torch.bfloat16, 96000 + trial * 10)) + on_time = [f for r, f in enumerate(fractions) if r != late] + assert all(f > 0.9 for f in on_time), ( + f"mixed trial {trial}: the stale read did not show on every on-time rank: {fractions}" + ) + _barrier() + + +def _sweep_table( + hidden: int, seed: int, dtype: torch.dtype, n_max: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """This rank's payload rows and their exact sum over the group, n_max + SWEEP_SLACK rows. + + The sum is exact in `dtype` (see `_payload`), so it is stored in `dtype`. + """ + shape = (n_max + SWEEP_SLACK, hidden) + return _payload(RANK, shape, seed, dtype), _sum_ref(shape, seed, dtype) + + +def _sweep_offsets(k: int) -> Tuple[int, int, int]: + """Row offsets of step k's three calls into their tables. + + Each call reads its rows at its own offset, so no two nearby calls hand the + op the same values at the same row: a result left over from an earlier + call (the stale read of the mixed suite) cannot pass for this one. + """ + return (k * 131) % SWEEP_SLACK, (k * 173 + 61) % SWEEP_SLACK, (k * 197 + 263) % SWEEP_SLACK + + +def check_core_tp_sums_at_every_row_count() -> None: + """Every T in 1..N_MAX, the core target's layer sequence, twice over. + + AUTO picks its transport from (T, hidden), and NCCL its algorithm and + protocol from the message size, so every row count is run, not sampled: + per step, the engram rows' bf16 [T, 6144] sum on NCCL (T up to + ENGRAM_N_MAX, the engine's max_num_tokens), the attention's fp32 [T, 5120] + sum on CORE_FP32, the MoE's fp32 [T, 5120] sum on CORE_FP32 -- the order a + tensor-parallel engram layer issues them in, the fp32 ones on one workspace + when the group has one. Above ENGRAM_N_MAX a step is a decoder-pass layer + (21-39): the two fp32 sums only, up to N_MAX. First with T ascending, then + in a seeded shuffle, so every band of AUTO_FP32_BANDS is entered from every + other band (a prefill chunk between decode steps). Every element of every + result is compared with the exact sum; the mismatch counts stay on the + device and are read once per pass. + + This is the driver of the entry's two core cells for this layout: their + strategy, shape and row range are asserted to be the ones swept here. + """ + core = {c.spec["dtype"]: c.spec for c in _layout_cells("core")} + assert set(core) == {torch.float32, torch.bfloat16}, list(core) + for dtype, strategy, trailing, n_max in ( + (torch.float32, CORE_FP32, (HIDDEN,), N_MAX), + (torch.bfloat16, NCCL, (ENGRAM_HIDDEN,), ENGRAM_N_MAX), + ): + spec = core[dtype] + assert spec["strategy"] == strategy, (dtype, spec["strategy"], strategy) + assert tuple(spec["trailing"]) == trailing, (dtype, spec["trailing"]) + assert list(spec["rows"]) == list(range(1, n_max + 1)), (dtype, spec["rows"]) + with validating(allreduce): + # The guard admits both core forms; the sweep itself runs unguarded. + _call(_payload(RANK, (1, HIDDEN), 1, torch.float32), CORE_FP32) + _call(_payload(RANK, (1, ENGRAM_HIDDEN), 1, torch.bfloat16), NCCL) + torch.cuda.synchronize() + _barrier() + + wo, wo_sum = _sweep_table(HIDDEN, 120001, torch.float32, N_MAX) + moe, moe_sum = _sweep_table(HIDDEN, 120002, torch.float32, N_MAX) + eng, eng_sum = _sweep_table(ENGRAM_HIDDEN, 120003, torch.bfloat16, ENGRAM_N_MAX) + + # The gate is not blind to a stale result: at every band edge, the right + # answer of a step differs on nearly every element from the answers of the + # calls before it on the same transport (the previous three fp32 sums on + # the workspace, the previous two bf16 sums). + k = 1000 + for t in sorted( + {1, 8, ENGRAM_N_MAX, N_MAX} | {edge for band in AUTO_FP32_BANDS for edge in band[:2]} + ): + o_wo = _sweep_offsets(k)[0] + o_eng = _sweep_offsets(k)[1] + wo_now = wo_sum[o_wo : o_wo + t] + earlier = [ + ("fp32", wo_now, moe_sum, _sweep_offsets(k - 1)[2]), + ("fp32", wo_now, wo_sum, _sweep_offsets(k - 1)[0]), + ("fp32", wo_now, moe_sum, _sweep_offsets(k - 2)[2]), + ] + if t <= ENGRAM_N_MAX: + earlier += [ + ("bf16", eng_sum[o_eng : o_eng + t], eng_sum, _sweep_offsets(k - 1)[1]), + ("bf16", eng_sum[o_eng : o_eng + t], eng_sum, _sweep_offsets(k - 2)[1]), + ] + for kind, now, table, o in earlier: + fraction = (now != table[o : o + t]).float().mean().item() + assert fraction > 0.9, f"{kind} T={t}: a stale result differs on only {fraction:.4f}" + + ascending = list(range(1, N_MAX + 1)) + shuffled = list(ascending) + random.Random(4242).shuffle(shuffled) + k = 0 + for name, order in (("ascending", ascending), ("shuffled", shuffled)): + bad = torch.zeros(3, N_MAX + 1, dtype=torch.int64, device="cuda") + calls = 0 + for t in order: + k += 1 + o_wo, o_eng, o_moe = _sweep_offsets(k) + if t <= ENGRAM_N_MAX: + e = _call(eng[o_eng : o_eng + t], NCCL) + assert e.shape == (t, ENGRAM_HIDDEN) and e.dtype == torch.bfloat16, ( + t, + e.shape, + e.dtype, + ) + bad[1, t] += (e != eng_sum[o_eng : o_eng + t]).sum() + calls += 1 + a = _call(wo[o_wo : o_wo + t], CORE_FP32) + b = _call(moe[o_moe : o_moe + t], CORE_FP32) + assert a.shape == (t, HIDDEN) and a.dtype == torch.float32, (t, a.shape, a.dtype) + assert b.shape == (t, HIDDEN) and b.dtype == torch.float32, (t, b.shape, b.dtype) + bad[0, t] += (a != wo_sum[o_wo : o_wo + t]).sum() + bad[2, t] += (b != moe_sum[o_moe : o_moe + t]).sum() + calls += 2 + torch.cuda.synchronize() + wrong = bad.nonzero().tolist() + fp32_name = STRATEGIES[CORE_FP32] + labels = (f"attention fp32 {fp32_name}", "engram bf16 NCCL", f"MoE fp32 {fp32_name}") + assert not wrong, f"{name}: wrong sums at " + ", ".join( + f"{labels[i]} T={t} ({bad[i, t].item()} elements)" for i, t in wrong[:20] + ) + if RANK == 0: + print( + f"[rank 0] core sweep {name}: fp32 {fp32_name} T=1..{N_MAX}, engram bf16 NCCL T=1..{ENGRAM_N_MAX}, " + f"{calls} calls, every element exact", + flush=True, + ) + _barrier() + del wo, wo_sum, moe, moe_sum, eng, eng_sum + torch.cuda.empty_cache() + _barrier() + + +def _core_forward( + f_in: List[torch.Tensor], e_in: List[torch.Tensor] +) -> List[Tuple[str, int, torch.Tensor]]: + """One decode forward's all-reduce sites, in the core target's order. + + The 80 fp32 sites cycle over the buffers of `f_in`; engram layer j reads + `e_in[j]`. Returns (kind, buffer index, result) per site. + """ + sites: List[Tuple[str, int, torch.Tensor]] = [] + fp32_site = 0 + for layer in range(CORE_LAYERS): + if layer in ENGRAM_LAYER_IDS: + j = ENGRAM_LAYER_IDS.index(layer) + sites.append(("bf16", j, _call(e_in[j], NCCL))) + for _ in range(2): # attention output projection, then routed experts + i = fp32_site % len(f_in) + sites.append(("fp32", i, _call(f_in[i], CORE_FP32))) + fp32_site += 1 + return sites + + +def check_core_tp_decode_graphs_replay_bitwise_equal_to_eager() -> None: + """One decode graph per CORE_CAPTURE_ROWS entry, the core target's 82 sites each. + + The engine's state after warm-up: every graph captured (largest first) into + one shared pool and kept alive. Then three passes over all of them in + shuffled order, each with new data copied into the captured inputs: + + 1. exact payloads, one graph at a time, with an eager prefill-sized sum + (PREFILL_BETWEEN_REPLAYS, rotating) on both call forms between replays: + every site bitwise equal to the exact sum; + 2. exact payloads, every graph replayed back to back with nothing + synchronizing, then all checked; + 3. unrestricted normal payloads, where the sum order shows in the bits: + every site bitwise equal to an eager call on the same buffer afterwards. + + Before every replay of passes 1 and 2, every site's output is checked to + differ from the new exact sum on more than 90% of its elements, so a replay + that did not run (a stale output) cannot pass. + """ + distinct = 8 # fp32 input buffers per graph, cycled over its 80 fp32 sites + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for t in CORE_CAPTURE_ROWS: + _call(_payload(RANK, (t, HIDDEN), 1, torch.float32), CORE_FP32) + _call(_payload(RANK, (t, ENGRAM_HIDDEN), 2, torch.bfloat16), NCCL) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + _barrier() + + pool = None + graphs: Dict[int, torch.cuda.CUDAGraph] = {} + f_in: Dict[int, List[torch.Tensor]] = {} + e_in: Dict[int, List[torch.Tensor]] = {} + sites: Dict[int, List[Tuple[str, int, torch.Tensor]]] = {} + for t in sorted(CORE_CAPTURE_ROWS, reverse=True): + f_in[t] = [_payload(RANK, (t, HIDDEN), 30 + i, torch.float32) for i in range(distinct)] + e_in[t] = [ + _payload(RANK, (t, ENGRAM_HIDDEN), 50 + j, torch.bfloat16) + for j in range(len(ENGRAM_LAYER_IDS)) + ] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + sites[t] = _core_forward(f_in[t], e_in[t]) + assert len(sites[t]) == 2 * CORE_LAYERS + len(ENGRAM_LAYER_IDS), len(sites[t]) + pool = graph.pool() if pool is None else pool + graphs[t] = graph + torch.cuda.synchronize() + _barrier() + + def stage_exact(t: int, seed: int) -> Dict[Tuple[str, int], torch.Tensor]: + refs = {} + for i, x in enumerate(f_in[t]): + x.copy_(_payload(RANK, (t, HIDDEN), seed + i, torch.float32)) + refs[("fp32", i)] = _sum_ref((t, HIDDEN), seed + i, torch.float32) + for j, x in enumerate(e_in[t]): + x.copy_(_payload(RANK, (t, ENGRAM_HIDDEN), seed + 100 + j, torch.bfloat16)) + refs[("bf16", j)] = _sum_ref((t, ENGRAM_HIDDEN), seed + 100 + j, torch.bfloat16) + return refs + + def verify(t: int, refs: Dict[Tuple[str, int], torch.Tensor], where: str) -> None: + bad = torch.zeros((), dtype=torch.int64, device="cuda") + for kind, i, out in sites[t]: + bad += (out != refs[(kind, i)]).sum() + assert bad.item() == 0, ( + f"{where} T={t}: {bad.item()} wrong elements over {len(sites[t])} sites" + ) + + stale_worst = [1.0] + + def assert_stale_is_visible( + t: int, refs: Dict[Tuple[str, int], torch.Tensor], where: str + ) -> None: + """Before the replay: what the outputs hold now must fail the gate.""" + fractions = torch.stack( + [(out != refs[(kind, i)]).double().mean() for kind, i, out in sites[t]] + ) + worst = fractions.min().item() + stale_worst[0] = min(stale_worst[0], worst) + assert worst > 0.9, ( + f"{where} T={t}: a stale output differs from the new sum on only {worst:.4f}" + ) + + rng = random.Random(777) + order = list(CORE_CAPTURE_ROWS) + # Pass 1: checked one graph at a time, prefill-sized eager traffic in between. + rng.shuffle(order) + for pos, t in enumerate(order): + refs = stage_exact(t, 200000 + 1000 * t) + assert_stale_is_visible(t, refs, "pass 1 (before replay)") + graphs[t].replay() + prefill = PREFILL_BETWEEN_REPLAYS[pos % len(PREFILL_BETWEEN_REPLAYS)] + seed = 300000 + pos + eager_f = _call(_payload(RANK, (prefill, HIDDEN), seed, torch.float32), CORE_FP32) + eager_e = None + if prefill <= ENGRAM_N_MAX: + eager_e = _call(_payload(RANK, (prefill, ENGRAM_HIDDEN), seed, torch.bfloat16), NCCL) + torch.cuda.synchronize() + verify(t, refs, "pass 1 (checked)") + _assert_bitwise( + eager_f, + _sum_ref((prefill, HIDDEN), seed, torch.float32), + f"pass 1 eager fp32 T={prefill}", + ) + if eager_e is not None: + _assert_bitwise( + eager_e, + _sum_ref((prefill, ENGRAM_HIDDEN), seed, torch.bfloat16), + f"pass 1 eager bf16 T={prefill}", + ) + del eager_f, eager_e + _barrier() + # Pass 2: every graph back to back, nothing synchronizing between replays. + rng.shuffle(order) + staged = {t: stage_exact(t, 400000 + 1000 * t) for t in order} + for t in order: + assert_stale_is_visible(t, staged[t], "pass 2 (before replay)") + for t in order: + graphs[t].replay() + torch.cuda.synchronize() + for t in order: + verify(t, staged[t], "pass 2 (burst)") + del staged + _barrier() + # Pass 3: unrestricted payloads; the replay must be bit for bit what an + # eager call on the same buffer returns. + rng.shuffle(order) + for pos, t in enumerate(order): + for i, x in enumerate(f_in[t]): + x.copy_(_random_parts((t, HIDDEN), 500000 + 1000 * t + i, torch.float32)[RANK]) + for j, x in enumerate(e_in[t]): + x.copy_(_random_parts((t, ENGRAM_HIDDEN), 600000 + 1000 * t + j, torch.bfloat16)[RANK]) + graphs[t].replay() + torch.cuda.synchronize() + eager = {("fp32", i): _call(x, CORE_FP32) for i, x in enumerate(f_in[t])} + eager.update({("bf16", j): _call(x, NCCL) for j, x in enumerate(e_in[t])}) + torch.cuda.synchronize() + for s, (kind, i, out) in enumerate(sites[t]): + assert torch.equal( + out.view(torch.int32 if kind == "fp32" else torch.int16), + eager[(kind, i)].view(torch.int32 if kind == "fp32" else torch.int16), + ), f"pass 3 T={t} site={s} ({kind}): replay differs from eager on the same data" + _assert_same_bits_on_every_rank(sites[t][0][2], f"pass 3 T={t} first site") + _barrier() + if RANK == 0: + print( + f"[rank 0] core graphs: {len(graphs)} captured at {list(CORE_CAPTURE_ROWS)}; before every replay of " + "passes 1-2 a stale output " + f"differed from the new reference on at least {stale_worst[0]:.4f} of its elements", + flush=True, + ) + for held in (graphs, sites, f_in, e_in): + held.clear() + torch.cuda.synchronize() + _barrier() + + +def check_engram_rows_on_auto_corrupt_the_next_fp32_sum() -> None: + """The wrong strategy for the engram rows, run: the control for the core rule. + + The core target's layer order is MoE fp32 sum, then on an engram layer the + engram rows' bf16 sum, the attention fp32 sum and the MoE fp32 sum. With + the engram sum on NCCL (the target's choice) all four are exact even with + one rank late to the last one. With the engram sum on AUTO -- the same + workspace as the fp32 sums -- the one-shot kernel's reset after the bf16 + call covers half of the fp32 buffer before it, and the last fp32 sum, which + reuses that buffer, returns the stale half on every on-time rank. + Runs in a process of its own: it leaves the workspace inconsistent. + """ + late = WORLD - 1 + + def layer(engram_strategy: int, t: int, seed: int) -> Tuple[List[float], float]: + prev = _call(_payload(RANK, (t, HIDDEN), seed, torch.float32), AUTO) + eng = _call(_payload(RANK, (t, ENGRAM_HIDDEN), seed + 1, torch.bfloat16), engram_strategy) + attn = _call(_payload(RANK, (t, HIDDEN), seed + 2, torch.float32), AUTO) + torch.cuda.synchronize() + head = [ + (prev != _sum_ref((t, HIDDEN), seed, torch.float32)).float().mean().item(), + (eng != _sum_ref((t, ENGRAM_HIDDEN), seed + 1, torch.bfloat16)).float().mean().item(), + (attn != _sum_ref((t, HIDDEN), seed + 2, torch.float32)).float().mean().item(), + ] + x = _payload(RANK, (t, HIDDEN), seed + 3, torch.float32) + torch.cuda.synchronize() + _barrier() + if RANK == late: + torch.cuda._sleep(200_000_000) # hold this rank's push back + moe = _call(x, AUTO) + torch.cuda.synchronize() + return head, (moe != _sum_ref((t, HIDDEN), seed + 3, torch.float32)).float().mean().item() + + rows = (1, 8) + for t in rows: + for trial in range(3): + head, last = layer(NCCL, t, 130000 + 10 * trial + 100 * t) + got = _host_allgather((head, last)) + assert all(h == [0.0, 0.0, 0.0] and tail == 0.0 for h, tail in got), ( + f"engram on NCCL, T={t} trial {trial}: expected every sum exact, got {got}" + ) + for t in rows: + for trial in range(3): + _, last = layer(AUTO, t, 140000 + 10 * trial + 100 * t) + got = _host_allgather(last) + on_time = [f for r, f in enumerate(got) if r != late] + assert all(f > 0.9 for f in on_time), ( + f"engram on AUTO, T={t} trial {trial}: the stale read did not show on every on-time rank: {got}" + ) + _barrier() + + +SUITE_CHECKS: Dict[str, Tuple[Callable[[], None], ...]] = { + "fp32": ( + check_certified_cells_fp32, + check_fp32_partial_sums_are_reduced_in_fp32, + check_trigger_completion_at_end_false_gives_the_same_bits, + check_engram_rows_owned_by_one_rank_on_the_workspace_free_strategies, + check_fp32_on_the_workspace_interleaved_with_bf16_on_nccl, + check_cuda_graph_decode_graphs_replay_in_any_order, + check_the_gate_discriminates_a_wrong_reduce, + check_wrapper_guards_a_non_contiguous_input, + check_wrapper_guards_float8, + ), + "bf16": ( + check_certified_cells_bf16, + check_engram_rows_owned_by_one_rank_on_every_strategy, + ), + # The DeepSeek-V4.1 core target's tensor-parallel surface: + # every row count, and every capture batch size, of its two call forms. + "core": ( + check_core_tp_sums_at_every_row_count, + check_core_tp_decode_graphs_replay_bitwise_equal_to_eager, + ), + # Each in a process of its own: they deliberately leave the workspace's + # ONESHOT state inconsistent. + "mixed": (check_oneshot_reads_stale_data_after_an_element_size_change,), + "mixed_core": (check_engram_rows_on_auto_corrupt_the_next_fp32_sum,), +} + + +# Checks and suites about the IPC workspace. A group that spans nodes has none +# (the target's `ar_one_node` is False there and it sums on NCCL), so for such a +# group these are not run: the checks are dropped from their suite, and the +# suites refuse to start. +WORKSPACE_CHECKS = ( + check_trigger_completion_at_end_false_gives_the_same_bits, + check_fp32_on_the_workspace_interleaved_with_bf16_on_nccl, +) +WORKSPACE_SUITES = ("mixed", "mixed_core") + + +def _suite_checks(suite: str) -> Tuple[Callable[[], None], ...]: + """The checks of `suite` that apply to this group.""" + if ONE_NODE: + return SUITE_CHECKS[suite] + assert suite not in WORKSPACE_SUITES, ( + f"suite {suite} is about the IPC workspace, which a group spanning nodes does not have" + ) + return tuple(c for c in SUITE_CHECKS[suite] if c not in WORKSPACE_CHECKS) + + +def _engine_dist(mapping: Any) -> Any: + """The engine's TorchDist for `mapping`, with its cluster probe read from hostnames. + + TorchDist asks Ray for each rank's node and GPU; these ranks are not Ray + actors, so the same (node, device) list comes from a host-side gather + instead. Everything else is the engine's own code: the device mesh the + Mapping's process groups come from, the per-node local group, `init_pg` + (which the `_pg` op resolves ranks and the node topology through), and + `tp_allgather` (which the IPC workspace exchanges its handles through). + """ + from tensorrt_llm._torch.distributed.communicator import TorchDist + + class _TorchDistOnHostnames(TorchDist): + def _get_cluster_info(self): + if self.cluster_info is None: + info: List[Any] = [None] * dist.get_world_size() + dist.all_gather_object(info, (socket.gethostname(), torch.cuda.current_device())) + self.cluster_info = info + return self.cluster_info + + return _TorchDistOnHostnames(mapping) + + +def _run_one_rank(suite: str, group_size: Optional[int]) -> int: + """Body of one rank: run the suite's checks; a failing rank exits nonzero. + + The world is cut into contiguous groups of `group_size` ranks (all of it + when None), as `Mapping.tp_group` cuts a pipeline-parallel world into its + stages' TP groups; every group runs the suite at once on its own ranks. + The rank's device is its node-local rank (LOCAL_RANK under torchrun, the + global rank under this file's own launcher), so the same body runs under + both. Rendezvous is torch.distributed's env:// (RANK, WORLD_SIZE, + MASTER_ADDR, MASTER_PORT), with MPI disabled. + """ + global \ + HOST_GROUP, \ + RANK, \ + WORLD, \ + GROUP, \ + WORKSPACE, \ + ONE_NODE, \ + STRATEGIES, \ + CORE_FP32, \ + AUTO_FP32_BANDS, \ + LAYOUT, \ + allreduce, \ + validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + world_rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ.get("LOCAL_RANK", world_rank)) + torch.cuda.set_device(local_rank) + # The backend pair the engine's Ray workers initialize: the `_pg` op's node + # topology probe exchanges host values over the world group, which needs a + # CPU backend beside NCCL. device_id keeps NCCL communicator teardown off the + # implicit-device path. + dist.init_process_group( + backend="cuda:nccl,cpu:gloo", + init_method="env://", + device_id=torch.device("cuda", local_rank), + ) + + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.comm import ( # noqa: E501 + allreduce as entry, + ) + from tensorrt_llm._torch.distributed import communicator + from tensorrt_llm._torch.distributed.ops import get_allreduce_workspace + from tensorrt_llm.mapping import Mapping + + # This file runs as a spawned rank process, not under pytest, so the repo's + # `__extra_import_path__` mechanism is not in play and the shared helper is + # reached the plain way. The mutation is confined to these rank processes. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + allreduce, validating = entry.allreduce, _validating_cm + WORLD = group_size or world_size + assert WORLD >= 2 and world_size % WORLD == 0, ( + f"a collective needs at least 2 ranks, and the world ({world_size}) must split into " + f"groups of {WORLD}" + ) + stage = world_rank // WORLD + RANK = world_rank - stage * WORLD + GROUP = list(range(stage * WORLD, (stage + 1) * WORLD)) + + hosts: List[Any] = [None] * world_size + dist.all_gather_object(hosts, socket.gethostname()) + nodes = len(set(hosts)) + assert world_size % nodes == 0, f"uneven ranks per node: {hosts}" + ONE_NODE = len({hosts[r] for r in GROUP}) == 1 + layouts = { + (c.spec["world_size"], c.spec["group_size"], c.spec["nodes"]): c.spec["layout"] + for c in allreduce.CELLS + } + assert (world_size, WORLD, nodes) in layouts, ( + f"no cells at world size {world_size}, group size {WORLD}, {nodes} node(s)" + ) + LAYOUT = layouts[(world_size, WORLD, nodes)] + + # One gloo group per block for the host-side checks. Every rank must enter + # new_group for every block, in the same order, or it deadlocks. + for first in range(0, world_size, WORLD): + block = dist.new_group(list(range(first, first + WORLD)), backend="gloo") + if first == stage * WORLD: + HOST_GROUP = block + + # The Mapping the target builds for its group (a pipeline-parallel world + # passes its pp_size), and the engine's TorchDist over it. The IPC + # workspace reaches the engine's communicator object through + # `Distributed.get`, which would build a Ray-backed one; it gets this one. + mapping = Mapping( + world_size=world_size, + rank=world_rank, + gpus_per_node=world_size // nodes, + tp_size=WORLD, + pp_size=world_size // WORLD, + ) + engine_dist = _engine_dist(mapping) + + def _the_engine_dist(m: Any) -> Any: + assert m == mapping, f"the workspace asked for another Mapping: {m}" + return engine_dist + + communicator.Distributed.get = staticmethod(_the_engine_dist) + assert list(mapping.tp_group) == GROUP, (mapping.tp_group, GROUP) + # One binding instead of a per-call argument: every check's call site + # keeps the certified shape, and the wrapper hands the group's boxed + # ProcessGroup to the `_pg` symbol -- the routing the product's non-MPI + # mode does in _torch/distributed/ops.py. + allreduce.bind_const(process_group=mapping.tp_group_pg.boxed()) + + if ONE_NODE: + assert WORLD in AUTO_FP32_BANDS_BY_GROUP, ( + f"AUTO's bands at group size {WORLD} were never read" + ) + STRATEGIES, CORE_FP32, AUTO_FP32_BANDS = ( + dict(ALL_STRATEGIES), + AUTO, + AUTO_FP32_BANDS_BY_GROUP[WORLD], + ) + # Real state: the IPC workspace the target builds for its group. + WORKSPACE = get_allreduce_workspace(mapping) + else: + STRATEGIES = {k: v for k, v in ALL_STRATEGIES.items() if k not in WORKSPACE_STRATEGIES} + CORE_FP32, AUTO_FP32_BANDS, WORKSPACE = NCCL, NCCL_BANDS, None + checks = _suite_checks(suite) + # One eager call before any check (and so before any capture). + _call(torch.zeros(1, HIDDEN, device="cuda"), 0) + torch.cuda.synchronize() + major, minor = torch.cuda.get_device_capability() + who: List[Any] = [None] * world_size + dist.all_gather_object( + who, + f"world rank {world_rank} host {socket.gethostname()} device {local_rank} " + f"({torch.cuda.get_device_name()} sm_{major}{minor}) group {GROUP} position {RANK} " + f"entry={entry.__file__}", + ) + if world_rank == 0: + print( + f"[rank 0] suite {suite} layout {LAYOUT}: sm_{major}{minor} " + f"{torch.cuda.get_device_name()} world_size={world_size} group_size={WORLD} " + f"groups={[list(range(g, g + WORLD)) for g in range(0, world_size, WORLD)]} " + f"one_node={ONE_NODE} strategies={list(STRATEGIES.values())} " + f"core fp32 on {STRATEGIES[CORE_FP32]} entry={entry.__file__}", + flush=True, + ) + print("\n".join(f"[rank 0] {line}" for line in who), flush=True) + dropped = [c.__name__ for c in SUITE_CHECKS[suite] if c not in checks] + if dropped: + print( + f"[rank 0] suite {suite}: not run for a group spanning nodes " + f"(no IPC workspace): {dropped}", + flush=True, + ) + + for check in checks: + try: + check() + except BaseException: + import traceback + + print(f"[rank {world_rank}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # Exit hard rather than return: a rank that leaves a collective + # early wedges every sibling in it, and the launcher kills the rest + # the moment it sees this nonzero exit. + os._exit(13) + if world_rank == 0: + print(f"[rank 0] passed {check.__name__}", flush=True) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {world_rank}] suite {suite}: {len(checks)} checks passed", flush=True) + # Exit without tearing the process groups down: destroying a communicator + # that captured collectives into CUDA graphs can wedge NCCL's watchdog, and + # these are throwaway processes the launcher owns. + sys.stdout.flush() + os._exit(0) + + +def _group_size_arg() -> Optional[int]: + """`--group-size G` from the command line, or None (the whole world).""" + if _GROUP_FLAG not in sys.argv: + return None + return int(sys.argv[sys.argv.index(_GROUP_FLAG) + 1]) + + +def _run_suite(suite: str, world_size: int, group_size: Optional[int]) -> None: + """One suite: `world_size` fresh worker processes, under the suite's deadline. + + The ranks rendezvous through torch.distributed's env:// method, so the only + thing this needs is a free port. Each child gets its own session; on the + first nonzero exit or on the deadline, every session is killed, which is + what keeps one wedged collective from outliving the run. + """ + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + env = dict( + os.environ, + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + WORLD_SIZE=str(world_size), + TLLM_DISABLE_MPI="1", + ) + # Drop the launching step's PMIx/PMI/OMPI/Slurm namespace from the + # children: they would otherwise register as participants of that step, + # and the first hard exit (a failing rank's) trips its errhandler, which + # kills the whole step before anything is reported. + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + group = [] if group_size is None else [_GROUP_FLAG, str(group_size)] + command = [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG, suite, *group] + print(f"[launcher] {world_size} x {' '.join(command)} (rendezvous on :{port})", flush=True) + children = [ + subprocess.Popen( + command, env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), start_new_session=True + ) + for r in range(world_size) + ] + + def _kill_all() -> None: + for p in children: + try: + os.killpg(p.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for p in children: + p.wait() + + per_suite_s = SUITE_DEADLINE_S[suite] + deadline = time.monotonic() + per_suite_s + live = set(range(world_size)) + failure = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {per_suite_s}s (wedged)" + break + for r in sorted(live): + code = children[r].poll() + if code is None: + continue + live.discard(r) + if code != 0: + failure = f"rank {r} exited {code}" + break + time.sleep(0.5) + if failure is not None: + _kill_all() + raise AssertionError(f"the {world_size}-rank {suite} suite failed: {failure}") + + +def _spawn_ranks(group_size: Optional[int]) -> None: + """Every suite in turn, one worker process per claimed device; no external launcher.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, ( + "set CUDA_VISIBLE_DEVICES to the devices this run owns, " + "e.g. export CUDA_VISIBLE_DEVICES=0,1,2,3" + ) + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size >= 2, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); a collective test needs at least 2" + ) + for suite in SUITES: + _run_suite(suite, world_size, group_size) + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank(sys.argv[sys.argv.index(_WORKER_FLAG) + 1], _group_size_arg())) + _spawn_ranks(_group_size_arg()) + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/_pp_send_recv_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/_pp_send_recv_op_matrix.py new file mode 100644 index 000000000000..d6dc5e78e5a7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/_pp_send_recv_op_matrix.py @@ -0,0 +1,714 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the DeepSeek-V4.1 target's pp_send_tensors / pp_recv_tensors. + +Its own launcher. Run plainly, it re-executes itself under ``mpirun`` with one rank per +device of the first 4 of CUDA_VISIBLE_DEVICES, under a deadline: world 4, PP2 x TP2 on one +node. With ``--two-node`` it starts the same rank body at world 8 (PP2 x TP4) through the +two-node launcher named by MODELING_V2_TWO_NODE_MPIRUN: a command prefix that starts 8 ranks +with ranks 0-3 on one node and 4-7 on another, so stage 0 is one node, stage 1 the other, +and every hand-over crosses between them. Each form drives the entries' cells of its world +size (`pp_send_tensors.CELLS`, which `pp_recv_tensors.CELLS` must equal). + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _pp_send_recv_op_matrix.py + MODELING_V2_TWO_NODE_MPIRUN="" python _pp_send_recv_op_matrix.py --two-node + +Unlike the other comm matrices this keeps ``mpirun``: the hand-over it certifies runs on the +engine's MPI-path PP communicator (``PPCommNCCL``), whose eager send copies each tensor and +sends the copy on its own stream -- which is what lets the sender overwrite its tensors +right after the call (phases 1, 2, 4, 5) -- and the no-copy control uses that +communicator's raw send. With MPI disabled the engine builds ``PPCommTorch`` instead, which +sends the caller's tensors with no copy, so those checks would not be about it. + +Stage-0 rank r sends to rank r + TP (``Mapping.next_pp_rank``); a sender and its receiver +are a *pair*, and pairs run independently. Every step hands over what the DeepSeek-V4.1 +target hands over at its stage boundary, in its order: the hyper-connection stream +h bf16 [T, 4, 5120], the last FFN mixes yF fp32 [T, 24] and their RMS statistic rF fp32 [T]. + +Delivery is checked bitwise, per tensor, on the receiver, and the payload is built so that +a stale or misrouted delivery cannot pass: every step's words are a fixed random base +(the same on every rank) plus, as int32 words, a tag unique to (phase, step, sender), so the +bytes of any other step, phase or sender differ from the expected ones in every 32-bit +word; and every receive buffer is filled with 0xFF bytes first, so an unwritten word shows. +The op moves bytes and the dtypes only size the transfer, so tagged words that are not +finite bf16 / fp32 values are fine. The gate is the entries' own `compare` against their +`reference` (the sender's bytes). + +Phases, on every pair (1-3 are the cells, run inside `validating` so the guards are shown +to admit every certified input): +1. eager, every T in 1..16384, the attention-TP layout: all pairs hand over the same T in + a step, T in a shuffled order; +2. eager, every T in 1..16384, the attention-DP layout (``enable_attention_dp``: each rank + holds its own rows): the pairs hand over different T in the same step; + in 1 and 2 the sender overwrites and drops its tensors right after each send; +3. CUDA graphs at every row count the target captures: T = G and 6 G (DSpark verify) over + the 46 graph batch sizes, 82 values in 1..1536, each after an eager hand-over at that T; + one graph per T on both stages, each replayed 3 times on new tags (ascending, descending, + and a per-pair shuffle with the sender overwriting its static tensors after each replay); +4. eager steps and graph replays interleaved at random T with the sender never waiting in + between (the engine's mix of eager context steps and captured decode steps); +5. a late receiver: its stream is held ~20 ms before the receive while the sender + overwrites, drops and reallocates its tensors (eager), or overwrites its static tensors + (graph); +6. controls: a send that does not copy (the communicator's raw send, held back, the tensors + overwritten meanwhile) must be seen as wrong; so must a missed receive, the previous + step's bytes and the other sender's bytes, by the entries' gates too; the guards + (`is_valid`, under `validating`) must refuse a strided view (what + the raw op does with one, eagerly and under capture, is recorded); +7. timing of one hand-over at T = 1, 6, 1536, 8192, 16384. +""" + +import os +import random +import shlex +import shutil +import signal +import socket +import statistics +import subprocess +import sys +import tempfile +import time +from pathlib import Path + +import torch + +assert torch.cuda.is_available(), "pp_send_tensors requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +_TWO_NODE_FLAG = "--two-node" +TWO_NODE_ENV = "MODELING_V2_TWO_NODE_MPIRUN" +_PID_DIR_ENV = "PP_SEND_RECV_MATRIX_PID_DIR" +DEADLINE_S = 1200 +LOCAL_WORLD = 4 +HIDDEN, HC, MIX = 5120, 4, 24 +NAMES = ("h", "yF", "rF") +DTYPES = (torch.bfloat16, torch.float32, torch.float32) +# max_num_tokens up to 16384 per rank. +N_MAX = 16384 +# The 46 CUDA-graph batch sizes: every G in 1..32 plus the padded default list up to 256. +GRAPH_BATCHES = list(range(1, 33)) + list(range(40, 129, 8)) + [192, 256] +# A captured step hands over T = G rows (DSpark off) or 6 G (DSpark on: 1 + 5 drafts). +GRAPH_T = sorted(set(GRAPH_BATCHES) | {6 * g for g in GRAPH_BATCHES}) +LATE_T = (1, 6, 1536, 8192, 16384) +TIMING_T = (1, 6, 1536, 8192, 16384) +INTERLEAVED_STEPS = 400 +SENTINEL = -1 # 0xFFFFFFFF: what a receive buffer holds before the receive +JUNK = 0x5A5A5A5A # what a sender writes over its tensors after handing them over +LATE_CYCLES = 40_000_000 # about 20 ms of GPU clock +# Tag phases. +PH_SWEEP_TP, PH_SWEEP_DP, PH_CAPTURE, PH_REPLAY, PH_MIX, PH_LATE, PH_CONTROL, PH_TIMING = ( + 1, + 2, + 3, + 4, + 8, + 9, + 10, + 11, +) + +assert len(GRAPH_BATCHES) == 46 and len(GRAPH_T) == 82 and GRAPH_T[-1] == 1536 + + +def _shapes(T): + return ((T, HC, HIDDEN), (T, MIX), (T,)) + + +def _tag(phase: int, step: int, sender: int) -> int: + assert 0 < phase < 256 and 0 <= step < (1 << 18) and 0 <= sender < 15 + return (phase << 22) | (step << 4) | (sender + 1) + + +def _untag(value: int) -> str: + value &= 0xFFFFFFFF + return f"phase {value >> 22} step {(value >> 4) & 0x3FFFF} sender {(value & 15) - 1}" + + +def _diagnose(got, base, tag: int) -> str: + """Say what a wrong delivery holds: unwritten words, or another step's / sender's words.""" + delta = got.reshape(-1) - base.reshape(-1) + wrong = delta != tag + n = int(wrong.sum()) + first = int(wrong.nonzero()[0, 0]) + unfilled = int((got.reshape(-1)[wrong] == SENTINEL).sum()) + values, counts = torch.unique(delta[wrong], return_counts=True) + top = sorted(zip(counts.tolist(), values.tolist()), reverse=True)[:3] + carried = "; ".join(f"{c} words carry +{v & 0xFFFFFFFF:#x} ({_untag(v)})" for c, v in top) + return ( + f"{n}/{delta.numel()} words wrong, first at word {first}, {unfilled} still 0xFF; {carried}" + ) + + +def _rank_body() -> int: + from mpi4py import MPI + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + # This node's ranks do not live in mpirun's process group: leave the pid where the + # launcher's deadline can find it. + pid_dir = os.environ.get(_PID_DIR_ENV) + if pid_dir: + Path(pid_dir, f"rank{rank}.pid").write_text(str(os.getpid())) + try: + return _run(comm) + except BaseException: + import traceback + + print(f"[rank {rank}] FAILED with an exception", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # Abort rather than return: a rank that leaves the hand-over wedges its peer. + comm.Abort(1) + return 1 + + +def _run(comm) -> int: + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.comm import ( # noqa: E501 + pp_recv_tensors as recv_entry, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.comm import ( # noqa: E501 + pp_send_tensors as send_entry, + ) + from tensorrt_llm._torch.distributed import communicator + from tensorrt_llm.mapping import Mapping + + # This file runs as an mpirun rank, not under pytest, so the repo's + # `__extra_import_path__` mechanism is not in play and the shared helper is + # reached the plain way. The mutation is confined to these rank processes. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating + + # The entries' instances: every call below goes through them. + pp_send_tensors, pp_recv_tensors = send_entry.pp_send_tensors, recv_entry.pp_recv_tensors + rank, world = comm.Get_rank(), comm.Get_size() + assert world in (4, 8), f"PP2 x TP2 (4 ranks) or PP2 x TP4 (8 ranks), got {world}" + tp = world // 2 + + # The cells of this world size, which phases 1-3 drive: one matrix certifies both + # halves, so the two entries must name the same cells. + assert [c.spec for c in pp_recv_tensors.CELLS] == [c.spec for c in pp_send_tensors.CELLS], ( + "pp_send_tensors and pp_recv_tensors name different cells" + ) + cells = { + (c.spec["attention_dp"], c.spec["mode"]): c + for c in pp_send_tensors.CELLS + if c.spec["world_size"] == world + } + assert set(cells) == {(False, "eager"), (True, "eager"), (False, "graph")}, list(cells) + handover = [(n, s[1:], d) for n, s, d in zip(NAMES, _shapes(1), DTYPES)] + for c in cells.values(): + assert c.spec["tp_size"] == tp and c.spec["pp_size"] == 2, c.why + assert [(n, tuple(t), d) for n, t, d in c.spec["tensors"]] == handover, c.why + for key in ((False, "eager"), (True, "eager")): + assert list(cells[key].spec["rows"]) == list(range(1, N_MAX + 1)), cells[key].why + assert list(cells[(False, "graph")].spec["rows"]) == GRAPH_T, cells[(False, "graph")].why + local = int(os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK", rank)) + per_node = int(os.environ.get("OMPI_COMM_WORLD_LOCAL_SIZE", world)) + torch.cuda.set_device(local) + dev = torch.device("cuda", local) + kw = dict(world_size=world, rank=rank, tp_size=tp, pp_size=2, gpus_per_node=per_node) + layouts = { + "attention-TP": Mapping(**kw), + "attention-DP": Mapping(**kw, moe_ep_size=tp, enable_attention_dp=True), + } + m = layouts["attention-TP"] + first = m.is_first_pp_rank() + pair = m.tp_rank + sender = rank if first else m.prev_pp_rank() + receiver = m.next_pp_rank() if first else rank + assert receiver == sender + tp and pair == sender, (rank, sender, receiver, pair) + for mm in layouts.values(): + assert (mm.next_pp_rank() if first else mm.prev_pp_rank()) == ( + receiver if first else sender + ) + communicator.init_pp_comm(layouts["attention-TP"]) + + major, minor = torch.cuda.get_device_capability() + who = comm.allgather( + ( + socket.gethostname(), + local, + torch.cuda.get_device_name(), + f"sm_{major}{minor}", + send_entry.__file__, + recv_entry.__file__, + ) + ) + hosts = [w[0] for w in who] + crosses = [hosts[s] != hosts[s + tp] for s in range(tp)] + # World 8 certifies the two-tray form: every pair must cross between the nodes. + assert all(crosses) if world == 8 else not any(crosses), (world, hosts) + if rank == 0: + print(f"[rank 0] world_size={world} PP2 x TP{tp}; pairs cross nodes: {crosses}", flush=True) + for r, (host, dvc, name, sm, fs, fr) in enumerate(who): + role = f"stage 0, sends to {r + tp}" if r < tp else f"stage 1, receives from {r - tp}" + print( + f"[rank 0] rank {r} host {host} device {dvc} {name} {sm} ({role}) entry={fs} {fr}", + flush=True, + ) + + g = torch.Generator(device=dev).manual_seed(0x5EED) + base = [ + torch.randn((N_MAX, HC, HIDDEN), generator=g, device=dev) + .to(torch.bfloat16) + .view(torch.int32), + torch.randn((N_MAX, MIX), generator=g, device=dev).view(torch.int32), + (torch.rand((N_MAX,), generator=g, device=dev) + 0.5).view(torch.int32), + ] + digest = [int(b.to(torch.int64).sum()) for b in base] + assert all(d == digest for d in comm.allgather(digest)), ( + "the base payload differs between ranks" + ) + + failures = [] + checks = 0 + + def payload(T, tag): + return [torch.add(b[:T], tag).view(d) for b, d in zip(base, DTYPES)] + + def fresh(T): + bufs = [torch.empty(s, dtype=d, device=dev) for s, d in zip(_shapes(T), DTYPES)] + for b in bufs: + b.view(torch.int32).fill_(SENTINEL) + return bufs + + def matches(bufs, T, tag): + return [torch.equal(x.view(torch.int32), torch.add(b[:T], tag)) for x, b in zip(bufs, base)] + + def verify(bufs, T, tag, what): + """The entries' own gate against their own reference (the sender's bytes, which + every rank rebuilds from the shared base and the tag); on a failure, say what the + wrong delivery holds instead.""" + nonlocal checks + checks += 1 + sent = payload(T, tag) + try: + pp_recv_tensors.compare(bufs, pp_recv_tensors.reference(sent)) + pp_send_tensors.compare(bufs, pp_send_tensors.reference(sent)) + except AssertionError as exc: + wrong = [k for k, ok in enumerate(matches(bufs, T, tag)) if not ok] + failures.extend( + f"[rank {rank}] {what} T={T} {NAMES[k]}: " + f"{_diagnose(bufs[k].view(torch.int32), base[k][:T], tag)}" + for k in wrong + ) + if not wrong: + failures.append(f"[rank {rank}] {what} T={T}: {exc}") + + def junk(bufs): + for b in bufs: + b.view(torch.int32).fill_(JUNK) + + def gate_rejects_each(bufs, T, tag): + """Whether both entries' gates fail each tensor of `bufs` on its own.""" + rejected = [] + for x, want in zip(bufs, payload(T, tag)): + refused = 0 + for op in (pp_recv_tensors, pp_send_tensors): + try: + op.compare([x], op.reference([want])) + except AssertionError: + refused += 1 + rejected.append(refused == 2) + return rejected + + def eager(T, phase, step, *, late=False, sync=True, what="eager"): + tag = _tag(phase, step, sender) + if first: + bufs = payload(T, tag) + pp_send_tensors(bufs) + # The eager send works on a copy: overwrite and drop the tensors at once. + junk(bufs) + del bufs + if late: + # Reallocate the freed sizes and fill them with junk: a copy that the + # caching allocator handed out again before the send read it would show. + again = [fresh(T) for _ in range(2)] + for bufs in again: + junk(bufs) + del again + if sync: + torch.cuda.synchronize() + else: + bufs = fresh(T) + if late: + torch.cuda._sleep(LATE_CYCLES) + got = pp_recv_tensors(bufs) + assert got is bufs, "pp_recv_tensors must return the list it filled" + verify(bufs, T, tag, f"{what} step {step}") + + def sweep(layout, phase, order): + assert sorted(order) == list(range(1, N_MAX + 1)) + communicator.init_pp_comm(layouts[layout]) + comm.Barrier() + t0 = time.perf_counter() + for step, T in enumerate(order): + eager(T, phase, step, what=f"{layout} eager sweep") + torch.cuda.synchronize() + comm.Barrier() + return time.perf_counter() - t0 + + # Phases 1-3 are the cells: the guards run on every call, which shows they admit + # every certified input. + with validating(pp_send_tensors, pp_recv_tensors): + # 1. eager, every T, the attention-TP layout: the same shuffled T on every pair. + shared = list(cells[(False, "eager")].spec["rows"]) + random.Random(1).shuffle(shared) + secs_tp = sweep("attention-TP", PH_SWEEP_TP, shared) + # 2. eager, every T, the attention-DP layout: a different T on every pair in a step. + per_pair = [ + list(range(1, N_MAX + 1)), + list(range(N_MAX, 0, -1)), + [(s * 7919) % N_MAX + 1 for s in range(N_MAX)], + [(s * 4099 + 5000) % N_MAX + 1 for s in range(N_MAX)], + ] + secs_dp = sweep("attention-DP", PH_SWEEP_DP, per_pair[pair]) + communicator.init_pp_comm(layouts["attention-TP"]) + checks_eager = checks + + # 3. CUDA graphs: an eager hand-over at T, then one graph per T on both stages. + graphs = {} + capture = torch.cuda.Stream() + for i, T in enumerate(GRAPH_T): + eager(T, PH_CAPTURE, i, what="eager before capture") + static = [torch.zeros(s, dtype=d, device=dev) for s, d in zip(_shapes(T), DTYPES)] + comm.Barrier() + capture.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture): + if first: + pp_send_tensors(static) + else: + pp_recv_tensors(static) + torch.cuda.current_stream().wait_stream(capture) + torch.cuda.synchronize() + graphs[T] = (graph, static) + comm.Barrier() + + def replay(T, phase, step, *, late=False, clobber=False, what="replay"): + graph, static = graphs[T] + tag = _tag(phase, step, sender) + if first: + for s, b in zip(static, base): + torch.add(b[:T], tag, out=s.view(torch.int32)) + graph.replay() + if clobber or late: + # Stream-ordered after the captured send: the replay must have sent the + # bytes the static tensors held when it ran. + junk(static) + else: + for s in static: + s.view(torch.int32).fill_(SENTINEL) + if late: + torch.cuda._sleep(LATE_CYCLES) + graph.replay() + verify(static, T, tag, f"{what} step {step}") + + shuffled = list(GRAPH_T) + random.Random(100 + pair).shuffle(shuffled) + rounds = ( + ("ascending", GRAPH_T, False), + ("descending", GRAPH_T[::-1], False), + ("per-pair shuffle, sender overwrites", shuffled, True), + ) + for r, (name, order, clobber) in enumerate(rounds): + for step, T in enumerate(order): + replay(T, PH_REPLAY + r, step, clobber=clobber, what=f"graph {name}") + torch.cuda.synchronize() + comm.Barrier() + checks_graph = checks - checks_eager + + # 4. eager steps and replays interleaved, the sender not waiting in between. + rng = random.Random(900 + pair) + mix = {"eager": 0, "replay": 0} + for step in range(INTERLEAVED_STEPS): + if rng.random() < 0.5: + replay(rng.choice(GRAPH_T), PH_MIX, step, what="interleaved replay") + mix["replay"] += 1 + else: + T = rng.choice(GRAPH_T) if rng.random() < 0.5 else int(round(2 ** rng.uniform(0, 14))) + eager(T, PH_MIX, step, sync=False, what="interleaved eager") + mix["eager"] += 1 + if first and step % 8 == 7: + torch.cuda.synchronize() # bound the copies in flight + torch.cuda.synchronize() + comm.Barrier() + + # 5. a late receiver. + for i, T in enumerate(LATE_T): + eager(T, PH_LATE, i, late=True, what="late receiver, eager") + for i, T in enumerate(t for t in LATE_T if t in graphs): + replay(T, PH_LATE, 100 + i, late=True, what="late receiver, replay") + torch.cuda.synchronize() + comm.Barrier() + + # 6. controls: the checks see what they guard against. + control = [] + side = torch.cuda.Stream() + for i, T in enumerate((1536, 16384)): + tag = _tag(PH_CONTROL, i, sender) + if first: + # The communicator's own send without the eager path's copy, held back ~20 ms + # on a side stream while the current stream overwrites the tensors. + bufs = payload(T, tag) + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + torch.cuda._sleep(LATE_CYCLES) + for t in bufs: + communicator._pp_comm.nccl_comm.send(t, receiver) + junk(bufs) + torch.cuda.synchronize() + del bufs + else: + bufs = fresh(T) + pp_recv_tensors(bufs) + ok = matches(bufs, T, tag) + overwritten = [bool((x.view(torch.int32) == JUNK).all()) for x in bufs] + rejected = gate_rejects_each(bufs, T, tag) + control.append( + f"send without a copy, T={T}: delivered == payload {ok}, == overwrite {overwritten}, " + f"gate rejects {rejected}" + ) + if any(ok) or not all(rejected): + failures.append( + f"[rank {rank}] control T={T}: an overwritten send passed the check" + ) + if not first: + T = 257 + now = _tag(PH_CONTROL, 10, sender) + wrong = { + "missed receive": fresh(T), + "previous step's bytes": payload(T, _tag(PH_CONTROL, 9, sender)), + "the other sender's bytes": payload(T, _tag(PH_CONTROL, 10, (sender + 1) % tp)), + } + for name, bufs in wrong.items(): + ok = matches(bufs, T, now) + differ = [ + bool((x.view(torch.int32) != torch.add(b[:T], now)).all()) + for x, b in zip(bufs, base) + ] + rejected = gate_rejects_each(bufs, T, now) + control.append( + f"{name}: passes {ok}, differs in every word {differ}, gate rejects {rejected}" + ) + if any(ok) or not all(differ) or not all(rejected): + failures.append(f"[rank {rank}] control {name}: not seen in every word") + # A strided view (every other hc slot, [T, 2, 5120] of [T, 4, 5120]): the guards must + # refuse it. What the unguarded op (`raw_call`, outside `validating`) does with one is + # recorded, eagerly and under capture. + T = 64 + for j, captured in enumerate((False, True)): + tag = _tag(PH_CONTROL, 20 + j, sender) + how = "captured" if captured else "eager" + if first: + view = payload(T, tag)[0][:, ::2, :] + if not captured: + try: + with validating(pp_send_tensors): + pp_send_tensors([view]) + failures.append(f"[rank {rank}] the send guard accepted a strided view") + except AssertionError as e: + control.append(f"send guard, strided view: refused, AssertionError({e})") + pp_send_tensors.raw_call([view]) + else: + graph = torch.cuda.CUDAGraph() + capture.wait_stream(torch.cuda.current_stream()) + with torch.cuda.graph(graph, stream=capture): + pp_send_tensors.raw_call([view]) + torch.cuda.current_stream().wait_stream(capture) + comm.Barrier() + graph.replay() + torch.cuda.synchronize() + else: + got = torch.empty((T, HC // 2, HIDDEN), dtype=torch.bfloat16, device=dev) + got.view(torch.int32).fill_(SENTINEL) + if not captured: + pp_recv_tensors([got]) + else: + graph = torch.cuda.CUDAGraph() + capture.wait_stream(torch.cuda.current_stream()) + with torch.cuda.graph(graph, stream=capture): + pp_recv_tensors([got]) + torch.cuda.current_stream().wait_stream(capture) + comm.Barrier() + graph.replay() + torch.cuda.synchronize() + whole = torch.add(base[0][:T], tag).view(torch.bfloat16) + as_view = torch.equal(got, whole[:, ::2, :].contiguous()) + as_storage = torch.equal(got.reshape(-1), whole.reshape(-1)[: got.numel()]) + control.append( + f"raw op send of a strided view, {how}: delivered the view's elements {as_view}, " + f"the first {got.numel() * 2} bytes of its storage {as_storage}" + ) + tag = _tag(PH_CONTROL, 22, sender) + if first: + pp_send_tensors([payload(T, tag)[0][:, : HC // 2, :].contiguous()]) + torch.cuda.synchronize() + else: + big = torch.empty((T, HC, HIDDEN), dtype=torch.bfloat16, device=dev) + big.view(torch.int32).fill_(SENTINEL) + try: + with validating(pp_recv_tensors): + pp_recv_tensors([big[:, ::2, :]]) + failures.append(f"[rank {rank}] the recv guard accepted a strided view") + except AssertionError as e: + control.append(f"recv guard, strided view: refused, AssertionError({e})") + pp_recv_tensors.raw_call([big[:, ::2, :]]) + torch.cuda.synchronize() + want = torch.add(base[0][:T], tag).view(torch.bfloat16)[:, : HC // 2, :].contiguous() + in_view = torch.equal(big[:, ::2, :], want) + in_storage = torch.equal(big.reshape(-1)[: want.numel()], want.reshape(-1)) + control.append( + f"raw op receive into a strided view, eager: the view holds the payload {in_view}, " + f"the first {want.numel() * 2} bytes of its storage hold it {in_storage}" + ) + comm.Barrier() + + # 7. timing of one hand-over (host wall time from a barrier to the data in place). + timing = {} + for i, T in enumerate(TIMING_T): + tag = _tag(PH_TIMING, i, sender) + bufs = payload(T, tag) if first else fresh(T) + samples = [] + for _ in range(10): + torch.cuda.synchronize() + comm.Barrier() + t0 = time.perf_counter() + if first: + pp_send_tensors(bufs) + else: + pp_recv_tensors(bufs) + torch.cuda.synchronize() + samples.append(time.perf_counter() - t0) + if not first: + verify(bufs, T, tag, "timing") + timing[T] = statistics.median(samples) * 1e6 + comm.Barrier() + + total_failures = comm.allreduce(len(failures)) + total_checks = comm.allreduce(checks) + report = comm.gather( + ( + rank, + first, + failures[:20], + control, + timing, + checks, + checks_eager, + checks_graph, + mix, + secs_tp, + secs_dp, + ), + root=0, + ) + if rank == 0: + for r, is_first, fails, ctrl, tim, n, ne, ng, mx, st, sd in report: + role = "sender" if is_first else "receiver" + print( + f"[rank {r}] {role}: timing us {', '.join(f'T={t}: {v:.0f}' for t, v in tim.items())}; " + f"sweeps {st:.1f} s / {sd:.1f} s", + flush=True, + ) + if not is_first: + print( + f"[rank {r}] checks {n} (eager sweeps {ne}; capture phase {ng}: 82 eager hand-overs + 246 " + f"replays); interleaved {mx}", + flush=True, + ) + for line in ctrl: + print(f"[rank {r}] control: {line}", flush=True) + for line in fails: + print(f"FAILURE {line}", flush=True) + status = "PP_SEND_RECV_OK" if total_failures == 0 else "PP_SEND_RECV_FAILED" + print( + f"{status} world_size={world} checks={total_checks} failures={total_failures} " + f"eager_T=1..{N_MAX} x2 layouts graph_T={len(GRAPH_T)} values 1..{GRAPH_T[-1]} " + f"x3 replays interleaved={INTERLEAVED_STEPS} late={list(LATE_T)}", + flush=True, + ) + return 0 if total_failures == 0 else 1 + + +def _scrubbed_env(**extra): + """The environment without the allocation's Slurm/PMI context (a rank attached to it + that dies abnormally takes the allocation down).""" + env = dict(os.environ, **extra) + for v in list(env): + if v.startswith( + ("PMI", "PMIX", "OMPI", "SLURM", "PRTE", "PALS", "MPI_LOCALRANKID", "MPI_LOCALNRANKS") + ): + env.pop(v) + return env + + +def _wait(cmd, env, pid_dir) -> int: + print(f"[launcher] {' '.join(cmd)}", flush=True) + # Own process group, so the deadline reaches every process it started. + proc = subprocess.Popen(cmd, env=env, start_new_session=True) + try: + return proc.wait(timeout=DEADLINE_S) + except subprocess.TimeoutExpired: + # TERM first: a launcher script gets to run its cleanup. + os.killpg(proc.pid, signal.SIGTERM) + try: + proc.wait(timeout=90) + except subprocess.TimeoutExpired: + os.killpg(proc.pid, signal.SIGKILL) + proc.wait() + if pid_dir: + for pid_file in Path(pid_dir).glob("rank*.pid"): + try: + os.kill(int(pid_file.read_text()), signal.SIGKILL) + except (ProcessLookupError, ValueError): + pass + print(f"pp_send_recv matrix wedged: stopped after {DEADLINE_S} s", flush=True) + return 124 + finally: + if pid_dir: + shutil.rmtree(pid_dir, ignore_errors=True) + + +def _launch_one_node() -> int: + visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") + devices = [d for d in visible.split(",") if d.strip()] or [ + str(i) for i in range(torch.cuda.device_count()) + ] + assert len(devices) >= LOCAL_WORLD, f"needs {LOCAL_WORLD} devices, sees {devices}" + pid_dir = tempfile.mkdtemp(prefix="pp-send-recv-matrix-") + env = _scrubbed_env( + CUDA_VISIBLE_DEVICES=",".join(devices[:LOCAL_WORLD]), **{_PID_DIR_ENV: pid_dir} + ) + cmd = [ + "mpirun", + "-np", + str(LOCAL_WORLD), + "--oversubscribe", + "--allow-run-as-root", + "-x", + _PID_DIR_ENV, + sys.executable, + str(Path(__file__).resolve()), + _WORKER_FLAG, + ] + return _wait(cmd, env, pid_dir) + + +def _launch_two_node() -> int: + prefix = os.environ.get(TWO_NODE_ENV, "").strip() + assert prefix, ( + f"{_TWO_NODE_FLAG} needs {TWO_NODE_ENV}: a command prefix that starts 8 ranks, " + f"ranks 0-3 on one node and 4-7 on another" + ) + cmd = shlex.split(prefix) + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG] + # The other node's ranks cannot be signalled from here; the prefix owns them. + return _wait(cmd, _scrubbed_env(), None) + + +if __name__ == "__main__": + if _WORKER_FLAG in sys.argv: + sys.exit(_rank_body()) + sys.exit(_launch_two_node() if _TWO_NODE_FLAG in sys.argv else _launch_one_node()) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/test_modeling_v2_allreduce.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/test_modeling_v2_allreduce.py new file mode 100644 index 000000000000..51bc4042e2d5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/test_modeling_v2_allreduce.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the allreduce op's certification matrix. + +Distinct from ``tests/unittest/_torch/multi_gpu/test_allreduce.py`` and not a +duplicate of it: that file covers the fusion patterns through the +``AllReduce`` module, while this covers the op itself -- the plain summing +all-reduce, strategy by strategy, over the payloads a DeepSeek-V4.1 +tensor-parallel forward hands it. + +The matrix is ``_allreduce_op_matrix.py`` beside this file, its own launcher: +it spawns one worker process per device for each suite, rendezvoused through +torch.distributed's env://, under a deadline it enforces by killing every +worker. One case per single-node layout of the entry's cells (world 4; world +2; world 4 cut into the pp2 stages' groups of 2), each over the first visible +devices. The two-node layouts are run by hand (see the entry's ``note``). + +Started by file path in a fresh interpreter: the launcher must carry no CUDA +or torch.distributed state of its own, and the ranks reach the entry by +absolute import, so this tree needs no package context. +""" + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.comm.allreduce import ( # noqa: E501 + allreduce as op, +) + +assert torch.cuda.is_available(), "allreduce requires CUDA devices" + +_MATRIX = Path(__file__).resolve().with_name("_allreduce_op_matrix.py") +_LAUNCHER_GRACE_S = 300 +"""Headroom over the matrix's own deadline, so its message wins the race.""" + +# (layout, world size, group size) of every cell that fits on one node. +LAYOUTS = sorted( + { + (c.spec["layout"], c.spec["world_size"], c.spec["group_size"]) + for c in op.CELLS + if c.spec["nodes"] == 1 + } +) + + +def _devices(world_size: int) -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + devices = [d for d in visible.split(",") if d.strip()] + else: + devices = [str(i) for i in range(torch.cuda.device_count())] + assert len(devices) >= world_size, ( + f"this case runs at world size {world_size}; only {len(devices)} device(s) are visible" + ) + return ",".join(devices[:world_size]) + + +def _matrix_constants(): + """The matrix's module-level budget, read without importing it by name + (tests/ is not a package); its tensorrt_llm imports live in the rank body.""" + spec = importlib.util.spec_from_file_location(_MATRIX.stem, _MATRIX) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +# Each case spawns its own ranks over the first 2 or 4 visible devices. Under +# xdist several workers would fight for the same GPUs, so this must run alone. +@pytest.mark.no_xdist +@pytest.mark.parametrize("layout", LAYOUTS, ids=[name for name, _, _ in LAYOUTS]) +def test_allreduce_op_matrix(layout) -> None: + name, world_size, group_size = layout + env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices(world_size)) + # The ranks import tensorrt_llm absolutely, and a source checkout is not + # necessarily installed. + # tests/unittest/_torch/modeling_v2/models/deepseek_v41//comm -> repo root. + repo_root = Path(__file__).resolve().parents[8] + assert (repo_root / "tensorrt_llm").is_dir(), ( + f"expected the repo root at {repo_root}, found no tensorrt_llm/ there; " + f"this file moved without its parents[] index following" + ) + env["PYTHONPATH"] = os.pathsep.join(p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p) + + completed = subprocess.run( + [sys.executable, str(_MATRIX), "--group-size", str(group_size)], + env=env, + capture_output=True, + text=True, + timeout=_matrix_constants().DEADLINE_S + _LAUNCHER_GRACE_S, + ) + if completed.returncode != 0: + raise AssertionError( + f"{_MATRIX.name} ({name}) exited {completed.returncode}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) + # What the run did, for `pytest -rA` to show beside the pass. + print(f"[{name}] world_size={world_size} group_size={group_size}") + print( + "\n".join( + line + for line in completed.stdout.splitlines() + if line.startswith(("[rank", "[launcher]", "OK")) + ) + ) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/test_modeling_v2_pp_recv_tensors.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/test_modeling_v2_pp_recv_tensors.py new file mode 100644 index 000000000000..546592ea2254 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/test_modeling_v2_pp_recv_tensors.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the pp_recv_tensors certification: the receiving half of +the hand-over matrix that ``test_modeling_v2_pp_send_tensors.py`` starts, in the same +one-node form (PP2 x TP2). The two-tray cells are run by hand (see the entry's +``note``).""" + +import pytest +from test_modeling_v2_pp_send_tensors import FORMS, run_pp_send_recv_matrix + + +@pytest.mark.no_xdist +@pytest.mark.parametrize("form", list(FORMS)) +def test_pp_recv_op_matrix(form: str) -> None: + run_pp_send_recv_matrix(form) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/test_modeling_v2_pp_send_tensors.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/test_modeling_v2_pp_send_tensors.py new file mode 100644 index 000000000000..89df53be58a4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/comm/test_modeling_v2_pp_send_tensors.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the pp_send_tensors certification matrix. + +The matrix is ``_pp_send_recv_op_matrix.py`` beside this file, its own launcher: +one 4-rank ``mpirun`` over the first 4 visible devices of this node (PP2 x TP2, +and the same world as PP2 x DEP2), under a deadline. It drives the one-node cells +of both ``pp_send_tensors`` and ``pp_recv_tensors``. The two-tray cells (8 ranks, +PP2 x TP4, one stage per node) cannot be started from one node: their hand-run +command is in the entry's ``note``. + +The same job certifies ``pp_recv_tensors`` (``test_modeling_v2_pp_recv_tensors.py``): +neither half of a hand-over runs without the other. + +The matrix keeps ``mpirun`` rather than a torch.distributed rendezvous: what it +certifies is the engine's MPI-path PP communicator (``PPCommNCCL``), whose eager +send copies the tensors onto its own stream; the ProcessGroup communicator the +engine uses with MPI disabled (``PPCommTorch``) sends the caller's tensors with no copy, +so the matrix's overwrite-after-send checks are not about it. +""" + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +assert torch.cuda.is_available(), "pp_send_tensors requires CUDA devices" + +_MATRIX = Path(__file__).resolve().with_name("_pp_send_recv_op_matrix.py") +FORMS = {"one_node_pp2xtp2": []} + + +def _matrix_constants(): + """The matrix's module-level constants, read without importing it by name (tests/ is + not a package) and without importing tensorrt_llm (the launcher must not).""" + spec = importlib.util.spec_from_file_location(_MATRIX.stem, _MATRIX) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def run_pp_send_recv_matrix(form: str) -> None: + """Start the matrix in `form` and assert it passed.""" + constants = _matrix_constants() + # tests/unittest/_torch/modeling_v2/models/deepseek_v41//comm -> repo root. + repo_root = Path(__file__).resolve().parents[8] + assert (repo_root / "tensorrt_llm").is_dir(), ( + f"expected the repo root at {repo_root}, found no tensorrt_llm/ there; " + f"this file moved without its parents[] index following" + ) + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join(p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p) + done = subprocess.run( + [sys.executable, str(_MATRIX), *FORMS[form]], + env=env, + capture_output=True, + text=True, + timeout=constants.DEADLINE_S + 300, + ) + # The matrix's own summary (placement, checks, controls, timing) for the run log. + print( + "\n".join( + line + for line in done.stdout.splitlines() + if "rank 0" in line + or "PP_SEND_RECV" in line + or line.startswith(("FAILURE", "[launcher]", "[rank")) + ) + ) + assert done.returncode == 0 and "PP_SEND_RECV_OK" in done.stdout, ( + f"{_MATRIX.name} ({form}) exited {done.returncode}\n--- stdout ---\n{done.stdout}\n" + f"--- stderr ---\n{done.stderr}" + ) + + +# The job owns four devices; under xdist workers would fight over them. +@pytest.mark.no_xdist +@pytest.mark.parametrize("form", list(FORMS)) +def test_pp_send_op_matrix(form: str) -> None: + run_pp_send_recv_matrix(form) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/embedding/test_modeling_v2_engram_gather.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/embedding/test_modeling_v2_engram_gather.py new file mode 100644 index 000000000000..f682739e11a6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/embedding/test_modeling_v2_engram_gather.py @@ -0,0 +1,2385 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the engram_gather entry of the DeepSeek-V4.1 target. + +What is certified: one Triton launch gathers rows of an FP8 E4M3 ``[rows, 256]`` +table with one E8M0 scale per 32 channels, dequantizes them in fp32 and rounds +once to bf16, writing zeros for ids outside the caller's row range -- the +DeepSeek-V4.1 Engram lookup (``ParallelEngramEmbedding.forward`` of the +checkpoint's ``inference/model.py``). Two residencies go through the same +kernel: a table in device memory and a table in pinned host memory read +through the address ``cudaHostGetDevicePointer`` returns (zero-copy). + +The entry owns what is certified: ``CELLS`` is the list, ``reference`` and +``compare`` are the gate. This file turns a cell's spec into tables and ids and +keeps what a cell cannot carry: the prefix sweeps over every row count, the +largest-shard and far-end runs, CUDA-graph capture, the layout guards, the +wrong-variant controls, the checkpoint's real rows and the real-size +node-shared host table (one process, four processes, one-GPU processes). + +Four sources of truth: + +* synthetic tables, against the entry's ``reference`` (torch ops in the + reference's own order: embedding lookup, fp32 multiply, one bf16 cast, + masked fill); +* REAL rows of the checkpoint's layer-1 table, read row by row with + ``safetensors`` ``get_slice`` at the ids the golden reference hashed for its + prompts (``golden//engram_hash.pt``) -- never the whole table; +* the golden layer dumps: the gathered rows pushed through the reference's own + next step (FP8 1x32 activation quantize, fp32 GEMM with ``wkv``) must + reproduce the golden ``engram_value``; +* the checkpoint's OWN two tables (202.76 GB) in the node's one shared host copy + of the target's ``engram_host_table.py`` (host_offload), through both of + its backends: ``vmm`` (the default: CUDA VMM host-NUMA allocations, CUDA views + handed to the gather) and ``shm`` (a ``/dev/shm`` file registered with + ``cudaHostRegister``, pinned CPU views), each checked against the checkpoint's + bytes read through an independent ``mmap`` of its safetensors shards, and each + gather held to a reference built from the table's bytes as the CPU reads them. +The last three need external directories and skip, naming them, when one is +absent. The real-size host tests need about 270 GB of available host memory +(the table plus the fill's headroom) unless the table is already there. The +four-rank case runs one process per GPU and test_vmm_table_created_by_a_one_gpu_process +runs processes that each see one GPU of each socket: both need four GPUs visible +on two NUMA nodes, and so does the check that the vmm slabs span both sockets. +Each skips, saying what it lacks, when the node or the visible GPUs fall short. + +The gate is bit equality. Every product of an E4M3 value and a power of two is +exact in fp32 (subnormals included, as the byte sweep shows), and the single +fp32 -> bf16 rounding is round-to-nearest-even on both sides, so any +difference is a defect. Deliberately wrong variants (scale per 16 channels, +out-of-shard rows not zeroed, E4M3 bytes read as E5M2) are run through the +same gate and must fail it. +""" + +import functools +import hashlib +import json +import os +import subprocess +import sys +import time +import warnings +from contextlib import contextmanager +from pathlib import Path + +import pytest +import torch +from utils.llm_data import llm_models_root + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.embedding.engram_gather import ( # noqa: E501 + engram_gather as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "engram_gather requires a CUDA device" + +DIM = 256 # engram_head_dim +BLOCK = 32 # one E8M0 scale per 32 channels +LOOKUPS = 24 # (engram_max_ngram_size - 1) * engram_n_heads +WORLD = 4 +REAL_ROWS_L1 = 384_006_168 # engram_num_embeddings[0], the layer-1 table +REAL_ROWS_L14 = 384_016_682 + +# Reference activations dumped from the checkpoint's own inference code; no default +# location, so the cases that compare against them skip unless this is set. +_GOLDEN_DIR = os.environ.get("MODELING_V2_DSV41_GOLDEN_DIR", "") +# The checkpoint resolves under LLM_MODELS_ROOT, as the whole-model gates' does; +# MODELING_V2_DSV41_CKPT_DIR overrides it. A case that needs it skips when it is absent. +_MODELS_ROOT = llm_models_root() +_CKPT_DIR = os.environ.get("MODELING_V2_DSV41_CKPT_DIR") or ( + str(_MODELS_ROOT / "DeepSeek-V4.1-Flash-NVFP4") if _MODELS_ROOT else "" +) +_REAL_CASES = ("capital_fr", "repeat_ngram") + +DEV = torch.device("cuda") + + +# ─── reference: the entry's own; controls: deliberately wrong variants of it ─── + + +# The three deliberately wrong variants every gate below must reject. + + +def _wrong_scale_per_16(ids, weight, scale, row_start, row_end): + """The same 8 scale bytes applied per 16 channels: channel c takes byte (c // 16) % 8.""" + ids = ids.to(DEV).long() + w8, s8 = weight.to(DEV).view(torch.uint8), scale.to(DEV).view(torch.uint8) + mask = (ids < row_start) | (ids >= row_end) + local = (ids - row_start).masked_fill(mask, 0) + s = s8[local].view(torch.float8_e8m0fnu).float() + s = s.repeat_interleave(16, -1).repeat(*([1] * (s.dim() - 1)), 2) # [.., 256]: (c // 16) % 8 + v = w8[local].view(torch.float8_e4m3fn).float() * s + return v.to(torch.bfloat16).masked_fill(mask.unsqueeze(-1), 0) + + +def _wrong_no_zero_fill(ids, weight, scale, row_start, row_end): + """Out-of-shard ids read local row 0 and are left in the output (no final masked fill).""" + ids = ids.to(DEV).long() + w8, s8 = weight.to(DEV).view(torch.uint8), scale.to(DEV).view(torch.uint8) + mask = (ids < row_start) | (ids >= row_end) + local = (ids - row_start).masked_fill(mask, 0) + v = w8[local].view(torch.float8_e4m3fn).float().unflatten(-1, (-1, BLOCK)) + v = v * s8[local].view(torch.float8_e8m0fnu).float().unsqueeze(-1) + return v.flatten(-2).to(torch.bfloat16) + + +def _wrong_e5m2(ids, weight, scale, row_start, row_end): + """The table bytes decoded as float8_e5m2 instead of float8_e4m3fn.""" + ids = ids.to(DEV).long() + w8, s8 = weight.to(DEV).view(torch.uint8), scale.to(DEV).view(torch.uint8) + mask = (ids < row_start) | (ids >= row_end) + local = (ids - row_start).masked_fill(mask, 0) + v = w8[local].view(torch.float8_e5m2).float().unflatten(-1, (-1, BLOCK)) + v = v * s8[local].view(torch.float8_e8m0fnu).float().unsqueeze(-1) + return v.flatten(-2).to(torch.bfloat16).masked_fill(mask.unsqueeze(-1), 0) + + +_WRONG = { + "scale_per_16": _wrong_scale_per_16, + "no_zero_fill": _wrong_no_zero_fill, + "e5m2": _wrong_e5m2, +} + + +# ─── gate ─── + + +def _mismatch(actual, expected) -> float: + """Fraction of elements whose bf16 bit pattern differs (NaN == NaN whatever the payload).""" + a, e = actual.view(torch.int16), expected.view(torch.int16) + same = (a == e) | (torch.isnan(actual) & torch.isnan(expected)) + return 1.0 - same.float().mean().item() + + +def _summed(expected): + """What a bf16 sum over ranks of ``expected`` (one owner, +0 elsewhere) yields: + identical except that -0.0 becomes +0.0 (IEEE -0 + +0 = +0). The reference's own + all-reduce does the same, so the rank sum is compared against this.""" + return torch.zeros_like(expected) + expected + + +def _assert_controls_fail(actual, ids, weight, scale, row_start, row_end, variants, label): + for name in variants: + wrong = _WRONG[name](ids, weight, scale, row_start, row_end) + with pytest.raises(AssertionError): + op.compare(actual, wrong) + print( + f" control {label} {name}: rejected, mismatch fraction {_mismatch(actual, wrong):.6f}" + ) + + +# ─── table builders ─── + + +def _gen(seed: int) -> torch.Generator: + g = torch.Generator(device=DEV) + g.manual_seed(seed) + return g + + +def _table(rows: int, seed: int): + """Device [rows, 256] E4M3 bytes (every byte except the two NaN codes) and + [rows, 8] E8M0 bytes in 110..134 (the checkpoint's layer-1 rows read here + use 117..121).""" + g = _gen(seed) + w = torch.randint(0, 256, (rows, DIM), dtype=torch.uint8, device=DEV, generator=g) + w[(w & 0x7F) == 0x7F] = 0x38 # 0x7F / 0xFF are E4M3 NaN; the byte sweep covers them + s = torch.randint(110, 135, (rows, DIM // BLOCK), dtype=torch.uint8, device=DEV, generator=g) + return w.view(torch.float8_e4m3fn), s.view(torch.float8_e8m0fnu) + + +def _pinned(t: torch.Tensor) -> torch.Tensor: + h = torch.empty(t.shape, dtype=torch.uint8, pin_memory=True) + h.copy_(t.view(torch.uint8)) + assert h.is_pinned() + return h.view(t.dtype) + + +def _ids(tokens: int, rows: int, seed: int, extra=()): + """[tokens, 24] int64 ids uniform over [0, rows), with the listed ids written + over the first positions (boundaries, padding).""" + ids = torch.randint( + 0, rows, (tokens, LOOKUPS), dtype=torch.int64, device=DEV, generator=_gen(seed) + ) + extra = torch.tensor(list(extra), dtype=torch.int64, device=DEV) + n = min(extra.numel(), ids.numel()) + ids.view(-1)[:n] = extra[:n] + return ids + + +def _shard_bounds(rows: int, rank: int): + part = (rows + WORLD - 1) // WORLD # ceil(rows / world_size), as the reference allocates + return rank * part, (rank + 1) * part, part + + +# ─── every byte pair ─── + + +def _drive_byte_pairs(spec): + """All 256 E4M3 bytes x all 256 E8M0 bytes: NaN codes, -0, subnormal products + (scale bytes 0..2), overflow to inf (scale >= 2^121 x 448), E8M0 255 = NaN.""" + residency = spec["residency"] + assert (spec["T"], spec["rows"]) == (11, 256) + w = torch.arange(256, dtype=torch.uint8, device=DEV).repeat(256, 1).view(torch.float8_e4m3fn) + s = ( + torch.arange(256, dtype=torch.uint8, device=DEV)[:, None] + .repeat(1, DIM // BLOCK) + .view(torch.float8_e8m0fnu) + ) + ids = torch.full((11, LOOKUPS), -1, dtype=torch.int64, device=DEV) # 264 slots, 8 padding + ids.view(-1)[:256] = torch.arange(256, device=DEV) + tw, ts = (w, s) if residency == "device" else (_pinned(w), _pinned(s)) + out = op(ids, tw, ts, 0, 256) + expected = op.reference(ids, w, s, 0, 256) + op.compare(out, expected) + flat = out.view(-1, DIM) + subnormal = (flat != 0) & (flat.abs() < torch.finfo(torch.bfloat16).tiny) + print( + f"\n {residency}: {torch.isnan(out).sum().item()} NaN, {torch.isinf(out).sum().item()} inf, " + f"{subnormal.sum().item()} bf16 subnormals, bit-exact" + ) + assert subnormal.any() and torch.isinf(out).any() and torch.isnan(out).any() + assert (flat[256:].view(torch.int16) == 0).all() + + +# ─── device table, row-sharded by rank ─── + + +def _drive_device_shards(spec): + """Four [ceil(rows/4), 256] shards (the last padded past the table's end with + rows no valid id names): each rank's output is its owned rows exactly and + +0 elsewhere, and the bf16 sum over ranks is the unsharded gather. The four + ranks' calls run in this process, on one device; the sum is the arithmetic the + all-reduce performs, not the collective.""" + tokens, rows = spec["T"], spec["rows"] + assert spec["world"] == WORLD + full_w, full_s = _table(rows, seed=tokens) + _, _, part = _shard_bounds(rows, 0) + boundaries = [ + 0, + part - 1, + part, + 2 * part - 1, + 2 * part, + 3 * part - 1, + 3 * part, + rows - 1, + -1, + -1, + ] + ids = _ids(tokens, rows, seed=100 + tokens, extra=boundaries) + total = torch.zeros(tokens, LOOKUPS, DIM, dtype=torch.bfloat16, device=DEV) + owners = torch.zeros(tokens, LOOKUPS, dtype=torch.int64, device=DEV) + for rank in range(WORLD): + lo, hi, _ = _shard_bounds(rows, rank) + w_r = torch.empty(part, DIM, dtype=torch.uint8, device=DEV) + s_r = torch.empty(part, DIM // BLOCK, dtype=torch.uint8, device=DEV) + w_r.copy_( + torch.randint(0, 0x70, w_r.shape, dtype=torch.uint8, device=DEV, generator=_gen(rank)) + ) # pad rows + s_r.fill_(127) + n = min(hi, rows) - lo + w_r[:n] = full_w.view(torch.uint8)[lo : lo + n] + s_r[:n] = full_s.view(torch.uint8)[lo : lo + n] + w_r, s_r = w_r.view(torch.float8_e4m3fn), s_r.view(torch.float8_e8m0fnu) + out = op(ids, w_r, s_r, lo, hi) + op.compare(out, op.reference(ids, w_r, s_r, lo, hi)) + owned = (ids >= lo) & (ids < hi) + assert (out[~owned].view(torch.int16) == 0).all(), "an out-of-shard row is not +0" + if rank == 1 and tokens == 32: + _assert_controls_fail(out, ids, w_r, s_r, lo, hi, _WRONG, "device rank-1 shard") + owners += owned.long() + total += out + assert ((owners == 1) | (ids < 0)).all() and (owners[ids < 0] == 0).all() + op.compare(total, _summed(op.reference(ids, full_w, full_s, 0, rows))) + + +# ─── pinned host table (host_offload) ─── + + +def _drive_host_shared(spec): + """host_offload's shared layout: one pinned copy of the whole table, row range + [0, rows), read through the mapped pointer. No device copy of the table is made: + device memory allocated across the call is unchanged with a caller-owned out.""" + tokens, rows = spec["T"], spec["rows"] + w, s = _table(rows, seed=7 + tokens) + hw, hs = _pinned(w), _pinned(s) + ids = _ids(tokens, rows, seed=tokens, extra=[0, rows - 1, -1, rows]) + out = torch.empty(tokens, LOOKUPS, DIM, dtype=torch.bfloat16, device=DEV) + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + op(ids, hw, hs, 0, rows, out=out) + torch.cuda.synchronize() + assert torch.cuda.memory_allocated() == before + assert hw.device.type == "cpu" and hw.is_pinned() and hs.is_pinned() + op.compare(out, op.reference(ids, w, s, 0, rows)) + + +def _drive_host_shards(spec): + """Per-rank host shards as row views (storage offsets) of one pinned buffer: + the mapped pointer of an interior address, zero-fill and the rank sum.""" + tokens, rows = spec["T"], spec["rows"] + assert spec["world"] == WORLD + w, s = _table(rows, seed=31) + _, _, part = _shard_bounds(rows, 0) + hw = torch.zeros(WORLD * part, DIM, dtype=torch.uint8, pin_memory=True) + hs = torch.full((WORLD * part, DIM // BLOCK), 127, dtype=torch.uint8, pin_memory=True) + hw[:rows].copy_(w.view(torch.uint8)) + hs[:rows].copy_(s.view(torch.uint8)) + hw, hs = hw.view(torch.float8_e4m3fn), hs.view(torch.float8_e8m0fnu) + ids = _ids(tokens, rows, seed=5, extra=[part - 1, part, 3 * part, rows - 1, -1]) + total = torch.zeros(tokens, LOOKUPS, DIM, dtype=torch.bfloat16, device=DEV) + for rank in range(WORLD): + lo, hi, _ = _shard_bounds(rows, rank) + out = op(ids, hw[lo:hi], hs[lo:hi], lo, hi) + op.compare(out, op.reference(ids, hw[lo:hi], hs[lo:hi], lo, hi)) + total += out + op.compare(total, _summed(op.reference(ids, w, s, 0, rows))) + + +# ─── byte offsets past 2 GiB (real tables are 98 GB) ─── + + +def _drive_offsets_2gib(spec): + """A 9,000,000-row shard (2.30 GB) whose first global row is rank 3's start of + the layer-14 table; ids at local rows above 2^31 / 256 = 8,388,608 are read at + byte offsets past 2 GiB. Ids just outside the shard on both sides give zeros.""" + residency, rows = spec["residency"], spec["rows"] + lo = 3 * ((REAL_ROWS_L14 + WORLD - 1) // WORLD) # 288,012,513 + assert lo == spec["row_start"] and spec["T"] == 32 + hi = lo + rows + w, s = _table(rows, seed=77) + tw, ts = (w, s) if residency == "device" else (_pinned(w), _pinned(s)) + local = torch.randint(8_388_608, rows, (32, LOOKUPS), device=DEV, generator=_gen(3)) + local[:, :4] = torch.randint(0, rows, (32, 4), device=DEV, generator=_gen(4)) + ids = lo + local + ids[0, :6] = torch.tensor([lo, hi - 1, lo - 1, hi, 0, REAL_ROWS_L14 - 1], device=DEV) + out = op(ids, tw, ts, lo, hi) + op.compare(out, op.reference(ids, w, s, lo, hi)) + assert (out[0, 2:6].view(torch.int16) == 0).all() + + +# ─── CUDA graphs ─── + + +@pytest.mark.parametrize("tokens", [1, 32, 4096]) +@pytest.mark.parametrize("residency", ["device", "host"]) +def test_cuda_graph_capture_replay(residency, tokens): + """Capture one call (the host pointer is resolved during capture), replay with + new ids, then rewrite the table in place and replay again: the replay reads the + live table (a captured staging copy would return stale rows).""" + rows = 50_000 + w, s = _table(rows, seed=11) + tw, ts = (w.clone(), s.clone()) if residency == "device" else (_pinned(w), _pinned(s)) + ids = _ids(tokens, rows, seed=12, extra=[-1, rows]) + out = torch.empty(tokens, LOOKUPS, DIM, dtype=torch.bfloat16, device=DEV) + op(ids, tw, ts, 0, rows, out=out) # compile outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(ids, tw, ts, 0, rows, out=out) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + op.compare(out, op.reference(ids, tw, ts, 0, rows)) + + ids.copy_(_ids(tokens, rows, seed=13, extra=[rows - 1, -5])) + graph.replay() + torch.cuda.synchronize() + op.compare(out, op.reference(ids, tw, ts, 0, rows)) + + new_w, new_s = _table(rows, seed=14) + tw.view(torch.uint8).copy_(new_w.view(torch.uint8)) + ts.view(torch.uint8).copy_(new_s.view(torch.uint8)) + torch.cuda.synchronize() + graph.replay() + torch.cuda.synchronize() + op.compare(out, op.reference(ids, new_w, new_s, 0, rows)) + eager = op(ids, tw, ts, 0, rows) + assert torch.equal(eager.view(torch.int16), out.view(torch.int16)) + + +# ─── every row count 1..N_max is the same per-row function ─── +# +# The core target calls this op on one rank's T tokens (tep4 / tp4: the +# [:, layer, :] view of the hash output, row stride 48; T = 1..max_num_tokens, up to +# 16384 per rank) and, under dep4, on the all-gathered ids of the four ranks (the +# contiguous int64 view of an int32 [4 * rows, 48] all-gather buffer; up to +# 4 x 16384 = 65536 rows). T reaches +# the kernel only as ``n_rows = T * 24``, which sizes the grid and masks the last +# program. The sweeps below check that every T from 1 to N_max returns exactly the +# first T rows of the N_max call, which is itself held bit-exact to the reference. + + +def _jit_variants(fn) -> int: + """Compiled variants Triton holds for ``fn`` on every device (Triton >= 3.2 keeps them + in ``device_caches[device][0]``, older releases in ``cache[device]``).""" + dc = getattr(fn, "device_caches", None) + if dc: + return sum(len(v[0]) for v in dc.values()) + return sum(len(v) for v in fn.cache.values()) + + +_POISON = ( + 0x7FA5 # a bf16 NaN bit pattern the kernel never produces (the tables here hold no NaN bytes) +) +_GUARD = 64 # rows past the prefix that must still hold the poison after each call + + +def _core_shard(full_w, full_s, rows: int, rank: int, seed: int): + """Rank ``rank``'s [ceil(rows / 4), 256] device shard, as the core holds it: bounds + ``row_start = rank * part``, ``row_end = min(row_start + part, rows)``; the last + shard's rows past the table's end hold bytes no valid id names.""" + part = (rows + WORLD - 1) // WORLD + lo, hi = rank * part, min(rank * part + part, rows) + w_r = torch.randint(0, 0x70, (part, DIM), dtype=torch.uint8, device=DEV, generator=_gen(seed)) + s_r = torch.full((part, DIM // BLOCK), 127, dtype=torch.uint8, device=DEV) + w_r[: hi - lo] = full_w.view(torch.uint8)[lo:hi] + s_r[: hi - lo] = full_s.view(torch.uint8)[lo:hi] + return w_r.view(torch.float8_e4m3fn), s_r.view(torch.float8_e8m0fnu), lo, hi + + +def _sprinkle(ids: torch.Tensor, values, seed: int) -> None: + """Write each of ``values`` at 64 random positions of ``ids`` (in place), and all of + them into the first token row, so even r = 1 sees every edge.""" + g = _gen(seed) + for v in values: + t = torch.randint(0, ids.shape[0], (64,), device=DEV, generator=g) + k = torch.randint(0, ids.shape[1], (64,), device=DEV, generator=g) + ids[t, k] = v + ids[0, : len(values)] = torch.tensor(list(values), dtype=ids.dtype, device=DEV) + + +@pytest.mark.parametrize( + "form,n_max,rank", [("tp_view", 16384, 1), ("dep4_gathered", 4 * 16384, 3)] +) +def test_every_row_count_is_a_prefix_of_n_max(form, n_max, rank): + """Every prefix ids[:r], r = 1..N_max, returns exactly the first r rows of the N_max call. + + ``tp_view``: rank 1's shard (row_start = part > 0) over a [:, 1, :] view of a + [N, 2, 24] int64 hash buffer (row stride 48), N_max 16384. ``dep4_gathered``: + rank 3's (padded, last) shard over the int64 view of an int32 [N, 48] all-gather + buffer, N_max 65536, with the zero padding rows the core's pad-to-max leaves. + Ids include -1, row_start - 1, row_start, row_end - 1, row_end, 0 and rows - 1. + Each call writes into a row-prefix view of one poisoned buffer: its r rows must + equal the N_max rows bit for bit and the next 64 rows must still hold the poison + (mismatches are summed on the device, one synchronization at the end). The Triton + kernel's compiled variants are counted: after r = 1 and r = 2 (n_rows = 24 r is + divisible by 16 exactly for even r, Triton's integer specialization) no further + variant may appear for any r.""" + rows = 100_003 + full_w, full_s = _table(rows, seed=4000 + rank) + w_r, s_r, lo, hi = _core_shard(full_w, full_s, rows, rank, seed=4100 + rank) + if form == "tp_view": + hashes = torch.randint( + 0, rows, (n_max, 2, LOOKUPS), dtype=torch.int64, device=DEV, generator=_gen(4200) + ) + ids = hashes[:, 1, :] + assert ids.stride() == (2 * LOOKUPS, 1) + else: + g32 = torch.empty(n_max, 2 * LOOKUPS, dtype=torch.int32, device=DEV) + ids = g32.view(torch.int64) + ids.copy_( + torch.randint( + 0, rows, (n_max, LOOKUPS), dtype=torch.int64, device=DEV, generator=_gen(4201) + ) + ) + per_rank = n_max // WORLD + for r_ in range( + WORLD + ): # pad-to-max rows at the end of each rank's block: all-zero int32 words + g32[(r_ + 1) * per_rank - 5 * (r_ + 1) : (r_ + 1) * per_rank] = 0 + assert ids.is_contiguous() and ids.shape == (n_max, LOOKUPS) + _sprinkle(ids, [-1, lo - 1, lo, hi - 1, hi, 0, rows - 1], seed=4300) + owned = (ids >= lo) & (ids < hi) + assert 0 < owned.float().mean().item() < 0.5 and (ids == -1).any() + + full = op(ids, w_r, s_r, lo, hi) + op.compare(full, op.reference(ids, w_r, s_r, lo, hi)) + full_bits = full.view(torch.int16) + buf = torch.empty(n_max + _GUARD, LOOKUPS, DIM, dtype=torch.bfloat16, device=DEV) + bits = buf.view(torch.int16) + bits.fill_(_POISON) + bad = torch.zeros((), dtype=torch.int64, device=DEV) + over = torch.zeros((), dtype=torch.int64, device=DEV) + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.embedding.engram_gather import ( # noqa: E501 + _engram_gather_kernel as fn, + ) + + before = _jit_variants(fn) + settled = None + for r in range(1, n_max + 1): + res = op(ids[:r], w_r, s_r, lo, hi, out=buf[:r]) + assert res.data_ptr() == buf.data_ptr() + bad += (bits[:r] != full_bits[:r]).sum() + over += (bits[r : r + _GUARD] != _POISON).sum() + if r == 2: + settled = _jit_variants(fn) + torch.cuda.synchronize() + after = _jit_variants(fn) + print( + f"\n {form}: r = 1..{n_max}, rank {rank} [{lo}, {hi}): {bad.item()} elements differ from the " + f"{n_max}-row result, {over.item()} guard elements written; kernel variants {before} before, " + f"{settled} after r = 2, {after} after r = {n_max}" + ) + assert bad.item() == 0, f"{bad.item()} elements differ from the {n_max}-row result's rows" + assert over.item() == 0, f"{over.item()} elements past a prefix were written" + assert after == settled and settled - before <= 2, (before, settled, after) + del buf, full + + +def test_row_sweep_harness_sees_a_one_row_shift(): + """Control for the sweeps above: the same bitwise per-row check rejects a call that + reads the ids one token row off, at every r tried, including r = 1.""" + rows = 100_003 + full_w, full_s = _table(rows, seed=4400) + w_r, s_r, lo, hi = _core_shard(full_w, full_s, rows, 1, seed=4401) + hashes = torch.randint( + 0, rows, (65, 2, LOOKUPS), dtype=torch.int64, device=DEV, generator=_gen(4402) + ) + hashes[:, 1, :8] = lo + torch.arange(65 * 8, device=DEV).view( + 65, 8 + ) # distinct owned rows per token + ids = hashes[:, 1, :] + full = op(ids[:64], w_r, s_r, lo, hi) + for r in (1, 2, 7, 63, 64): + shifted = op(ids[1 : r + 1], w_r, s_r, lo, hi) + n_bad = int((shifted.view(torch.int16) != full.view(torch.int16)[:r]).sum().item()) + print(f" [control row shift] r={r}: {n_bad} / {r * LOOKUPS * DIM} elements differ") + assert n_bad > 0 + + +# ─── table size: any per-rank shard up to the largest the target allocates ─── + +# The target allocates every rank's shard as [ceil(rows / 4), 256] (+ [.., 8] scale bytes): +# 96,001,542 rows for layer 1 and 96,004,171 for layer 14 (the largest), whatever the +# free memory; rank 3's layer-14 shard holds 96,004,169 loaded rows (row_end = rows). +_PART_L1 = (REAL_ROWS_L1 + WORLD - 1) // WORLD +_PART_L14 = (REAL_ROWS_L14 + WORLD - 1) // WORLD +_PART_MAX = max(_PART_L1, _PART_L14) + + +def _fill_rows(buf_w, buf_s, local: torch.Tensor, seed: int) -> None: + """Random E4M3 bytes (no NaN codes) and E8M0 bytes 110..134 into the listed local rows.""" + g = _gen(seed) + w = torch.randint(0, 256, (local.numel(), DIM), dtype=torch.uint8, device=DEV, generator=g) + w[(w & 0x7F) == 0x7F] = 0x38 + buf_w[local] = w + buf_s[local] = torch.randint( + 110, 135, (local.numel(), DIM // BLOCK), dtype=torch.uint8, device=DEV, generator=g + ) + + +def _wrong_wrap32(ids, weight, scale, row_start, row_end): + """The gather with a 32-bit byte offset: local row r read at (r * 256) mod 2^32.""" + ids = ids.to(DEV).long() + wrapped = row_start + (ids - row_start) % ((1 << 32) // DIM) + return op.reference( + torch.where((ids >= row_start) & (ids < row_end), wrapped, ids), + weight, + scale, + row_start, + row_end, + ) + + +def test_table_size_invariance_and_far_end_at_the_largest_shard(): + """The table-row axis. The table is a checkpoint constant, not a pool the + runtime sizes from free memory, and the kernel receives it as a raw address plus + (row_start, row_end); this runs that instead of arguing it. + + (1) Invariance: one device buffer of the largest per-rank shard the target allocates + (layer 14, 96,004,171 rows, 24.6 GB + 0.77 GB scale bytes), rows 0..4095 random and + the rest a constant row. The same 16384-token call (the tp [:, 1, :] view, rank 3's + layer-14 row_start 288,012,513; ids in the first 4096 local rows plus -1, ids below + row_start and at row_end) on views [:R] for R = 4096 (smallest holding them), + 2^24 (mid) and 96,004,171, and on a separate exact 4096-row allocation, is bit-exact + to the reference and bitwise equal across all four -- once with row_end fixed at + row_start + 4096 and once with row_end = row_start + R. + (2) Far end: for every rank of both layers, with the target's own bounds + (row_end = min(row_start + part, rows)), ids among the last 2^20 loaded rows of the + shard plus the last loaded row, row_end - 1, row_end, row_start, row_start - 1 and -1, + in the tp view (T 16384) and, for rank 3, the dep4 all-gathered form (T 65536): bit-exact. + Control: a gather with a 32-bit byte offset (row r read at (256 r) mod 2^32) fails.""" + buf_w = torch.full((_PART_MAX, DIM), 0x38, dtype=torch.uint8, device=DEV) + buf_s = torch.full((_PART_MAX, DIM // BLOCK), 127, dtype=torch.uint8, device=DEV) + try: + small = 4096 + _fill_rows(buf_w, buf_s, torch.arange(small, device=DEV), seed=9100) + lo = 3 * _PART_L14 # 288,012,513 + hashes = torch.randint( + lo, lo + small, (16384, 2, LOOKUPS), dtype=torch.int64, device=DEV, generator=_gen(9101) + ) + ids = hashes[:, 1, :] + _sprinkle(ids, [-1, lo - 1, lo, lo + small - 1, 0], seed=9102) + exact_w, exact_s = buf_w[:small].clone(), buf_s[:small].clone() + views = [ + (small, buf_w[:small], buf_s[:small]), + (1 << 24, buf_w[: 1 << 24], buf_s[: 1 << 24]), + (_PART_MAX, buf_w, buf_s), + ("exact", exact_w, exact_s), + ] + for fixed_end in (True, False): + outs = [] + for size, w, s in views: + hi = lo + small if fixed_end or size == "exact" else lo + w.shape[0] + tw, ts = w.view(torch.float8_e4m3fn), s.view(torch.float8_e8m0fnu) + out = op(ids, tw, ts, lo, hi) + op.compare(out, op.reference(ids, tw, ts, lo, hi)) + outs.append(out) + for out in outs[1:]: + assert torch.equal(out.view(torch.int16), outs[0].view(torch.int16)) + print( + f"\n invariance (row_end {'fixed' if fixed_end else '= row_start + R'}): R = " + f"{[v[0] for v in views]}: bit-exact to the reference and bitwise equal" + ) + # zero-fill: ids below row_start and -1 read nothing + assert (outs[0].view(torch.int16)[0, 0] == 0).all() and ( + outs[0].view(torch.int16)[0, 1] == 0 + ).all() + + buf_w[:small] = 0x38 + buf_s[:small] = 127 + for layer, rows, part in ( + ("L1", REAL_ROWS_L1, _PART_L1), + ("L14", REAL_ROWS_L14, _PART_L14), + ): + tw, ts = buf_w[:part].view(torch.float8_e4m3fn), buf_s[:part].view(torch.float8_e8m0fnu) + for rank in range(WORLD): + lo = rank * part + hi = min(lo + part, rows) + n_loaded = hi - lo + tail = torch.randint( + n_loaded - (1 << 20), n_loaded, (4096,), device=DEV, generator=_gen(9200 + rank) + ) + local = torch.unique(torch.cat([tail, torch.tensor([n_loaded - 1, 0], device=DEV)])) + _fill_rows(buf_w, buf_s, local, seed=9300 + rank) + forms = [("tp_view", 16384)] + ( + [("dep4_gathered", 4 * 16384)] if rank == WORLD - 1 else [] + ) + for form, t in forms: + pick = torch.randint( + 0, local.numel(), (t, LOOKUPS), device=DEV, generator=_gen(9400 + rank) + ) + if form == "tp_view": + hashes = torch.empty(t, 2, LOOKUPS, dtype=torch.int64, device=DEV) + hashes[:, 1, :] = lo + local[pick] + ids = hashes[:, 1, :] + else: + g32 = torch.empty(t, 2 * LOOKUPS, dtype=torch.int32, device=DEV) + ids = g32.view(torch.int64) + ids.copy_(lo + local[pick]) + _sprinkle( + ids, [-1, lo - 1, lo, hi - 1, hi, lo + n_loaded - 1], seed=9500 + rank + ) + out = op(ids, tw, ts, lo, hi) + op.compare(out, op.reference(ids, tw, ts, lo, hi)) + far = int((ids - lo).clamp(max=n_loaded - 1).max().item()) + print( + f" far end {layer} rank {rank} {form} T {t}: [{lo}, {hi}) of a {part}-row shard, " + f"max local row {far} (byte offset {far * DIM}): bit-exact" + ) + if rank == 0 and form == "tp_view": + with pytest.raises(AssertionError): + op.compare(out, _wrong_wrap32(ids, tw, ts, lo, hi)) + print( + f" control {layer} 32-bit byte offset: rejected, mismatch fraction " + f"{_mismatch(out, _wrong_wrap32(ids, tw, ts, lo, hi)):.6f}" + ) + buf_w[local] = 0x38 + buf_s[local] = 127 + finally: + del buf_w, buf_s + torch.cuda.empty_cache() + + +# ─── capture at every graph batch the gates capture ─── + +# The generation-only CUDA-graph batches an engine captures per rank, Gset: every batch +# 1..32 plus TRT-LLM's padded default list for cuda_graph_config {enable_padding: true, +# max_batch_size: 256} (1, 2, 4, 8, 16..128 step 8, 192, 256; the padding-off list at 256, +# 1..32, 64, 128, 256, is a subset). A decode step carries G tokens and a DSpark +# target-verify step 6 G (1 + 5 drafts): T = G and 6 G, 82 counts up to 1536. tep4 / tp4 +# and host_offload gather the rank's T tokens through a [:, layer, :] view of the hash +# output (row stride 48); dep4 with device tables gathers the all-gathered ids of the four +# ranks, 4 x T rows (82 counts up to 6144), contiguous. +_GRAPH_BATCHES = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_CAPTURE_T = sorted(set(_GRAPH_BATCHES) | {6 * b for b in _GRAPH_BATCHES}) +assert len(_GRAPH_BATCHES) == 46 and len(_CAPTURE_T) == 82 and _CAPTURE_T[-1] == 1536 +_CAPTURE_CASES = [("tp_view", t) for t in _CAPTURE_T] + [ + ("dep4_gathered", 4 * t) for t in _CAPTURE_T +] + + +def _capture_ids(form: str, tokens: int, rows: int, seed: int): + """[tokens, 2, 24] ids for ``tp_view`` (both layers' views are gathered) or + [tokens, 24] for ``dep4_gathered`` (each rank's block of tokens / 4 rows ends in + zero padding rows, as the core's pad-to-max before the all-gather leaves it).""" + part = (rows + WORLD - 1) // WORLD + edge = [0, part - 1, part, 2 * part - 1, 2 * part, 3 * part - 1, 3 * part, rows - 1, -1] + if form == "tp_view": + ids = torch.randint( + 0, rows, (tokens, 2, LOOKUPS), dtype=torch.int64, device=DEV, generator=_gen(seed) + ) + ids[0, 0, : len(edge)] = torch.tensor(edge, device=DEV) + ids[-1, 1, : len(edge)] = torch.tensor(edge[::-1], device=DEV) + return ids + ids = torch.randint( + 0, rows, (tokens, LOOKUPS), dtype=torch.int64, device=DEV, generator=_gen(seed) + ) + ids.view(-1)[: len(edge)] = torch.tensor(edge, device=DEV) + per_rank = tokens // WORLD + for rank in range(WORLD): + pad = ( + (seed + rank) % 3 if per_rank > 2 else 0 + ) # 0..2 padding rows at the end of each rank's block + if pad: + ids[(rank + 1) * per_rank - pad : (rank + 1) * per_rank] = 0 + return ids + + +@pytest.mark.parametrize("form,tokens", _CAPTURE_CASES) +def test_cuda_graph_at_every_capture_batch(form, tokens): + """The core's calls, captured per rank with ``out`` allocated inside the capture + (``engram_gather(ids, shard_w, shard_s, row_start, row_end)``), after an eager + warm-up at the same shape. ``tp_view``: the two Engram layers' ``hashes[:, l, :]`` + views of one static ``[T, 2, 24]`` int64 buffer; ``dep4_gathered``: the int64 + view of a static ``[4 * rows, 48]`` int32 all-gather buffer. Three replays on new + ids written in place: every rank's replay equals the same eager call bit for bit + and the reference, out-of-shard rows are +0, and the bf16 sum over the four ranks + is the unsharded gather.""" + rows = 100_003 + full_w, full_s = _table(rows, seed=900 + tokens) + if form == "tp_view": + static = _capture_ids(form, tokens, rows, seed=0) + views = [static[:, 0, :], static[:, 1, :]] + assert all(v.stride() == (2 * LOOKUPS, 1) for v in views) + else: + g32 = torch.empty(tokens, 2 * LOOKUPS, dtype=torch.int32, device=DEV) + g32.view(torch.int64).copy_(_capture_ids(form, tokens, rows, seed=0)) + views = [g32.view(torch.int64)] + assert views[0].shape == (tokens, LOOKUPS) and views[0].is_contiguous() + capture_ids = [v.clone() for v in views] + shards = [_core_shard(full_w, full_s, rows, rank, seed=50 + rank) for rank in range(WORLD)] + graphs, outs = [], [] + for w_r, s_r, lo, hi in shards: + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for v in views: + op(v, w_r, s_r, lo, hi) # engine warm-up at this shape + torch.cuda.current_stream().wait_stream(s) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs.append([op(v, w_r, s_r, lo, hi) for v in views]) + graphs.append(graph) + for rep in (1, 2, 3): + new = _capture_ids(form, tokens, rows, seed=rep) + if form == "tp_view": + static.copy_(new) + else: + g32.view(torch.int64).copy_(new) + for rank, graph in enumerate(graphs): + for o in outs[rank]: + o.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + for li, v in enumerate(views): + total = torch.zeros(tokens, LOOKUPS, DIM, dtype=torch.bfloat16, device=DEV) + owners = torch.zeros(tokens, LOOKUPS, dtype=torch.int64, device=DEV) + for rank, (w_r, s_r, lo, hi) in enumerate(shards): + out = outs[rank][li] + eager = op(v, w_r, s_r, lo, hi) + assert torch.equal(out.view(torch.int16), eager.view(torch.int16)), ( + f"rank {rank}: replay != eager" + ) + expected = op.reference(v, w_r, s_r, lo, hi) + op.compare(out, expected) + owned = (v >= lo) & (v < hi) + assert (out[~owned].view(torch.int16) == 0).all(), "an out-of-shard row is not +0" + stale = op.reference(capture_ids[li], w_r, s_r, lo, hi) + if owned.any() or ((capture_ids[li] >= lo) & (capture_ids[li] < hi)).any(): + # the replay read the new ids: the capture-time ids' rows would differ + assert not torch.equal(expected.view(torch.int16), stale.view(torch.int16)), ( + f"rank {rank}" + ) + if rep == 1 and li == 0 and rank == 1: + print( + f"\n {form} T={tokens}: 4 ranks x {len(views)} calls x 3 replays, replay == eager == " + f"reference bit for bit; rank 1 differs from the capture-time ids' rows on " + f"{_mismatch(out, stale):.6f}" + ) + _assert_controls_fail( + out, v, w_r, s_r, lo, hi, _WRONG, f"graph {form} T={tokens} rank 1" + ) + owners += owned.long() + total += out + assert ((owners == 1) | (v < 0)).all() and (owners[v < 0] == 0).all() + op.compare(total, _summed(op.reference(v, full_w, full_s, 0, rows))) + + +# ─── capture through the node's shared pinned host table (host_offload) ─── + + +def _write_synthetic_checkpoint(ckpt_dir: str, tables) -> None: + """A two-layer 'checkpoint' holding only ``layers.{1,14}.engram.embed.{weight,scale}``, + in one safetensors shard written by hand (header length, JSON header, raw bytes), + plus its ``model.safetensors.index.json``: what ``engram_host_table`` reads.""" + import struct + + header, blobs, at = {}, [], 0 + for layer, (w, s) in tables.items(): + for kind, t, dtype in (("weight", w, "F8_E4M3"), ("scale", s, "U8")): + b = t.view(torch.uint8).cpu().contiguous().numpy().tobytes() + header[f"layers.{layer}.engram.embed.{kind}"] = { + "dtype": dtype, + "shape": list(t.shape), + "data_offsets": [at, at + len(b)], + } + blobs.append(b) + at += len(b) + raw = json.dumps(header).encode() + raw += b" " * (-len(raw) % 8) + shard = "model-00001-of-00001.safetensors" + with open(os.path.join(ckpt_dir, shard), "wb") as f: + f.write(struct.pack(" None: # noqa: F811 — pytest resolves the fixture by this name + """Every cell through ``op`` (guarded by ``is_valid``), ``op.reference`` and ``op.compare``, + with the structural checks (zero fill, rank sum, untouched memory) its configuration carries.""" + certifying(op) + spec = cell.spec + assert (spec["D"], spec["B"], spec["K"]) == (DIM, BLOCK, LOOKUPS) + _DRIVERS[spec["kind"]](spec) + + +def test_is_valid_rejects_layouts_the_kernel_misreads(): + """Each case ``is_valid`` refuses (inside ``validating``) was run through the launch + without the guard on sm_103 and returned wrong values (or wrote outside ``out``) without + an error; the measured results are in ``is_valid``'s docstring.""" + rows = 1000 + w, s = _table(rows, seed=41) + ids = _ids(7, rows, seed=42) + out = torch.empty(7, LOOKUPS, DIM, dtype=torch.bfloat16, device=DEV) + bad = { + "ids inner stride": lambda: op(ids.t().contiguous().t(), w, s, 0, rows), + "ids float32": lambda: op(ids.float(), w, s, 0, rows), + "weight float8_e5m2": lambda: op( + ids, w.view(torch.uint8).view(torch.float8_e5m2), s, 0, rows + ), + "scale float32": lambda: op(ids, w, s.float(), 0, rows), + "weight column slice": lambda: op( + ids, + torch.cat([w.view(torch.uint8)] * 2, 1).view(torch.float8_e4m3fn)[:, :DIM], + s, + 0, + rows, + ), + "scale column slice": lambda: op( + ids, w, torch.cat([s.view(torch.uint8)] * 2, 1)[:, : DIM // BLOCK], 0, rows + ), + "scale with fewer rows than weight": lambda: op(ids, w, s[: rows // 2], 0, rows), + "scale width not dividing D": lambda: op( + ids, w, s.view(torch.uint8)[:, :7].contiguous(), 0, rows + ), + "row range beyond the rows held": lambda: op(ids, w[:100], s[:100], 0, 200), + "out last-dim slice": lambda: op( + ids, + w, + s, + 0, + rows, + out=torch.empty(7, LOOKUPS, 2 * DIM, dtype=torch.bfloat16, device=DEV)[..., :DIM], + ), + "out too few token rows": lambda: op(ids, w, s, 0, rows, out=out[:3]), + } + match = { + "ids inner stride": r"ids must be \[T, K\]", + "ids float32": "ids must be int64 or int32", + "weight float8_e5m2": "weight must be float8_e4m3fn", + "scale float32": "scale must be float8_e8m0fnu or uint8", + "weight column slice": "weight must be a contiguous", + "scale column slice": "scale must be a contiguous", + "scale with fewer rows than weight": "does not match weight", + "scale width not dividing D": "does not match weight", + "row range beyond the rows held": "exceeds the", + "out last-dim slice": "out must be a contiguous", + "out too few token rows": "out must be a contiguous", + } + assert set(match) == set(bad) + for name, call in bad.items(): + with pytest.raises(AssertionError, match=match[name]): + with validating(op): + call() + print(f" rejected: {name}") + # CPU ids are refused by Triton's own launcher, loudly + with pytest.raises(ValueError, match="cannot be accessed from Triton"): + op(ids.cpu(), w, s, 0, rows) + + +# ─── REAL rows of the checkpoint's layer-1 table ─── + + +def _real_available() -> bool: + return all( + os.path.isfile(os.path.join(_GOLDEN_DIR, c, "engram_hash.pt")) for c in _REAL_CASES + ) and os.path.isfile(os.path.join(_CKPT_DIR, "model.safetensors.index.json")) + + +def _skip_without_real(): + if not _real_available(): + pytest.skip(f"golden dumps ({_GOLDEN_DIR}) or checkpoint ({_CKPT_DIR}) not present") + + +@functools.lru_cache(maxsize=None) +def _weight_file(name: str) -> str: + with open(os.path.join(_CKPT_DIR, "model.safetensors.index.json")) as f: + return os.path.join(_CKPT_DIR, json.load(f)["weight_map"][name]) + + +@functools.lru_cache(maxsize=None) +def _real_rows(cases: tuple): + """Layer-1 hash ids of the golden cases ([T, 24] each, offsets applied), and the + rows they name, read one row at a time through get_slice (never the table).""" + from safetensors import safe_open + + hids = [ + torch.load(os.path.join(_GOLDEN_DIR, c, "engram_hash.pt"), weights_only=True)["hash_ids"][ + :, 0, : + ] + for c in cases + ] + uniq = torch.unique(torch.cat([h.reshape(-1) for h in hids])) + got = [] + for name in ("layers.1.engram.embed.weight", "layers.1.engram.embed.scale"): + with safe_open(_weight_file(name), framework="pt", device="cpu") as f: + sl = f.get_slice(name) + assert sl.get_shape()[0] == REAL_ROWS_L1 + got.append(torch.cat([sl[r : r + 1] for r in uniq.tolist()])) + w, s = got + assert w.dtype == torch.float8_e4m3fn and s.dtype == torch.float8_e8m0fnu, (w.dtype, s.dtype) + return hids, uniq, w, s + + +@pytest.mark.parametrize("residency", ["device", "host"]) +@pytest.mark.parametrize("case", _REAL_CASES) +def test_real_rows_compact_table(case, residency): + """The golden prompt's real rows in a compact table (ids remapped): the device + and the pinned host path both reproduce the reference dequantization bit for bit.""" + _skip_without_real() + (hid,), uniq, w, s = _real_rows((case,)) + ids = torch.searchsorted(uniq, hid.contiguous()).to(DEV) + tw, ts = (w.to(DEV), s.to(DEV)) if residency == "device" else (_pinned(w), _pinned(s)) + out = op(ids, tw, ts, 0, uniq.numel()) + op.compare(out, op.reference(ids, w, s, 0, uniq.numel())) + print( + f"\n {case} {residency}: {tuple(hid.shape)} ids, {uniq.numel()} distinct real rows, bit-exact" + ) + _assert_controls_fail(out, ids, w, s, 0, uniq.numel(), ("scale_per_16", "e5m2"), f"real {case}") + + +def test_real_ids_at_the_real_shard_size(): + """The real global ids against a device buffer of the real per-rank shard size + (ceil(384,006,168 / 4) = 96,001,542 rows, 24.6 GB): each rank's call reads its + own rows at their true local offsets (up to 24.6 GB into the buffer) and zeroes + the rest, and the rank sum equals the unsharded gather.""" + _skip_without_real() + hids, uniq, w, s = _real_rows(_REAL_CASES) + ids = torch.cat(hids).to(DEV) + part = (REAL_ROWS_L1 + WORLD - 1) // WORLD + buf_w = torch.empty(part, DIM, dtype=torch.uint8, device=DEV) + buf_s = torch.empty(part, DIM // BLOCK, dtype=torch.uint8, device=DEV) + # local row 0 of every rank is a known nonzero row unless a real id lands there, + # so the no-zero-fill control (out-of-shard ids read local row 0) has something to show + buf_w[0].fill_(0x38) + buf_s[0].fill_(127) + uw, us = w.view(torch.uint8).to(DEV), s.view(torch.uint8).to(DEV) + uniq_d = uniq.to(DEV) + expected = op.reference(torch.searchsorted(uniq_d, ids), w, s, 0, uniq.numel()) + total = torch.zeros_like(expected) + for rank in range(WORLD): + lo, hi = rank * part, (rank + 1) * part + mine = (uniq_d >= lo) & (uniq_d < hi) + buf_w[uniq_d[mine] - lo] = uw[mine] + buf_s[uniq_d[mine] - lo] = us[mine] + tw, ts = buf_w.view(torch.float8_e4m3fn), buf_s.view(torch.float8_e8m0fnu) + out = op(ids, tw, ts, lo, hi) + owned = (ids >= lo) & (ids < hi) + want = expected.masked_fill(~owned.unsqueeze(-1), 0) + op.compare(out, want) + print( + f"\n rank {rank}: [{lo}, {hi}) owns {owned.sum().item()} of {ids.numel()} lookups " + f"(max local row {(uniq_d[mine] - lo).max().item() if mine.any() else -1})" + ) + if 0 < owned.sum().item() < ids.numel(): + for name in _WRONG: + wrong = _WRONG[name](ids, tw, ts, lo, hi) + with pytest.raises(AssertionError): + op.compare(out, wrong) + print( + f" control real rank-{rank} {name}: rejected, mismatch fraction {_mismatch(out, wrong):.6f}" + ) + total += out + op.compare(total, _summed(expected)) + del buf_w, buf_s + torch.cuda.empty_cache() + + +@contextmanager +def _ieee_fp32_matmul(): + """The golden reference runs its fp32 GEMMs with TF32 off.""" + prev_allow = torch.backends.cuda.matmul.allow_tf32 + prev_prec = getattr(torch.backends.cuda.matmul, "fp32_precision", None) + torch.backends.cuda.matmul.allow_tf32 = False + if prev_prec is not None: + torch.backends.cuda.matmul.fp32_precision = "ieee" + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = prev_allow + if prev_prec is not None: + torch.backends.cuda.matmul.fp32_precision = prev_prec + + +def _fp8_act_qdq(x): + """The reference's FP8 1x32 activation quantize-dequantize: amax floored at 1e-4, + scale 2^ceil(log2(amax / 448)) from the fp32 bits, E4M3 of the clamped quotient.""" + xf = x.float().unflatten(-1, (-1, 32)) + v = xf.abs().amax(-1).clamp_min(1e-4) * (1.0 / 448.0) + bits = v.contiguous().view(torch.int32) + e = ((bits >> 23) & 0xFF) - 127 + ((bits & 0x7FFFFF) != 0).to(torch.int32) + sc = ((e + 127) << 23).view(torch.float32) + q = (xf / sc.unsqueeze(-1)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + return (q.float() * sc.unsqueeze(-1)).flatten(-2) + + +@pytest.mark.parametrize("case", _REAL_CASES) +def test_real_rows_reproduce_golden_engram_value(case): + """Gathered rows -> the reference's next step (FP8 1x32 activation quantize, + IEEE fp32 GEMM with the dequantized 32x32-block ``wkv``, bf16) -> the value + slice must match the golden ``engram_value`` of layer 1.""" + _skip_without_real() + from safetensors import safe_open + + (hid,), uniq, w, s = _real_rows((case,)) + ids = torch.searchsorted(uniq, hid.contiguous()).to(DEV) + out = op(ids, w.to(DEV), s.to(DEV), 0, uniq.numel()) + wkv = {} + for name in ("layers.1.engram.wkv.weight", "layers.1.engram.wkv.scale"): + with safe_open(_weight_file(name), framework="pt", device="cpu") as f: + wkv[name] = f.get_tensor(name).to(DEV) + ww, ws = wkv["layers.1.engram.wkv.weight"], wkv["layers.1.engram.wkv.scale"] + n, k = ww.shape + sc = ws.float().repeat_interleave(32, 0)[:n].repeat_interleave(32, 1)[:, :k] + wkv_f32 = ww.float() * sc + golden = torch.load( + os.path.join(_GOLDEN_DIR, case, "layers", "L01.pt"), + map_location="cpu", + mmap=True, + weights_only=True, + ) + golden_value = golden["engram_value"].to(DEV) + + def value_of(rows_bf16): + with _ieee_fp32_matmul(): + kv = torch.matmul(_fp8_act_qdq(rows_bf16.flatten(-2)), wkv_f32.t()).to(torch.bfloat16) + return kv[:, n - golden_value.shape[-1] :] + + value = value_of(out) + exact = 1.0 - _mismatch(value, golden_value) + print(f"\n {case}: engram_value {tuple(value.shape)} bit-exact vs golden {exact:.6f}") + torch.testing.assert_close(value, golden_value) + for name in ("scale_per_16", "e5m2"): + wrong = value_of(_WRONG[name](ids, w, s, 0, uniq.numel())) + with pytest.raises(AssertionError): + torch.testing.assert_close(wrong, golden_value) + rel = ((wrong.float() - golden_value.float()).norm() / golden_value.float().norm()).item() + print( + f" control golden {case} {name}: rejected, rel L2 {rel:.4g}, " + f"bit-exact {1.0 - _mismatch(wrong, golden_value):.6f}" + ) + + +# ─── host_offload at the real size: the checkpoint's own tables in the node-shared pinned file ─── +# +# host_offload keeps ONE copy of both Engram tables per node -- [384,006,168, 256] and +# [384,016,682, 256] float8_e4m3fn plus [rows, 8] E8M0 scale bytes, 202.76 GB -- that +# The target's engram_host_table.py fills from the checkpoint, in one of two backends +# ($TRTLLM_DSV41_ENGRAM_HOST_TABLE_BACKEND): ``vmm`` (the default) allocates CUDA VMM host-NUMA +# slabs, 1 GiB each alternating between the sockets, shares them between the node's processes as +# file descriptors, and hands out CUDA views of that host memory; ``shm`` fills a /dev/shm file +# that every rank maps (MAP_SHARED) and cudaHostRegisters, and hands out pinned CPU views. Either +# way every rank gathers its OWN tokens over ALL rows (row_start 0, row_end rows; no zero fill, no +# collective) in the target's form: the [:, layer, :] view of the [T, 2, 24] hash output (row +# stride 48), with the scale table passed as uint8 bytes. One call carries at most max_num_tokens = 16384 tokens per +# rank (this path has no all-gather, on every member). Weight rows reach byte offsets past 2^36 of their table, scale +# rows past 2^31. + +_HOST_LAYERS = ((1, REAL_ROWS_L1), (14, REAL_ROWS_L14)) +_HOST_T_MAX = 16384 +_COMMON_SEED = 31337 # the ids every rank of the four-rank case gathers +_RANK_DEADLINE_S = 1800 # a rank that has not exited by then is killed (by PID) and the case fails + + +def _offset_line_rows(rows: int): + """Rows on both sides of each 2^k byte line (k = 31..36) of the 256-byte-row weight table and + of the 2^31 line of the 8-byte-row scale table (row 2^28, also the weight table's 2^36 line).""" + lines = [(1 << k) // DIM for k in range(31, 37)] + [(1 << 31) // (DIM // BLOCK)] + return sorted({r + d for r in lines for d in (-2, -1, 0, 1) if 0 <= r + d < rows}) + + +def _token0_rows(rows: int): + """24 edge rows for token 0: the last two rows, row 0, the rows just below and at each + 2^31..2^36 weight byte line, and the next-to-last rows to fill the 24 lookups.""" + edge = [rows - 1, rows - 2, 0] + [ + x for k in range(31, 37) for x in ((1 << k) // DIM - 1, (1 << k) // DIM) + ] + edge += list(range(rows - 3, rows - 3 - (LOOKUPS - len(edge)), -1)) + assert len(edge) == LOOKUPS + return edge + + +def _host_ids(tokens: int, seed: int) -> torch.Tensor: + """[T, 2, 24] int64 ids in the hash output's layout (layer 1 in column 0, layer 14 in column 1), + uniform over each table's rows. Token 0 holds ``_token0_rows`` in a seed-dependent order (so + even T = 1 reaches the last row and every offset line, and consecutive seeds differ), and the + last 4096 rows, the first 1024 and the offset-line rows are also written at random positions. + Drawn on the CPU from a seeded generator, so every process draws the same ids.""" + g = torch.Generator().manual_seed(seed) + ids = torch.empty(tokens, 2, LOOKUPS, dtype=torch.int64) + for li, (_, rows) in enumerate(_HOST_LAYERS): + col = torch.randint(0, rows, (tokens * LOOKUPS,), generator=g) + special = torch.cat( + [ + torch.arange(rows - 4096, rows), + torch.arange(1024), + torch.tensor(_offset_line_rows(rows)), + ] + ) + k = min(col.numel() // 4, special.numel()) + col[torch.randperm(col.numel(), generator=g)[:k]] = special[ + torch.randperm(special.numel(), generator=g)[:k] + ] + col[:LOOKUPS] = torch.tensor(_token0_rows(rows))[torch.randperm(LOOKUPS, generator=g)] + ids[:, li, :] = col.view(tokens, LOOKUPS) + return ids.to(DEV) + + +def _host_ref(ids: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """The reference on a host table without copying the table to the device: the rows ``ids`` + name are read on the CPU through this process's own mapping (``weight`` / ``scale`` are the + CPU byte views of ``_cpu_bytes``) and dequantized by ``_ref``.""" + flat = ids.reshape(-1).cpu() + local = torch.arange(flat.numel(), device=DEV).view(ids.shape) + return op.reference( + local, weight.view(torch.uint8)[flat], scale.view(torch.uint8)[flat], 0, flat.numel() + ) + + +def _host_wrong_wrap( + ids: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, weight_mod: int, scale_mod: int +): + """The gather with truncated byte offsets: weight row r read at byte (256 r) mod ``weight_mod`` + and scale row r at byte (8 r) mod ``scale_mod`` -- what a 32-bit (or signed 32-bit) offset reads.""" + flat = ids.reshape(-1).cpu() + local = torch.arange(flat.numel(), device=DEV).view(ids.shape) + w = weight.view(torch.uint8)[flat % (weight_mod // DIM)] + s = scale.view(torch.uint8)[flat % (scale_mod // (DIM // BLOCK))] + return op.reference(local, w, s, 0, flat.numel()) + + +def _digest(out: torch.Tensor) -> str: + return hashlib.sha256(out.view(torch.int16).cpu().numpy().tobytes()).hexdigest() + + +_CKPT_TABLES = {} + + +def _checkpoint_table(layer: int): + """(weight [rows, 256], scale [rows, 8]) uint8 CPU views of layer ``layer``'s Engram table in + the checkpoint's own safetensors shard, mapped read-only with ``mmap`` at the offsets the + shard's header gives -- independent of engram_host_table's fill and of ``safetensors``.""" + import mmap + import struct + import warnings + + if layer not in _CKPT_TABLES: + views = [] + for kind, dtype in (("weight", "F8_E4M3"), ("scale", "F8_E8M0")): + name = f"layers.{layer}.engram.embed.{kind}" + with open(_weight_file(name), "rb") as f: + (length,) = struct.unpack(" str: + """``vmm`` for a VmmEngramHostTable, ``shm`` for the /dev/shm file's SharedEngramHostTable.""" + return "vmm" if type(table).__name__ == "VmmEngramHostTable" else "shm" + + +def _gather_views(table, layer: int): + """(weight, scale bytes) of ``layer`` exactly as the target hands them to engram_gather: + ``table.weight(layer)`` and ``table.scale(layer)`` viewed as uint8 (CUDA views of host memory + under vmm, pinned CPU views under shm).""" + return table.weight(layer), table.scale(layer).view(torch.uint8) + + +def _cpu_bytes(table, layer: int): + """uint8 CPU tensors ``[rows, 256]`` / ``[rows, 8]`` over ``layer``'s host bytes, read by the CPU + through this process's own host mapping (never through the GPU): the shm views themselves, or + under vmm the host view ``host_buffer`` exposes of the same VMM allocations.""" + if _table_backend(table) == "shm": + return table.weight(layer).view(torch.uint8), table.scale(layer).view(torch.uint8) + out = [] + for kind, cols in (("weight", DIM), ("scale", DIM // BLOCK)): + buf = table.host_buffer(f"layers.{layer}.engram.embed.{kind}") + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) # torch warns that the buffer is read-only + out.append(torch.frombuffer(buf, dtype=torch.uint8).view(-1, cols)) + return tuple(out) + + +def _check_views(table, device) -> None: + """Each layer's gather views are what the backend promises: CUDA views on ``device`` (vmm; + ``check_gather_views`` too) or pinned CPU tensors (shm).""" + device = torch.device(device) + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + for layer, rows in _HOST_LAYERS: + hw, hs = table.weight(layer), table.scale(layer) + assert hw.shape == (rows, DIM) and hw.dtype == torch.float8_e4m3fn, (hw.shape, hw.dtype) + assert hs.shape == (rows, DIM // BLOCK) and hs.dtype == torch.float8_e8m0fnu, ( + hs.shape, + hs.dtype, + ) + if _table_backend(table) == "vmm": + assert hw.is_cuda and hs.is_cuda and hw.device == hs.device == device, ( + hw.device, + device, + ) + else: + assert hw.device.type == "cpu" and hw.is_pinned() and hs.is_pinned() + if _table_backend(table) == "vmm": + table.check_gather_views(device) + + +@contextmanager +def _backend_env(backend: str): + """``$TRTLLM_DSV41_ENGRAM_HOST_TABLE_BACKEND`` for the duration: unset for ``vmm`` (the default + the target runs with), ``shm`` for the file backend; the previous value restored after.""" + key = "TRTLLM_DSV41_ENGRAM_HOST_TABLE_BACKEND" + saved = os.environ.get(key) + if backend == "vmm": + os.environ.pop(key, None) + else: + os.environ[key] = backend + try: + yield + finally: + if saved is None: + os.environ.pop(key, None) + else: + os.environ[key] = saved + + +@pytest.fixture(scope="module", params=["vmm", "shm"]) +def real_host_table(request): + """The checkpoint's own node-shared table, opened through the target's entry point + ``open_shared_engram_table`` as local rank 0 on this process's device, under each backend: + + * ``vmm`` (``$TRTLLM_DSV41_ENGRAM_HOST_TABLE_BACKEND`` unset, the default): import the table + from a process on this node that already serves it, or create the VMM host-NUMA slabs, fill + them from the checkpoint and serve them on the node's abstract socket. Closed afterwards; + the host memory goes when the last process mapping it closes. + * ``shm``: fill, or reuse a complete file whose manifest names the same checkpoint bytes, at + the module's default path; then map and ``cudaHostRegister``. A file this fixture filled is + removed afterwards if no other process holds it; a file that was already there is left in + place, as the module's reuse contract has it. + + The backend variable stays set for the fixture's lifetime, so the four-rank case's processes + inherit it.""" + if not os.path.isfile(os.path.join(_CKPT_DIR, "model.safetensors.index.json")): + pytest.skip(f"checkpoint ({_CKPT_DIR}) not present") + backend = request.param + with _backend_env(backend): + eht = _eht() + assert eht.host_table_backend() == backend + path = eht.default_table_path() + existed = backend == "shm" and eht.describe_shared_engram_table(path)["state"] == "ready" + torch.cuda.synchronize() + alloc0, free0 = torch.cuda.memory_allocated(), torch.cuda.mem_get_info()[0] + t0 = time.monotonic() + table = eht.open_shared_engram_table(_CKPT_DIR, local_rank=0, path=path, device=DEV) + seconds = time.monotonic() - t0 + torch.cuda.synchronize() + assert _table_backend(table) == backend, type(table).__name__ + info = { + "backend": backend, + "path": table.path, + "alloc_delta": torch.cuda.memory_allocated() - alloc0, + "free_delta": free0 - torch.cuda.mem_get_info()[0], + "fill": table.manifest.get("fill") or {}, + } + fill = info["fill"] + if backend == "shm": + info["existed"] = existed + info["inode"] = os.stat(path).st_ino + print( + f"\n shm host table {path} (inode {info['inode']}): {'reused' if existed else 'filled'}, opened in " + f"{seconds:.2f} s (registration {table.register_seconds:.2f} s, read_only {table.read_only}); " + f"{table.nbytes} table bytes; fill record: {fill.get('seconds', 0):.1f} s, numa {fill.get('numa')}, " + f"huge-page bytes {fill.get('huge_page_bytes')}; device memory_allocated delta " + f"{info['alloc_delta']} B, " + f"cudaMemGetInfo delta {info['free_delta'] / 2**20:.1f} MiB" + ) + else: + info["existed"] = not table.created + print( + f"\n vmm host table: {table.describe()}; opened in {seconds:.2f} s; {table.nbytes} table bytes in " + f"{table.allocated_bytes} allocated; fill record: pid {fill.get('pid')}, alloc " + f"{fill.get('alloc_seconds', 0):.1f} s, " + f"fill {fill.get('seconds', 0):.1f} s, granularity {fill.get('granularity')}, numa bytes " + f"{fill.get('numa_bytes')}; device memory_allocated delta {info['alloc_delta']} B, cudaMemGetInfo " + f"delta {info['free_delta'] / 2**20:.1f} MiB" + ) + try: + yield table, info + finally: + torch.cuda.synchronize() + table.close() + if backend == "shm" and not existed: + eht.remove_shared_engram_table(path) + + +# ─── the vmm table created by a process that sees ONE GPU ─── +# +# The driver refuses cuMemCreate(HOST_NUMA, node) with CUDA_ERROR_INVALID_VALUE for a CPU socket that +# none of the process's visible GPUs is attached to. The vmm backend +# probes the sockets first (_vmm_usable_nodes) and places every slab on the ones it may allocate +# on: both with GPUs of both sockets visible, else its GPUs' socket; a process that sees only a GPU +# of the other socket imports the same slabs and reads them over the socket link. The cases below +# run that in separate processes, each with exactly one GPU visible (CUDA_VISIBLE_DEVICES), because +# visibility is fixed when CUDA initializes. + +_ONE_GPU_CAPTURE_T = ( + 1, + 6, + 8, + 32, + 48, + 128, + 192, + 256, + 768, + 1536, +) # G and 6 G for G in 1, 8, 32, 128, 256 +_ONE_GPU_WAIT_S = 900 # how long a creator keeps serving while the importer runs + + +def _vmm_create_probe() -> dict: + """In this process (its current device, with whatever GPUs are visible): per CPU NUMA node, + whether ``cuMemCreate`` of one allocation-granularity host-NUMA slab, exactly as the vmm backend + asks for it (``_vmm_prop``), succeeds; the device's own NUMA node; and the sockets the backend's + own ``_vmm_usable_nodes`` finds usable.""" + from cuda.bindings import driver as cu + + eht = _eht() + torch.zeros(1, device="cuda") + torch.cuda.synchronize() # the primary context is current + dev = torch.cuda.current_device() + err, numa = cu.cuDeviceGetAttribute(cu.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_HOST_NUMA_ID, dev) + assert err == cu.CUresult.CUDA_SUCCESS, err + out = { + "visible": os.environ.get("CUDA_VISIBLE_DEVICES"), + "device_numa": int(numa), + "create": {}, + } + nodes = eht._cpu_memory_nodes() + for node in nodes: + gran = eht._vmm_granularity([node]) + err, handle = cu.cuMemCreate(gran, eht._vmm_prop(node), 0) + out["create"][str(node)] = err.name + if err == cu.CUresult.CUDA_SUCCESS: + cu.cuMemRelease(handle) + usable, refused = eht._vmm_usable_nodes(nodes) + out["usable"] = usable + out["refused"] = {str(n): why for n, why in refused.items()} + return out + + +def _one_gpu_create_main(name: str, out_dir: str) -> None: + """The creator, run as ``python --vmm-one-gpu create --name --out `` with one + GPU visible: create the table under ``name`` (no process serves it), check the slab placement and + the checkpoint's bytes, run every T 1..16384 through both layers against the reference and the + captured steps, write ``create.json``, then keep serving until ``/done`` appears (the + importer's turn) and close. Any failed check raises (exit code 1).""" + eht = _eht() + dev = torch.device("cuda", 0) + torch.cuda.set_device(dev) + torch.zeros(1, device=dev) + torch.cuda.synchronize() + alloc0 = torch.cuda.memory_allocated() + t0 = time.monotonic() + table = eht.open_vmm_engram_table(_CKPT_DIR, local_rank=0, device=dev, wait_s=300.0, name=name) + open_s = time.monotonic() - t0 + torch.cuda.synchronize() + alloc_delta = torch.cuda.memory_allocated() - alloc0 + probe = _vmm_create_probe() + own = probe["device_numa"] + fill = table.manifest.get("fill") or {} + numa = table.numa_bytes() + assert table.created and table.path == name and alloc_delta == 0, ( + table.created, + table.path, + alloc_delta, + ) + assert numa == {own: table.allocated_bytes}, (numa, own) # every slab on this GPU's socket + assert fill.get("device_host_nodes") == {"0": own}, fill.get("device_host_nodes") + assert fill.get("numa_nodes_refused") == { + str(n): "cuMemCreate: CUDA_ERROR_INVALID_VALUE" for n in (0, 1) if n != own + }, fill.get("numa_nodes_refused") + _check_views(table, dev) + checked, bad = table.verify(full=False) + assert checked > 0 and not bad, (checked, bad) + result = { + "visible": os.environ.get("CUDA_VISIBLE_DEVICES"), + "pid": os.getpid(), + "device_numa": own, + "open_s": round(open_s, 2), + "alloc_s": round(fill.get("alloc_seconds", 0.0), 2), + "fill_s": round(fill.get("seconds", 0.0), 2), + "numa_bytes": {str(k): v for k, v in numa.items()}, + "allocated_bytes": table.allocated_bytes, + "slabs": len(table._slabs), + "refused": fill.get("numa_nodes_refused"), + "crc_chunks_checked": checked, + "sweep_s": {}, + "digest": {}, + } + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.embedding.engram_gather import ( # noqa: E501 + _engram_gather_kernel as fn, + ) + + tabs = [_gather_views(table, layer) + (rows,) for layer, rows in _HOST_LAYERS] + cpu = [_cpu_bytes(table, layer) for layer, _ in _HOST_LAYERS] + # the table's bytes on far, edge and random rows equal the checkpoint's own (independent mmap) + for li, (layer, rows) in enumerate(_HOST_LAYERS): + ck_w, ck_s = _checkpoint_table(layer) + g = torch.Generator().manual_seed(7300 + layer) + picked = torch.cat( + [ + torch.arange(rows - 4096, rows), + torch.arange(1024), + torch.tensor(_offset_line_rows(rows)), + torch.randint(0, rows, (8192,), generator=g), + ] + ) + assert torch.equal(cpu[li][0][picked], ck_w[picked]) and torch.equal( + cpu[li][1][picked], ck_s[picked] + ), layer + # every T 1..16384: the 16384-token call == reference, every prefix its first rows, the guard rows untouched + hashes = _host_ids(_HOST_T_MAX, seed=8100) + buf = torch.empty(_HOST_T_MAX + _GUARD, LOOKUPS, DIM, dtype=torch.bfloat16, device=dev) + bits = buf.view(torch.int16) + for li, (w, s, rows) in enumerate(tabs): + ids = hashes[:, li, :] + full = op(ids, w, s, 0, rows) + op.compare(full, _host_ref(ids, *cpu[li])) + wrong = _host_wrong_wrap(ids, *cpu[li], 1 << 32, 1 << 64) + assert _mismatch(full, wrong) > 0.5, "the 32-bit weight-offset control was not rejected" + result["digest"][str(_HOST_LAYERS[li][0])] = _digest(full) + full_bits = full.view(torch.int16) + bits.fill_(_POISON) + bad_t = torch.zeros((), dtype=torch.int64, device=dev) + over = torch.zeros((), dtype=torch.int64, device=dev) + t1 = time.monotonic() + before, settled = _jit_variants(fn), None + for r in range(1, _HOST_T_MAX + 1): + op(ids[:r], w, s, 0, rows, out=buf[:r]) + bad_t += (bits[:r] != full_bits[:r]).sum() + over += (bits[r : r + _GUARD] != _POISON).sum() + if r == 2: + settled = _jit_variants(fn) + torch.cuda.synchronize() + assert bad_t.item() == 0 and over.item() == 0, (bad_t.item(), over.item()) + assert _jit_variants(fn) == settled and settled - before <= 2, ( + before, + settled, + _jit_variants(fn), + ) + result["sweep_s"][str(_HOST_LAYERS[li][0])] = round(time.monotonic() - t1, 1) + del buf + # capture: both layers in one graph at G and 6 G counts, 3 replays on new ids each + stale_min = 1.0 + for tokens in _ONE_GPU_CAPTURE_T: + static = _host_ids(tokens, seed=0) + views = [static[:, 0, :], static[:, 1, :]] + capture_ids = static.clone() + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for v, (w, s, rows) in zip(views, tabs): + op(v, w, s, 0, rows) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = [op(v, w, s, 0, rows) for v, (w, s, rows) in zip(views, tabs)] + for rep in (1, 2, 3): + static.copy_(_host_ids(tokens, seed=rep)) + for o in outs: + o.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + for li, (v, (w, s, rows)) in enumerate(zip(views, tabs)): + assert torch.equal( + outs[li].view(torch.int16), op(v, w, s, 0, rows).view(torch.int16) + ) + expected = _host_ref(v, *cpu[li]) + op.compare(outs[li], expected) + stale_min = min( + stale_min, _mismatch(expected, _host_ref(capture_ids[:, li, :], *cpu[li])) + ) + del graph, outs + assert stale_min > 0.5, stale_min + result["captured_T"] = list(_ONE_GPU_CAPTURE_T) + result["stale_min"] = round(stale_min, 6) + with open(os.path.join(out_dir, "create.json"), "w") as f: + json.dump(result, f) + with open(os.path.join(out_dir, "serving"), "w") as f: + f.write(str(os.getpid())) + deadline = time.monotonic() + _ONE_GPU_WAIT_S + while not os.path.exists(os.path.join(out_dir, "done")) and time.monotonic() < deadline: + time.sleep(0.5) + del tabs, cpu + torch.cuda.synchronize() + table.close() + assert os.path.exists(os.path.join(out_dir, "done")), ( + "the importer did not finish while this process served" + ) + + +def _one_gpu_import_main(name: str, out_dir: str) -> None: + """The importer, run as ``python --vmm-one-gpu import --name --out `` with one GPU + of the OTHER socket visible, while the creator serves: import the table, check that no slab is on this + GPU's socket (it reads every row over the socket link), and gather the creator's 16384 tokens on both + layers (bit-exact to the reference and to the creator's digest) and the captured steps.""" + eht = _eht() + dev = torch.device("cuda", 0) + torch.cuda.set_device(dev) + torch.zeros(1, device=dev) + torch.cuda.synchronize() + alloc0 = torch.cuda.memory_allocated() + t0 = time.monotonic() + table = eht.open_vmm_engram_table(_CKPT_DIR, local_rank=1, device=dev, wait_s=120.0, name=name) + open_s = time.monotonic() - t0 + torch.cuda.synchronize() + probe = _vmm_create_probe() + numa = table.numa_bytes() + assert not table.created and table.path == name and torch.cuda.memory_allocated() == alloc0 + assert probe["device_numa"] not in numa and len(numa) == 1, (probe["device_numa"], numa) + _check_views(table, dev) + tabs = [_gather_views(table, layer) + (rows,) for layer, rows in _HOST_LAYERS] + cpu = [_cpu_bytes(table, layer) for layer, _ in _HOST_LAYERS] + hashes = _host_ids(_HOST_T_MAX, seed=8100) + digest = {} + for li, (w, s, rows) in enumerate(tabs): + out = op(hashes[:, li, :], w, s, 0, rows) + op.compare(out, _host_ref(hashes[:, li, :], *cpu[li])) + digest[str(_HOST_LAYERS[li][0])] = _digest(out) + for tokens in (1, 48, 1536): + static = _host_ids(tokens, seed=100) + views = [static[:, 0, :], static[:, 1, :]] + for v, (w, s, rows) in zip(views, tabs): + op(v, w, s, 0, rows) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = [op(v, w, s, 0, rows) for v, (w, s, rows) in zip(views, tabs)] + for rep in (1, 2): + static.copy_(_host_ids(tokens, seed=1000 * rep)) + graph.replay() + torch.cuda.synchronize() + for li, (v, (w, s, rows)) in enumerate(zip(views, tabs)): + assert torch.equal( + outs[li].view(torch.int16), op(v, w, s, 0, rows).view(torch.int16) + ) + op.compare(outs[li], _host_ref(v, *cpu[li])) + del graph, outs + result = { + "visible": os.environ.get("CUDA_VISIBLE_DEVICES"), + "pid": os.getpid(), + "device_numa": probe["device_numa"], + "numa_bytes": {str(k): v for k, v in numa.items()}, + "open_s": round(open_s, 2), + "digest": digest, + "captured_T": [1, 48, 1536], + } + del tabs, cpu + torch.cuda.synchronize() + table.close() + with open(os.path.join(out_dir, "import.json"), "w") as f: + json.dump(result, f) + + +def _spawn(args, env, log_path): + with open(log_path, "w") as log: + return subprocess.Popen( + [sys.executable, str(Path(__file__).resolve())] + args, + env=env, + stdout=log, + stderr=subprocess.STDOUT, + ) + + +def _wait_all(procs, logs, deadline_s): + """Wait for ``procs`` up to ``deadline_s``; kill (by PID) any still running; fail with the logs' tails.""" + deadline = time.monotonic() + deadline_s + while any(p.poll() is None for p in procs) and time.monotonic() < deadline: + time.sleep(1) + wedged = [i for i, p in enumerate(procs) if p.poll() is None] + for p in procs: + if p.poll() is None: + p.kill() + p.wait() + tails = "".join(f"\n--- {log.name} ---\n{log.read_text()[-3000:]}" for log in logs) + assert not wedged, f"processes {wedged} did not exit within {deadline_s} s{tails}" + assert [p.returncode for p in procs] == [0] * len(procs), ( + f"exit codes {[p.returncode for p in procs]}{tails}" + ) + + +def test_vmm_table_created_by_a_one_gpu_process(tmp_path): + """The vmm table with one GPU visible. + + (1) Probe. In this process (four GPUs visible) ``cuMemCreate`` of a host-NUMA slab, as the backend + asks for it, succeeds on both CPU NUMA nodes, and ``_vmm_usable_nodes`` finds both usable. In a + process that sees one GPU, run once per GPU: it succeeds on that GPU's own node, returns + CUDA_ERROR_INVALID_VALUE on the other socket's node, and ``_vmm_usable_nodes`` finds only the own + node usable (the other refused with that error). + (2) Create and read, once with a GPU of each socket visible (the first and the last visible GPU): + a process that sees only that GPU creates the table under a name no process serves; it succeeds, + every slab lands on that GPU's socket, the manifest records the other socket as refused, the views + are CUDA views with no device memory behind them, sampled chunk CRCs re-verify, and the bytes on + far, edge and random rows equal the checkpoint's. Through that table, both layers: the 16384-token + call equals the reference (and rejects the 32-bit offset control), every prefix r = 1..16384 returns + exactly its first r rows with the guard rows untouched and no new Triton variant after r = 2, and + both layers in one graph captured at T = G and 6 G for G in 1, 8, 32, 128, 256 replay 3 times on new + ids equal to eager and to the reference. + (3) While it serves, a process that sees only a GPU of the OTHER socket imports the table (no slab + on its own socket) and gathers the same 16384 tokens: bit-exact to the reference and the same + SHA-256 as the creator; captured at T 1, 48, 1536, replay == eager == reference. + A process that has not exited within the deadline is killed by PID and the case fails.""" + if not os.path.isfile(os.path.join(_CKPT_DIR, "model.safetensors.index.json")): + pytest.skip(f"checkpoint ({_CKPT_DIR}) not present") + visible = [d for d in os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",") if d.strip()] + visible = visible or [str(i) for i in range(torch.cuda.device_count())] + if len(visible) < WORLD: + pytest.skip(f"needs {WORLD} visible GPUs, one process per GPU; {len(visible)} visible") + here = _vmm_create_probe() + if len(here["usable"]) < 2: + pytest.skip( + f"needs visible GPUs on two NUMA nodes (a creator on one socket, an importer on the " + f"other); the visible GPUs reach NUMA nodes {here['usable']} only" + ) + print( + f"\n {len(visible)} GPUs visible (current device NUMA node {here['device_numa']}): cuMemCreate per node " + f"{here['create']}; usable {here['usable']}, refused {here['refused']}" + ) + assert set(here["create"].values()) == {"CUDA_SUCCESS"} and len(here["create"]) == 2, here + assert here["usable"] == [0, 1] and not here["refused"], here + repo_root = Path(__file__).resolve().parents[8] + base_env = dict(os.environ) + # the repo root (tensorrt_llm), tests/unittest (the utils package this file imports), tests + # (test_common, which utils imports; the unittest conftest adds it for pytest) and + # tests/unittest/_torch/modeling_v2 (_validating, which pytest reaches via __extra_import_path__) + base_env["PYTHONPATH"] = os.pathsep.join( + p + for p in ( + str(repo_root), + str(Path(__file__).resolve().parents[6]), + str(Path(__file__).resolve().parents[7]), + str(Path(__file__).resolve().parents[4]), + base_env.get("PYTHONPATH", ""), + ) + if p + ) + base_env.pop("TRTLLM_DSV41_ENGRAM_HOST_TABLE_BACKEND", None) + # (1) the probe, one process per GPU + procs, logs = [], [] + for i, dev in enumerate(visible[:WORLD]): + logs.append(tmp_path / f"probe{i}.log") + procs.append( + _spawn( + ["--vmm-create-probe", str(tmp_path / f"probe{i}.json")], + dict(base_env, CUDA_VISIBLE_DEVICES=dev), + logs[-1], + ) + ) + _wait_all(procs, logs, _RANK_DEADLINE_S) + node_of = {} + for i, dev in enumerate(visible[:WORLD]): + r = json.loads((tmp_path / f"probe{i}.json").read_text()) + own = r["device_numa"] + node_of[dev] = own + print( + f" one GPU visible ({r['visible']}, NUMA node {own}): cuMemCreate per node {r['create']}; usable " + f"{r['usable']}, refused {r['refused']}" + ) + assert r["create"][str(own)] == "CUDA_SUCCESS", r + assert all( + v == "CUDA_ERROR_INVALID_VALUE" for n, v in r["create"].items() if n != str(own) + ), r + assert r["usable"] == [own] and r["refused"] == { + str(1 - own): "cuMemCreate: CUDA_ERROR_INVALID_VALUE" + }, r + # (2) + (3): a creator on each socket, an importer on the other + creators = [visible[0], visible[WORLD - 1]] + assert node_of[creators[0]] != node_of[creators[1]], node_of + for k, cdev in enumerate(creators): + idev = next(d for d in visible[:WORLD] if node_of[d] != node_of[cdev]) + name = f"cg-engram-gather-onegpu-{os.getpid()}-{k}" + run = tmp_path / f"run{k}" + run.mkdir() + clog, ilog = run / "create.log", run / "import.log" + creator = _spawn( + ["--vmm-one-gpu", "create", "--name", name, "--out", str(run)], + dict(base_env, CUDA_VISIBLE_DEVICES=cdev), + clog, + ) + deadline = time.monotonic() + _RANK_DEADLINE_S + while ( + not (run / "serving").exists() + and creator.poll() is None + and time.monotonic() < deadline + ): + time.sleep(1) + try: + if (run / "serving").exists(): + importer = _spawn( + ["--vmm-one-gpu", "import", "--name", name, "--out", str(run)], + dict(base_env, CUDA_VISIBLE_DEVICES=idev), + ilog, + ) + _wait_all([importer], [ilog], 600) + finally: + (run / "done").write_text("done") # the creator stops serving and exits either way + _wait_all( + [creator], + [clog] + ([ilog] if ilog.exists() else []), + max(60.0, deadline - time.monotonic()), + ) + c = json.loads((run / "create.json").read_text()) + i_ = json.loads((run / "import.json").read_text()) + print( + f" one GPU visible ({cdev}, NUMA node {c['device_numa']}) created the table in {c['open_s']} s (alloc " + f"{c['alloc_s']} s, fill {c['fill_s']} s): {c['slabs']} slabs, {c['allocated_bytes']} bytes, NUMA bytes " + f"{c['numa_bytes']}, refused {c['refused']}; {c['crc_chunks_checked']} chunk CRCs re-verified; every T " + f"1..{_HOST_T_MAX} bit-exact (sweep {c['sweep_s']} s per layer); captured at T {c['captured_T']}, replay " + f"== eager == reference (min fraction differing from the capture-time ids {c['stale_min']})" + ) + print( + f" one GPU visible ({idev}, NUMA node {i_['device_numa']}) imported it in {i_['open_s']} s (slabs on " + f"{list(i_['numa_bytes'])}): {_HOST_T_MAX} tokens sha256 L1 {i_['digest']['1'][:16]} L14 " + f"{i_['digest']['14'][:16]} == the creator's, bit-exact; captured at T {i_['captured_T']}" + ) + assert c["numa_bytes"] == {str(c["device_numa"]): c["allocated_bytes"]}, c + assert i_["device_numa"] != c["device_numa"] and list(i_["numa_bytes"]) == [ + str(c["device_numa"]) + ], (c, i_) + assert i_["digest"] == c["digest"], (i_["digest"], c["digest"]) + + +def test_real_host_table_far_and_random_rows_equal_the_checkpoint(real_host_table): + """Both of the checkpoint's tables through the node's shared host copy (each backend), in the + target's call form (``[:, layer, :]`` view of a ``[T, 2, 24]`` hash buffer, uint8 scale view, + rows ``[0, rows)``): the last 4096 rows (the deepest, 384,006,167 / 384,016,681, start 98.3 GB + into the weight table), the first 1024, the rows on both sides of the 2^31 .. 2^36 byte lines of + the weight table and of the 2^31 line of the scale table, the four-rank shard boundaries and 8192 + random rows are gathered and compared bit for bit with the dequantization of the SAME rows read + from the checkpoint's safetensors shards through an independent mmap; the table's bytes on those + rows, read by the CPU, equal the checkpoint's. The views are CUDA views on this device (vmm) or + pinned CPU tensors (shm); no device memory holds the table. Controls, each rejected by the same + gate: a 32-bit weight byte offset, a signed 32-bit scale byte offset, the scale applied per 16 + channels, E5M2 decoding.""" + table, info = real_host_table + assert table.layers == [1, 14] + assert ( + table.nbytes + == sum(rows * (DIM + DIM // BLOCK) for _, rows in _HOST_LAYERS) + == 202_758_032_400 + ) + assert info["alloc_delta"] == 0 and info["free_delta"] < (1 << 30), ( + info + ) # no device copy of 202.76 GB + _check_views(table, DEV) + if info["backend"] == "vmm": + print( + f"\n vmm: {table.allocated_bytes} bytes allocated over NUMA nodes {table.numa_bytes()}" + ) + for li, (layer, rows) in enumerate(_HOST_LAYERS): + hw, hs = _gather_views(table, layer) + cw, cs = _cpu_bytes(table, layer) + ck_w, ck_s = _checkpoint_table(layer) + assert ck_w.shape == (rows, DIM) and ck_s.shape == (rows, DIM // BLOCK) + part = (rows + WORLD - 1) // WORLD + g = torch.Generator().manual_seed(7100 + layer) + picked = torch.cat( + [ + torch.arange(rows - 4096, rows), + torch.arange(1024), + torch.tensor(_offset_line_rows(rows)), + torch.tensor([b + d for b in (part, 2 * part, 3 * part) for d in (-2, -1, 0, 1)]), + torch.randint(0, rows, (8192,), generator=g), + ] + ) + picked = torch.cat( + [picked, torch.randint(0, rows, (-picked.numel() % LOOKUPS,), generator=g)] + ) + tokens = picked.numel() // LOOKUPS + hashes = torch.randint( + 0, + REAL_ROWS_L1, + (tokens, 2, LOOKUPS), + dtype=torch.int64, + device=DEV, + generator=_gen(7200 + layer), + ) + hashes[:, li, :] = picked.view(tokens, LOOKUPS).to(DEV) + ids = hashes[:, li, :] + assert ids.stride() == (2 * LOOKUPS, 1) + out = op(ids, hw, hs, 0, rows) + t0 = time.monotonic() + ck_rows_w, ck_rows_s = ck_w[picked], ck_s[picked] # page faults into the checkpoint shards + read_s = time.monotonic() - t0 + local = torch.arange(picked.numel(), device=DEV).view(tokens, LOOKUPS) + expected = op.reference(local, ck_rows_w, ck_rows_s, 0, picked.numel()) + op.compare(out, expected) + assert torch.equal(cw[picked], ck_rows_w) and torch.equal(cs[picked], ck_rows_s) + deepest = int(picked.max()) + print( + f"\n {info['backend']} layer {layer}: {picked.numel()} rows ({tokens} tokens x 24) incl. the last " + "4096, the 2^31..2^36 " + f"byte lines and 8192 random rows == checkpoint (mmap, {read_s:.1f} s), bit-exact; deepest row " + f"{deepest}: weight byte offset {deepest * DIM}, scale byte offset {deepest * (DIM // BLOCK)}" + ) + for label, weight_mod, scale_mod in ( + ("32-bit weight byte offset", 1 << 32, 1 << 64), + ("signed 32-bit scale byte offset", 1 << 64, 1 << 31), + ): + wrong = _host_wrong_wrap(ids, cw, cs, weight_mod, scale_mod) + with pytest.raises(AssertionError): + op.compare(out, wrong) + print( + f" control {info['backend']} L{layer} {label}: rejected, mismatch fraction {_mismatch(out, wrong):.6f}" + ) + _assert_controls_fail( + out, + local, + ck_rows_w, + ck_rows_s, + 0, + picked.numel(), + ("scale_per_16", "e5m2"), + f"real host {info['backend']} L{layer}", + ) + + +def test_real_host_table_every_row_count_is_a_prefix_of_16384(real_host_table): + """The prefix rule through the real mapping, per layer, in the target's form: the 16384-token call + (``_host_ids``: far rows and every offset line included) is bit-exact to the reference built from + the table's own rows (read on the CPU through the mapping; the previous case ties those bytes to + the checkpoint) and rejects the 32-bit offset control, under each backend; then every prefix ``ids[:r]``, + r = 1..16384, returns exactly the first r rows of that call into a row-prefix view of one + poisoned buffer whose 64 rows past the prefix keep the poison. No Triton variant appears + after r = 2 (n_rows = 24 r is divisible by 16 exactly for even r).""" + table, info = real_host_table + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.embedding.engram_gather import ( # noqa: E501 + _engram_gather_kernel as fn, + ) + + hashes = _host_ids(_HOST_T_MAX, seed=8100) + buf = torch.empty(_HOST_T_MAX + _GUARD, LOOKUPS, DIM, dtype=torch.bfloat16, device=DEV) + bits = buf.view(torch.int16) + for li, (layer, rows) in enumerate(_HOST_LAYERS): + ids = hashes[:, li, :] + hw, hs = _gather_views(table, layer) + cw, cs = _cpu_bytes(table, layer) + full = op(ids, hw, hs, 0, rows) + op.compare(full, _host_ref(ids, cw, cs)) + with pytest.raises(AssertionError): + op.compare(full, _host_wrong_wrap(ids, cw, cs, 1 << 32, 1 << 64)) + full_bits = full.view(torch.int16) + bits.fill_(_POISON) + bad = torch.zeros((), dtype=torch.int64, device=DEV) + over = torch.zeros((), dtype=torch.int64, device=DEV) + before = _jit_variants(fn) + settled = None + t0 = time.monotonic() + for r in range(1, _HOST_T_MAX + 1): + res = op(ids[:r], hw, hs, 0, rows, out=buf[:r]) + assert res.data_ptr() == buf.data_ptr() + bad += (bits[:r] != full_bits[:r]).sum() + over += (bits[r : r + _GUARD] != _POISON).sum() + if r == 2: + settled = _jit_variants(fn) + torch.cuda.synchronize() + after = _jit_variants(fn) + far = int(ids.max().item()) + print( + f"\n {info['backend']} layer {layer}: T {_HOST_T_MAX} bit-exact (deepest row {far}, byte offset " + f"{far * DIM}); r = 1..{_HOST_T_MAX} " + f"in {time.monotonic() - t0:.1f} s: {bad.item()} elements differ from the {_HOST_T_MAX}-row result, " + f"{over.item()} guard elements written; kernel variants {before} before, {settled} after r = 2, {after} " + f"after r = {_HOST_T_MAX}" + ) + assert bad.item() == 0 and over.item() == 0, (bad.item(), over.item()) + assert after == settled and settled - before <= 2, (before, settled, after) + del buf + + +@pytest.mark.parametrize("tokens", _CAPTURE_T) +def test_real_host_table_cuda_graph_at_every_capture_batch(real_host_table, tokens): + """The target's captured Engram step at every T of the capture set (Gset graph batches G and 6 G + for the DSpark verify step, 82 counts up to 1536), under each backend: both layers' gathers through the real + mapping in one graph, ``out`` allocated inside the capture, after an eager warm-up on a side stream. Three replays + on new ids (far rows and every offset line in token 0) written in place into the static + ``[T, 2, 24]`` buffer: each replay equals the eager call bit for bit and the reference, and + differs from the capture-time ids' rows (it read the new ids).""" + table, info = real_host_table + tabs = [_gather_views(table, layer) + (rows,) for layer, rows in _HOST_LAYERS] + cpu = [_cpu_bytes(table, layer) for layer, _ in _HOST_LAYERS] + static = _host_ids(tokens, seed=0) + views = [static[:, 0, :], static[:, 1, :]] + capture_ids = static.clone() + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for v, (w, s, rows) in zip(views, tabs): + op(v, w, s, 0, rows) # engine warm-up at this shape + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = [op(v, w, s, 0, rows) for v, (w, s, rows) in zip(views, tabs)] + for rep in (1, 2, 3): + static.copy_(_host_ids(tokens, seed=rep)) + for o in outs: + o.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + for li, (v, (w, s, rows)) in enumerate(zip(views, tabs)): + eager = op(v, w, s, 0, rows) + assert torch.equal(outs[li].view(torch.int16), eager.view(torch.int16)), ( + f"layer column {li}: replay != eager" + ) + expected = _host_ref(v, *cpu[li]) + op.compare(outs[li], expected) + stale = _host_ref(capture_ids[:, li, :], *cpu[li]) + assert not torch.equal(expected.view(torch.int16), stale.view(torch.int16)), ( + "the replay did not read new ids" + ) + if rep == 1 and li == 0: + print( + f"\n {info['backend']} T={tokens}: 2 layers x 3 replays == eager == reference bit for bit; " + "layer 1 differs from " + f"the capture-time ids' rows on {_mismatch(outs[li], stale):.6f}" + ) + del graph, outs + + +def _host_rank_main(rank: int, out_dir: str) -> None: + """One engine-like rank of the four-rank case, run as ``python --engram-host-rank + --out ``: attach through ``open_shared_engram_table`` on its own GPU, gather, and + write what it saw to ``/rank.json``. Any failed check raises (exit code 1).""" + eht = _eht() + torch.cuda.set_device(rank) + dev = torch.device("cuda", rank) + torch.zeros(1, device=dev) + torch.cuda.synchronize() + alloc0 = torch.cuda.memory_allocated() + path = eht.default_table_path() + backend = eht.host_table_backend() + t0 = time.monotonic() + table = eht.open_shared_engram_table( + _CKPT_DIR, local_rank=rank, path=path, device=dev, wait_s=600.0 + ) + open_s = time.monotonic() - t0 + torch.cuda.synchronize() + alloc_delta = torch.cuda.memory_allocated() - alloc0 + assert _table_backend(table) == backend, (type(table).__name__, backend) + _check_views(table, dev) + assert alloc_delta == 0, alloc_delta + tabs = [_gather_views(table, layer) + (rows,) for layer, rows in _HOST_LAYERS] + cpu = [_cpu_bytes(table, layer) for layer, _ in _HOST_LAYERS] + common = _host_ids(_HOST_T_MAX, seed=_COMMON_SEED) + digest = {} + for li, (w, s, rows) in enumerate(tabs): + out = op(common[:, li, :], w, s, 0, rows) + op.compare(out, _host_ref(common[:, li, :], *cpu[li])) + digest[str(_HOST_LAYERS[li][0])] = _digest(out) + own_tokens = 4096 - 7 * rank # the ranks' own, unequal token counts + own = _host_ids(own_tokens, seed=9000 + rank) + for li, (w, s, rows) in enumerate(tabs): + op.compare(op(own[:, li, :], w, s, 0, rows), _host_ref(own[:, li, :], *cpu[li])) + captured = [] + for tokens in (1, 48, 192, 1536): + static = _host_ids(tokens, seed=100 + rank) + views = [static[:, 0, :], static[:, 1, :]] + for v, (w, s, rows) in zip(views, tabs): + op(v, w, s, 0, rows) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = [op(v, w, s, 0, rows) for v, (w, s, rows) in zip(views, tabs)] + for rep in (1, 2): + static.copy_(_host_ids(tokens, seed=1000 * rep + rank)) + graph.replay() + torch.cuda.synchronize() + for li, (v, (w, s, rows)) in enumerate(zip(views, tabs)): + assert torch.equal( + outs[li].view(torch.int16), op(v, w, s, 0, rows).view(torch.int16) + ) + op.compare(outs[li], _host_ref(v, *cpu[li])) + captured.append(tokens) + del graph, outs + props = torch.cuda.get_device_properties(dev) + fill = table.manifest.get("fill") or {} + result = { + "rank": rank, + "pid": os.getpid(), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "device_uuid": str(getattr(props, "uuid", rank)), + "backend": backend, + "path": table.path, + "inode": os.stat(table.path).st_ino if backend == "shm" else None, + "created": bool(getattr(table, "created", False)), + "fill_pid": fill.get("pid"), + "fill_time": fill.get("unix_time"), + "weight_l1_address": tabs[0][0].data_ptr(), + "read_only": table.read_only, + "open_s": round(open_s, 2), + "register_s": round(table.register_seconds, 2), + "alloc_delta": alloc_delta, + "digest": digest, + "own_tokens": own_tokens, + "captured_T": captured, + } + del tabs, cpu, w, s + torch.cuda.synchronize() + table.close() # the file stays: the pytest process that launched this rank still holds it + with open(os.path.join(out_dir, f"rank{rank}.json"), "w") as f: + json.dump(result, f) + + +def test_real_host_table_vmm_slabs_span_both_sockets(real_host_table): + """Under vmm the table's creator places its slabs on every CPU NUMA node its visible GPUs let it + allocate on (the driver refuses a host-NUMA allocation on a socket none of them is attached + to). A creator that saw GPUs of both sockets therefore splits the table about evenly over both, + which this checks; a table created from one socket cannot, and the case skips.""" + table, info = real_host_table + if info["backend"] != "vmm": + pytest.skip("the shm backend's placement is the kernel's page policy, not per-socket slabs") + eht = _eht() + nodes = list(eht._cpu_memory_nodes()) + refused = info["fill"].get("numa_nodes_refused") or {} + usable = [n for n in nodes if str(n) not in refused] + if len(usable) < 2: + pytest.skip( + f"the table's creator could allocate on NUMA nodes {usable} of {nodes} only (refused " + f"{refused}); spanning both sockets needs visible GPUs on two NUMA nodes" + ) + numa = table.numa_bytes() + assert len(numa) == 2 and min(numa.values()) > 0.45 * table.allocated_bytes, numa + + +def test_real_host_table_four_ranks_read_the_same_mapping(real_host_table, tmp_path): + """One process per GPU on four GPUs, as the engine's ranks run, under each backend: each opens the + SAME node copy through ``open_shared_engram_table`` (local rank 0 fills or reuses it, the others + attach; here every rank reuses the table this process holds: under vmm each imports the + fds this process serves, under shm each maps the same /dev/shm file), with no device memory for + the table, and gathers + (1) the same 16384-token ids on both layers -- every rank's result has the same SHA-256 as this + process's own call, which is bit-exact to the reference; (2) its own 4096 - 7 rank tokens, + bit-exact; (3) both layers captured in one graph at T 1, 48, 192 and 1536, two replays on new ids, + replay == eager == reference. A rank that has not exited within the deadline is killed by PID + and the case fails.""" + table, info = real_host_table + visible = [d for d in os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",") if d.strip()] + visible = visible or [str(i) for i in range(torch.cuda.device_count())] + if len(visible) < WORLD: + pytest.skip(f"runs one process per GPU on {WORLD} GPUs; {len(visible)} visible") + common = _host_ids(_HOST_T_MAX, seed=_COMMON_SEED) + want = {} + for li, (layer, rows) in enumerate(_HOST_LAYERS): + w, s = _gather_views(table, layer) + out = op(common[:, li, :], w, s, 0, rows) + op.compare(out, _host_ref(common[:, li, :], *_cpu_bytes(table, layer))) + want[str(layer)] = _digest(out) + repo_root = Path(__file__).resolve().parents[8] + assert (repo_root / "tensorrt_llm").is_dir(), repo_root + env = dict(os.environ, CUDA_VISIBLE_DEVICES=",".join(visible[:WORLD])) + # the repo root (tensorrt_llm), tests/unittest (the utils package this file imports), tests + # (test_common, which utils imports; the unittest conftest adds it for pytest) and + # tests/unittest/_torch/modeling_v2 (_validating, which pytest reaches via __extra_import_path__) + env["PYTHONPATH"] = os.pathsep.join( + p + for p in ( + str(repo_root), + str(Path(__file__).resolve().parents[6]), + str(Path(__file__).resolve().parents[7]), + str(Path(__file__).resolve().parents[4]), + env.get("PYTHONPATH", ""), + ) + if p + ) + procs = [] + for rank in range(WORLD): + with open(tmp_path / f"rank{rank}.log", "w") as log: + procs.append( + subprocess.Popen( + [ + sys.executable, + str(Path(__file__).resolve()), + "--engram-host-rank", + str(rank), + "--out", + str(tmp_path), + ], + env=env, + stdout=log, + stderr=subprocess.STDOUT, + ) + ) + deadline = time.monotonic() + _RANK_DEADLINE_S + while any(p.poll() is None for p in procs) and time.monotonic() < deadline: + time.sleep(1) + wedged = [rank for rank, p in enumerate(procs) if p.poll() is None] + for p in procs: + if p.poll() is None: + p.kill() + p.wait() + tails = "".join( + f"\n--- rank {r} ---\n{(tmp_path / f'rank{r}.log').read_text()[-3000:]}" + for r in range(WORLD) + ) + assert not wedged, f"ranks {wedged} did not exit within {_RANK_DEADLINE_S} s{tails}" + codes = [p.returncode for p in procs] + assert codes == [0] * WORLD, f"rank exit codes {codes}{tails}" + results = [json.loads((tmp_path / f"rank{r}.json").read_text()) for r in range(WORLD)] + print() + for r in results: + print( + f" {r['backend']} rank {r['rank']} pid {r['pid']} device {r['device_uuid']}: {r['path']} " + + ( + f"inode {r['inode']} " + if r["backend"] == "shm" + else f"imported (created {r['created']}; fill pid {r['fill_pid']}) " + ) + + f"at 0x{r['weight_l1_address']:x}, opened {r['open_s']} s (map/registration {r['register_s']} s), " + f"device delta {r['alloc_delta']} B; common {_HOST_T_MAX} tokens sha256 L1 {r['digest']['1'][:16]} " + f"L14 {r['digest']['14'][:16]}; own {r['own_tokens']} tokens and captures at T {r['captured_T']} bit-exact" + ) + assert r["backend"] == info["backend"] and r["path"] == info["path"], (r, info) + if info["backend"] == "shm": + assert r["inode"] == info["inode"], (r, info) + else: # the same allocations: imported, from the fill this process's fixture made or imported + assert not r["created"] and (r["fill_pid"], r["fill_time"]) == ( + info["fill"].get("pid"), + info["fill"].get("unix_time"), + ) + assert r["alloc_delta"] == 0 + assert r["digest"] == want, (r["rank"], r["digest"], want) + assert ( + len({r["pid"] for r in results}) == WORLD + and len({r["device_uuid"] for r in results}) == WORLD + ) + print( + f" this process (cuda:0): sha256 L1 {want['1'][:16]} L14 {want['14'][:16]}, bit-exact to the reference" + ) + + +if __name__ == "__main__": + import argparse + + _parser = argparse.ArgumentParser( + description="a rank of test_real_host_table_four_ranks_read_the_same_mapping, " + "or a process of test_vmm_table_created_by_a_one_gpu_process" + ) + _parser.add_argument("--engram-host-rank", type=int) + _parser.add_argument("--out") + _parser.add_argument("--vmm-create-probe") + _parser.add_argument("--vmm-one-gpu", choices=("create", "import")) + _parser.add_argument("--name") + _args = _parser.parse_args() + if _args.vmm_create_probe: + with open(_args.vmm_create_probe, "w") as _f: + json.dump(_vmm_create_probe(), _f) + elif _args.vmm_one_gpu == "create": + _one_gpu_create_main(_args.name, _args.out) + elif _args.vmm_one_gpu == "import": + _one_gpu_import_main(_args.name, _args.out) + else: + _host_rank_main(_args.engram_host_rank, _args.out) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/embedding/test_modeling_v2_engram_hash.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/embedding/test_modeling_v2_engram_hash.py new file mode 100644 index 000000000000..61547b40c14f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/embedding/test_modeling_v2_engram_hash.py @@ -0,0 +1,4527 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the engram_hash entry of the DeepSeek-V4.1 target. + +The op computes DeepSeek-V4.1's Engram n-gram hash ids on the device and keeps +each request's token history in a paged, position-indexed int32 pool, so a +step's ids need no host-side history. What is certified is bit-exactness +against the checkpoint's ``inference/engram.py`` ``NgramHashState`` semantics +across every way a request's tokens reach the op: one-shot prefill, chunked +prefill, decode, and speculative verify runs whose rejected tail is rewound. + +The entry owns what is certified: ``CELLS`` is the list, ``reference`` and +``compare`` are the per-call gate (the hash ids, given the pool as it was before +the call). This file turns a cell's spec into a schedule of calls and keeps what +a cell cannot carry: the sweeps over every token count, the engine envelope, +the engine's own pools, block-table widths, the encoder-replay history call, +CUDA-graph capture, the rejections, the wrong-variant controls and the +checkpoint / golden comparisons. + +Four sources of truth: + +* the entry's ``reference``, one call at a time, reading the history from the + pool it is given; +* ``_RefState``: the checkpoint's ``NgramHashState.forward`` rewritten here in + torch ops for one request (int64 end to end, position-indexed cache, the + cumulative DEAD / sequence-start blocking), driven with the same run schedule + as the op. It carries each request's history itself, independently of the + op's pool, so it is what certifies the pool writes across calls. The hash + constants are rebuilt here with the checkpoint's recipe (numpy PCG64 + multipliers, sympy primes) and anchored to the config; +* the checkpoint's own ``inference/engram.py`` module, imported from the + checkpoint directory with the checkpoint's tokenizer (real 99092-class map); +* the golden dumps (``golden//engram_hash.pt`` under ``$MODELING_V2_DSV41_GOLDEN_DIR``). + These last two need external directories and skip, naming them, when absent. + +Every comparison is exact (``torch.equal``) -- the output is integer ids. The +deliberately wrong variants (``pad_id`` 2 instead of ``token_map[2]``, int32 or +float64 products, non-cumulative blocking, multipliers seeded by layer index, +history not carried) run through the same comparison and must fail it. +""" + +import functools +import importlib.util +import os +import types + +import numpy as np +import pytest +import torch +from utils.llm_data import llm_models_root + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.embedding.engram_hash import ( # noqa: E501 + engram_hash as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "engram_hash requires a CUDA device" + +# DeepSeek-V4.1-Flash text_config (config.json) +LAYER_IDS = (1, 14) +MAX_NGRAM = 4 +N_HEADS = 8 +ENGRAM_VOCAB = 16_000_000 +COMPRESSED_VOCAB = 99_092 +NUM_EMBEDDINGS = (384_006_168, 384_016_682) +PAD_TOKEN_ID = 2 +VOCAB = 129_280 +COLS = (MAX_NGRAM - 1) * N_HEADS +DEAD = -1 +# Written into every pool slot up front. It is a plausible-looking positive id, +# so a read of a slot nobody wrote hashes to garbage instead of hiding as DEAD. +SENTINEL = 0x5A5A5A5A +DEV = "cuda" + +# Per-request length L (max_seq_len) and DSpark: with DSpark the engine hands the manager +# Lm = L + 14 (max_draft_len 5), so positions run to Lm - 1; without it Lm = L. The block table the +# target passes, metadata.engram_token_block_table, is max_blocks_per_seq columns wide, which +# KVCacheManagerV2 sets to ceil((Lm + headroom) / 128) rounded up to a multiple of 4 (headroom 1, +# or 10 with DSpark): 132 / 260 / 516 / 1028 / 2052 / 8196 at the L below, DSpark on and off alike +# (test_engine_width_and_capture_at_every_L builds the manager at each and asserts it; the cg +# engine dumps show 8196 at L 1M, 1028 at L 128k and 260 at L 32k). +_L_SET = (16384, 32768, 65536, 131072, 262144, 1048576) +_L_CONFIGS = [(L, ds) for L in _L_SET for ds in (False, True)] + + +def _lm(L: int, ds: bool) -> int: + return L + 14 if ds else L + + +def _engine_width(L: int, ds: bool) -> int: + blocks = -(-(_lm(L, ds) + (10 if ds else 1)) // 128) + return -(-blocks // 4) * 4 + + +assert [_engine_width(L, ds) for L, ds in _L_CONFIGS] == [ + 132, + 132, + 260, + 260, + 516, + 516, + 1028, + 1028, + 2052, + 2052, + 8196, + 8196, +] + + +# Reference activations dumped from the checkpoint's own inference code; no default +# location, so the cases that compare against them skip unless this is set. +_GOLDEN_DIR = os.environ.get("MODELING_V2_DSV41_GOLDEN_DIR", "") +# The checkpoint resolves under LLM_MODELS_ROOT, as the whole-model gates' does; +# MODELING_V2_DSV41_CKPT_DIR overrides it. A case that needs it skips when it is absent. +_MODELS_ROOT = llm_models_root() +_CKPT_DIR = os.environ.get("MODELING_V2_DSV41_CKPT_DIR") or ( + str(_MODELS_ROOT / "DeepSeek-V4.1-Flash-NVFP4") if _MODELS_ROOT else "" +) +_GOLDEN_CASES = ( + "capital_fr", + "arith_chat", + "repeat_ngram", + "long_retrieval", + "long_topk", + "boot_natalia", + "boot_capital", + "boot_passkey", +) + + +# ─── hash constants, rebuilt with the checkpoint's recipe ────────────────────── + + +def _multipliers(layer_ids=LAYER_IDS) -> torch.Tensor: + """engram.py compute_hash_multipliers: odd, per-layer numpy PCG64 seeded 10007 * layer_id.""" + bound = max(1, (np.iinfo(np.int64).max // COMPRESSED_VOCAB) // 2) + rows = [] + for layer_id in layer_ids: + g = np.random.default_rng(10007 * layer_id) + rows.append(g.integers(low=0, high=bound, size=(MAX_NGRAM,), dtype=np.int64) * 2 + 1) + return torch.tensor(np.stack(rows)) + + +@functools.lru_cache(maxsize=None) +def _primes() -> torch.Tensor: + """engram.py EngramLayout.from_args: layer -> n-gram size -> head, next unused prime above 15999999.""" + from sympy import isprime + + seen, layers = set(), [] + for _ in LAYER_IDS: + per_ngram = [] + for _ in range(MAX_NGRAM - 1): + cur, sizes = ENGRAM_VOCAB - 1, [] + for _ in range(N_HEADS): + cur += 1 + while not isprime(cur) or cur in seen: + cur += 1 + seen.add(cur) + sizes.append(cur) + per_ngram.append(sizes) + layers.append(per_ngram) + return torch.tensor(layers, dtype=torch.int64) + + +def _offsets(primes: torch.Tensor) -> torch.Tensor: + flat = primes.flatten(1) + return torch.cat([torch.zeros_like(flat[:, :1]), flat.cumsum(1)[:, :-1]], dim=1) + + +@functools.lru_cache(maxsize=None) +def _synthetic_token_map() -> torch.Tensor: + """A seeded stand-in for the tokenizer-derived map: every id -> [0, 99092), max class present.""" + g = torch.Generator().manual_seed(1234) + m = torch.randint(0, COMPRESSED_VOCAB, (VOCAB,), generator=g, dtype=torch.int32) + m[7] = COMPRESSED_VOCAB - 1 # the largest class, so products come within 1.1e15 of 2^63 + m[8] = COMPRESSED_VOCAB - 1 + assert int(m[PAD_TOKEN_ID]) != PAD_TOKEN_ID + return m + + +class _Consts: + def __init__(self, token_map: torch.Tensor, layer_ids=LAYER_IDS): + self.token_map = token_map + self.multipliers = _multipliers(layer_ids) + self.primes = _primes() + self.offsets = _offsets(self.primes) + self.pad_id = int(token_map[PAD_TOKEN_ID]) + self.dev = types.SimpleNamespace( + token_map=token_map.to(DEV), + multipliers=self.multipliers.to(DEV), + primes=self.primes.to(DEV), + offsets=self.offsets.to(DEV), + ) + + +# ─── reference: the checkpoint's NgramHashState.forward, one request ──────────── + + +class _RefState: + """engram.py NgramHashState for one request, in torch ops on the CPU. + + ``variant`` selects a deliberately wrong implementation for the controls: + ``int32`` / ``float64`` products, ``noncumulative`` blocking. + """ + + def __init__( + self, + consts: _Consts, + max_len: int, + variant=None, + pad_id=None, + multipliers=None, + base: int = 0, + ): + """``base``: the cache holds positions [base, max_len) only (a request whose calls never + reach below ``base``: every run starts at or after base + N - 1, or base is 0).""" + self.c = consts + self.variant = variant + self.pad_id = consts.pad_id if pad_id is None else pad_id + self.multipliers = consts.multipliers if multipliers is None else multipliers + self.base = base + self.cache = torch.full((max_len - base,), SENTINEL, dtype=torch.int64) + + def forward(self, ids: torch.Tensor, start_pos: int, mask=None) -> torch.Tensor: + ids = ids.to(torch.int64) + n = ids.numel() + assert n == 0 or self.base == 0 or start_pos - (MAX_NGRAM - 1) >= self.base, ( + start_pos, + self.base, + ) + in_vocab = (ids >= 0) & (ids < self.c.token_map.numel()) + compressed = torch.where( + in_vocab, self.c.token_map[ids.clamp(0, self.c.token_map.numel() - 1)].long(), DEAD + ) + if mask is not None: + compressed = torch.where(mask, compressed, DEAD) + self.cache[start_pos - self.base : start_pos - self.base + n] = compressed + positions = torch.arange(start_pos, start_pos + n) + tokens, blocked = [], torch.zeros(n, dtype=torch.bool) + for shift in range(MAX_NGRAM): + source = self.cache[(positions - shift).clamp_min(0) - self.base] + hit = (positions < shift) | (source == DEAD) + blocked = hit if self.variant == "noncumulative" else blocked | hit + tokens.append(torch.where(blocked, self.pad_id, source)) + tokens = torch.stack(tokens, dim=-1) # [n, MAX_NGRAM] + if self.variant == "int32": + products = tokens.int().unsqueeze(1) * self.multipliers.int() + elif self.variant == "float64": + products = (tokens.double().unsqueeze(1) * self.multipliers.double()).long() + else: + products = tokens.unsqueeze(1) * self.multipliers # [n, L, MAX_NGRAM] + rolling, hashes = products[..., 0], [] + for i in range(1, MAX_NGRAM): + rolling = torch.bitwise_xor(rolling, products[..., i]) + hashes.append(rolling.unsqueeze(-1) % self.c.primes[:, i - 1]) + return torch.cat(hashes, dim=-1) + self.c.offsets + + +# ─── op harness: paged pool, per-step block tables in batch order ─────────────── + + +class _Pool: + """A paged int32 token pool, sized independently of any batch, and each request's pages. + + ``num_pages`` is chosen by this test to be plainly larger than the requests + need; every request owns a random, disjoint set of physical pages, and each + step's block table is rebuilt in that step's batch order. ``lows[r]``, when + given, is the lowest position request r's calls touch: its logical pages below + that are unmapped (-1), as in the engine's windowed ENGRAM_TOKEN role, so a + request a million positions long needs only the pages of its window. + """ + + def __init__( + self, max_lens: dict, page_size: int, num_pages: int = 4096, seed: int = 0, lows=None + ): + self.page_size = page_size + blocks = {r: -(-n // page_size) for r, n in max_lens.items()} + first = {r: (lows or {}).get(r, 0) // page_size for r in max_lens} + need = sum(blocks[r] - first[r] for r in blocks) + num_pages = max(num_pages, 2 * need + 1) + perm = torch.randperm(num_pages, generator=torch.Generator().manual_seed(seed)) + self.pages, at = {}, 0 + for r, nb in blocks.items(): + own = perm[at : at + nb - first[r]].to(torch.int32) + self.pages[r] = ( + torch.cat([torch.full((first[r],), -1, dtype=torch.int32), own]) + if first[r] + else own + ) + at += nb - first[r] + self.spare = perm[at:] # never owned by a request + self.max_blocks = max(blocks.values()) + self.pool = torch.full((num_pages, page_size), SENTINEL, dtype=torch.int32, device=DEV) + + def table(self, req_ids, override=None) -> torch.Tensor: + """[B, max_blocks] int32; ``override[r]`` replaces request r's row (windowing tests).""" + rows = [] + for r in req_ids: + row = torch.full((self.max_blocks,), -1, dtype=torch.int32) + src = self.pages[r] if override is None or r not in override else override[r] + row[: src.numel()] = src + rows.append(row) + return torch.stack(rows).to(DEV) + + def slot(self, r: int, p: int): + return int(self.pages[r][p // self.page_size]), p % self.page_size + + +def _call( + consts: _Consts, pool: _Pool, runs, *, table=None, pad_id=None, dtypes=None, token_map=None +): + """One op call over ``runs`` = [(req, ids[int64 cpu], start_pos, mask or None)], in batch order.""" + args, kwargs = _call_args( + consts, pool, runs, table=table, pad_id=pad_id, dtypes=dtypes, token_map=token_map + ) + return op(*args, **kwargs).cpu() + + +def _call_args( + consts: _Consts, pool: _Pool, runs, *, table=None, pad_id=None, dtypes=None, token_map=None +): + """The op's (positional args, kwargs) for one call over ``runs``, in batch order.""" + dtypes = dtypes or {} + lens = [ids.numel() for _, ids, _, _ in runs] + input_ids = torch.cat([ids for _, ids, _, _ in runs]).to(dtypes.get("ids", torch.int32)) + req_idx = torch.repeat_interleave(torch.arange(len(runs)), torch.tensor(lens)).to( + dtypes.get("meta", torch.int32) + ) + cu = torch.tensor([0] + list(np.cumsum(lens)), dtype=dtypes.get("meta", torch.int32)) + start = torch.tensor([s for _, _, s, _ in runs], dtype=dtypes.get("meta", torch.int32)) + masks = [m for _, _, _, m in runs] + token_mask = None + if any(m is not None for m in masks): + token_mask = torch.cat( + [torch.ones(n, dtype=torch.bool) if m is None else m for n, m in zip(lens, masks)] + ).to(DEV) + if table is None: + table = pool.table([r for r, _, _, _ in runs]) + args = ( + input_ids.to(DEV), + req_idx.to(DEV), + cu.to(DEV), + start.to(DEV), + pool.pool, + table.to(DEV).to(dtypes.get("table", torch.int32)), + consts.dev.token_map if token_map is None else token_map, + consts.dev.multipliers, + consts.dev.primes, + consts.dev.offsets, + consts.pad_id if pad_id is None else pad_id, + ) + return args, dict(token_mask=token_mask) + + +def _certify_schedule(consts, pool, steps, refs, *, tables=None, **kw): + """Every call of ``steps`` through ``op`` (its ``is_valid`` armed by ``certifying``), gated + by ``op.compare`` against ``op.reference`` on the pool as it was before the call, and every + row also against the per-request ``_RefState`` (which certifies the pool writes, since it + carries the history itself). ``tables[k]``, when given, is call k's block table. Returns + the rows that differ from ``_RefState``.""" + bad = 0 + for k, runs in enumerate(steps): + table = None if tables is None else tables[k] + args, kwargs = _call_args(consts, pool, runs, table=table, **kw) + before = pool.pool.clone() + got = op(*args, **kwargs) + op.compare(got, op.reference(*args[:4], before, *args[5:], **kwargs)) + got = got.cpu() + at = 0 + for r, ids, start, mask in runs: + want = refs[r].forward(ids, start, mask) + bad += int((got[at : at + ids.numel()] != want).flatten(1).any(1).sum()) + at += ids.numel() + return bad + + +def _run_schedule(consts, pool, steps, refs, **kw): + """Run ``steps`` (lists of runs) through the op and the references; return the mismatching row count.""" + bad = 0 + for runs in steps: + got = _call(consts, pool, runs, **kw) + at = 0 + for r, ids, start, mask in runs: + want = refs[r].forward(ids, start, mask) + bad += int((got[at : at + ids.numel()] != want).flatten(1).any(1).sum()) + at += ids.numel() + return bad + + +def _ids(n: int, g: torch.Generator) -> torch.Tensor: + return torch.randint(0, VOCAB, (n,), generator=g, dtype=torch.int64) + + +# ─── constants ───────────────────────────────────────────────────────────────── + + +def test_constants_match_checkpoint_layout(): + """The rebuilt constants are the checkpoint's: multipliers, 48 distinct primes, table sizes.""" + m = _multipliers() + assert m.tolist() == [ + [76632096046245, 4839876093313, 35959672319349, 73987337458391], + [67716810739261, 51510806800915, 30921347202721, 82619226485591], + ] + p = _primes() + assert tuple(p.shape) == (2, MAX_NGRAM - 1, N_HEADS) + assert len(set(p.flatten().tolist())) == 48 + assert int(p[0, 0, 0]) == 16000057 and int(p[0, -1, -1]) == 16000463 + assert int(p[1, 0, 0]) == 16000477 and int(p[1, -1, -1]) == 16000889 + assert tuple(p.flatten(1).sum(1).tolist()) == NUM_EMBEDDINGS + # the largest product fits int64, and nothing narrower + assert (COMPRESSED_VOCAB - 1) * int(m.max()) < 2**63 + assert (COMPRESSED_VOCAB - 1) * int(m.max()) > 2**62 + + +# ─── one-shot prefill ────────────────────────────────────────────────────────── + + +def _drive_prefill(spec): + """One-shot prefill runs from position 0 in one call.""" + lens, page_size = list(spec["lens"]), spec["page_size"] + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(sum(lens) + page_size) + reqs = {r: _ids(n, g) for r, n in enumerate(lens)} + reqs[0][: min(4, lens[0])] = 7 # the largest compressed class at the sequence start + pool = _Pool({r: n for r, n in enumerate(lens)}, page_size, seed=len(lens)) + refs = {r: _RefState(consts, n) for r, n in enumerate(lens)} + runs = [(r, reqs[r], 0, None) for r in reqs] + assert _certify_schedule(consts, pool, [runs], refs) == 0 + # the pool holds each token's compressed id at its position, and nothing else was written + host = pool.pool.cpu() + for r, ids in reqs.items(): + for p in (0, ids.numel() // 2, ids.numel() - 1): + page, col = pool.slot(r, p) + assert int(host[page, col]) == int(consts.token_map[ids[p]]) + assert int((host != SENTINEL).sum()) == sum(lens) + + +# ─── chunked prefill + decode, mixed batches ──────────────────────────────────── + +# request: (prompt length, chunk sizes -- the last takes the remainder, decode steps) +_CHUNKED = { + 0: (301, [128, 1, 2, 125], 6), + 1: (130, [127, 3], 5), + 2: (1027, [255, 1, 256], 4), + 3: (5, [1, 1, 1, 1], 7), + 4: (520, [259], 3), + 5: (2, [1], 9), +} + + +def _chunked_seqs(g: torch.Generator): + return {r: _ids(n + d, g) for r, (n, _, d) in _CHUNKED.items()} + + +def _chunked_steps(seqs, masks=None): + """Mixed steps: every live request contributes its next chunk or one decode token.""" + plan = {} + for r, (n, chunks, d) in _CHUNKED.items(): + runs, at = [], 0 + for c in chunks: + runs.append((at, at + c)) + at += c + if at < n: + runs.append((at, n)) + runs += [(n + i, n + i + 1) for i in range(d)] + plan[r] = runs + steps, k = [], 0 + while any(plan.values()): + order = [r for r in plan if plan[r]] + order = order[k % len(order) :] + order[: k % len(order)] + ctx = [r for r in order if plan[r][0][0] < _CHUNKED[r][0]] + gen = [r for r in order if r not in ctx] + step = [] + for r in ctx + gen: + a, b = plan[r].pop(0) + m = None if masks is None else masks[r][a:b] + step.append((r, seqs[r][a:b], a, m)) + steps.append(step) + k += 1 + return steps + + +def _history_rows(steps) -> int: + """Rows that read a predecessor from the pool: the first min(len, N - 1) rows of every run past position 0.""" + return sum(min(ids.numel(), MAX_NGRAM - 1) for runs in steps for _, ids, a, _ in runs if a > 0) + + +def _chunked_schedule(g: torch.Generator, masks=None): + seqs = _chunked_seqs(g) + return seqs, _chunked_steps(seqs, masks) + + +def _drive_chunked(spec): + """The mixed chunked-prefill + decode schedule of six requests. With int32 indices: the + history-not-carried control. With int64 ids, run indices, block table and token map: the + same schedule (measured on another seed).""" + page_size = spec["page_size"] + consts = _Consts(_synthetic_token_map()) + if spec["index_dtype"] == torch.int64: + seqs, steps = _chunked_schedule(torch.Generator().manual_seed(13)) + max_lens = {r: s.numel() for r, s in seqs.items()} + pool = _Pool(max_lens, page_size, seed=1) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + dtypes = {"ids": torch.int64, "meta": torch.int64, "table": torch.int64} + tm64 = consts.dev.token_map.to(torch.int64) + assert _certify_schedule(consts, pool, steps, refs, dtypes=dtypes, token_map=tm64) == 0 + return + seqs, steps = _chunked_schedule(torch.Generator().manual_seed(7)) + max_lens = {r: s.numel() for r, s in seqs.items()} + pool = _Pool(max_lens, page_size, seed=page_size) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + assert len(steps) >= 10 + assert _certify_schedule(consts, pool, steps, refs) == 0 + # control: the same schedule with the pool cleared (DEAD) before every call, i.e. history not carried + pool2 = _Pool(max_lens, page_size, seed=page_size) + refs2 = {r: _RefState(consts, n) for r, n in max_lens.items()} + bad = 0 + for runs in steps: + pool2.pool.fill_(DEAD) + bad += _run_schedule(consts, pool2, [runs], refs2) + total = sum(n for n in max_lens.values()) + print(f"[control] history not carried, page_size={page_size}: {bad} of {total} rows differ") + assert bad == _history_rows(steps) + + +# ─── speculative verify with rewind ───────────────────────────────────────────── + + +def _drive_verify(spec): + """Runs of 1 + drafts tokens; only a random prefix is accepted and the next run starts after it.""" + drafts = spec["drafts"] + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(drafts) + prompts = {0: 3, 1: 40, 2: 129, 3: 2, 4: 700, 5: 1} + n_steps = 12 + max_lens = {r: n + n_steps * (drafts + 1) + 1 for r, n in prompts.items()} + pool = _Pool(max_lens, 128, seed=11) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + steps = [[(r, _ids(n, g), 0, None) for r, n in prompts.items()]] + start = dict(prompts) + for k in range(n_steps): + order = list(prompts)[k % 6 :] + list(prompts)[: k % 6] + step = [] + for r in order: + step.append((r, _ids(drafts + 1, g), start[r], None)) + start[r] += 1 + int( + torch.randint(0, drafts + 1, (1,), generator=g) + ) # anchor + accepted drafts + steps.append(step) + assert _certify_schedule(consts, pool, steps, refs) == 0 + + +# ─── pad id ──────────────────────────────────────────────────────────────────── + + +def test_pad_id_is_the_compressed_pad_token(): + """pad_id is token_map[2], a compressed id; passing the raw id 2 is caught wherever padding occurs.""" + consts = _Consts(_synthetic_token_map()) + assert consts.pad_id != PAD_TOKEN_ID + seqs, steps = _chunked_schedule(torch.Generator().manual_seed(17)) + max_lens = {r: s.numel() for r, s in seqs.items()} + pool = _Pool(max_lens, 128, seed=8) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + assert _run_schedule(consts, pool, steps, refs) == 0 + pool = _Pool(max_lens, 128, seed=8) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + bad = _run_schedule(consts, pool, steps, refs, pad_id=PAD_TOKEN_ID) + # every request's first min(len, 3) positions read past the sequence start + expect = sum(min(n, MAX_NGRAM - 1) for n in max_lens.values()) + print( + f"[control] pad_id=2 instead of token_map[2]={consts.pad_id}: {bad} of {sum(max_lens.values())} rows differ" + ) + assert bad == expect + + +# ─── DEAD tokens (image spans) and out-of-range ids ───────────────────────────── + + +def _drive_dead(spec): + """token_mask False and ids outside [0, V) are DEAD; blocking is cumulative and crosses calls.""" + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(3) + masks = {r: torch.ones(n + d, dtype=torch.bool) for r, (n, _, d) in _CHUNKED.items()} + masks[0][0:1] = False # first token + masks[0][126:130] = False # a span containing the 128 chunk boundary + masks[0][131] = False # right at a chunk start + masks[1][60:61] = False # a single dead token mid-chunk + masks[1][126] = False # 1 before a chunk start + masks[2][250:255] = False # ends exactly at the 255 chunk boundary + masks[2][1027:1028] = False # the first decode token + masks[3][2] = False + masks[4][100:200] = False + seqs = _chunked_seqs(g) + seqs[4][300] = VOCAB # out of range -> DEAD + seqs[4][302] = -5 # negative -> DEAD + seqs[2][10] = VOCAB + 12345 + steps = _chunked_steps(seqs, masks) + max_lens = {r: s.numel() for r, s in seqs.items()} + pool = _Pool(max_lens, 128, seed=5) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + assert _certify_schedule(consts, pool, steps, refs) == 0 + # DEAD is stored as -1 in the pool + page, col = pool.slot(4, 300) + assert int(pool.pool[page, col]) == DEAD + # control: the same schedule against non-cumulative blocking + pool2 = _Pool(max_lens, 128, seed=5) + wrong = {r: _RefState(consts, n, variant="noncumulative") for r, n in max_lens.items()} + bad = _run_schedule(consts, pool2, steps, wrong) + print(f"[control] non-cumulative blocking: {bad} of {sum(max_lens.values())} rows differ") + assert bad > 0 + + +# ─── windowed pool: only the pages holding [start - 3, start) must be real ──────── + + +def _drive_windowed(spec): + """Pages outside the read/write window may be unmapped (-1: writes skipped) or a shared scratch page.""" + consts = _Consts(_synthetic_token_map()) + seqs, steps = _chunked_schedule(torch.Generator().manual_seed(9)) + ps = 128 + max_lens = {r: s.numel() for r, s in seqs.items()} + pool = _Pool(max_lens, ps, seed=2) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + scratch = int(pool.spare[0]) + tables = [] + for k, runs in enumerate(steps): + override = {} + for r, ids, a, _ in runs: + b = a + ids.numel() + keep = {p // ps for p in range(max(0, a - 3), a)} | { + p // ps for p in range(max(0, b - 3), b) + } + row = pool.pages[r].clone() + for blk in range(row.numel()): + if blk not in keep: + row[blk] = -1 if (blk + k) % 2 else scratch + override[r] = row + tables.append(pool.table([r for r, *_ in runs], override)) + assert _certify_schedule(consts, pool, steps, refs, tables=tables) == 0 + # control: unmap the page holding start - 1 on a continuation call -> history rows are wrong + pool3 = _Pool(max_lens, ps, seed=2) + refs3 = {r: _RefState(consts, n) for r, n in max_lens.items()} + bad = 0 + for runs in steps: + override = {} + for r, ids, a, _ in runs: + row = pool3.pages[r].clone() + if a > 0: + row[(a - 1) // ps] = -1 + override[r] = row + bad += _run_schedule( + consts, pool3, [runs], refs3, table=pool3.table([r for r, *_ in runs], override) + ) + print(f"[control] page of start-1 unmapped: {bad} rows differ") + assert bad == _history_rows(steps) + + +# ─── CUDA graph: one decode step captured, replayed with new ids / positions ───── + + +def test_cuda_graph_decode_replay(): + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(21) + B, n_steps = 8, 20 + prompts = {r: int(torch.randint(1, 300, (1,), generator=g)) for r in range(B)} + max_lens = {r: n + n_steps + 1 for r, n in prompts.items()} + pool = _Pool(max_lens, 128, seed=4) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + assert ( + _run_schedule(consts, pool, [[(r, _ids(n, g), 0, None) for r, n in prompts.items()]], refs) + == 0 + ) + ids = torch.zeros(B, dtype=torch.int32, device=DEV) + req = torch.arange(B, dtype=torch.int32, device=DEV) + cu = torch.arange(B + 1, dtype=torch.int32, device=DEV) + start = torch.tensor([prompts[r] for r in range(B)], dtype=torch.int32, device=DEV) + table = pool.table(list(range(B))) + out = torch.empty(B, len(LAYER_IDS), COLS, dtype=torch.int64, device=DEV) + c = consts.dev + + def step(): + op( + ids, + req, + cu, + start, + pool.pool, + table, + c.token_map, + c.multipliers, + c.primes, + c.offsets, + consts.pad_id, + out=out, + ) + + # warm up (compiles) on a copy of the pool, then restore it + saved = pool.pool.clone() + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + step() + torch.cuda.current_stream().wait_stream(s) + pool.pool.copy_(saved) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + pool.pool.copy_(saved) + pos = dict(prompts) + for _ in range(n_steps): + new = _ids(B, g) + ids.copy_(new.to(torch.int32)) + start.copy_(torch.tensor([pos[r] for r in range(B)], dtype=torch.int32)) + graph.replay() + got = out.cpu() + for r in range(B): + assert torch.equal(got[r : r + 1], refs[r].forward(new[r : r + 1], pos[r])) + pos[r] += 1 + + +# ─── index dtypes and the out= path ───────────────────────────────────────────── + + +def _drive_empty(spec): + """T = 0, B = 0: an empty int64 result, no launch, nothing written.""" + consts = _Consts(_synthetic_token_map()) + pool = _Pool({0: 8}, spec["page_size"]) + c = consts.dev + z = torch.empty(0, dtype=torch.int32, device=DEV) + table = torch.empty(0, 1, dtype=torch.int32, device=DEV) + cu = torch.zeros(1, dtype=torch.int32, device=DEV) + args = (z, z, cu, z, pool.pool, table, c.token_map, c.multipliers, c.primes, c.offsets) + before = pool.pool.clone() + out = op(*args, consts.pad_id) + assert tuple(out.shape) == (0, len(LAYER_IDS), COLS) + assert int((pool.pool != SENTINEL).sum()) == 0 + op.compare(out, op.reference(*args[:4], before, *args[5:], consts.pad_id)) + + +_DRIVERS = { + "prefill": _drive_prefill, + "chunked": _drive_chunked, + "verify": _drive_verify, + "dead": _drive_dead, + "windowed": _drive_windowed, + "empty": _drive_empty, +} + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Every call of the cell's schedule through ``op``, ``op.reference`` and ``op.compare``, + every row against the cross-call ``_RefState``, and the checks the configuration carries + (pool contents, history-not-carried / non-cumulative / unmapped-page controls).""" + certifying(op) + spec = cell.spec + assert (spec["L"], spec["N"], spec["H"]) == (len(LAYER_IDS), MAX_NGRAM, N_HEADS) + _DRIVERS[spec["kind"]](spec) + + +def test_rejects_malformed_inputs(): + consts = _Consts(_synthetic_token_map()) + pool = _Pool({0: 8}, 128) + c = consts.dev + ids = torch.tensor([1, 2, 3], dtype=torch.int32, device=DEV) + req = torch.zeros(3, dtype=torch.int32, device=DEV) + cu = torch.tensor([0, 3], dtype=torch.int32, device=DEV) + st = torch.zeros(1, dtype=torch.int32, device=DEV) + tb = pool.table([0]) + good = dict( + input_ids=ids, + req_idx_per_token=req, + cu_seq_lens=cu, + start_pos=st, + token_pool=pool.pool, + block_table=tb, + token_map=c.token_map, + multipliers=c.multipliers, + primes=c.primes, + offsets=c.offsets, + pad_id=consts.pad_id, + ) + bad = { + "ids int16": dict(input_ids=ids.to(torch.int16)), + "ids on cpu": dict(input_ids=ids.cpu()), + "strided ids": dict(input_ids=torch.arange(6, dtype=torch.int32, device=DEV)[::2]), + "req short": dict(req_idx_per_token=req[:2]), + "cu wrong length": dict(cu_seq_lens=torch.tensor([0, 3, 3], dtype=torch.int32, device=DEV)), + "pool int64": dict(token_pool=pool.pool.to(torch.int64)), + "multipliers int32": dict(multipliers=c.multipliers.int()), + "primes 7 heads": dict(primes=c.primes[..., :7].contiguous()), + "offsets short": dict(offsets=c.offsets[:, :-1].contiguous()), + "mask not bool": dict(token_mask=torch.ones(3, dtype=torch.int32, device=DEV)), + "strided mask": dict(token_mask=torch.ones(6, dtype=torch.bool, device=DEV)[::2]), + } + for name, patch in bad.items(): + with pytest.raises(ValueError, match="^engram_hash: ") as err: + op(**{**good, **patch}) + print(f"[rejected] {name}: {err.value}") + assert int((pool.pool != SENTINEL).sum()) == 0 + + +# ─── every token count 1..16384 is the same per-token function ─── +# +# The core target makes one call per step over all of the rank's tokens: context +# runs first (fresh prompts and chunks continuing from history), then generation +# runs (1-token decode or 1 + 5-token DSpark verify). A token's ids depend only on +# its own run's row, start and table, its predecessors and the pool; T reaches the +# kernel only as ``num_tokens``, which sizes the grid (64 tokens per program) and +# masks the last program. So every prefix of a mixed 16384-token step -- a valid call +# over the runs it touches, the last one cut short -- must return exactly the first r +# rows of the 16384-token call and leave the pool exactly as the first r tokens' +# writes do. + + +def _jit_variants(fn) -> int: + """Compiled variants Triton holds for ``fn`` on every device (Triton >= 3.2 keeps them + in ``device_caches[device][0]``, older releases in ``cache[device]``).""" + dc = getattr(fn, "device_caches", None) + if dc: + return sum(len(v[0]) for v in dc.values()) + return sum(len(v) for v in fn.cache.values()) + + +def _plant_history(consts, pool, refs, req, start, g, dead_every=7): + """Store compressed ids for positions [start - 3, start) of ``req`` -- what earlier + calls would have left -- in the pool and in the reference's cache; every + ``dead_every``-th one is DEAD (an image token).""" + for p in range(max(0, start - (MAX_NGRAM - 1)), start): + c = int(consts.token_map[int(torch.randint(0, VOCAB, (1,), generator=g))]) + if int(torch.randint(0, dead_every, (1,), generator=g)) == 0: + c = DEAD + page, col = pool.slot(req, p) + pool.pool[page, col] = c + refs[req].cache[p - refs[req].base] = c + + +def _step_tensors(consts, pool, runs): + """Device tensors of one call over ``runs`` = [(req, ids, start)] in batch order, plus + each token's pool slot (flat index) and the compressed id it stores.""" + lens = [ids.numel() for _, ids, _ in runs] + input_ids = torch.cat([ids for _, ids, _ in runs]).to(torch.int32) + req_idx = torch.repeat_interleave(torch.arange(len(runs)), torch.tensor(lens)).to(torch.int32) + cu = torch.tensor([0] + list(np.cumsum(lens)), dtype=torch.int32) + start = torch.tensor([s for _, _, s in runs], dtype=torch.int32) + table = pool.table([r for r, _, _ in runs]) + slots, vals = [], [] + assert table.shape[1] == pool.max_blocks + for r, ids, s in runs: + for o in range(ids.numel()): + page, col = pool.slot(r, s + o) + slots.append(page * pool.page_size + col) + vals.append(consts.token_map[ids].to(torch.int32)) + return ( + input_ids.to(DEV), + req_idx, + cu.to(DEV), + start.to(DEV), + table, + torch.tensor(slots, dtype=torch.int64, device=DEV), + torch.cat(vals).to(DEV), + ) + + +# contexts (start, tokens): fresh prompts, chunks continuing from history, a 1-token context, chunks at +# L 32768 .. 1048576 and one ending at Lm = 1,048,590 (L 1M with DSpark, the last position 1,048,589) +_SWEEP_LM = _lm(1048576, True) +_SWEEP_CONTEXTS = [ + (0, 3001), + (2048, 1500), + (131, 1203), + (5, 900), + (0, 1), + (777, 1531), + (1_000_000, 2500), + (262_000, 1700), + (65_530, 1111), + (_SWEEP_LM - 2000, 2000), + (131_071, 344), + (32_767, 29), +] +# 180 decode starts: page edges, the gates' and every L's last positions, and a spread +_SWEEP_DECODE_STARTS = ( + [1, 2, 3, 127, 128, 129, 16383, 4095, 4096] + + [37 * k + 11 for k in range(11)] + + [ + 16397, + 32767, + 32781, + 65535, + 65549, + 131071, + 131085, + 262143, + 262157, + 1048575, + _SWEEP_LM - 1, + _SWEEP_LM - 2, + 1048448, + 1048447, + 524288, + 524287, + ] + + [6553 * k + 1021 for k in range(144)] +) +# 64 verify-run starts (runs of 1 + 5): page edges, runs ending at every L's Lm, and a spread +_SWEEP_VERIFY_STARTS = [ + 1, + 2, + 126, + 127, + 16378, + 3000, + 16392, + 32776, + 65544, + 131080, + 262152, + _SWEEP_LM - 6, + _SWEEP_LM - 7, + 1048570, +] + [20959 * k + 777 for k in range(50)] + + +def test_every_token_count_is_a_prefix_of_16384(): + """A mixed 16384-token step in the engine's order at max_seq_len 1,048,576 with DSpark (Lm + 1,048,590): 12 contexts (10 continuing from history, one ending at Lm), 180 decode runs and 64 + verify runs of 1 + 5 tokens (positions up to 1,048,589, runs ending at every L's Lm); B = + 256; a block table of the engine's 8196 columns at that L, each request's pages below its + window unmapped; DEAD tokens in the planted histories. The 16384-token call equals the + reference. Then every prefix r = 1..16384 (the runs it touches, the last cut at r) is + called into a row-prefix view of one poisoned output buffer: its r rows must equal the + full call's rows, the next 64 rows keep the poison, and the whole pool must equal the + planted histories plus exactly the first r tokens' writes. Mismatches are summed on the + device. The kernel's compiled variants are counted: after r = 1, 2 and 16 (``num_tokens`` + == 1, not divisible by 16, divisible by 16: Triton's integer specialization) no further + variant may appear.""" + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(8192) + runs = [(r, _ids(n, g), st) for r, (st, n) in enumerate(_SWEEP_CONTEXTS)] + nc = len(runs) + runs += [(nc + k, _ids(1, g), st) for k, st in enumerate(_SWEEP_DECODE_STARTS)] + nd = len(runs) + runs += [(nd + k, _ids(6, g), st) for k, st in enumerate(_SWEEP_VERIFY_STARTS)] + n_max = sum(ids.numel() for _, ids, _ in runs) + assert n_max == 16384 and len(runs) == 256, (n_max, len(runs)) + max_lens = {r: st + ids.numel() for r, ids, st in runs} + assert max(max_lens.values()) == _SWEEP_LM + lows = {r: max(0, st - (MAX_NGRAM - 1)) for r, _, st in runs} + pool = _Pool(max_lens, 128, seed=8192, lows=lows) + assert pool.max_blocks <= _engine_width(1048576, True) + pool.max_blocks = _engine_width(1048576, True) + refs = {r: _RefState(consts, n, base=lows[r]) for r, n in max_lens.items()} + for r, _, s in runs: + _plant_history(consts, pool, refs, r, s, g) + planted = pool.pool.clone() + ids_d, req_idx, cu, start, table, slots, vals = _step_tensors(consts, pool, runs) + c = consts.dev + + full = op( + ids_d, + req_idx.to(DEV), + cu, + start, + pool.pool, + table, + c.token_map, + c.multipliers, + c.primes, + c.offsets, + consts.pad_id, + ) + at = 0 + for r, ids, s in runs: + assert torch.equal(full[at : at + ids.numel()].cpu(), refs[r].forward(ids, s)), (r, s) + at += ids.numel() + expected_full = planted.clone() + expected_full.view(-1)[slots] = vals + assert torch.equal(pool.pool, expected_full) + + pool.pool.copy_(planted) + expected = planted.clone() + poison = -7 + obuf = torch.full((n_max + 64, len(LAYER_IDS), COLS), poison, dtype=torch.int64, device=DEV) + bad_out = torch.zeros((), dtype=torch.int64, device=DEV) + bad_pool = torch.zeros((), dtype=torch.int64, device=DEV) + over = torch.zeros((), dtype=torch.int64, device=DEV) + runs_of = req_idx.tolist() + req_d = req_idx.to(DEV) + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.embedding.engram_hash import ( # noqa: E501 + _engram_hash_kernel as fn, + ) + + before = _jit_variants(fn) + settled = None + for r in range(1, n_max + 1): + b = runs_of[r - 1] + 1 + res = op( + ids_d[:r], + req_d[:r], + cu[: b + 1], + start[:b], + pool.pool, + table, + c.token_map, + c.multipliers, + c.primes, + c.offsets, + consts.pad_id, + out=obuf[:r], + ) + assert res.data_ptr() == obuf.data_ptr() + expected.view(-1).index_copy_(0, slots[r - 1 : r], vals[r - 1 : r]) + bad_out += (obuf[:r] != full[:r]).sum() + over += (obuf[r : r + 64] != poison).sum() + bad_pool += (pool.pool != expected).sum() + if r == 16: + settled = _jit_variants(fn) + torch.cuda.synchronize() + after = _jit_variants(fn) + print( + f"\n r = 1..{n_max} over {len(runs)} runs (positions up to {max(max_lens.values()) - 1}, table " + f"{table.shape[1]} columns): {bad_out.item()} output elements differ from the {n_max}-token call, " + f"{over.item()} guard elements written, {bad_pool.item()} pool slot-checks differ; kernel variants " + f"{before} before, {settled} after r = 16, {after} after r = {n_max}" + ) + assert bad_out.item() == 0 and over.item() == 0 and bad_pool.item() == 0 + assert after == settled and settled - before <= 3, (before, settled, after) + + +def test_prefix_sweep_harness_sees_a_one_token_shift(): + """Control for the sweep above: its row check rejects the outputs of a call whose + tokens are one row off (ids shifted by one within a 64-token context), at r = 1..64.""" + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(64) + ids = _ids(65, g) + pool = _Pool({0: 65, 1: 65}, 128, seed=64) + full = _call(consts, pool, [(0, ids[:64], 0, None)]) + shifted = _call(consts, pool, [(1, ids[1:65], 0, None)]) + differ = [int((shifted[:r] != full[:r]).flatten(1).any(1).sum()) for r in range(1, 65)] + print( + f" [control one-token shift] rows that differ at r = 1, 2, 7, 64: {differ[0]}/1, {differ[1]}/2, " + f"{differ[6]}/7, {differ[63]}/64" + ) + assert all(d == r for r, d in zip(range(1, 65), differ)) + + +# ─── the (T, B) points of the engine's runtime-value dump, against the reference ─── + +# (contexts as (start, tokens), decode starts): the dumped (T, B) of the boot gates +_DUMP_STEPS = { + "T6_B6": ([], [5, 9, 12, 20, 33, 71]), + "T7_B7": ([], [5, 9, 12, 20, 33, 71, 130]), + "T48_B1": ([(0, 48)], []), + "T888_B1": ([(1024, 888)], []), + "T1018_B3": ([(0, 1016)], [17, 2047]), + "T1935_B3": ([(0, 1000), (2048, 934)], [129]), + "T2048_B1": ([(0, 2048)], []), +} + + +@pytest.mark.parametrize("name", list(_DUMP_STEPS)) +def test_dumped_steps_against_reference(name): + ctx, dec = _DUMP_STEPS[name] + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(len(name) * 101) + runs = [(k, _ids(n, g), s) for k, (s, n) in enumerate(ctx)] + runs += [(len(ctx) + k, _ids(1, g), s) for k, s in enumerate(dec)] + T, B = sum(ids.numel() for _, ids, _ in runs), len(runs) + assert f"T{T}_B{B}" == name + max_lens = {r: s + ids.numel() for r, ids, s in runs} + pool = _Pool(max_lens, 128, seed=T) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + for r, _, s in runs: + _plant_history(consts, pool, refs, r, s, g) + got = _call(consts, pool, [(r, ids, s, None) for r, ids, s in runs]) + at = 0 + for r, ids, s in runs: + assert torch.equal(got[at : at + ids.numel()], refs[r].forward(ids, s)), (name, r) + at += ids.numel() + + +# ─── the engine's own ENGRAM_TOKEN pool, far end ─── +# +# The pool the core hashes into is the V4.1 cache manager's ENGRAM_TOKEN role, +# ``get_buffers(engram_token_owner, ENGRAM_TOKEN).squeeze(-1)``. Its page count comes +# from the manager's capacity planning, not from any batch. The manager below is built +# at the accuracy gate's settings (tep4 rank 0: max_seq_len 16384, max_batch_size +# 32, max_num_tokens 8192, tokens_per_block 128, block reuse on, device quota +# 72.68959879875183 GiB), and its COMPRESS / SWA pools are checked against the gate's +# runtime-value dump, so its ENGRAM_TOKEN pool is the engine's. + +_ENGINE_QUOTA = int(round(72.68959879875183 * (1 << 30))) +_ENGINE_COMPRESS_R2_PAGES = ( + 453118, + 453119, + 453120, +) # the dump's ratio-2 COMPRESS views (per owner) +_ENGINE_SWA_PAGES = 137440 +# The largest quota any engine at kv_cache_fraction 0.5 can hand the manager on this GPU: +# half of the GB300's total device memory (297,020,948,480 B from torch.cuda.mem_get_info; +# the engine takes 0.5 of the FREE memory after the weights, which is always less). +_MAX_QUOTA = 148_510_474_240 +# ENGRAM_TOKEN pages per (max_batch_size, quota, max_num_tokens, KV storage, DSpark +# max_draft_len). The manager splits its quota among the pools by a typical step (one +# max_num_tokens context plus max_batch_size - 1 generation requests), so the ENGRAM_TOKEN share +# moves with the batch, with max_num_tokens, with DSpark and with the other roles' bytes: native +# Native KV storage (TRTLLM_DSV41_KV_STORAGE=native: SWA 528 B, main KV 288 B per entry) leaves +# more of the quota to it. The certified configuration set is the one the gates and witnesses +# launch: boot (2048, batch 8) and accuracy (8192, batch 32), DSpark off / 5, bf16 / +# native, and the native-storage witness's two legs, max_num_tokens 2048, batch 1, DSpark 5, native and +# bf16. All were probed at the _MAX_QUOTA ceiling (block reuse and encoder replay change +# nothing). +# Batch 1 leaves the most to this role: among unsplit managers the largest is the native-storage witness's +# native leg, 36,241,408 pages (18.56 GB, 4.64e9 elements, past 2^31 elements and 2^34 bytes); its +# bf16 contrast leg gives 32,243,712, the boot gate's native settings 26,345,472 and its bf16 ones +# 15,269,888. Every role's pages grow monotonically with the quota (14 quotas from 1 GiB to the +# ceiling at every probed configuration, both witness legs included). +# A pipeline-parallel (PP2) stage's manager holds half the layers, so at the same quota its +# ENGRAM_TOKEN share is larger: stage 0 (layers 0-19; the Engram history's owner, layer 1, is on +# it) built alone gives 57,999,360 pages native and 48,365,568 bf16 at (max_num_tokens 2048, batch 1, +# no DSpark, max_seq_len 16384), the maxima over the scanned engine envelope including PP2 +# stages (every max_seq_len gives the same). Stage alone is the upper bound (an engine's cross-rank allreduce(MIN) +# can only cut a stage's capacity), and tp does not change the pools. B_max is that stage-0 pool. +_ENGINE_TOKEN_PAGES = { + (32, _ENGINE_QUOTA, 8192, "bf16", 0): 2_359_296, + (8, _ENGINE_QUOTA, 8192, "bf16", 0): 5_242_880, + (32, _MAX_QUOTA, 8192, "bf16", 0): 4_456_448, + (8, _MAX_QUOTA, 8192, "bf16", 0): 10_027_008, + (8, _MAX_QUOTA, 2048, "bf16", 0): 15_269_888, + (8, _MAX_QUOTA, 2048, "native", 0): 26_345_472, + (1, _MAX_QUOTA, 2048, "bf16", 5): 32_243_712, + (1, _MAX_QUOTA, 2048, "native", 5): 36_241_408, + (1, _MAX_QUOTA, 2048, "bf16", 0, "pp2-stage0"): 48_365_568, + (1, _MAX_QUOTA, 2048, "native", 0, "pp2-stage0"): 57_999_360, +} +_UNSPLIT_B_MAX = 36_241_408 # largest unsplit manager (native-storage witness, native leg) +_B_MAX_CONFIG = (1, _MAX_QUOTA, 2048, "native", 0) +_B_MAX_STAGE = 0 +_B_MAX = _ENGINE_TOKEN_PAGES[_B_MAX_CONFIG + ("pp2-stage0",)] +# The block-table width the target passes: metadata.engram_token_block_table is a +# [max_num_sequences, max_blocks_per_seq] view (row stride = width) of the sliding block tables, +# and KVCacheManagerV2 sets max_blocks_per_seq = ceil((max_seq_len + headroom) / 128) rounded up +# to a multiple of 4: (16384 + 1) -> 129 -> 132 without DSpark, (16398 + 10) -> 129 -> 132 with +# it. 132 at every probed configuration. +_ENGINE_WIDTH = 132 + + +def _engine_manager( + max_batch_size: int = 32, + quota: int = _ENGINE_QUOTA, + max_num_tokens: int = 8192, + storage: str = "bf16", + draft: int = 0, + L: int = 16384, + pp_stage=None, +): + """The V4.1 cache manager at these settings; ``pp_stage`` builds that stage's manager of a PP2 + mapping instead (tp 1: ``Mapping(world_size=2, rank=stage, tp_size=1, pp_size=2)``, built + alone).""" + from types import SimpleNamespace + + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as T, + ) + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, + KvCacheConfig, + ) + from tensorrt_llm.mapping import Mapping + + text = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=list(LAYER_IDS), + engram_max_ngram_size=MAX_NGRAM, + sliding_window=128, + head_dim=512, + index_head_dim=128, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, + ) + config = SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **{k: v for k, v in text.items() if k != "model_type"}, + ) + sparse = DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ) + # The layout reads the storage mode once, at construction. + saved_storage = os.environ.get("TRTLLM_DSV41_KV_STORAGE") + os.environ["TRTLLM_DSV41_KV_STORAGE"] = storage + try: + mgr = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + max_gpu_total_bytes=quota, + host_cache_size=0, + event_buffer_max_size=0, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=40, + num_kv_heads=1, + head_dim=512, + tokens_per_block=128, + # With DSpark the engine hands the manager max_seq_len L + 14 (16398 at 16384: 16384 + 5 + 4 + 5). + max_seq_len=L + 14 if draft else L, + max_batch_size=max_batch_size, + mapping=( + Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + if pp_stage is None + else Mapping(world_size=2, rank=pp_stage, tp_size=1, pp_size=2, gpus_per_node=4) + ), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=VOCAB, + max_num_tokens=max_num_tokens, + sparse_attn_config=sparse, + pretrained_config=config, + spec_config=DSparkDecodingConfig(max_draft_len=draft, block_size=draft) + if draft + else None, + ) + finally: + if saved_storage is None: + os.environ.pop("TRTLLM_DSV41_KV_STORAGE", None) + else: + os.environ["TRTLLM_DSV41_KV_STORAGE"] = saved_storage + assert mgr.v41_layout.kv_storage == storage and mgr.max_seq_len == (L + 14 if draft else L), ( + mgr.max_seq_len + ) + assert mgr.max_blocks_per_seq == ( + _ENGINE_WIDTH if L == 16384 else _engine_width(L, bool(draft)) + ), (mgr.max_blocks_per_seq, L, draft) + if (max_batch_size, quota, max_num_tokens, storage, draft, L) == ( + 32, + _ENGINE_QUOTA, + 8192, + "bf16", + 0, + 16384, + ): + compress = [mgr.get_buffers(o, T.COMPRESS).shape[0] for o in (2, 8, 14)] + swa = tuple(mgr.get_buffers(0, T.SWA).shape) + assert sorted(compress) == sorted(_ENGINE_COMPRESS_R2_PAGES) and swa == ( + _ENGINE_SWA_PAGES, + 128, + 512, + ), ( + f"this manager does not reproduce the gate's engine: COMPRESS pages {compress}, SWA {swa}" + ) + raw = mgr.get_buffers(mgr.v41_layout.engram_token_owner, T.ENGRAM_TOKEN) + token_pool = raw.squeeze(-1) + assert token_pool.data_ptr() == raw.data_ptr() and token_pool.dtype == torch.int32 + key = (max_batch_size, quota, max_num_tokens, storage, draft) + ( + () if pp_stage is None else (f"pp2-stage{pp_stage}",) + ) + want = _ENGINE_TOKEN_PAGES.get(key) if L == 16384 else None + if pp_stage is not None: + assert want is not None and mgr.v41_layout.engram_token_owner == LAYER_IDS[0], ( + key, + mgr.v41_layout, + ) + assert token_pool.stride() == (128, 1) and token_pool.shape[1] == 128, ( + tuple(token_pool.shape), + token_pool.stride(), + ) + assert want is None or tuple(token_pool.shape) == (want, 128), (tuple(token_pool.shape), want) + return mgr, token_pool + + +class _FarEndPool(_Pool): + """Requests' pages taken from the LAST pages of a given pool tensor (the very last + page included), in random order; the whole tensor is set to SENTINEL first.""" + + def __init__(self, pool: torch.Tensor, max_lens: dict, seed: int): + self.page_size = pool.shape[1] + n = pool.shape[0] + blocks = {r: -(-m // self.page_size) for r, m in max_lens.items()} + need = sum(blocks.values()) + assert 2 * need < n, (need, n) + tail = torch.arange(n - 2 * need, n) + perm = tail[torch.randperm(tail.numel(), generator=torch.Generator().manual_seed(seed))] + perm = torch.cat([perm[perm == n - 1], perm[perm != n - 1]]) # the last page is owned + self.pages, at = {}, 0 + for r, nb in blocks.items(): + self.pages[r] = perm[at : at + nb].to(torch.int32) + at += nb + self.spare = perm[at:] + self.max_blocks = max(blocks.values()) + self.pool = pool + self.pool.fill_(SENTINEL) + + +def _expected_pool(consts, pool, seqs, lens): + """SENTINEL everywhere except each request's positions [0, lens[r]), which hold their compressed ids.""" + exp = torch.full_like(pool.pool, SENTINEL) + for r, n in lens.items(): + for p in range(n): + page, col = pool.slot(r, p) + exp[page, col] = int(consts.token_map[seqs[r][p]]) + return exp + + +@pytest.mark.parametrize( + "max_batch_size, quota, max_num_tokens, storage, draft, pp_stage", + [ + (32, _ENGINE_QUOTA, 8192, "bf16", 0, None), + (8, _ENGINE_QUOTA, 8192, "bf16", 0, None), + (32, _MAX_QUOTA, 8192, "bf16", 0, None), + (8, _MAX_QUOTA, 2048, "bf16", 0, None), + (1, _MAX_QUOTA, 2048, "bf16", 5, None), + (1, _MAX_QUOTA, 2048, "native", 5, None), + (1, _MAX_QUOTA, 2048, "bf16", 0, 0), + ], + ids=[ + "bs32-gate-quota", + "bs8-gate-quota", + "bs32-max-quota", + "boot-bs8-mnt2048-max-quota", + "kvquant-witness-bf16-bs1-mnt2048-dspark5-max-quota", + "kvquant-witness-native-bs1-mnt2048-dspark5-max-quota", + "pp2-stage0-bf16-bs1-mnt2048-max-quota", + ], +) +def test_engine_token_pool_far_end_chunked_and_decode( + max_batch_size, quota, max_num_tokens, storage, draft, pp_stage +): + """The chunked-prefill + decode schedule of six requests (the same as + test_chunked_prefill_and_decode_continuation) on the engine's own ENGRAM_TOKEN pool, + every request's pages among the pool's last pages (the last page owned): every call + equals the reference, and the whole pool afterwards holds exactly each request's + compressed ids at its positions and SENTINEL everywhere else. At max_batch_size 8 + the pool is [5242880, 128] int32 (2.68 GB), so its last pages sit past 2^31 bytes; at the + largest quota (0.5 x device memory) and batch 32 it is [4456448, 128] (2.28 GB); at the boot + gate's settings (max_num_tokens 2048, batch 8) and that quota it is [15269888, 128] (7.82 GB); + at the native-storage witness's bf16 contrast leg (max_num_tokens 2048, batch 1, DSpark 5, + TRTLLM_DSV41_KV_STORAGE=bf16) and that quota it is [32243712, 128] (16.51 GB); at its native leg + [36241408, 128] (the unsplit managers' largest); on a PP2 stage-0 manager at (2048, batch 1, bf16) + [48365568, 128] (24.76 GB). The largest configuration, the PP2 stage-0 manager in native storage + [57999360, 128], is test_history_pool_size_invariance_b_max's.""" + mgr, token_pool = _engine_manager( + max_batch_size, quota, max_num_tokens, storage, draft, pp_stage=pp_stage + ) + try: + consts = _Consts(_synthetic_token_map()) + seqs, steps = _chunked_schedule(torch.Generator().manual_seed(77)) + max_lens = {r: s.numel() for r, s in seqs.items()} + pool = _FarEndPool(token_pool, max_lens, seed=77) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + assert _run_schedule(consts, pool, steps, refs) == 0 + assert torch.equal(pool.pool, _expected_pool(consts, pool, seqs, max_lens)) + lowest = min(int(p.min()) for p in pool.pages.values()) + print( + f"\n engine ENGRAM_TOKEN pool {tuple(token_pool.shape)} int32" + f"{'' if pp_stage is None else f' (PP2 stage {pp_stage} manager)'}: requests on pages " + f"{lowest}..{token_pool.shape[0] - 1}; {len(steps)} calls == reference; pool exact, rest SENTINEL" + ) + finally: + mgr.shutdown() + + +# ─── capture at every graph batch, at every L's engine width, on that engine's pool ─── + +# The generation-only CUDA-graph batches an engine captures per rank, Gset: every batch 1..32 plus +# TRT-LLM's padded default list for cuda_graph_config {enable_padding: true, max_batch_size: 256} +# (1, 2, 4, 8, 16..128 step 8, 192, 256; the padding-off list at 256, 1..32, 64, 128, 256, is a +# subset). A decode step (DSpark off) hashes G tokens of G runs, a DSpark target-verify step 6 G +# tokens of G runs (next_n 6): T = G and 6 G, up to 1536. Under dep4 each rank hashes its own +# requests, so no gathered count reaches this op. The captured block table is the metadata's +# engram_token_block_table, max_blocks_per_seq columns (132 .. 8196 as L grows). +_GRAPH_BATCHES = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +assert len(_GRAPH_BATCHES) == 46 + + +class _WindowFarEndPool(_Pool): + """Each request's window of logical pages [lo // P, ceil(hi / P)) mapped to pages among the LAST + pages of the given pool tensor (the very last page included), in random order; the request's + other columns unmapped (-1), as the engine's windowed ENGRAM_TOKEN role leaves them. The whole + tensor is set to SENTINEL first.""" + + def __init__(self, pool: torch.Tensor, windows: dict, seed: int): + self.page_size = P = pool.shape[1] + n = pool.shape[0] + first = {r: lo // P for r, (lo, hi) in windows.items()} + blocks = {r: -(-hi // P) for r, (lo, hi) in windows.items()} + need = sum(blocks[r] - first[r] for r in windows) + assert 2 * need < n, (need, n) + tail = torch.arange(n - 2 * need, n) + perm = tail[torch.randperm(tail.numel(), generator=torch.Generator().manual_seed(seed))] + perm = torch.cat([perm[perm == n - 1], perm[perm != n - 1]]) # the last page is owned + self.pages, at = {}, 0 + for r in windows: + k = blocks[r] - first[r] + self.pages[r] = torch.cat( + [torch.full((first[r],), -1, dtype=torch.int32), perm[at : at + k].to(torch.int32)] + ) + at += k + self.spare = perm[at:] + self.max_blocks = max(blocks.values()) + self.pool = pool + self.pool.fill_(SENTINEL) + + +@pytest.mark.parametrize( + "L,dspark", _L_CONFIGS, ids=[f"L{L}-{'dspark' if ds else 'nodspark'}" for L, ds in _L_CONFIGS] +) +def test_engine_width_and_capture_at_every_L(L, dspark): + """The engine at max_seq_len L (the manager gets Lm = L + 14 with DSpark), max_batch_size 256, + max_num_tokens 16384: the V4.1 cache manager is built, its max_blocks_per_seq -- the width of + metadata.engram_token_block_table -- must equal _engine_width (132 .. 8196) with max_seq_len + not clamped, and on its own ENGRAM_TOKEN pool every graph batch G of Gset is captured: G + requests, request 0 whose last replayed run ends exactly at Lm (block-table column + (Lm - 1) // 128), request 1 whose first generated position opens a page, request 2 at the + sequence start, the rest at random positions below Lm; histories planted before each (DEAD + every ~7th); each request's window mapped to pages among the pool's last pages (the last + page owned), the rest of its row -1. One generation step (G runs of next_n tokens: 1 without + DSpark, 6 with it) is warmed up eagerly, captured, and replayed for 3 steps with ids and + start_pos updated in place (verify: the next start is 1 + a random accepted count, request 0 + all accepted, so rejected positions are rewritten). Every replay equals an eager call of the + same op on the same inputs and the reference, bit for bit, and differs from the replay + before it; the whole pool afterwards equals the planted histories plus exactly the writes.""" + lm, width, next_n, n_steps = _lm(L, dspark), _engine_width(L, dspark), 6 if dspark else 1, 3 + mgr, token_pool = _engine_manager( + max_batch_size=256, max_num_tokens=16384, draft=5 if dspark else 0, L=L + ) + try: + assert mgr.max_blocks_per_seq == width and mgr.max_seq_len == lm + consts = _Consts(_synthetic_token_map()) + c = consts.dev + replays = 0 + for G in _GRAPH_BATCHES: + g = torch.Generator().manual_seed(7 * L + 1000 * next_n + G) + span = (n_steps + 1) * next_n + pos0 = {r: _ri(g, 3, lm - span) for r in range(G)} + pos0[0] = lm - n_steps * next_n # all drafts accepted: the last replayed run ends at Lm + if G > 1: + pos0[1] = 128 * _ri(g, 1, (lm - span) // 128) # the first position opens a page + if G > 2: + pos0[2] = _ri(g, 0, 2) # predecessors before the sequence start + windows = { + r: (max(0, p0 - (MAX_NGRAM - 1)), min(lm, p0 + span)) for r, p0 in pos0.items() + } + pool = _WindowFarEndPool(token_pool, windows, seed=G) + pool.max_blocks = width + refs = {r: _RefState(consts, hi + 8, base=lo) for r, (lo, hi) in windows.items()} + for r in range(G): + _plant_history(consts, pool, refs, r, pos0[r], g) + expected = pool.pool.clone() + T = G * next_n + ids = torch.zeros(T, dtype=torch.int32, device=DEV) + req = torch.arange(G, dtype=torch.int32, device=DEV).repeat_interleave(next_n) + cu = torch.arange(0, T + 1, next_n, dtype=torch.int32, device=DEV) + start = torch.tensor([pos0[r] for r in range(G)], dtype=torch.int32, device=DEV) + table = pool.table(list(range(G))) + assert tuple(table.shape) == (G, width) and table.stride() == (width, 1) + out = torch.empty(T, len(LAYER_IDS), COLS, dtype=torch.int64, device=DEV) + + def step(): + op( + ids, + req, + cu, + start, + pool.pool, + table, + c.token_map, + c.multipliers, + c.primes, + c.offsets, + consts.pad_id, + out=out, + ) + + saved = pool.pool.clone() + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + step() # engine warm-up at this shape + torch.cuda.current_stream().wait_stream(side) + pool.pool.copy_(saved) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + pool.pool.copy_(saved) + del saved + pos, prev = dict(pos0), None + for _ in range(n_steps): + new = {r: _ids(next_n, g) for r in range(G)} + ids.copy_(torch.cat([new[r] for r in range(G)]).to(torch.int32)) + start.copy_(torch.tensor([pos[r] for r in range(G)], dtype=torch.int32)) + out.fill_(-7) + graph.replay() + got = out.clone() + step() # the same call eagerly: rewrites the same slots with the same ids + assert torch.equal(got, out), f"L={L} G={G} next_n={next_n}: replay != eager" + assert prev is None or not torch.equal(got, prev), ( + "a replay returned the previous step's ids" + ) + prev = got + got = got.cpu() + runs = [(r, new[r], pos[r], None) for r in range(G)] + for r in range(G): + want = refs[r].forward(new[r], pos[r]) + assert torch.equal(got[r * next_n : (r + 1) * next_n], want), ( + L, + G, + next_n, + r, + pos[r], + ) + _apply_writes(consts, pool, expected, runs) + for r in range(G): + adv = next_n if r == 0 else 1 + (_ri(g, 0, next_n - 1) if next_n > 1 else 0) + pos[r] += adv + replays += 1 + assert pos[0] == lm and max(pos0[r] + n_steps * next_n for r in range(G)) <= lm + assert torch.equal(pool.pool, expected), (L, G, next_n) + del graph, expected + print( + f"\n L={L} DSpark={'on' if dspark else 'off'}: manager max_seq_len {mgr.max_seq_len}, " + f"max_blocks_per_seq {mgr.max_blocks_per_seq}, ENGRAM_TOKEN pool {tuple(token_pool.shape)}; captured at " + f"all {len(_GRAPH_BATCHES)} G (T = G x {next_n}, up to {max(_GRAPH_BATCHES) * next_n}), {replays} replays " + f"== eager == reference, request 0 ending at Lm {lm} (column {(lm - 1) // 128} of {width}); pool exact" + ) + finally: + mgr.shutdown() + + +# ─── the multi-request envelope (B <= 256 runs, T <= 16384 rows, positions < Lm <= 1,048,590) ─── +# +# The engine's configuration bounds one call: max_num_tokens <= 16384 rows per rank, +# max_batch_size <= 256 runs, max_seq_len L <= 1,048,576 positions per request (Lm = L + 14 +# with DSpark). Inside that envelope a call is a +# partition of T rows into B runs: context chunks (fresh or continuing from history, with +# image spans DEAD through token_mask) and generation runs of 1 + D rows (D = 0 decode, +# D = 5 DSpark target verify, whose next run rewinds over the rejected tail). The runs are +# random here in number, kind, length, start, order, page placement and mask, and every +# call is checked against the reference, with the whole pool checked after every call. + +_ENV_MAX_POS = 16384 # max_seq_len of the gates; the default position bound of the helpers below +_ENV_MAX_T = 16384 # max_num_tokens: 2048, 4096, 8192 or 16384 per rank +_ENV_MAX_B = 256 # max_batch_size + + +def _jit_keys(fn) -> set: + """The compiled-variant keys Triton holds for ``fn`` (see _jit_variants).""" + dc = getattr(fn, "device_caches", None) + if dc: + return {(str(d), k) for d, v in dc.items() for k in v[0].keys()} + return {(str(d), k) for d, v in fn.cache.items() for k in v.keys()} + + +def _jit_clear(fn) -> None: + """Drop Triton's in-memory compiled variants of ``fn`` (the next calls re-bind, from the disk cache).""" + dc = getattr(fn, "device_caches", None) + if dc: + for v in dc.values(): + v[0].clear() + else: + for v in fn.cache.values(): + v.clear() + + +def _ri(g: torch.Generator, lo: int, hi: int) -> int: + """Uniform int in [lo, hi].""" + return lo + int(torch.randint(0, hi - lo + 1, (1,), generator=g)) + + +def _composition(g: torch.Generator, total: int, parts: int) -> list: + """``total`` split into ``parts`` positive lengths, uniformly over compositions.""" + if parts == 1: + return [total] + cuts = sorted((torch.randperm(total - 1, generator=g)[: parts - 1] + 1).tolist()) + edges = [0] + cuts + [total] + return [b - a for a, b in zip(edges[:-1], edges[1:])] + + +def _random_T(g: torch.Generator) -> int: + u = _ri(g, 0, 3) + if u == 0: + return _ri(g, 1, 64) + if u == 1: + return _ri(g, 65, 1024) + return _ri(g, 1025, _ENV_MAX_T) + + +def _feasible_G(T: int, B: int, D: int) -> list: + """Generation-run counts G that make a valid step of T rows and B runs (B - G contexts of >= 1 row).""" + out = [] + for G in range(0, B + 1): + rest = T - G * (1 + D) + if rest < 0: + continue + if (B - G == 0 and rest == 0) or (B - G > 0 and rest >= B - G): + out.append(G) + return out + + +def _image_mask(g: torch.Generator, spans_near: list, size: int = _ENV_MAX_POS + 8) -> torch.Tensor: + """bool over positions [0, size): False on 1..3 image spans, each starting near one of ``spans_near``.""" + m = torch.ones(size, dtype=torch.bool) + for _ in range(_ri(g, 1, 3)): + lo, hi = spans_near[_ri(g, 0, len(spans_near) - 1)] + a = _ri(g, max(0, lo), max(max(0, lo), hi)) + m[a : a + _ri(g, 1, 300)] = False + return m + + +def _plan_pair( + g: torch.Generator, + T: int, + B: int, + D: int, + G: int, + masked: bool, + shuffle: bool, + mask_all=False, + mask_fn=None, + max_pos: int = _ENV_MAX_POS, +): + """Two consecutive steps of the same B requests, positions below ``max_pos`` (the manager + length Lm). Step 1: B - G contexts (lengths a random composition of T - G (1 + D), start 0, + a random cached length, or the last start that leaves 8 positions) and G generation runs of + 1 + D rows at random positions (one in eight at the last start whose continuation can + still reach Lm). Step 2 continues every request that can: a generation run restarts at + start + 1 + accepted (a random accepted count for D = 5: the rejected tail is rewritten), a + context either continues with a random chunk (one in eight to the end of the budget or of + Lm) or turns into a generation run at its next position. ``mask_fn(g, s, n, size)``, when + given, builds a masked context's position mask (``size`` = max_pos + 8 positions) instead of + ``_image_mask``. Returns (steps, max position + 1 per request).""" + nc = B - G + ctx_lens = _composition(g, T - G * (1 + D), nc) if nc else [] + reqs = [] + for r in range(B): + if r < nc: + n = ctx_lens[r] + room = max_pos - 8 - n + u = _ri(g, 0, 9) + s = 0 if (room < 1 or u < 3) else (room if u == 9 else _ri(g, 1, room)) + reqs.append(("ctx", s, n)) + else: + top = max_pos - 2 * (1 + D) + reqs.append(("gen", top if _ri(g, 0, 7) == 0 else _ri(g, 1, top - 1), 1 + D)) + masks = {} + if masked: + for r, (kind, s, n) in enumerate(reqs): + if kind == "ctx" and (mask_all or _ri(g, 0, 1)): + if mask_fn is not None: + masks[r] = mask_fn(g, s, n, max_pos + 8) + else: + masks[r] = _image_mask( + g, [(s - 20, s + n - 1), (s + n - 6, s + n - 1)], max_pos + 8 + ) + step2, ctx2 = [], [] + for r, (kind, s, n) in enumerate(reqs): + if kind == "gen": + step2.append((r, "gen", s + 1 + (_ri(g, 0, D) if D else 0), 1 + D)) + elif s + n + 1 + D <= max_pos and _ri(g, 0, 1): + step2.append((r, "gen", s + n, 1 + D)) + else: + ctx2.append(r) + budget = _ENV_MAX_T - sum(n for _, _, _, n in step2) + for i, r in enumerate(ctx2): + _, s, n = reqs[r] + e = s + n + cap = min(max_pos - e, budget - (len(ctx2) - i - 1)) + if cap < 1: # at the end of Lm, or no rows left in the step: this request sits the step out + continue + L = ( + cap + if _ri(g, 0, 7) == 0 + else (_ri(g, 1, cap) if _ri(g, 0, 1) else _ri(g, 1, min(cap, 64))) + ) + step2.append((r, "ctx", e, L)) + budget -= L + step1 = [(r, kind, s, n) for r, (kind, s, n) in enumerate(reqs)] + + def order(step): + if shuffle: + perm = torch.randperm(len(step), generator=g).tolist() + return [step[i] for i in perm] + return [x for x in step if x[1] == "ctx"] + [x for x in step if x[1] == "gen"] + + steps = [] + for step in (step1, step2) if step2 else (step1,): + runs = [] + for r, kind, s, n in order(step): + m = masks[r][s : s + n].clone() if (kind == "ctx" and r in masks) else None + runs.append((r, _ids(n, g), s, m)) + steps.append(runs) + ends = {} + for runs in steps: + for r, ids, s, _ in runs: + ends[r] = max(ends.get(r, 0), s + ids.numel()) + return steps, ends + + +def _apply_writes(consts, pool, expected, runs): + """What one call must leave in the pool: every run's positions hold their compressed ids (-1 = DEAD).""" + P = pool.page_size + for r, ids, s, m in runs: + pos = torch.arange(s, s + ids.numel()) + slots = pool.pages[r][pos // P].long() * P + pos % P + vals = consts.token_map[ids].to(torch.int32) + if m is not None: + vals = torch.where(m, vals, torch.full_like(vals, DEAD)) + expected.view(-1)[slots.to(DEV)] = vals.to(DEV) + + +def _run_checked(consts, pool, refs, expected, steps, **kw): + """Run ``steps`` through the op; return (rows that differ from the reference, calls whose + pool afterwards differs from ``expected`` updated with that call's writes).""" + bad_rows, bad_pools = 0, 0 + for runs in steps: + got = _call(consts, pool, runs, **kw) + at = 0 + for r, ids, s, m in runs: + want = refs[r].forward(ids, s, m) + bad_rows += int((got[at : at + ids.numel()] != want).flatten(1).any(1).sum()) + at += ids.numel() + _apply_writes(consts, pool, expected, runs) + bad_pools += int(not torch.equal(pool.pool, expected)) + return bad_rows, bad_pools + + +def _lows(steps) -> dict: + """The lowest position each request's calls touch: its first run's start - (N - 1), >= 0.""" + lows = {} + for runs in steps: + for r, _, s, _ in runs: + lows[r] = min(lows.get(r, s), max(0, s - (MAX_NGRAM - 1))) + return lows + + +def _env_case( + consts, + g, + T, + B, + D, + G, + masked, + shuffle, + width, + two_steps=True, + mask_all=False, + mask_fn=None, + max_pos=_ENV_MAX_POS, + **kw, +): + """One random partition (and its continuation) on a fresh pool, positions below ``max_pos``: + histories planted before every run that starts past 0 (DEAD every ~7th), block tables + ``width`` columns wide (the engine's max_blocks_per_seq for the partition's L, _engine_width; + test_block_table_width_only_sizes_indexing covers 128..132 and wider). Each request's pages + below its lowest touched position are unmapped (-1), as in the engine's windowed pool.""" + steps, ends = _plan_pair(g, T, B, D, G, masked, shuffle, mask_all, mask_fn, max_pos) + if not two_steps: + steps = steps[:1] + lows = _lows(steps) + pool = _Pool(ends, 128, num_pages=4096, seed=_ri(g, 0, 10**6), lows=lows) + assert pool.max_blocks <= width, (pool.max_blocks, width) + pool.max_blocks = width + refs = {r: _RefState(consts, ends[r] + 8, base=lows[r]) for r in ends} + for r, _, s, _ in steps[0]: + _plant_history(consts, pool, refs, r, s, g) + expected = pool.pool.clone() + bad = _run_checked(consts, pool, refs, expected, steps, **kw) + return bad, steps + + +_ENV_EDGES = [ # (T, B, D, G, L, DSpark): the envelope's corners, run before the random partitions + (1, 1, 0, 0, 16384, False), # one fresh 1-token context + (1, 1, 0, 1, 1048576, False), # one decode row + (8192, 1, 0, 0, 16384, False), # one 8192-token chunk + (8192, 32, 0, 0, 16384, True), # 32 contexts filling the gates' max_num_tokens + (8192, 32, 5, 31, 16384, True), # 31 verify runs + one 8006-token context + (8192, 32, 0, 31, 16384, False), # 31 decode rows + one 8161-token context + (32, 32, 0, 32, 16384, False), # a full decode batch of the gates + (192, 32, 5, 32, 16384, True), # a full DSpark verify batch of the gates + (33, 32, 0, 31, 16384, False), # 31 decode rows + a 2-token context + (16384, 1, 0, 0, 1048576, True), # one 16384-token chunk + (16384, 1, 0, 0, 16384, False), # one 16384-token chunk that is the whole request + (16384, 256, 0, 0, 1048576, False), # 256 contexts filling max_num_tokens + (16384, 256, 5, 255, 1048576, True), # 255 verify runs + one 14854-token context + (16384, 256, 0, 255, 262144, False), # 255 decode rows + one 16129-token context + (256, 256, 0, 256, 1048576, False), # a full decode batch of 256 + (1536, 256, 5, 256, 1048576, True), # a full DSpark verify batch of 256 + (257, 256, 0, 255, 32768, False), # 255 decode rows + a 2-token context + (1531, 256, 5, 255, 65536, True), # 255 verify runs + a 1-token context +] +_ENV_PARTITIONS = 264 # every B in 1..256 at least once, every (L, DSpark) at least 22 times + + +def _env_shapes(g: torch.Generator, count: int, need_context: bool = False) -> list: + """``count`` random (T, B, D, G, L, DSpark): B = k % 256 + 1 for the k-th (every B once in + 256), (L, DSpark) cycling over the 12 configurations, D = 5 (DSpark verify runs) or 0 + (decode rows) with DSpark, 0 without, T from _random_T (1..16384) with T >= B, G a random + feasible generation-run count (< B when ``need_context``).""" + shapes = [] + while len(shapes) < count: + k = len(shapes) + B = k % _ENV_MAX_B + 1 + L, ds = _L_CONFIGS[k % len(_L_CONFIGS)] + T, D = _random_T(g), (5 * _ri(g, 0, 1) if ds else 0) + if T < B: + continue + feas = [G for G in _feasible_G(T, B, D) if G < B or not need_context] + if feas: + shapes.append((T, B, D, feas[_ri(g, 0, len(feas) - 1)], L, ds)) + return shapes + + +def test_envelope_random_partitions(): + """>= 200 seeded random partitions of T in [1, 16384] rows into B in [1, 256] runs (the + edges above first; then 264 random, every B in 1..256 at least once), each at a per-request + length L in {16384, 32768, 65536, 131072, 262144, 1048576} with DSpark on (positions below + Lm = L + 14) or off, and each followed by its continuation step: random context / generation + mix, D in {0, 5} with DSpark, random cached starts incl. odd ones up to Lm - 1, runs ending + at the far end Lm, random run order, random page placement, image spans DEAD through + token_mask on half of the partitions (spans inside chunks and across the step boundary). + Every call (both steps) equals the reference, and after every call the whole pool equals + the planted histories plus exactly the writes so far. Block tables are the engine's width + for the partition's L (_engine_width: 132 .. 8196 columns); pages below a request's window + are unmapped.""" + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(23_2023) + shapes = [] + for T, B, D, G, L, ds in _ENV_EDGES: + assert G in _feasible_G(T, B, D), (T, B, D, G) + shapes.append((T, B, D, G, L, ds)) + shapes += _env_shapes(g, _ENV_PARTITIONS) + bad_rows = bad_pools = calls = rows = masked_calls = 0 + seen_T2, seen_B, per_cfg, far = [], set(), {}, {} + for k, (T, B, D, G, L, ds) in enumerate(shapes): + masked, shuffle = bool(k % 2), bool((k // 2) % 2) + lm, width = _lm(L, ds), _engine_width(L, ds) + (br, bp), steps = _env_case(consts, g, T, B, D, G, masked, shuffle, width, max_pos=lm) + bad_rows, bad_pools = bad_rows + br, bad_pools + bp + for runs in steps: + n = sum(ids.numel() for _, ids, _, _ in runs) + assert 1 <= n <= _ENV_MAX_T and 1 <= len(runs) <= _ENV_MAX_B, (n, len(runs)) + last = max(s + ids.numel() for _, ids, s, _ in runs) + assert last <= lm, (last, lm) + far[(L, ds)] = far.get((L, ds), 0) + sum( + int(s + ids.numel() == lm) for _, ids, s, _ in runs + ) + calls, rows = calls + 1, rows + n + masked_calls += int(any(m is not None and not bool(m.all()) for *_, m in runs)) + seen_T2.append(n) + seen_B.add(B) + per_cfg[(L, ds)] = per_cfg.get((L, ds), 0) + 1 + print( + f"\n {len(shapes)} partitions ({calls} calls, {rows} rows; T {min(seen_T2)}..{max(seen_T2)}, " + f"B {min(seen_B)}..{max(seen_B)} ({len(seen_B)} distinct), {masked_calls} calls with DEAD image rows): " + f"{bad_rows} rows differ from the reference, {bad_pools} calls left the pool inexact" + ) + print( + " per (L, DSpark): partitions / runs ending at Lm (the far end, block-table column " + "(Lm - 1) // 128 of a table of _engine_width columns): " + + "; ".join( + f"L {L} {'on' if ds else 'off'} (Lm {_lm(L, ds)}, width {_engine_width(L, ds)}): " + f"{per_cfg[(L, ds)]} / {far[(L, ds)]}" + for L, ds in _L_CONFIGS + ) + ) + assert ( + len(shapes) >= 200 + and seen_B == set(range(1, _ENV_MAX_B + 1)) + and max(seen_T2) == _ENV_MAX_T + ) + assert all(per_cfg[c] >= 20 and far[c] > 0 for c in _L_CONFIGS), (per_cfg, far) + assert bad_rows == 0 and bad_pools == 0 + + +def test_envelope_harness_sees_a_rewind_skipped_and_a_wrong_mask(): + """Controls for the envelope check: the same partitions judged against a reference that + (a) ignores token_mask or (b) never rewinds (every verify run's next start is taken as + start + 1 + D) must fail it.""" + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(4242) + # (a) masks dropped from the reference only + steps, ends = _plan_pair(g, 3000, 8, 0, 0, masked=True, shuffle=False, mask_all=True) + assert any(m is not None for runs in steps for *_, m in runs) + pool = _Pool(ends, 128, seed=1) + refs = {r: _RefState(consts, _ENV_MAX_POS + 8) for r in ends} + for r, _, s, _ in steps[0]: + _plant_history(consts, pool, refs, r, s, g) + unmasked = [[(r, ids, s, None) for r, ids, s, _ in runs] for runs in steps] + bad_a = 0 + for runs, plain in zip(steps, unmasked): + got = _call(consts, pool, runs) + at = 0 + for r, ids, s, _ in plain: + bad_a += int( + (got[at : at + ids.numel()] != refs[r].forward(ids, s)).flatten(1).any(1).sum() + ) + at += ids.numel() + # (b) verify: the op rewinds (the next run starts after the accepted tokens); a reference + # fed the rejected drafts as history instead differs on the rows that read them + D, B = 5, 16 + pool = _Pool({r: 400 for r in range(B)}, 128, seed=2) + refs = {r: _RefState(consts, 400) for r in range(B)} + starts = {r: _ri(g, 10, 300) for r in range(B)} + for r in range(B): + _plant_history(consts, pool, refs, r, starts[r], g) + step1 = [(r, _ids(1 + D, g), starts[r], None) for r in range(B)] + _call(consts, pool, step1) + for r, ids, s, _ in step1: + refs[r].forward(ids, s) + acc = {r: _ri(g, 0, D - 1) for r in range(B)} # at least one draft rejected + step2 = [(r, _ids(1 + D, g), starts[r] + 1 + acc[r], None) for r in range(B)] + got = _call(consts, pool, step2) + bad_b, at = 0, 0 + for r, ids, s, _ in step2: + # the no-rewind reference: the rejected drafts stay in place and this run lands after them + wrong = _RefState(consts, 400) + wrong.cache.copy_(refs[r].cache) + bad_b += int( + (got[at : at + ids.numel()] != wrong.forward(ids, starts[r] + 1 + D)) + .flatten(1) + .any(1) + .sum() + ) + at += ids.numel() + print( + f" [control] reference without token_mask: {bad_a} rows differ; reference without rewind: " + f"{bad_b} of {B * (1 + D)} rows differ" + ) + assert bad_a > 0 and bad_b > 0 + + +# ─── the multimodal target's masked call over the envelope ─── +# +# The multimodal target passes ``token_mask`` = keep, a contiguous bool [T]: False on the +# image-span rows of its context chunks (DEAD), True on text rows and on every generation +# row. It passes the mask only in a step that holds at least one context chunk of an image +# request (eager; captured steps are generation-only and pass None). Such a step can be mixed +# with decode (D = 0) or DSpark verify (D = 5) runs, a chunk may lie wholly inside a span (all +# DEAD) or hold no span row (mask present, all True), and a span of one image continues +# across chunked-prefill steps. Every partition below puts a mask on every context run. + +_MM_SPAN_MAX = ( + 4096 # span length drawn log-uniform in [1, 4096] (MMMU spans run to ~1.1k rows per image) +) + + +def _log_len(g: torch.Generator, hi: int) -> int: + """Length in [1, hi], log-uniform (short and long spans both common).""" + if hi <= 1: + return 1 + u = float(torch.rand((), generator=g)) + return max(1, min(hi, int(round(hi**u)))) + + +def _mm_mask(kind: str): + """A ``mask_fn`` for ``_plan_pair``: a position mask over [0, size) for a context + chunk at [s, s + n). ``text``: all True (mask present, no span); ``dead``: all False (the + whole request is a span); ``one``: one span inside the chunk; ``multi``: 3..6 spans; + ``cross``: a span starting in the chunk's last rows that runs past the step boundary; + ``ends``: a span from position 0 and one ending exactly at the chunk end; ``mixed``: a + random pick of the others per request (the sweep's kind).""" + + def fn(g, s, n, size=_ENV_MAX_POS + 8): + def span(m, a, ln): + a = max(0, min(a, size - 1)) + m[a : a + ln] = False + + k = kind + if k == "mixed": + k = ("text", "dead", "one", "multi", "multi", "cross", "cross", "ends", "rand", "rand")[ + _ri(g, 0, 9) + ] + m = torch.ones(size, dtype=torch.bool) + if k == "text": + return m + if k == "dead": + return torch.zeros(size, dtype=torch.bool) + if k == "one": + ln = _log_len(g, n) + span(m, s + _ri(g, 0, n - ln), ln) + elif k == "multi": + for _ in range(_ri(g, 3, 6)): + ln = _log_len(g, max(1, n // 3)) + span(m, s + _ri(g, 0, max(0, n - ln)), ln) + elif k == "cross": + a = s + n - _ri(g, 1, min(n, 64)) + span(m, a, (s + n - a) + _ri(g, 1, _MM_SPAN_MAX)) + elif k == "ends": + span(m, 0, _log_len(g, s + n)) + ln = _log_len(g, n) + span(m, s + n - ln, ln) + else: # rand: 1..6 spans anywhere around the chunk, page edges and history included + for _ in range(_ri(g, 1, 6)): + ln = _log_len(g, _MM_SPAN_MAX) + anchor = ( + s - _ri(g, 0, 300), + s + _ri(g, 0, max(0, n - 1)), + s + n - _ri(g, 0, 300), + 128 * _ri(g, 0, (s + n) // 128) - _ri(g, 0, 8), + 0, + )[_ri(g, 0, 4)] + span(m, anchor, ln) + return m + + return fn + + +# (name, T, B, D, G, mask kind, L, DSpark): the masked call's corners, run before the random partitions +_MM_EDGES = [ + ("one-row DEAD context", 1, 1, 0, 0, "dead", 16384, False), + ("one context wholly inside a span", 300, 1, 0, 0, "dead", 1048576, False), + ("32 contexts, every row DEAD", 8192, 32, 0, 0, "dead", 16384, False), + ("DEAD 2-row context + 31 decode rows", 33, 32, 0, 31, "dead", 16384, False), + ("mask present, all True, one row", 1, 1, 0, 0, "text", 16384, False), + ("mask present, all True, with verify runs", 4000, 8, 5, 3, "text", 131072, True), + ("one span", 2048, 1, 0, 0, "one", 16384, False), + ("several spans per context", 8192, 4, 0, 0, "multi", 16384, False), + ("spans across the step boundary", 6000, 3, 0, 0, "cross", 16384, False), + ("spans across the step boundary, 32 contexts", 8192, 32, 0, 0, "cross", 16384, False), + ("spans from position 0 and to the chunk end", 8192, 32, 0, 0, "ends", 16384, False), + ("31 verify runs + a masked 8006-row context", 8192, 32, 5, 31, "mixed", 16384, True), + ("31 decode rows + a masked 8161-row context", 8192, 32, 0, 31, "mixed", 16384, False), + ("full DSpark verify batch + a 1-row DEAD context", 187, 32, 5, 31, "dead", 16384, True), + ("256 contexts, every row DEAD", 16384, 256, 0, 0, "dead", 1048576, False), + ("one 16384-row context, several spans", 16384, 1, 0, 0, "multi", 1048576, True), + ("spans across the step boundary, 256 contexts", 16384, 256, 0, 0, "cross", 262144, False), + ( + "spans from position 0 and to the chunk end, 256 contexts", + 16384, + 256, + 0, + 0, + "ends", + 65536, + False, + ), + ("255 verify runs + a masked 14854-row context", 16384, 256, 5, 255, "mixed", 1048576, True), + ("255 decode rows + a masked 16129-row context", 16384, 256, 0, 255, "mixed", 32768, False), + ("256-request verify batch + a 1-row DEAD context", 1531, 256, 5, 255, "dead", 1048576, True), +] + + +def _mm_shapes(g: torch.Generator, count: int) -> list: + """``count`` random (T, B, D, G, L, DSpark) with at least one context run (G < B), every B in + 1..256 once per 256, the 12 (L, DSpark) configurations in turn.""" + return _env_shapes(g, count, need_context=True) + + +def _mm_stats(steps, st: dict) -> None: + """Tally what the masked calls of one partition covered.""" + for runs in steps: + n = sum(ids.numel() for _, ids, _, _ in runs) + st["calls"] += 1 + st["rows"] += n + masks = [m for *_, m in runs] + if not any(m is not None for m in masks): + st["unmasked_calls"] += 1 + continue + keep = torch.cat( + [torch.ones(ids.numel(), dtype=torch.bool) if m is None else m for _, ids, _, m in runs] + ) + dead = int((~keep).sum()) + gen_lens = {ids.numel() for _, ids, _, m in runs if m is None} + st["masked_calls"] += 1 + st["masked_rows"] += n + st["dead_rows"] += dead + st["T_max"] = max(st["T_max"], n) + st["B_seen"].add(len(runs)) + st["all_dead_calls"] += int(dead == n) + st["all_true_calls"] += int(dead == 0) + st["mixed_decode_calls"] += int(1 in gen_lens and dead > 0) + st["mixed_verify_calls"] += int(6 in gen_lens and dead > 0) + for _, ids, s, m in runs: + if m is None or bool(m.all()): + continue + st["span_runs"] += 1 + st["pos0_dead"] += int(s == 0 and not bool(m[0])) + st["dead_at_run_end"] += int(not bool(m[-1])) + st["dead_at_run_start_after_history"] += int(s > 0 and not bool(m[0])) + + +def _mm_new_stats() -> dict: + keys = ( + "calls", + "rows", + "unmasked_calls", + "masked_calls", + "masked_rows", + "dead_rows", + "T_max", + "all_dead_calls", + "all_true_calls", + "mixed_decode_calls", + "mixed_verify_calls", + "span_runs", + "pos0_dead", + "dead_at_run_end", + "dead_at_run_start_after_history", + ) + st = {k: 0 for k in keys} + st["B_seen"] = set() + return st + + +def test_multimodal_masked_calls_over_the_envelope(): + """The multimodal target's masked call (``token_mask`` = keep, bool [T]) over the + envelope: the 21 corners above, then 264 seeded random partitions of T in [1, 16384] + rows into B in [1, 256] runs (every B at least once), at every (L, DSpark) of _L_CONFIGS + (positions below Lm, block tables of the engine width for L), D in {0, 5} with DSpark, at + least one context run each, and a mask on + EVERY context run (random kind per request: no span, all DEAD, one span, several spans, + spans across the step boundary, from position 0 / to the chunk end, 1..6 random spans of + 1..4096 rows around the chunk, page edges and history included); generation rows True. + Each partition's continuation step follows (verify rewinds, contexts continuing inside a + span). Every call equals the reference, and after every call the whole pool equals the + planted histories plus exactly the writes so far (DEAD stored as -1). The kernel's + compiled variants are pinned: after one masked and one unmasked call per ``num_tokens`` + class at the engine's 132-column tables, the whole sweep (tables 132 .. 8196 columns wide) + compiles nothing.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.embedding.engram_hash import ( # noqa: E501 + _engram_hash_kernel as fn, + ) + + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(13_2026) + # pin the variants: T == 1, T % 16 == 0, other; masked and unmasked; int32, width 132 + _jit_clear(fn) + for T, masked in ((1, True), (16, True), (37, True), (1, False), (16, False), (37, False)): + (br, bp), _ = _env_case( + consts, + g, + T, + 1, + 0, + 0, + masked, + False, + _ENGINE_WIDTH, + two_steps=False, + mask_all=masked, + mask_fn=_mm_mask("one") if masked else None, + ) + assert br == 0 and bp == 0, (T, masked, br, bp) + pinned = _jit_keys(fn) + assert len(pinned) == 6, f"expected one variant per (num_tokens class, mask), got {len(pinned)}" + + st = _mm_new_stats() + bad_rows = bad_pools = 0 + for name, T, B, D, G, kind, L, ds in _MM_EDGES: + assert G in _feasible_G(T, B, D) and G < B, (name, T, B, D, G) + (br, bp), steps = _env_case( + consts, + g, + T, + B, + D, + G, + True, + False, + _engine_width(L, ds), + mask_all=True, + mask_fn=_mm_mask(kind), + max_pos=_lm(L, ds), + ) + assert br == 0 and bp == 0, f"{name}: {br} rows differ, {bp} pools inexact" + assert _call_masked(steps[0]), name + _mm_stats(steps, st) + edges = {k: (set(v) if isinstance(v, set) else v) for k, v in st.items()} + shapes = _mm_shapes(g, _ENV_PARTITIONS) + cfgs = set() + for k, (T, B, D, G, L, ds) in enumerate(shapes): + (br, bp), steps = _env_case( + consts, + g, + T, + B, + D, + G, + True, + bool(k % 2), + _engine_width(L, ds), + mask_all=True, + mask_fn=_mm_mask("mixed"), + max_pos=_lm(L, ds), + ) + bad_rows, bad_pools = bad_rows + br, bad_pools + bp + assert _call_masked(steps[0]), (k, T, B, D, G) + _mm_stats(steps, st) + cfgs.add((L, ds)) + compiled = _jit_keys(fn) - pinned + print( + f"\n masked envelope: {len(_MM_EDGES)} corners + {len(shapes)} random partitions, every one's first call " + f"masked; {st['calls']} calls ({st['rows']} rows), {st['masked_calls']} with a token_mask " + f"({st['masked_rows']} rows, {st['dead_rows']} DEAD; T up to {st['T_max']}, " + f"{len(st['B_seen'])} distinct B), {st['unmasked_calls']} continuation calls without one; " + f"all-DEAD calls {st['all_dead_calls']}, all-True masked calls {st['all_true_calls']}, " + f"DEAD rows beside decode runs in {st['mixed_decode_calls']} calls and beside verify runs in " + f"{st['mixed_verify_calls']}; {st['span_runs']} runs with spans: DEAD at position 0 in " + f"{st['pos0_dead']}, DEAD first row after history in {st['dead_at_run_start_after_history']}, " + f"DEAD last row in {st['dead_at_run_end']}; {bad_rows} rows differ from the reference, " + f"{bad_pools} calls left the pool inexact; {len(pinned)} pinned variants, {len(compiled)} compiled after" + ) + assert bad_rows == 0 and bad_pools == 0 + assert not compiled, ( + f"the masked sweep compiled {len(compiled)} variant(s) beyond the pinned ones" + ) + assert ( + len(shapes) >= 200 + and st["B_seen"] == set(range(1, _ENV_MAX_B + 1)) + and st["T_max"] == _ENV_MAX_T + ) + assert cfgs == set(_L_CONFIGS) + assert st["masked_calls"] >= len(shapes) + len(_MM_EDGES) + # An all-DEAD call needs every context of the step to draw the all-DEAD kind, which the random + # partitions (B up to 256) rarely do; the corners hold such calls at B 1, 32 and 256. + assert edges["all_dead_calls"] >= 3 + for key in ( + "all_true_calls", + "mixed_decode_calls", + "mixed_verify_calls", + "pos0_dead", + "dead_at_run_end", + "dead_at_run_start_after_history", + ): + assert edges[key] >= 1 and st[key] > edges[key], (key, edges[key], st[key]) + + +def _call_masked(runs) -> bool: + """Whether the op receives a token_mask for this step (some run carries a mask).""" + return any(m is not None for *_, m in runs) + + +def test_multimodal_masked_harness_sees_a_wrong_mask(): + """Controls for the masked sweep: 24 of its partitions (same planner, fresh seed) judged + against a reference that (a) ignores ``token_mask`` or (b) applies it one row late. Every + partition where the wrong mask differs from the true one must fail.""" + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(13_2027) + fails = {"ignored": [], "one row late": []} + for T, B, D, G, _, _ in _mm_shapes(g, 24): + steps, ends = _plan_pair(g, T, B, D, G, True, True, True, _mm_mask("mixed")) + seed = _ri(g, 0, 10**6) + for ctl, wrong in ( + ("ignored", lambda m: None), + ("one row late", lambda m: None if m is None else torch.cat([m[:1], m[:-1]])), + ): + # the wrong mask changes a row: a DEAD row it drops, or a row it shifts + differs = any( + m is not None + and (not bool(m.all()) if wrong(m) is None else not torch.equal(wrong(m), m)) + for runs in steps + for *_, m in runs + ) + pool = _Pool(ends, 128, num_pages=4096, seed=seed) + pool.max_blocks = _ENGINE_WIDTH + refs = {r: _RefState(consts, _ENV_MAX_POS + 8) for r in ends} + gh = torch.Generator().manual_seed(seed) + for r, _, s, _ in steps[0]: + _plant_history(consts, pool, refs, r, s, gh) + bad = 0 + for runs in steps: + got = _call(consts, pool, runs) + at = 0 + for r, ids, s, m in runs: + bad += int( + (got[at : at + ids.numel()] != refs[r].forward(ids, s, wrong(m))) + .flatten(1) + .any(1) + .sum() + ) + at += ids.numel() + if differs: + fails[ctl].append(bad) + print( + f" [control] partitions where the wrong mask differs: ignored {len(fails['ignored'])}, one row late " + f"{len(fails['one row late'])}; rows differing per partition, mask ignored: min {min(fails['ignored'])} " + f"max {max(fails['ignored'])}; one row late: min {min(fails['one row late'])} " + f"max {max(fails['one row late'])}" + ) + assert len(fails["ignored"]) >= 20 and len(fails["one row late"]) >= 15 + assert all(b > 0 for b in fails["ignored"]) and all(b > 0 for b in fails["one row late"]) + + +# ─── the kernel-variant map over the envelope; B is not a launch argument ─── +# +# The launch is ``_engram_hash_kernel[(ceil(T / 64),)]`` with T as ``num_tokens``; B reaches +# the kernel only through the contents of req_idx_per_token / cu_seq_lens / start_pos / +# block_table, never as an argument. Triton specializes integer arguments (== 1, divisible by +# 16) and dtypes, and HAS_MASK is a constexpr, so the variants over the envelope are indexed +# by (num_tokens class, block-table width class, mask, index dtypes); the map below compiles +# and checks each, then pins that a different B (1..256) or T within a class compiles nothing. +# Widths: 128 (divisible by 16), the engine's 132 at L 16384, 129 (neither), and the engine's +# widths at L 32768 .. 1048576 (260, 516, 1028, 2052, 8196: 4 mod 16, neither) with positions up +# to that L's Lm; a class whose Triton specialization an earlier class already compiled must +# compile nothing. +_MAP_WIDTHS = ( + (128, 16384, False), + (_ENGINE_WIDTH, 16384, False), + (129, 16384, False), + (260, 32768, True), + (516, 65536, False), + (1028, 131072, True), + (2052, 262144, False), + (8196, 1048576, True), +) + + +def test_kernel_variant_map_over_the_envelope(): + import inspect + + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.embedding.engram_hash import ( # noqa: E501 + _engram_hash_kernel as fn, + ) + + arg_names = list(inspect.signature(fn.fn).parameters) + assert "num_tokens" in arg_names and not any("batch" in a for a in arg_names), arg_names + consts = _Consts(_synthetic_token_map()) + tm64 = consts.dev.token_map.to(torch.int64) + i64 = {"ids": torch.int64, "meta": torch.int64, "table": torch.int64} + g = torch.Generator().manual_seed(64) + t_classes = (("T==1", (1,)), ("T%16==0", (16, 4096, 16384)), ("T%16!=0", (37, 1001, 16383))) + variant_of, bad_rows, bad_pools, calls, n_classes, shared = {}, 0, 0, 0, 0, [] + _jit_clear(fn) # the map starts from no variant, so every class compiles its own + assert not _jit_keys(fn) + for dname, kw in (("int32", {}), ("int64", {"dtypes": i64, "token_map": tm64})): + for masked in (False, True): + for width, L, ds in _MAP_WIDTHS: + assert width in (128, 129) or width == _engine_width(L, ds) + for tname, Ts in t_classes: + cls = f"{tname} width={width} mask={'yes' if masked else 'no'} {dname}" + first = True + for T in Ts: + for B in sorted({1, min(T, _ENV_MAX_B), _ri(g, 1, min(T, _ENV_MAX_B))}): + D = 0 + feas = _feasible_G(T, B, D) + G = ( + 0 if masked else feas[_ri(g, 0, len(feas) - 1)] + ) # masked: every run a context + before = _jit_keys(fn) + (br, bp), steps = _env_case( + consts, + g, + T, + B, + D, + G, + masked, + False, + width, + two_steps=False, + mask_all=masked, + max_pos=_lm(L, ds), + **kw, + ) + has_mask = any(m is not None for *_, m in steps[0]) + assert has_mask == masked, (cls, T, B) + new = _jit_keys(fn) - before + bad_rows, bad_pools, calls = bad_rows + br, bad_pools + bp, calls + 1 + if first: + assert len(new) <= 1, ( + f"{cls}: its first call compiled {len(new)} variants" + ) + for key in new: + variant_of[key] = f"{cls} (first at T={T} B={B})" + if not new: + shared.append(cls) + first = False + n_classes += 1 + else: + assert not new, ( + f"{cls}: T={T} B={B} compiled {len(new)} new variant(s)" + ) + # closure: random envelope partitions in the target's int32 form compile nothing further + before = _jit_keys(fn) + for k in range(40): + T, D = _random_T(g), 5 * (k % 2) + B = _ri(g, 1, min(_ENV_MAX_B, T)) + L, ds = ( + _L_CONFIGS[(2 * k + 1) % len(_L_CONFIGS)] + if D + else _L_CONFIGS[(2 * k) % len(_L_CONFIGS)] + ) + feas = _feasible_G(T, B, D) + if not feas: + continue + (br, bp), _ = _env_case( + consts, + g, + T, + B, + D, + feas[-1], + bool(k % 3), + True, + _engine_width(L, ds), + max_pos=_lm(L, ds), + ) + bad_rows, bad_pools, calls = bad_rows + br, bad_pools + bp, calls + 1 + after = _jit_keys(fn) + print( + f"\n variant map: {len(variant_of)} variant(s) compiled by this map over {calls} checked calls:" + ) + for key, cls in sorted(variant_of.items(), key=lambda kv: kv[1]): + print(f" {cls}") + print( + f" {len(shared)} class(es) reused a variant an earlier class compiled (same Triton specialization):" + ) + for cls in shared: + print(f" {cls}") + print( + f" 40 random envelope partitions after the map compiled {len(after - before)} further variant(s); " + f"{bad_rows} rows differ from the reference, {bad_pools} pools inexact" + ) + assert after == before + assert bad_rows == 0 and bad_pools == 0 + assert n_classes == 3 * 2 * 2 * len(_MAP_WIDTHS) and len(variant_of) + len(shared) == n_classes + # 129 and the engine's 132 .. 8196 are neither 1 nor a multiple of 16: one specialization + assert all(" width=128 " not in cls and f" width={_ENGINE_WIDTH} " not in cls for cls in shared) + assert len(variant_of) == 2 * 2 * 2 * 3 + + +# ─── every (T, B) the gates' runtime-value dumps show ─── + +# (T, B) of the engram_hash call in every boot-gate log and the three accuracy-preview logs, +# parsed one line per tuple with an arity check and the out-of-envelope garbles dropped (B 375). +_DUMP_TB = [ + (1, 1), + (2, 1), + (2, 2), + (3, 3), + (4, 4), + (5, 5), + (6, 6), + (7, 7), + (8, 8), + (14, 1), + (16, 1), + (16, 16), + (24, 24), + (32, 32), + (48, 1), + (524, 32), + (525, 32), + (534, 32), + (560, 32), + (580, 32), + (581, 32), + (586, 32), + (587, 32), + (593, 32), + (600, 32), + (602, 32), + (604, 32), + (609, 32), + (610, 32), + (612, 32), + (617, 32), + (621, 32), + (624, 32), + (625, 32), + (626, 32), + (628, 32), + (629, 32), + (630, 32), + (632, 32), + (634, 32), + (637, 32), + (640, 32), + (643, 32), + (645, 32), + (647, 32), + (648, 32), + (650, 32), + (654, 32), + (656, 32), + (657, 32), + (658, 32), + (660, 32), + (661, 32), + (662, 32), + (663, 32), + (665, 32), + (666, 32), + (668, 32), + (671, 32), + (672, 32), + (673, 32), + (674, 32), + (675, 32), + (677, 32), + (680, 32), + (681, 32), + (682, 32), + (683, 32), + (684, 32), + (685, 32), + (687, 32), + (691, 32), + (692, 32), + (693, 32), + (694, 32), + (696, 32), + (697, 32), + (698, 32), + (699, 32), + (701, 32), + (703, 32), + (705, 32), + (706, 32), + (707, 32), + (710, 32), + (711, 32), + (712, 32), + (713, 32), + (714, 32), + (715, 32), + (716, 32), + (718, 32), + (719, 32), + (720, 32), + (721, 32), + (722, 32), + (723, 32), + (724, 32), + (725, 32), + (727, 32), + (728, 32), + (729, 32), + (730, 32), + (731, 32), + (732, 32), + (733, 32), + (734, 32), + (735, 32), + (737, 32), + (739, 32), + (741, 32), + (742, 32), + (743, 32), + (744, 32), + (745, 32), + (746, 32), + (749, 32), + (750, 32), + (751, 32), + (752, 32), + (753, 32), + (754, 32), + (755, 32), + (756, 32), + (757, 32), + (758, 32), + (759, 32), + (760, 32), + (761, 32), + (763, 32), + (764, 32), + (765, 32), + (766, 32), + (767, 32), + (768, 32), + (769, 32), + (770, 32), + (771, 32), + (772, 32), + (773, 32), + (774, 32), + (776, 32), + (777, 32), + (778, 32), + (779, 32), + (780, 32), + (782, 32), + (783, 32), + (784, 32), + (785, 32), + (786, 32), + (787, 32), + (788, 32), + (789, 32), + (790, 32), + (791, 32), + (792, 32), + (793, 32), + (794, 32), + (795, 32), + (796, 32), + (797, 32), + (798, 32), + (799, 32), + (800, 32), + (801, 32), + (802, 32), + (803, 32), + (804, 32), + (805, 32), + (806, 32), + (807, 32), + (808, 32), + (809, 32), + (810, 32), + (811, 32), + (812, 32), + (813, 32), + (814, 32), + (815, 32), + (816, 32), + (817, 32), + (818, 32), + (819, 32), + (821, 32), + (822, 32), + (823, 32), + (824, 32), + (825, 32), + (826, 32), + (827, 32), + (828, 32), + (829, 32), + (830, 32), + (831, 32), + (832, 32), + (833, 32), + (834, 32), + (835, 32), + (836, 32), + (837, 32), + (838, 32), + (839, 32), + (840, 32), + (841, 32), + (842, 32), + (843, 32), + (844, 32), + (845, 32), + (846, 32), + (847, 32), + (848, 32), + (849, 32), + (850, 32), + (851, 32), + (853, 32), + (854, 32), + (855, 32), + (856, 32), + (857, 32), + (858, 32), + (859, 32), + (860, 32), + (861, 32), + (862, 32), + (863, 32), + (864, 32), + (865, 32), + (866, 32), + (867, 32), + (868, 32), + (869, 32), + (870, 32), + (871, 32), + (872, 32), + (874, 32), + (875, 32), + (876, 32), + (877, 32), + (878, 32), + (879, 32), + (880, 32), + (881, 32), + (882, 32), + (883, 32), + (884, 32), + (885, 32), + (886, 32), + (887, 32), + (888, 1), + (888, 32), + (889, 32), + (890, 1), + (890, 32), + (891, 32), + (892, 32), + (893, 32), + (894, 32), + (895, 32), + (897, 32), + (898, 32), + (899, 32), + (900, 32), + (901, 32), + (902, 32), + (903, 32), + (904, 32), + (905, 32), + (906, 32), + (907, 32), + (908, 32), + (909, 32), + (910, 32), + (911, 32), + (912, 32), + (913, 32), + (914, 32), + (915, 32), + (916, 32), + (917, 32), + (918, 32), + (919, 32), + (920, 32), + (921, 32), + (922, 32), + (923, 32), + (924, 32), + (925, 32), + (926, 32), + (927, 32), + (928, 32), + (929, 32), + (930, 32), + (931, 31), + (931, 32), + (932, 32), + (933, 32), + (934, 32), + (935, 32), + (937, 32), + (938, 32), + (939, 32), + (940, 32), + (941, 32), + (942, 32), + (943, 32), + (944, 32), + (945, 32), + (946, 32), + (947, 32), + (948, 32), + (949, 32), + (950, 32), + (951, 32), + (952, 32), + (953, 32), + (954, 32), + (955, 32), + (956, 32), + (957, 32), + (958, 32), + (959, 32), + (960, 32), + (961, 32), + (962, 32), + (963, 32), + (964, 32), + (965, 32), + (966, 32), + (967, 32), + (968, 32), + (969, 32), + (970, 32), + (971, 32), + (972, 32), + (973, 32), + (974, 32), + (975, 32), + (976, 32), + (977, 32), + (978, 32), + (980, 32), + (981, 32), + (982, 32), + (983, 32), + (984, 32), + (985, 32), + (986, 32), + (987, 32), + (988, 32), + (989, 32), + (990, 32), + (991, 32), + (992, 32), + (993, 32), + (994, 32), + (995, 32), + (996, 32), + (997, 32), + (998, 32), + (999, 32), + (1000, 32), + (1001, 32), + (1003, 32), + (1004, 32), + (1005, 32), + (1006, 32), + (1007, 32), + (1008, 32), + (1010, 32), + (1011, 32), + (1013, 32), + (1014, 32), + (1015, 32), + (1017, 32), + (1018, 3), + (1018, 32), + (1019, 32), + (1020, 32), + (1021, 32), + (1022, 32), + (1023, 32), + (1024, 32), + (1025, 32), + (1027, 32), + (1028, 32), + (1029, 32), + (1030, 32), + (1031, 32), + (1033, 32), + (1034, 32), + (1035, 32), + (1036, 32), + (1038, 32), + (1039, 32), + (1042, 32), + (1043, 32), + (1044, 32), + (1045, 32), + (1046, 32), + (1047, 32), + (1049, 32), + (1050, 32), + (1051, 32), + (1052, 32), + (1053, 32), + (1054, 32), + (1056, 32), + (1057, 32), + (1059, 32), + (1062, 32), + (1063, 32), + (1065, 32), + (1067, 32), + (1069, 32), + (1070, 32), + (1071, 32), + (1073, 32), + (1074, 32), + (1076, 32), + (1077, 32), + (1079, 32), + (1081, 32), + (1086, 32), + (1088, 32), + (1089, 32), + (1093, 32), + (1095, 32), + (1096, 32), + (1097, 32), + (1098, 32), + (1099, 32), + (1101, 32), + (1103, 32), + (1107, 32), + (1110, 32), + (1111, 32), + (1112, 32), + (1113, 32), + (1115, 32), + (1116, 32), + (1117, 32), + (1118, 32), + (1120, 32), + (1121, 32), + (1122, 32), + (1123, 32), + (1124, 32), + (1125, 32), + (1126, 32), + (1129, 32), + (1130, 32), + (1131, 32), + (1132, 32), + (1135, 32), + (1137, 32), + (1140, 32), + (1142, 32), + (1146, 32), + (1149, 32), + (1150, 32), + (1151, 32), + (1153, 32), + (1156, 32), + (1157, 32), + (1159, 32), + (1160, 32), + (1161, 32), + (1166, 32), + (1167, 32), + (1168, 32), + (1170, 32), + (1171, 32), + (1172, 32), + (1173, 32), + (1174, 32), + (1176, 32), + (1177, 32), + (1180, 32), + (1181, 32), + (1185, 32), + (1187, 1), + (1188, 32), + (1191, 32), + (1201, 32), + (1202, 32), + (1203, 32), + (1208, 32), + (1209, 32), + (1213, 32), + (1216, 32), + (1229, 32), + (1231, 32), + (1232, 32), + (1235, 32), + (1242, 32), + (1248, 32), + (1257, 32), + (1267, 32), + (1268, 32), + (1271, 32), + (1292, 32), + (1304, 32), + (1305, 32), + (1313, 32), + (1329, 32), + (1332, 32), + (1343, 32), + (1363, 32), + (1374, 32), + (1375, 32), + (1385, 32), + (1404, 32), + (1418, 32), + (1421, 32), + (1425, 32), + (1427, 32), + (1433, 32), + (1437, 32), + (1438, 32), + (1440, 32), + (1449, 32), + (1453, 32), + (1454, 32), + (1462, 32), + (1463, 32), + (1469, 32), + (1478, 32), + (1492, 32), + (1494, 32), + (1496, 32), + (1497, 32), + (1502, 32), + (1503, 32), + (1507, 32), + (1511, 32), + (1512, 32), + (1514, 32), + (1524, 32), + (1526, 32), + (1527, 32), + (1528, 32), + (1532, 32), + (1535, 32), + (1537, 32), + (1538, 32), + (1540, 32), + (1541, 32), + (1542, 32), + (1543, 32), + (1544, 32), + (1548, 32), + (1549, 32), + (1550, 32), + (1553, 32), + (1554, 32), + (1555, 32), + (1558, 32), + (1563, 32), + (1565, 32), + (1567, 32), + (1568, 32), + (1569, 32), + (1572, 32), + (1573, 32), + (1576, 32), + (1577, 32), + (1580, 32), + (1581, 32), + (1584, 32), + (1591, 32), + (1592, 32), + (1596, 32), + (1597, 32), + (1598, 32), + (1599, 32), + (1600, 32), + (1601, 32), + (1604, 32), + (1605, 32), + (1609, 32), + (1610, 32), + (1613, 32), + (1614, 32), + (1615, 32), + (1621, 32), + (1623, 32), + (1624, 32), + (1625, 32), + (1626, 32), + (1631, 32), + (1635, 32), + (1636, 32), + (1638, 32), + (1640, 32), + (1644, 32), + (1645, 32), + (1646, 32), + (1650, 32), + (1651, 32), + (1654, 32), + (1655, 32), + (1656, 32), + (1657, 32), + (1659, 32), + (1660, 32), + (1666, 32), + (1667, 32), + (1669, 32), + (1670, 32), + (1671, 32), + (1672, 32), + (1674, 32), + (1676, 32), + (1677, 32), + (1679, 32), + (1684, 32), + (1685, 32), + (1686, 32), + (1691, 32), + (1692, 32), + (1693, 32), + (1694, 32), + (1695, 32), + (1696, 32), + (1697, 32), + (1698, 32), + (1699, 32), + (1700, 32), + (1702, 32), + (1705, 32), + (1706, 32), + (1708, 32), + (1710, 32), + (1712, 32), + (1713, 32), + (1714, 32), + (1715, 32), + (1716, 32), + (1717, 32), + (1718, 32), + (1719, 32), + (1721, 32), + (1723, 32), + (1724, 32), + (1725, 32), + (1726, 32), + (1729, 32), + (1730, 32), + (1731, 32), + (1732, 32), + (1733, 32), + (1734, 32), + (1735, 32), + (1736, 32), + (1737, 32), + (1739, 32), + (1740, 32), + (1742, 32), + (1743, 32), + (1744, 32), + (1746, 32), + (1749, 32), + (1750, 32), + (1753, 32), + (1754, 32), + (1755, 32), + (1756, 32), + (1758, 32), + (1761, 32), + (1764, 32), + (1765, 32), + (1767, 32), + (1768, 32), + (1770, 32), + (1771, 32), + (1772, 32), + (1773, 32), + (1774, 32), + (1777, 32), + (1778, 32), + (1779, 32), + (1780, 32), + (1782, 32), + (1783, 32), + (1784, 32), + (1785, 32), + (1786, 32), + (1787, 32), + (1788, 32), + (1789, 32), + (1790, 32), + (1794, 32), + (1796, 32), + (1799, 32), + (1800, 32), + (1801, 32), + (1802, 32), + (1803, 32), + (1805, 32), + (1806, 32), + (1810, 32), + (1811, 32), + (1812, 32), + (1813, 32), + (1815, 32), + (1820, 32), + (1821, 32), + (1823, 32), + (1824, 32), + (1827, 32), + (1829, 32), + (1831, 32), + (1832, 32), + (1833, 32), + (1834, 32), + (1835, 32), + (1838, 32), + (1840, 32), + (1842, 32), + (1844, 32), + (1847, 32), + (1848, 32), + (1849, 32), + (1850, 32), + (1851, 32), + (1853, 32), + (1854, 32), + (1855, 32), + (1856, 32), + (1858, 32), + (1862, 32), + (1866, 32), + (1867, 32), + (1868, 32), + (1873, 32), + (1874, 32), + (1880, 32), + (1882, 32), + (1885, 32), + (1886, 32), + (1887, 32), + (1889, 32), + (1891, 32), + (1893, 32), + (1894, 32), + (1898, 32), + (1899, 32), + (1901, 32), + (1903, 32), + (1905, 32), + (1906, 32), + (1908, 32), + (1912, 32), + (1915, 32), + (1917, 32), + (1921, 32), + (1922, 32), + (1925, 32), + (1931, 32), + (1932, 32), + (1935, 3), + (1940, 32), + (1941, 32), + (1942, 32), + (1943, 32), + (1944, 32), + (1950, 32), + (1951, 32), + (1954, 32), + (1955, 32), + (1956, 32), + (1961, 32), + (1965, 32), + (1967, 32), + (1969, 32), + (1977, 32), + (1982, 32), + (1983, 32), + (1984, 32), + (1986, 32), + (1994, 32), + (1998, 32), + (1999, 32), + (2002, 32), + (2003, 32), + (2005, 32), + (2006, 32), + (2007, 32), + (2009, 32), + (2010, 32), + (2012, 32), + (2014, 32), + (2017, 32), + (2018, 32), + (2023, 32), + (2027, 32), + (2028, 32), + (2034, 32), + (2037, 32), + (2038, 32), + (2039, 32), + (2043, 32), + (2046, 32), + (2047, 32), + (2048, 1), + (2049, 32), + (2050, 32), + (2055, 32), + (2056, 32), + (2064, 32), + (2075, 32), + (2081, 32), + (2082, 32), + (2084, 32), + (2102, 32), + (2103, 32), + (2107, 32), + (2112, 32), + (2114, 32), + (2122, 32), + (2124, 32), + (2135, 32), + (2137, 32), + (2147, 32), + (2151, 32), + (2165, 32), + (2170, 32), + (2176, 32), + (2184, 32), + (2190, 32), + (2193, 32), + (2194, 32), + (2195, 32), + (2209, 32), + (2223, 32), + (2242, 32), + (2258, 32), + (2261, 32), + (2290, 32), + (2342, 32), + (2362, 32), + (2379, 32), + (2396, 32), + (2398, 32), + (2430, 32), + (2440, 32), + (2443, 32), + (2448, 32), + (2462, 32), + (2486, 32), + (2494, 32), + (2495, 32), + (2509, 32), + (2518, 32), + (2527, 32), + (2535, 32), + (2540, 32), + (2552, 32), + (2570, 32), + (2585, 32), + (2634, 32), + (2655, 32), + (2658, 32), + (2672, 32), + (2673, 32), + (2683, 32), + (2716, 32), + (2722, 32), + (2727, 32), + (2742, 32), + (2744, 32), + (2750, 32), + (2757, 32), + (2778, 32), + (2798, 32), + (2820, 32), + (2842, 32), + (2848, 32), + (2880, 32), + (2893, 32), + (2895, 32), + (2939, 32), + (2997, 32), + (3141, 32), + (3189, 32), + (3234, 32), + (3278, 32), + (3305, 32), + (3317, 32), + (3365, 32), + (3629, 32), + (8087, 28), + (8104, 19), + (8112, 20), + (8118, 10), + (8126, 17), + (8142, 10), + (8142, 17), + (8143, 17), + (8150, 8), + (8172, 29), + (8173, 28), + (8177, 28), + (8178, 8), + (8179, 29), + (8192, 1), +] +# (T, B) of the engram_hash call in further engine dumps (max_num_tokens 16384, max_batch_size +# 256 / 32 / 8, max_seq_len 32768 / 131072 / 1048576), the values not already listed above. +_DUMP_TB_ENVELOPE = [ + (5, 1), + (6, 1), + (7, 2), + (12, 1), + (12, 2), + (18, 3), + (24, 4), + (29, 3), + (34, 3), + (48, 8), + (60, 10), + (96, 16), + (108, 18), + (144, 24), + (168, 28), + (192, 32), + (228, 38), + (240, 40), + (288, 48), + (330, 55), + (336, 56), + (384, 64), + (432, 72), + (462, 77), + (480, 80), + (528, 88), + (546, 91), + (576, 96), + (597, 1), + (610, 1), + (624, 104), + (672, 112), + (720, 120), + (768, 128), + (780, 130), + (852, 142), + (1152, 192), + (1536, 256), + (8467, 50), + (9193, 50), + (9240, 50), + (9364, 158), + (9370, 159), + (9524, 50), + (10441, 1), + (10444, 1), + (10447, 2), + (10451, 1), + (16298, 25), + (16308, 74), + (16316, 145), + (16320, 133), + (16322, 146), + (16327, 23), + (16328, 23), + (16329, 158), + (16333, 105), + (16340, 40), + (16378, 24), + (16384, 1), +] + + +def test_every_dumped_step_against_reference(): + """Every dumped (T, B) built as one call with exactly those values (B - G contexts, + G decode rows; histories planted, random starts, half with image spans), against the + reference, the pool exact afterwards. The gates' points run at L 16384 (132-column tables); + the envelope engines' points at the (L, DSpark) configurations in turn (positions below + Lm, tables of the engine width for L).""" + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(907) + bad_rows = bad_pools = 0 + points = [(T, B, 16384, False) for T, B in _DUMP_TB] + points += [ + (T, B) + _L_CONFIGS[k % len(_L_CONFIGS)] for k, (T, B) in enumerate(_DUMP_TB_ENVELOPE) + ] + assert not set(_DUMP_TB) & set(_DUMP_TB_ENVELOPE) + for k, (T, B, L, ds) in enumerate(points): + feas = _feasible_G(T, B, 0) + G = feas[_ri(g, 0, len(feas) - 1)] + (br, bp), steps = _env_case( + consts, + g, + T, + B, + 0, + G, + bool(k % 2), + bool(k % 3 == 0), + _engine_width(L, ds), + two_steps=False, + max_pos=_lm(L, ds), + ) + runs = steps[0] + assert (sum(ids.numel() for _, ids, _, _ in runs), len(runs)) == (T, B) + bad_rows, bad_pools = bad_rows + br, bad_pools + bp + print( + f"\n {len(points)} dumped (T, B) steps ({len(_DUMP_TB)} gate points, {len(_DUMP_TB_ENVELOPE)} envelope " + f"points; B up to {max(p[1] for p in points)}, T up to {max(p[0] for p in points)}): {bad_rows} rows " + f"differ, {bad_pools} pools inexact" + ) + assert bad_rows == 0 and bad_pools == 0 + + +# ─── history-pool size invariance up to the largest ENGRAM_TOKEN pool ─── + + +def _masked_chunked_schedule(g: torch.Generator): + masks = {r: torch.ones(n + d, dtype=torch.bool) for r, (n, _, d) in _CHUNKED.items()} + masks[0][126:130] = False + masks[2][250:255] = False + masks[4][100:200] = False + return _chunked_schedule(g, masks) + + +class _FixedPool(_Pool): + """The requests' pages fixed at given page indices of whatever pool tensor is attached.""" + + def __init__(self, pages: dict, page_size: int, pool: torch.Tensor): + self.page_size = page_size + self.pages = pages + self.max_blocks = max(p.numel() for p in pages.values()) + self.pool = pool + + +def _run_outputs(consts, pool, steps, refs): + outs = [_call(consts, pool, runs) for runs in steps] + bad = 0 + for got, runs in zip(outs, steps): + at = 0 + for r, ids, s, m in runs: + bad += int( + (got[at : at + ids.numel()] != refs[r].forward(ids, s, m)).flatten(1).any(1).sum() + ) + at += ids.numel() + return torch.cat(outs), bad + + +def _count_not(t: torch.Tensor, value: int, chunk: int = 1 << 20) -> int: + return sum(int((t[i : i + chunk] != value).sum()) for i in range(0, t.shape[0], chunk)) + + +def _far_end_expected(consts, pool, seqs, max_lens, steps): + """SENTINEL everywhere except each request's positions, which hold its compressed ids + (-1 where the schedule's mask made the token DEAD).""" + masks = {r: torch.ones(n, dtype=torch.bool) for r, n in max_lens.items()} + for runs in steps: + for r, ids, s, m in runs: + if m is not None: + masks[r][s : s + ids.numel()] = m + exp = torch.full_like(pool.pool, SENTINEL) + for r, n in max_lens.items(): + pos = torch.arange(n) + slots = pool.pages[r][pos // 128].long() * 128 + pos % 128 + vals = consts.token_map[seqs[r]].to(torch.int32) + exp.view(-1)[slots.to(DEV)] = torch.where(masks[r], vals, torch.full_like(vals, DEAD)).to( + DEV + ) + return exp + + +def _boundary_pages(max_lens: dict, lines: list, seed: int) -> dict: + """Each request's pages drawn, shuffled, from windows around the given page indices (so + the schedule's slots sit on both sides of each line).""" + blocks = {r: -(-n // 128) for r, n in max_lens.items()} + need = sum(blocks.values()) + half = -(-need // (2 * len(lines))) + 1 + cand = torch.cat([torch.arange(p - half, p + half) for p in lines]) + cand = cand[torch.randperm(cand.numel(), generator=torch.Generator().manual_seed(seed))] + pages, at = {}, 0 + for r, nb in blocks.items(): + pages[r] = cand[at : at + nb].to(torch.int32) + at += nb + return pages + + +def test_history_pool_size_invariance_b_max(): + """The largest ENGRAM_TOKEN pool measured over the engine envelope including PP2 stages (stage + alone), at 0.5 x device memory: the PP2 stage-0 manager, native storage, max_num_tokens 2048, + max_batch_size 1, no DSpark (_B_MAX_CONFIG) -> [57999360, 128] int32 (29.70 GB, 7.42e9 elements), + the live manager's own buffer. + (1) Invariance: the masked chunked-prefill + decode schedule of six requests, pages at fixed + random indices below 2 x their need, run on views pool[:B] for B = the smallest pool that + holds them, 2^23 (the first page at byte 2^32), the old bound 10,027,008, the bf16 boot + gate's 15,269,888, 2^24 + 1 (the first page past element 2^31), the native boot gate's + 26,345,472, B_max // 2, the witness's bf16 leg's 32,243,712, 2^25 + 1 (the first page past + byte 2^34), the unsplit managers' largest 36,241,408, the PP2 stage-0 bf16 pool 48,365,568 and + B_max, and + on a separate allocation of exactly the smallest size: all give bit-identical hashes + (= the reference) and bit-identical pages, and nothing past the smallest size is written. + (2) Boundaries: the same schedule with its pages on both sides of page 2^23 (byte offset + 2^32), page 2^24 (element offset 2^31, byte 2^33) and page 2^25 (element offset 2^32, byte + 2^34): = reference, pool exact. + (3) Far end: every page among the pool's last pages, the last page owned: = reference, pool + exact, SENTINEL elsewhere. Controls: the far-end pages wrapped as a 32-bit byte offset + (page mod 2^23), a signed 32-bit element offset (page mod 2^24) or an unsigned 32-bit element + offset (page mod 2^25) would wrap them fail it.""" + mgr, token_pool = _engine_manager(*_B_MAX_CONFIG, pp_stage=_B_MAX_STAGE) + try: + assert tuple(token_pool.shape) == (_B_MAX, 128) and _B_MAX * 128 > 2**32 + consts = _Consts(_synthetic_token_map()) + seqs, steps = _masked_chunked_schedule(torch.Generator().manual_seed(2323)) + max_lens = {r: s.numel() for r, s in seqs.items()} + ref_pool = _Pool(max_lens, 128, num_pages=1, seed=2323) # pages below 2 x need + 1 + pages = ref_pool.pages + b_small = max(int(p.max()) for p in pages.values()) + 1 + sizes = [ + ("smallest", b_small), + ("2^23", 1 << 23), + ("old bound", 10_027_008), + ("bf16 boot", 15_269_888), + ("2^24 + 1", (1 << 24) + 1), + ("native boot", 26_345_472), + ("B_max // 2", _B_MAX // 2), + ("bf16 witness leg", 32_243_712), + ("2^25 + 1", (1 << 25) + 1), + ("unsplit B_max", _UNSPLIT_B_MAX), + ("bf16 PP2 stage 0", 48_365_568), + ("B_max", _B_MAX), + ] + assert [b for _, b in sizes] == sorted(b for _, b in sizes) + token_pool.fill_(SENTINEL) + results = {} + for name, B in sizes: + token_pool[:b_small].fill_(SENTINEL) + pool = _FixedPool(pages, 128, token_pool[:B]) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + out, bad = _run_outputs(consts, pool, steps, refs) + assert bad == 0, (name, bad) + results[name] = (out, token_pool[:b_small].clone()) + assert _count_not(token_pool[b_small:], SENTINEL) == 0, name + own = torch.full((b_small, 128), SENTINEL, dtype=torch.int32, device=DEV) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + out, bad = _run_outputs(consts, _FixedPool(pages, 128, own), steps, refs) + assert bad == 0 + results["own allocation"] = (out, own) + base_out, base_pool = results["smallest"] + for name, (o, p) in results.items(): + assert torch.equal(o, base_out) and torch.equal(p, base_pool), name + print( + f"\n invariance: pages < {b_small}; B = {', '.join(str(b) for _, b in sizes)} (views of the manager's " + f"[{_B_MAX}, 128] buffer: PP2 stage-0 manager, native storage, max_num_tokens 2048, batch 1, no DSpark) " + f"and an own " + f"[{b_small}, 128] " + f"allocation: {len(steps)} calls, hashes and pages bit-identical, = reference; nothing written past " + f"page {b_small - 1}" + ) + + lines = [1 << 23, 1 << 24, 1 << 25] + assert _B_MAX > lines[-1] + bnd = _FixedPool(_boundary_pages(max_lens, lines, 2323), 128, token_pool) + bnd.pool.fill_(SENTINEL) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + _, bad = _run_outputs(consts, bnd, steps, refs) + assert bad == 0 + exp = _far_end_expected(consts, bnd, seqs, max_lens, steps) + assert torch.equal(bnd.pool, exp) + del exp + used = torch.cat(list(bnd.pages.values())) + straddle = [(int((used < p).sum()), int((used >= p).sum())) for p in lines] + print( + f" boundaries: pages on both sides of 2^23 (byte 2^32) {straddle[0]}, 2^24 (element 2^31, byte " + f"2^33) {straddle[1]} and 2^25 (element 2^32, byte 2^34) {straddle[2]} (below, at-or-above): " + f"= reference, pool exact" + ) + assert all(lo > 0 and hi > 0 for lo, hi in straddle) + + far = _FarEndPool(token_pool, max_lens, seed=2323) + assert int(max(int(p.max()) for p in far.pages.values())) == _B_MAX - 1 + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + _, bad = _run_outputs(consts, far, steps, refs) + assert bad == 0 + exp = _far_end_expected(consts, far, seqs, max_lens, steps) + assert torch.equal(far.pool, exp) + lowest = min(int(p.min()) for p in far.pages.values()) + wrap_differs = {} + for bits in (23, 24, 25): + far.pool.fill_(SENTINEL) + wrapped = _FixedPool({r: p % (1 << bits) for r, p in far.pages.items()}, 128, far.pool) + refs = {r: _RefState(consts, n) for r, n in max_lens.items()} + _run_outputs(consts, wrapped, steps, refs) + wrap_differs[bits] = not torch.equal(far.pool, exp) + print( + f" far end: pages {lowest}..{_B_MAX - 1} (element offsets past 2^32, byte offsets past 2^34), " + f"= reference, pool exact, rest SENTINEL; [control] pages wrapped mod 2^23 (32-bit byte offset) -> pool " + f"check fails: {wrap_differs[23]}; mod 2^24 (signed 32-bit element offset): {wrap_differs[24]}; " + f"mod 2^25 (unsigned 32-bit element offset): {wrap_differs[25]}" + ) + assert lowest > (1 << 25) and wrap_differs[23] and wrap_differs[24] and wrap_differs[25] + del exp + finally: + mgr.shutdown() + + +# ─── the block-table width only sizes indexing ─── +# +# The target passes metadata.engram_token_block_table: [max_num_sequences, max_blocks_per_seq], +# row stride = width; the width is 132 at max_seq_len 16384 and grows with L to 8196 at 1,048,576 +# (_engine_width). The kernel reads row b at bt_row_stride * b and column p // P, masked by +# p // P < max_blocks. Without DSpark positions stay below L, so column L / 128 - 1 is the last +# one a call addresses; with DSpark the manager's max_seq_len is Lm = L + 14 and a verify run at +# the end of a full-length request writes positions up to Lm - 1, i.e. column L / 128. + +_WIDTHS = (128, 129, 130, 131, 132, 136, 260, 516, 1028, 2052, 8196) + + +class _WidePool(_Pool): + """_Pool whose per-step table is ``width`` columns wide, a column-prefix view of a + [B + 3, width + pad] allocation (row stride width + pad). The columns past each request's + own pages, the pad columns and the 3 rows past B hold OTHER requests' page ids, so a read or + a write through any of them would corrupt a checked row or a checked page. ``corrupt`` = + (req, col) replaces one of the request's own columns with a foreign page (the control).""" + + def __init__(self, max_lens: dict, seed: int, width: int, pad: int, corrupt=None): + super().__init__(max_lens, 128, num_pages=4096, seed=seed) + self.width, self.pad, self.corrupt = width, pad, corrupt + self.junk = torch.Generator().manual_seed(1000 * width + pad) + self.all_pages = torch.cat(list(self.pages.values())) + self.all_pages = self.all_pages[self.all_pages >= 0] + + def table(self, req_ids, override=None) -> torch.Tensor: + n = len(req_ids) + full = self.all_pages[ + torch.randint( + 0, self.all_pages.numel(), (n + 3, self.width + self.pad), generator=self.junk + ) + ] + for i, r in enumerate(req_ids): + own = self.pages[r] + full[i, : own.numel()] = own + if self.corrupt is not None and self.corrupt[0] == r: + others = torch.cat([p for q, p in self.pages.items() if q != r]) + full[i, self.corrupt[1]] = others[0] + view = full.to(DEV)[:, : self.width] + assert view.stride() == (self.width + self.pad, 1) + return view + + +def _width_run(consts, steps, ends, seed, width, pad, corrupt=None): + """``steps`` on a fresh _WidePool (pages fixed by ``seed``) with histories planted from + ``seed``: (all outputs, final pool, rows != reference, calls whose pool != expected).""" + pool = _WidePool(ends, seed, width, pad, corrupt) + refs = {r: _RefState(consts, n + 2) for r, n in ends.items()} + g = torch.Generator().manual_seed(seed) + for r, _, s, _ in steps[0]: + _plant_history(consts, pool, refs, r, s, g) + expected = pool.pool.clone() + outs, bad_rows, bad_pools = [], 0, 0 + for runs in steps: + got = _call(consts, pool, runs) + at = 0 + for r, ids, s, m in runs: + bad_rows += int( + (got[at : at + ids.numel()] != refs[r].forward(ids, s, m)).flatten(1).any(1).sum() + ) + at += ids.numel() + _apply_writes(consts, pool, expected, runs) + bad_pools += int(not torch.equal(pool.pool, expected)) + outs.append(got) + return torch.cat(outs), pool.pool.clone(), bad_rows, bad_pools + + +_DSPARK_MAX_SEQ_LEN = 16398 # the manager's max_seq_len with DSpark (max_draft_len 5, overlap on) + + +def _dspark_tail_steps(g: torch.Generator, L: int = 16384): + """DSpark verify runs of 6 at the end of full-length requests of max_seq_len L, into logical + page L / 128: from L - 2 (the last real position) to L + 3, from Lm - 6 to Lm - 1 (the + manager's last position, Lm = L + 14), beside a decode row at L - 1; then the rewound + continuation from L.""" + lm = _lm(L, True) + step1 = [ + (0, _ids(6, g), L - 2, None), + (1, _ids(6, g), lm - 6, None), + (2, _ids(1, g), L - 1, None), + ] + step2 = [(0, _ids(6, g), L, None), (2, _ids(6, g), L, None)] + ends = {0: L + 6, 1: lm, 2: L + 6} + return [step1, step2], ends + + +def _plain_tail_steps(g: torch.Generator, L: int): + """Without DSpark, runs ending at max_seq_len L (logical page L / 128 - 1, the last column a + call can address): a 200-row chunk from L - 400, a decode row at L - 2 and a 128-row chunk + ending at L; then the chunk continued to L and the decode row at L - 1.""" + step1 = [ + (0, _ids(200, g), L - 400, None), + (1, _ids(1, g), L - 2, None), + (2, _ids(128, g), L - 128, None), + ] + step2 = [(0, _ids(200, g), L - 200, None), (1, _ids(1, g), L - 1, None)] + ends = {0: L, 1: L, 2: L} + return [step1, step2], ends + + +def _last_page_steps(g: torch.Generator): + """Three requests whose runs end at position 16383 (logical page 127, the last column a call + can address): a 128-token chunk from 16256, a decode row at 16383 and a 6-row verify run from + 16378; then a second step on page 127 again (the verify rewound by 2).""" + step1 = [ + (0, _ids(128, g), 16256, None), + (1, _ids(1, g), 16383, None), + (2, _ids(6, g), 16378, None), + ] + step2 = [(2, _ids(4, g), 16380, None), (1, _ids(1, g), 16383, None)] + ends = {0: 16384, 1: 16384, 2: 16384} + return [step1, step2], ends + + +def test_block_table_width_only_sizes_indexing(): + """The same steps with block tables 128..132, 136, 260, 516, 1028, 2052 and 8196 columns wide + (the engine's widths at L 16384 .. 1048576 among them), each as a column-prefix view whose + row stride is the width (the engine's form) or the width + 5, with every column and row the + call must not read holding another request's pages: the hashes and the whole pool are + bit-identical at every width and stride, and equal the reference. The steps are the 9 + gate-envelope corners (positions below 16384), 15 random envelope partitions (D 0 / 5, image + spans, shuffled order, each with its continuation step), three requests ending on logical + page 127, DSpark verify runs into logical page 128 (positions up to 16397, run at widths >= + 129 only), and at every L of _L_SET the far end with and without DSpark: runs ending at L + (last column L / 128 - 1) and DSpark verify runs ending at Lm = L + 14 (column L / 128), + each run at every listed width above its last page (at L 1048576: 8196 only). Control: the + same last-page steps at width 132 with the decode request's own column 127 (its history and + its write) pointing at a foreign page fail the check.""" + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(132) + cases = [] + for T, B, D, G, _, _ in _ENV_EDGES[:9]: + steps, ends = _plan_pair( + g, T, B, D, G, masked=bool(len(cases) % 2), shuffle=bool(len(cases) % 3 == 0) + ) + cases.append((f"edge T={T} B={B} D={D} G={G}", steps, ends)) + while len(cases) < 24: + B = _ri(g, 1, _ENV_MAX_B) + T, D = _random_T(g), 5 * _ri(g, 0, 1) + feas = _feasible_G(T, B, D) if T >= B else [] + if not feas: + continue + G = feas[_ri(g, 0, len(feas) - 1)] + steps, ends = _plan_pair( + g, T, B, D, G, masked=bool(len(cases) % 2), shuffle=bool(len(cases) % 3 == 0) + ) + cases.append((f"random T={T} B={B} D={D} G={G}", steps, ends)) + steps, ends = _last_page_steps(g) + cases.append(("last page", steps, ends)) + steps, ends = _dspark_tail_steps(g) + cases.append(("DSpark tail", steps, ends)) + for L in _L_SET: + steps, ends = _plain_tail_steps(g, L) + cases.append((f"L {L} tail", steps, ends)) + steps, ends = _dspark_tail_steps(g, L) + cases.append((f"L {L} DSpark tail", steps, ends)) + bad_rows = bad_pools = runs_done = 0 + max_page, widths_run = {}, {} + for k, (name, steps, ends) in enumerate(cases): + top = max((s + ids.numel() - 1) // 128 for runs in steps for _, ids, s, _ in runs) + max_page[name] = top + base = None + widths_run[name] = [w for w in _WIDTHS if w > top] + for width in widths_run[name]: + for pad in (0, 5): + out, pool, br, bp = _width_run(consts, steps, ends, 7000 + k, width, pad) + bad_rows, bad_pools, runs_done = bad_rows + br, bad_pools + bp, runs_done + 1 + if base is None: + base = (out, pool) + else: + assert torch.equal(out, base[0]) and torch.equal(pool, base[1]), ( + name, + width, + pad, + ) + k_last = [name for name, _, _ in cases].index("last page") + steps, ends = cases[k_last][1], cases[k_last][2] + _, _, cr, cp = _width_run( + consts, steps, ends, 7000 + k_last, _ENGINE_WIDTH, 0, corrupt=(1, 127) + ) + print( + f"\n {len(cases)} step pairs x widths {_WIDTHS} (those > the case's last logical page) x row stride " + f"(width, width + 5) = {runs_done} runs; last logical page per case: max {max(max_page.values())} " + f"(last page {max_page['last page']}, DSpark tail {max_page['DSpark tail']}): hashes and pools bit-identical " + f"across widths and strides, {bad_rows} rows differ from the reference, {bad_pools} calls left the pool " + f"inexact; [control] column 127 of a request pointing at a foreign page: {cr} rows differ, {cp} calls left " + f"the pool inexact" + ) + print( + " far end per L (last logical page, widths run): " + + "; ".join( + f"L {L}: {max_page[f'L {L} tail']} {widths_run[f'L {L} tail']}, DSpark {max_page[f'L {L} DSpark tail']} " + f"{widths_run[f'L {L} DSpark tail']}" + for L in _L_SET + ) + ) + assert max_page["last page"] == 127 and max_page["DSpark tail"] == 128 + for L in _L_SET: + assert ( + max_page[f"L {L} tail"] == L // 128 - 1 and max_page[f"L {L} DSpark tail"] == L // 128 + ) + assert _engine_width(L, False) in widths_run[f"L {L} tail"] + assert _engine_width(L, True) in widths_run[f"L {L} DSpark tail"] + assert bad_rows == 0 and bad_pools == 0 + assert cr > 0 and cp > 0 + + +# ─── encoder replay: the history call and its empty runs ─────────────────────── +# +# Under encoder SWA bounded replay the target calls the op twice in a step that holds a +# prefix hit. First the HISTORY call, over the metadata's v41_replay_history_* fields: +# input_ids v41_replay_history_token_ids_cuda[:hist] int32 prompt ids +# req_idx_per_token v41_replay_history_req_idx_cuda[:hist] int32 context row b +# cu_seq_lens v41_replay_history_cu_seq_lens_cuda[:nc + 1] int32 +# start_pos v41_replay_history_start_pos_cuda[:nc] int32 +# block_table metadata.engram_token_block_table (_engine_width columns; rows 0..nc - 1 read) +# token_mask None +# with ONE RUN PER CONTEXT REQUEST in batch order (B = num_contexts, 1..256). A request whose +# first chunk replays from s = reused - min(128, reused) > 0 has min(3, s) rows, the prompt +# ids of positions [s - len, s), at start_pos s - len. Every other context request (no hit, +# a later chunk, or a replay from s = 0 after a hit of <= 128 tokens) has an EMPTY run: +# cu_seq_lens[b] == cu_seq_lens[b + 1], start_pos 0, no row. reused is a multiple of 2 (the +# compress ratio), so a non-empty run is 3 rows at an odd start_pos >= 1 (s >= 4) or 2 rows +# at start_pos 0 (s = 2). The call is made only when some run is non-empty; its hashes are +# discarded, its pool writes are what matters: the step's main call (all context chunks, +# then the generation runs) reads [s - 3, s) of every replay request, whose chunk starts at +# s with R = min(128, reused) replay rows and possibly nothing after them. + +_REPLAY_WINDOW = 128 # replay_window: a hit of reused tokens re-forwards min(128, reused) +_HISTORY_LEN = MAX_NGRAM - 1 # engram_history_len +_REPLAY_KINDS = ("r3", "r2", "r0", "plain") + + +def _replay_step( + g: torch.Generator, kinds, D: int = 0, G: int = 0, forced=None, L: int = 16384, ds: bool = False +) -> dict: + """One encoder-replay step: len(kinds) context requests in batch order, then G generation + runs of 1 + D rows. kinds[b]: 'r3' a replay from an even s >= 4 (3 history rows), 'r2' a + replay from s = 2 (2 history rows at 0), 'r0' a replay from s = 0 after a hit of 2..128 + tokens (empty history run), 'plain' no replay (empty history run; the chunk continues + from a cached c of any parity, 0 included). A replay chunk holds its R replay rows plus + 0 or more suffix rows. ``forced[b] = (start, rows)`` pins a context. Returns the history + runs (the engine's form: every context request, the empty runs included), the main + runs, each request's end and each context's start and R. Positions stay below the manager + length Lm of (L, ds); the step's block tables are the engine width for L.""" + forced = forced or {} + max_pos = _lm(L, ds) + nc = len(kinds) + mins, reused = [], {} + for b, kind in enumerate(kinds): + if kind in ("r3", "r2"): + reused[b] = None + mins.append(_REPLAY_WINDOW) + elif kind == "r0": + reused[b] = 2 * _ri(g, 1, _REPLAY_WINDOW // 2) + mins.append(reused[b]) + else: + mins.append(1) + budget = _ENV_MAX_T - sum(mins) - G * (1 + D) + assert budget >= 0 and nc + G <= _ENV_MAX_B, (kinds, D, G) + total = _ri(g, 0, budget) if _ri(g, 0, 2) else _ri(g, 0, min(budget, 64 * nc)) + cuts = sorted(_ri(g, 0, total) for _ in range(nc - 1)) + extra = [b - a for a, b in zip([0] + cuts, cuts + [total])] + extra = [ + 0 if _ri(g, 0, 2) == 0 else e for e in extra + ] # a third: a chunk of exactly its minimum + hist, main, starts, replay = [], [], [], [] + for b, kind in enumerate(kinds): + n = mins[b] + extra[b] + if kind == "r3": + s = 2 * _ri(g, 2, (max_pos - n) // 2) + elif kind == "r2": + s = 2 + elif kind == "r0": + s = 0 + else: + s = 0 if _ri(g, 0, 9) < 3 else _ri(g, 1, max_pos - n) + if b in forced: + s, n = forced[b] + assert n >= mins[b] and s + n <= max_pos and (kind != "r3" or (s >= 4 and s % 2 == 0)) + n_hist = min(_HISTORY_LEN, s) if kind in ("r3", "r2") else 0 + hist.append((b, _ids(n_hist, g), s - n_hist if n_hist else 0, None)) + main.append((b, _ids(n, g), s, None)) + starts.append(s) + replay.append(mins[b] if kind != "plain" else 0) + for j in range(G): + main.append((nc + j, _ids(1 + D, g), _ri(g, 1, max_pos - (1 + D)), None)) + assert sum(ids.numel() for _, ids, _, _ in main) <= _ENV_MAX_T and len(main) <= _ENV_MAX_B + ends = {r: s + ids.numel() for r, ids, s, _ in main} + return { + "kinds": list(kinds), + "D": D, + "G": G, + "hist": hist, + "main": main, + "ends": ends, + "start": starts, + "replay": replay, + "max_pos": max_pos, + "width": _engine_width(L, ds), + } + + +def _engine_args(pool: _Pool, runs): + """The op's index arguments for ``runs`` as the metadata builds them (int32, batch order; + an empty run keeps its cu_seq_lens entry, start_pos and block-table row).""" + lens = [ids.numel() for _, ids, _, _ in runs] + ids = torch.cat([x for _, x, _, _ in runs]).to(torch.int32) + req = torch.repeat_interleave(torch.arange(len(runs)), torch.tensor(lens)).to(torch.int32) + cu = torch.tensor([0] + list(np.cumsum(lens)), dtype=torch.int32) + start = torch.tensor([s for _, _, s, _ in runs], dtype=torch.int32) + return [ids, req, cu, start, pool.table([r for r, _, _, _ in runs])] + + +def _call_engine_args(consts: _Consts, pool: _Pool, args) -> torch.Tensor: + """One op call with the index arguments ``_engine_args`` built.""" + ids, req, cu, start, table = args + return op( + ids.to(DEV), + req.to(DEV), + cu.to(DEV), + start.to(DEV), + pool.pool, + table.to(DEV), + consts.dev.token_map, + consts.dev.multipliers, + consts.dev.primes, + consts.dev.offsets, + consts.pad_id, + ).cpu() + + +def _replay_case(consts: _Consts, step: dict, page_size: int, seed: int) -> dict: + """Run one replay step the way the target does -- the history call, then the main call -- + on a fresh pool, and measure (asserting nothing): + + * the history call (A, the engine's form) against the reference, and the whole pool + against the planted state plus exactly its writes; + * B: the same call without its empty runs, from the same pool: rows and pool bitwise; + * C: the empty runs given a random start_pos and another run's block-table row: rows and + pool bitwise (an empty run's start_pos and table row are never read); + * every empty-run request's pages unchanged by the call; + * the main call after it against the reference (replay chunks read the history the + history call stored), the whole pool exact; + * [control] W: req_idx_per_token numbered over the non-empty runs only; + * [control] the main call with the history call skipped.""" + hist, main, ends = step["hist"], step["main"], step["ends"] + g = torch.Generator().manual_seed(seed) + lows = {} + for r, ids, s, _ in list(hist) + list(main): + if ids.numel() or r >= len(hist): + lows[r] = min(lows.get(r, s), max(0, s - _HISTORY_LEN)) + for b, kind in enumerate( + step["kinds"] + ): # the replay's stale history below its start is planted too + lows[b] = min(lows[b], max(0, step["start"][b] - _HISTORY_LEN)) + pool = _Pool(ends, page_size, num_pages=4096, seed=seed, lows=lows) + pool.max_blocks = max(pool.max_blocks, step["width"]) + refs = {r: _RefState(consts, ends[r] + 8, base=lows[r]) for r in ends} + for b, kind in enumerate(step["kinds"]): + s = step["start"][b] + if kind == "r3": + _plant_history(consts, pool, refs, b, s - _HISTORY_LEN, g) # what the history rows read + if s > 0: + # plain: the chunk's real history; replay: stale ids the history call must overwrite + _plant_history(consts, pool, refs, b, s, g) + for r, _, s, _ in main[len(hist) :]: + _plant_history(consts, pool, refs, r, s, g) + pre = pool.pool.clone() + nonempty = [run for run in hist if run[1].numel()] + empty = [b for b, (_, ids, _, _) in enumerate(hist) if not ids.numel()] + + def run(args): + pool.pool.copy_(pre) + out = _call_engine_args(consts, pool, args) + return out, pool.pool.clone() + + args_a = _engine_args(pool, hist) + out_a, pool_a = run(args_a) + at, bad_hist = 0, 0 + for b, ids, s, _ in hist: + want = refs[b].forward(ids, s) + bad_hist += int((out_a[at : at + ids.numel()] != want).flatten(1).any(1).sum()) + at += ids.numel() + expected = pre.clone() + _apply_writes(consts, pool, expected, hist) + pool_ok = torch.equal(pool_a, expected) + out_b, pool_b = run(_engine_args(pool, nonempty)) + same_b = torch.equal(out_a, out_b) and torch.equal(pool_a, pool_b) + args_c = [x.clone() for x in args_a] + donor = nonempty[0][0] + for b in empty: + args_c[3][b] = _ri(g, 0, step["max_pos"] - 1) + args_c[4][b] = args_a[4][donor] + out_c, pool_c = run(args_c) + same_c = torch.equal(out_a, out_c) and torch.equal(pool_a, pool_c) + own = {b: pool.pages[b][pool.pages[b] >= 0].long().to(DEV) for b in empty} + untouched = all(torch.equal(pool_a[own[b]], pre[own[b]]) for b in empty) + args_w = list(args_a) + lens = [ids.numel() for _, ids, _, _ in nonempty] + args_w[1] = torch.repeat_interleave(torch.arange(len(nonempty)), torch.tensor(lens)).to( + torch.int32 + ) + out_w, pool_w = run(args_w) + caught_w = not (torch.equal(out_w, out_a) and torch.equal(pool_w, pool_a)) + empty_before = any( + b < max(r for r, (_, ids, _, _) in enumerate(hist) if ids.numel()) for b in empty + ) + # the main call on the history call's pool + pool.pool.copy_(pool_a) + got = _call(consts, pool, main) + wants = [refs[r].forward(ids, s, m) for r, ids, s, m in main] + want = torch.cat(wants) + bad_main = int((got != want).flatten(1).any(1).sum()) + _apply_writes(consts, pool, pool_a, main) + main_pool_ok = torch.equal(pool.pool, pool_a) + # control: no history call + pool.pool.copy_(pre) + bad_skip = int((_call(consts, pool, main) != want).flatten(1).any(1).sum()) + return { + "bad_hist": bad_hist, + "pool_ok": pool_ok, + "same_b": same_b, + "same_c": same_c, + "untouched": untouched, + "caught_w": caught_w, + "empty_before": empty_before, + "bad_main": bad_main, + "main_pool_ok": main_pool_ok, + "bad_skip": bad_skip, + "hist_rows": sum(ids.numel() for _, ids, _, _ in hist), + "main_rows": int(want.shape[0]), + "empty": empty, + } + + +# (kinds in batch order, D, G, forced {b: (start, rows)}[, L, DSpark]): the corners, each at page 128 +# and 256; L 16384 without DSpark unless given +_REPLAY_LM_1M = _lm(1048576, True) +_REPLAY_CORNERS = [ + ( + ["r3", "plain"], + 0, + 0, + {0: (3970, 129), 1: (1000, 64)}, + ), # the witness's (3, 2): replay first, empty last + ( + ["plain", "r3"], + 0, + 0, + {1: (3970, 128)}, + ), # empty first; the replay chunk is its 128 replay rows only + (["r3", "plain", "r3"], 0, 0, {}), # empty in the middle + (["plain", "plain", "r3", "r0", "plain"], 0, 0, {}), # several empties around one replay + (["r2", "plain"], 0, 0, {}), # a 2-row run at start_pos 0 beside an empty + (["plain", "r2"], 0, 0, {}), + (["r0", "r2", "r0"], 0, 0, {}), # the only non-empty run is 2 rows + (["r3"] + ["plain", "r0"] * 15 + ["plain"], 0, 0, {}), # B 32: 31 empties after the replay + (["plain", "r0"] * 15 + ["plain", "r3"], 0, 0, {}), # B 32: 31 empties before it + (["plain"] * 16 + ["r3"] + ["r0"] * 15, 0, 0, {}), # B 32: the replay in the middle + (["r3"] * 32, 0, 0, {}), # B 32, no empty run: 96 history rows + ( + ["r3"], + 0, + 0, + {0: (_ENV_MAX_POS - 128, 128)}, + ), # B 1 at the last replay start (start_pos 16253) + (["r3"], 0, 0, {0: (4, 128)}), # the first replay start with 3 history rows (start_pos 1) + (["r3", "plain"], 5, 6, {}), # beside six DSpark verify runs in the main call + (["plain", "r2", "r0", "r3"], 0, 10, {}), # beside ten decode rows + ( + ["plain"] * 8 + ["r3"] * 2 + ["plain"] * 8, + 5, + 14, + {}, + ), # 16 empties, 2 replays, 14 verify runs: B 32 + (["r3"] + ["plain"] * 255, 0, 0, {}, 1048576, False), # B 256: 255 empties after the replay + (["plain"] * 255 + ["r3"], 0, 0, {}, 262144, True), # B 256: 255 empties before it + ( + ["plain"] * 128 + ["r3"] + ["r0"] * 20 + ["plain"] * 107, + 0, + 0, + {}, + 65536, + False, + ), # B 256: in the middle + ( + ["r3"] * 127, + 0, + 0, + {}, + 1048576, + True, + ), # B 127, no empty run: 381 history rows, 16256 replay rows + ( + ["r3"], + 0, + 0, + {0: (_REPLAY_LM_1M - 128, 128)}, + 1048576, + True, + ), # the last replay start at Lm 1,048,590 + ( + ["r3"], + 0, + 0, + {0: (1048576 - 128, 128)}, + 1048576, + False, + ), # the last replay start at L 1,048,576 + (["r3", "plain"], 5, 250, {}, 1048576, True), # beside 250 DSpark verify runs: B 252 + ( + ["plain", "r2", "r0", "r3"] * 4, + 0, + 240, + {}, + 131072, + False, + ), # 16 contexts beside 240 decode rows: B 256 +] +_REPLAY_PARTITIONS = 256 + + +def _replay_cases(g: torch.Generator) -> list: + """(step, page size) for the corners (both page sizes) and then the seeded random + partitions (page 128, the engine's): every B 1..32 several times, kinds drawn per + context request (r3 0.35, r2 0.1, r0 0.15, plain 0.4; at least one r3 / r2), D in {0, 5}, + generation runs on half of them.""" + cases = [] + for corner in _REPLAY_CORNERS: + kinds, D, G, forced = corner[:4] + L, ds = corner[4:] if len(corner) > 4 else (16384, D > 0) + for ps in (128, 256): + cases.append((_replay_step(g, kinds, D, G, forced, L, ds), ps)) + k = 0 + while k < _REPLAY_PARTITIONS: + nc = k % _ENV_MAX_B + 1 + L, ds = _L_CONFIGS[k % len(_L_CONFIGS)] + k += 1 + kinds = [] + for _ in range(nc): + u = _ri(g, 0, 19) + kinds.append("r3" if u < 7 else "r2" if u < 9 else "r0" if u < 12 else "plain") + if not any(x in ("r3", "r2") for x in kinds): + kinds[_ri(g, 0, nc - 1)] = "r3" + # the step's rows fit max_num_tokens: a replay chunk holds its 128 (r3 / r2) or up to 128 (r0) + # replay rows, so past about 120 replays a random one turns into a plain chunk + while _REPLAY_WINDOW * sum(x != "plain" for x in kinds) + nc > _ENV_MAX_T: + idx = [i for i, x in enumerate(kinds) if x != "plain"] + kinds[idx[_ri(g, 0, len(idx) - 1)]] = "plain" + if not any(x in ("r3", "r2") for x in kinds): + kinds[kinds.index("plain")] = "r3" + D = 5 * _ri(g, 0, 1) if ds else 0 + room = _ENV_MAX_T - _REPLAY_WINDOW * sum(x != "plain" for x in kinds) - nc + G = min(_ri(g, 0, _ENV_MAX_B - nc), room // (1 + D)) if _ri(g, 0, 1) else 0 + cases.append((_replay_step(g, kinds, D, G, None, L, ds), 128)) + return cases + + +def test_replay_history_call_with_empty_runs(): + """The encoder-replay target's history call, empty runs included, over the envelope: + the 24 corners at page 128 and 256, then 256 seeded random partitions at page 128 + (B 1..256, every value; the 12 (L, DSpark) configurations in turn, positions below Lm up to + 1,048,590; empty runs anywhere, several per call, up to 255; 2- and 3-row runs; replay + starts up to Lm - 128; the main call's chunks of exactly their replay rows or longer, plain + chunks from any cached start, decode or DSpark verify runs after them; tables of the engine + width for L, 132 .. 8196 columns). Every history call equals the reference with the whole pool exact, + is bitwise equal (rows and pool) to the same call without its empty runs, and to the + call with the empty runs' start_pos and block-table rows changed; no empty-run request's + page changes; the main call after it equals the reference with the pool exact. Controls: + req_idx_per_token numbered over the non-empty runs fails exactly the calls where an + empty run precedes a non-empty one; skipping the history call fails the main call of + every step. No Triton variant beyond one per num_tokens class is compiled.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.embedding.engram_hash import ( # noqa: E501 + _engram_hash_kernel as fn, + ) + + consts = _Consts(_synthetic_token_map()) + g = torch.Generator().manual_seed(19_066) + warm = _Pool({0: 64}, 128, seed=3) + warm.max_blocks = _ENGINE_WIDTH + for T in (1, 16, 37): # one call per num_tokens class, in the target's int32 form + _call(consts, warm, [(0, _ids(T, g), 0, None)]) + variants = _jit_keys(fn) + cases = _replay_cases(g) + tot = { + "bad_hist": 0, + "bad_main": 0, + "pool_bad": 0, + "diff_b": 0, + "diff_c": 0, + "touched": 0, + "w_wrong": 0, + "skip_missed": 0, + "hist_rows": 0, + "main_rows": 0, + "w_caught": 0, + } + seen = { + "B": set(), + "first": 0, + "middle": 0, + "last": 0, + "multi": 0, + "all_but_one": 0, + "D": set(), + "ps": set(), + "r2": 0, + "r0": 0, + "zero_compute": 0, + "max_empty": 0, + "max_start": 0, + "calls_with_empty": 0, + "widths": set(), + } + for i, (step, ps) in enumerate(cases): + r = _replay_case(consts, step, ps, 50_000 + i) + tot["bad_hist"] += r["bad_hist"] + tot["bad_main"] += r["bad_main"] + tot["pool_bad"] += int(not r["pool_ok"]) + int(not r["main_pool_ok"]) + tot["diff_b"] += int(not r["same_b"]) + tot["diff_c"] += int(not r["same_c"]) + tot["touched"] += int(not r["untouched"]) + tot["w_wrong"] += int(r["caught_w"] != r["empty_before"]) + tot["w_caught"] += int(r["caught_w"]) + tot["skip_missed"] += int(r["bad_skip"] == 0) + tot["hist_rows"] += r["hist_rows"] + tot["main_rows"] += r["main_rows"] + nc, empty = len(step["kinds"]), r["empty"] + seen["B"].add(nc) + seen["ps"].add(ps) + if step["G"]: + seen["D"].add(step["D"]) + if empty: + seen["calls_with_empty"] += 1 + seen["first"] += int(0 in empty) + seen["last"] += int(nc - 1 in empty) + seen["middle"] += int(any(0 < b < nc - 1 for b in empty)) + seen["multi"] += int(len(empty) >= 2) + seen["all_but_one"] += int(nc >= 2 and len(empty) == nc - 1) + seen["max_empty"] = max(seen["max_empty"], len(empty)) + seen["r2"] += step["kinds"].count("r2") + seen["r0"] += step["kinds"].count("r0") + seen["zero_compute"] += sum( + 1 + for b, (_, ids, _, _) in enumerate(step["main"][:nc]) + if step["replay"][b] and ids.numel() == step["replay"][b] + ) + seen["max_start"] = max( + [seen["max_start"]] + [s for _, ids, s, _ in step["hist"] if ids.numel()] + ) + seen["widths"].add(step["width"]) + compiled = _jit_keys(fn) - variants + print( + f"\n {len(cases)} replay steps ({len(_REPLAY_CORNERS)} corners x pages 128 / 256, then " + f"{_REPLAY_PARTITIONS} random at 128): history calls {tot['hist_rows']} rows, main calls " + f"{tot['main_rows']} rows; B {min(seen['B'])}..{max(seen['B'])} ({len(seen['B'])} values); " + f"{seen['calls_with_empty']} history calls with empty runs (first {seen['first']}, middle " + f"{seen['middle']}, last {seen['last']}; >= 2 empties {seen['multi']}, all but one empty " + f"{seen['all_but_one']}, max {seen['max_empty']}); {seen['r2']} 2-row runs, {seen['r0']} replays from 0, " + f"{seen['zero_compute']} replay chunks of only their replay rows; history start_pos up to " + f"{seen['max_start']}; generation D {sorted(seen['D'])}; table widths {sorted(seen['widths'])}" + ) + print( + f" {tot['bad_hist']} history rows and {tot['bad_main']} main rows differ from the reference, " + f"{tot['pool_bad']} calls left the pool inexact; without the empty runs: {tot['diff_b']} calls differ; " + f"empty runs' start_pos / table rows changed: {tot['diff_c']} differ; {tot['touched']} calls changed an " + f"empty-run request's page; {len(compiled)} new Triton variant(s)" + ) + print( + f" [control] req_idx over the non-empty runs only: caught in {tot['w_caught']} calls, " + f"{tot['w_wrong']} calls where caught != (an empty run precedes a non-empty one); history call skipped: " + f"{tot['skip_missed']} steps where the main call still matched" + ) + assert tot["bad_hist"] == 0 and tot["bad_main"] == 0 and tot["pool_bad"] == 0 + assert tot["diff_b"] == 0 and tot["diff_c"] == 0 and tot["touched"] == 0 + assert tot["w_wrong"] == 0 and tot["w_caught"] > 0 and tot["skip_missed"] == 0 + assert not compiled + assert ( + seen["B"] == set(range(1, _ENV_MAX_B + 1)) + and seen["ps"] == {128, 256} + and seen["D"] == {0, 5} + ) + assert min(seen["first"], seen["middle"], seen["last"], seen["multi"], seen["all_but_one"]) > 0 + assert ( + seen["r2"] > 0 + and seen["r0"] > 0 + and seen["zero_compute"] > 0 + and seen["max_empty"] == _ENV_MAX_B - 1 + ) + assert seen["widths"] == {_engine_width(L, ds) for L, ds in _L_CONFIGS} + assert seen["max_start"] == _REPLAY_LM_1M - 128 - _HISTORY_LEN + + +# ─── the checkpoint's own engram.py and the golden dumps ──────────────────────── + + +def _need_external(): + missing = [d for d in (_GOLDEN_DIR, _CKPT_DIR) if not os.path.isdir(d)] + if missing: + pytest.skip(f"golden dumps / checkpoint not present: {missing}") + + +@functools.lru_cache(maxsize=None) +def _vendor(): + """The checkpoint's inference/engram.py, its NgramHashState built with the checkpoint tokenizer.""" + from transformers import PreTrainedTokenizerFast + + spec = importlib.util.spec_from_file_location( + "dsv41_vendor_engram", os.path.join(_CKPT_DIR, "inference", "engram.py") + ) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + tok = PreTrainedTokenizerFast(tokenizer_file=os.path.join(_CKPT_DIR, "tokenizer.json")) + assert len(tok) == VOCAB + args = types.SimpleNamespace( + engram_layer_ids=list(LAYER_IDS), + engram_max_ngram_size=MAX_NGRAM, + engram_n_heads=N_HEADS, + engram_vocab_size=ENGRAM_VOCAB, + engram_num_embeddings=list(NUM_EMBEDDINGS), + engram_head_dim=256, + engram_compressed_vocab_size=COMPRESSED_VOCAB, + engram_pad_id=PAD_TOKEN_ID, + max_batch_size=1, + max_seq_len=4096, + ) + layout = mod.EngramLayout.from_args(args) + return mod, layout, args, tok + + +@functools.lru_cache(maxsize=None) +def _vendor_state(): + """One vendor NgramHashState. Reusable across schedules: its cache is position-indexed and every + schedule here starts at position 0, so each position is written before it is read.""" + mod, layout, args, tok = _vendor() + return mod.NgramHashState(args, layout, tok) + + +@functools.lru_cache(maxsize=None) +def _real_consts() -> _Consts: + state = _vendor_state() + consts = _Consts(state.token_map.to(torch.int32)) + # the vendor's own buffers equal the constants rebuilt here + assert torch.equal(state.multipliers, consts.multipliers) + assert torch.equal(state.primes, consts.primes) + assert torch.equal(state.offsets, consts.offsets) + assert state.pad_id == consts.pad_id + return consts + + +def _golden(case): + tokens = torch.load(os.path.join(_GOLDEN_DIR, case, "tokens.pt"), weights_only=True) + gold = torch.load(os.path.join(_GOLDEN_DIR, case, "engram_hash.pt"), weights_only=True) + return ( + tokens["input_ids"].to(torch.int64), + int(tokens["n_prompt"]), + gold["hash_ids"], + gold["compressed_ids"], + ) + + +def _golden_runs(n_prompt: int, total: int): + """Chunked prompt (7, 1, 2, 64, rest) then one decode token per generated position.""" + runs, at = [], 0 + for c in (7, 1, 2, 64): + if at + c > n_prompt: + break + runs.append((at, at + c)) + at += c + if at < n_prompt: + runs.append((at, n_prompt)) + return runs + [(p, p + 1) for p in range(n_prompt, total)] + + +@pytest.mark.parametrize("case", _GOLDEN_CASES) +def test_golden_and_vendor_chunked_and_decode(case): + _need_external() + consts = _real_consts() + ids, n_prompt, gold, gold_comp = _golden(case) + T = ids.numel() + assert torch.equal(consts.token_map[ids].long(), gold_comp.long()) + spans = _golden_runs(n_prompt, T) + + # one-shot over every dumped row + pool = _Pool({0: T}, 128, seed=T) + one = _call(consts, pool, [(0, ids, 0, None)]) + assert torch.equal(one, gold) + + # chunked prefill + decode, alone and in a batch with an unrelated request + pool = _Pool({0: T, 1: 64 + T}, 128, seed=T + 1) + other = _ids(64 + T, torch.Generator().manual_seed(T)) + vendor = _vendor_state() + got = [] + for k, (a, b) in enumerate(spans): + mine_run = (0, ids[a:b], a, None) + runs = [mine_run] + if k % 2: + extra = (1, other[a:b], a, None) + runs = [extra, mine_run] if b - a == 1 else [mine_run, extra] + out = _call(consts, pool, runs) + at = 0 + for r, x, _, _ in runs: + if r == 0: + mine = out[at : at + x.numel()] + at += x.numel() + got.append(mine) + assert torch.equal(mine, vendor(ids[None, a:b], a)[0]), (case, a, b) + got = torch.cat(got) + assert torch.equal(got, gold) + + # controls, each against the golden ids + def frac(x): + return f"{int((x != gold).flatten(1).any(1).sum())}/{T} rows" + + # This tokenizer maps token 2 to class 2, so pad_id = token_map[2] happens to equal 2 here and a + # pad_id-2 mistake is invisible on real data; test_pad_id_is_the_compressed_pad_token shows it on a + # map where the two differ. The raw-id control below is the one a missing token map fails. + assert consts.pad_id == PAD_TOKEN_ID + pool = _Pool({0: T}, 128, seed=3) + raw_map = torch.arange(VOCAB, dtype=torch.int32, device=DEV) + raw = torch.cat( + [_call(consts, pool, [(0, ids[a:b], a, None)], token_map=raw_map) for a, b in spans] + ) + ctl = {"raw ids (token map skipped)": raw} + for variant in ("int32", "float64"): + ref = _RefState(consts, T, variant=variant) + ctl[variant] = torch.cat([ref.forward(ids[a:b], a) for a, b in spans]) + ref = _RefState(consts, T, multipliers=_multipliers((0, 1))) + ctl["seeds from layer index 0/1"] = torch.cat([ref.forward(ids[a:b], a) for a, b in spans]) + pool = _Pool({0: T}, 128, seed=3) + nohist = [] + for a, b in spans: + pool.pool.fill_(DEAD) + nohist.append(_call(consts, pool, [(0, ids[a:b], a, None)])) + ctl["history not carried"] = torch.cat(nohist) + ref = _RefState(consts, T) + assert torch.equal(torch.cat([ref.forward(ids[a:b], a) for a, b in spans]), gold) + print( + f"[control] {case} (T={T}, pad_id={consts.pad_id}): " + + ", ".join(f"{k}: {frac(v)}" for k, v in ctl.items()) + ) + for name, x in ctl.items(): + assert not torch.equal(x, gold), name + + +def test_vendor_image_spans_chunked(): + """Real token map, DEAD spans (vendor token_mask) across chunk and decode boundaries.""" + _need_external() + consts = _real_consts() + ids, n_prompt, _, _ = _golden("long_retrieval") + T = ids.numel() + mask = torch.ones(T, dtype=torch.bool) + for a, b in ( + (0, 1), + (5, 9), + (70, 76), + (80, 81), + (140, 200), + (n_prompt - 2, n_prompt + 1), + (T - 3, T - 2), + ): + mask[a:b] = False + spans = _golden_runs(n_prompt, T) + pool = _Pool({0: T}, 128, seed=6) + vendor = _vendor_state() + wrong = _RefState(consts, T, variant="noncumulative") + got, ctl = [], [] + for a, b in spans: + out = _call(consts, pool, [(0, ids[a:b], a, mask[a:b])]) + assert torch.equal(out, vendor(ids[None, a:b], a, mask[None, a:b])[0]), (a, b) + got.append(out) + ctl.append(wrong.forward(ids[a:b], a, mask[a:b])) + got, ctl = torch.cat(got), torch.cat(ctl) + print( + "[control] non-cumulative blocking vs vendor, image spans: " + f"{int((ctl != got).flatten(1).any(1).sum())}/{T} rows differ" + ) + assert not torch.equal(ctl, got) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/gemm/test_modeling_v2_mhc_gemm_sqrsum_fma.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/gemm/test_modeling_v2_mhc_gemm_sqrsum_fma.py new file mode 100644 index 000000000000..464957ed1223 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/gemm/test_modeling_v2_mhc_gemm_sqrsum_fma.py @@ -0,0 +1,438 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the mhc_gemm_sqrsum_fma entry. + +The entry owns what is certified: `CELLS` is the list, `reference` (a float64 +product built on the spot) and `compare` (a rigorous fp32 rounding-chain bound +plus a fidelity gate on `y`, torch's fp32 defaults on `r`) are the gate. This +file turns a cell's spec into tensors, runs every call into NaN-filled buffers +with guard rows, and keeps the checks a cell cannot carry: every tactic +bit-identical, CUDA-graph replay, M = 0, subnormal flushing, the wrong-variant +controls, the op's own rejections, the layout guard and the silent misreads it +exists for, the every-row-count prefix sweep with its one-row-shift control, +one compiled kernel per tile band, and capture at the engine's capture union. +""" + +import zlib + +import pytest +import torch +from torch.profiler import ProfilerActivity, profile + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.gemm.mhc_gemm_sqrsum_fma import ( # noqa: E501 + mhc_gemm_sqrsum_fma as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "mhc_gemm_sqrsum_fma requires a CUDA device" + +HC_MULT, HIDDEN = 4, 5120 +K_HC = HC_MULT * HIDDEN # 20480: the flattened [hc_mult, hidden] stream +N_HC = (2 + HC_MULT) * HC_MULT # 24 mixes: 4 pre, 4 post, 16 comb +TILE_NS = (1, 2, 3, 4, 6, 8, 12, 24) +# The entry's fidelity gate, in whose units the controls report. +_FIDELITY = 2.0**-17 + + +def _ref(x: torch.Tensor, w: torch.Tensor): + """The entry's reference for x [M, K] and w [N, K]: (y64, r64, bound, l2).""" + m, k = x.shape + return op.reference(x, w, None, None, m, w.shape[0], k) + + +def _y_errors(y: torch.Tensor, x: torch.Tensor, w: torch.Tensor): + """(max of err / bound, median of err / l2) for y against the entry's float64 reference.""" + y64, _, bound, l2 = _ref(x, w) + err = (y.double() - y64).abs() + return (err / bound).max().item(), (err / l2).median().item() + + +def _call(x: torch.Tensor, w: torch.Tensor, tile_n: int = 0, tile_m: int = 0): + """Run the op into NaN-filled buffers with one guard row on each side. + + Checks both sides of the in-place contract: every element of y and r is + written, the guard rows around them are left bit-for-bit untouched, and the + op returns the buffers it filled. + """ + m, k = x.shape + n = w.shape[0] + y_buf = torch.full((m + 2, n), float("nan"), device="cuda") + r_buf = torch.full((m + 2,), float("nan"), device="cuda") + y, r = y_buf[1 : m + 1], r_buf[1 : m + 1] + result = op(x, w, y, r, m, n, k, tile_n, tile_m) + assert result[0] is y and result[1] is r, "the op writes in place and returns (y, r)" + torch.cuda.synchronize() + assert y_buf[0].isnan().all() and y_buf[-1].isnan().all(), "y guard rows were written" + assert r_buf[0].isnan() and r_buf[-1].isnan(), "r guard slots were written" + assert not y.isnan().any(), "an element of y was not written" + assert not r.isnan().any(), "an element of r was not written" + return y, r + + +def _check(y: torch.Tensor, r: torch.Tensor, x: torch.Tensor, w: torch.Tensor) -> None: + op.compare((y, r), _ref(x, w)) + + +def _make(m: int, k: int, seed: int, flavor: str): + """x bf16 [m, k] and w fp32 [N_HC, k]. + + ``gauss``: iid normal operands. + ``hc_like``: shaped like the real streams measured on this checkpoint -- + four correlated hc copies per token, one massive attention-sink row + (~150x the others' norm), and heavy-tailed weights (a few entries ~100x + the weight rms). + """ + g = torch.Generator(device="cuda").manual_seed(seed) + if flavor == "gauss": + x = torch.randn(m, k, device="cuda", generator=g) + w = torch.randn(N_HC, k, device="cuda", generator=g) * 0.05 + else: + copies = k // HIDDEN if k % HIDDEN == 0 else 1 + base = torch.randn(m, 1, k // copies, device="cuda", generator=g) + x = (base + 0.3 * torch.randn(m, copies, k // copies, device="cuda", generator=g)).reshape( + m, k + ) * 0.08 + x[0] *= 150.0 + w = torch.randn(N_HC, k, device="cuda", generator=g) * 0.05 + spikes = torch.rand(N_HC, k, device="cuda", generator=g) < 1e-3 + w = torch.where(spikes, w * 100.0, w) + return x.to(torch.bfloat16), w.contiguous() + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + m, k = spec["M"], spec["K"] + x, w = _make(m, k, seed=zlib.crc32(cell.why.encode()) % 2**31, flavor=spec["flavor"]) + y, r = _call(x, w, spec["tile_n"], spec["tile_m"]) + op.compare((y, r), op.reference(x, w, y, r, m, N_HC, k, spec["tile_n"], spec["tile_m"])) + + +@pytest.mark.parametrize("m", (1, 8, 32, 48, 128, 129, 192, 2048, 2049, 4096)) +def test_every_tactic_is_bit_identical(m: int) -> None: + """tile_n only re-tiles the grid and tile_m is not read: results never change.""" + x, w = _make(m, K_HC, seed=100 + m, flavor="hc_like") + y0, r0 = _call(x, w) + _check(y0, r0, x, w) + for tile_n in TILE_NS: + for tile_m in (0, 1, 2): + y, r = _call(x, w, tile_n, tile_m) + assert torch.equal(y, y0) and torch.equal(r, r0), (tile_n, tile_m) + + +def test_cuda_graph_replay() -> None: + """Captured once on static buffers, replayed on new stream contents: bit-identical to eager.""" + m = 48 + x_static, w = _make(m, K_HC, seed=7, flavor="hc_like") + y = torch.empty(m, N_HC, device="cuda") + r = torch.empty(m, device="cuda") + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(x_static, w, y, r, m, N_HC, K_HC, 1, 0) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(x_static, w, y, r, m, N_HC, K_HC, 1, 0) + x_new, _ = _make(m, K_HC, seed=8, flavor="hc_like") + x_static.copy_(x_new) + graph.replay() + torch.cuda.synchronize() + y_eager, r_eager = _call(x_new, w, 1, 0) + assert torch.equal(y, y_eager) and torch.equal(r, r_eager) + + +def test_zero_tokens_is_a_no_op() -> None: + x, w = _make(1, K_HC, seed=9, flavor="gauss") + y = torch.empty(0, N_HC, device="cuda") + r = torch.empty(0, device="cuda") + op(x[:0], w, y, r, 0, N_HC, K_HC) + torch.cuda.synchronize() + + +def test_subnormal_inputs_are_flushed() -> None: + """The kernel is built with fast math: a subnormal x element contributes 0.""" + x = torch.zeros(1, 2048, device="cuda", dtype=torch.bfloat16) + x[0, 0] = 1e-39 # subnormal in bf16 and fp32 + w = torch.zeros(N_HC, 2048, device="cuda") + w[:, 0] = 1e10 # the exact product, ~1e-29, is a normal fp32 value + y, r = _call(x, w) + assert (y == 0).all() and (r == 0).all() + + +# ─── deliberately-wrong controls: each must land outside the gate ─────────── + + +def _tf32_round(w: torch.Tensor) -> torch.Tensor: + """Round fp32 to the 10-bit TF32 mantissa (round half away from zero).""" + bits = w.view(torch.int32) + return ((bits + 0x1000) & ~0x1FFF).view(torch.float32) + + +@pytest.mark.parametrize("m", (1, 32, 4096)) +def test_wrong_controls_are_rejected(m: int) -> None: + x, w = _make(m, K_HC, seed=200 + m, flavor="hc_like") + y, r = _call(x, w) + good_bound, good = _y_errors(y, x, w) + + # TF32-rounded weights -- the fidelity that disqualified the tensor-core path. + y_tf32, _ = _call(x, _tf32_round(w.clone())) + _, tf32 = _y_errors(y_tf32, x, w) + # bf16-rounded weights. + y_bf16, _ = _call(x, w.bfloat16().float()) + _, bf16 = _y_errors(y_bf16, x, w) + # The stream flattened hidden-major instead of copy-major -- a caller layout bug. + x_perm = x.view(m, HC_MULT, HIDDEN).transpose(1, 2).contiguous().view(m, K_HC) + y_perm, _ = _call(x_perm, w) + perm_bound, perm = _y_errors(y_perm, x, w) + print( + f"\nM={m}: median err/||x*w||_2 in units of the gate: correct {good / _FIDELITY:.3g} " + f"(rigorous-bound ratio {good_bound:.3g}), " + f"TF32 w {tf32 / _FIDELITY:.3g}, bf16 w {bf16 / _FIDELITY:.3g}, " + f"permuted stream {perm / _FIDELITY:.3g} (rigorous-bound ratio {perm_bound:.3g})" + ) + assert good <= _FIDELITY + assert tf32 > 4 * _FIDELITY, f"TF32 control only {tf32 / _FIDELITY:.3g}x the gate" + assert bf16 > 4 * _FIDELITY + assert perm > 4 * _FIDELITY and perm_bound > 1.0 + + # sum(x^2) accumulated from bf16-rounded squares. + r_bf16 = (x.float() * x.float()).bfloat16().float().sum(-1) + ref = _ref(x, w) + with pytest.raises(AssertionError): + op.compare((y, r_bf16), ref) # y is the correct one: only r can fail + op.compare((y, r), ref) + + +# ─── rejections, and the silent failures the guard's asserts exist for ────── + + +def _buffers(m: int, k: int): + x, w = _make(m, k, seed=11, flavor="gauss") + return x, w, torch.empty(m, N_HC, device="cuda"), torch.empty(m, device="cuda") + + +def test_tile_n_must_divide_n() -> None: + x, w, y, r = _buffers(4, 2048) + for tile_n in (5, 16): + with pytest.raises(RuntimeError, match=f"not divisible by tile_n={tile_n}"): + op(x, w, y, r, 4, N_HC, 2048, tile_n) + + +def test_dtypes_are_checked_by_the_op() -> None: + x, w, y, r = _buffers(4, 2048) + with pytest.raises(RuntimeError, match="expected scalar type BFloat16 but found Float"): + op(x.float(), w, y, r, 4, N_HC, 2048) + with pytest.raises(RuntimeError, match="expected scalar type Float but found BFloat16"): + op(x, w.bfloat16(), y, r, 4, N_HC, 2048) + with pytest.raises(RuntimeError, match="expected scalar type Float but found BFloat16"): + op(x, w, y.bfloat16(), r, 4, N_HC, 2048) + + +def test_raw_op_silently_misreads_strides_and_sizes() -> None: + """Why `is_valid` asserts layout and sizes: the op itself accepts these and returns wrong numbers.""" + m = 8 + big, w = _make(m, K_HC + 1024, seed=12, flavor="gauss") + w = w[:, :K_HC].contiguous() + x_strided = big[:, :K_HC] + y = torch.empty(m, N_HC, device="cuda") + r = torch.empty(m, device="cuda") + op.raw_call(x_strided, w, y, r, m, N_HC, K_HC, 0, 0) + torch.cuda.synchronize() + bound, _ = _y_errors(y, x_strided, w) + assert bound > 1.0, ( + "a row-strided x was read correctly; the contiguity assert may be unnecessary" + ) + + x = x_strided.contiguous() + op.raw_call(x, w, y, r, m, N_HC, K_HC - 1024, 0, 0) + torch.cuda.synchronize() + bound, _ = _y_errors(y, x, w) + assert bound > 1.0, "a K smaller than x's width was read correctly" + + +def test_guard_rejects_layouts_the_op_misreads() -> None: + m, k = 8, 2048 + x, w, y, r = _buffers(m, k) + wide = torch.empty(m, k + 64, device="cuda", dtype=torch.bfloat16) + cases = ( + ("x must be", (wide[:, :k], w, y, r, m, N_HC, k)), + ("w must be", (x, w.t().contiguous().t(), y, r, m, N_HC, k)), + ("y must be", (x, w, torch.empty(m, 2 * N_HC, device="cuda")[:, :N_HC], r, m, N_HC, k)), + ("r must be", (x, w, y, torch.empty(2 * m, device="cuda")[::2], m, N_HC, k)), + ("x must be", (x, w, y, r, m, N_HC, k - 1024)), + ("x must be", (x, w, y, r, m + 4, N_HC, k)), + ) + for match, args in cases: + with pytest.raises(AssertionError, match=match): + with validating(op): + op(*args) + + +# ─── every row count 1..32768 in the target's tile bands, and the capture union ── +# +# DeepSeek-V4.1's core target calls this op twice per layer with M = the rows +# of the pass: the tokens the rank holds this step (any integer up to the +# engine's max_num_tokens, 16384 at most in the envelope) or, in layers 21-39 +# of a step that runs the decoder bounded-replay pass, that pass's rows (the +# finishing requests' 128-row tails plus the generation rows, up to 32768). +# The DSpark drafter calls it at N = 5G rows (G draft requests, up to 256). It +# picks tile_n by M: 1 for M <= 128, 6 for 129..2048 and 24 above (`_tile_n` +# in the target), so the tile-24 band reaches 32768. M only sizes grid +# axis 0, where block (m, j) reads row m of x and writes row m of y and r. +# tile_n selects one of three compiled kernels. The tests below check both +# facts on the GPU, at every M from 1 to 32768, each M run with the tile its +# band gets; the entry's cells hold every band's edges and every dumped M +# against the float64 reference. + +_N_MAX = 32768 +_BANDS = ((1, 128, 1), (129, 2048, 6), (2049, _N_MAX, 24)) # (first M, last M, tile_n) +# The band points and runtime-dump row counts the entry certifies at the target's tile. +_CELL_BAND_M = tuple(sorted({c.spec["M"] for c in op.CELLS if c.spec["tile_n"] > 0})) +# The capture union per rank. G ranges over every CUDA-graph batch size an +# envelope engine captures: every G 1..32, plus TRT-LLM's padded default list +# for cuda_graph_config {enable_padding: true, max_batch_size: 256} +# (40..128 step 8, 192, 256); the padding-off list at 256 (1..32, 64, 128, 256) +# is a subset. A captured step is generation-only, so the trunk runs at G (a +# decode step) and 6G (the DSpark target-verify step, 1 + 5 draft tokens per +# request), and the drafter at 5G. The largest is 1536, so a captured call +# only ever takes tile 1 or 6. +_GSET = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_CAPTURE_M = tuple(sorted({k * g for g in _GSET for k in (1, 5, 6)})) +assert len(_GSET) == 46 and len(_CAPTURE_M) == 114 and max(_CAPTURE_M) == 1536 + + +def _target_tile(m: int) -> int: + for lo, hi, tile in _BANDS: + if lo <= m <= hi: + return tile + raise ValueError(m) + + +def test_every_row_count_is_the_same_per_row_function() -> None: + """Every prefix x[:m], m = 1..32768, at its band's tile_n, returns exactly the + first m rows of the 32768-row result, which is held to the float64 reference. + + The output buffers are NaN-filled before every call, and the row after the + last one is a guard row that must stay NaN. So a row the kernel skipped, or + a write past row m - 1, shows up as a mismatch. + """ + x, w = _make(_N_MAX, K_HC, seed=4242, flavor="hc_like") + full_y, full_r = _call(x, w, 24) + _check(full_y, full_r, x, w) + for tile in (1, 6): # the 32768-row result is the same bits at every band's tile + y_t, r_t = _call(x, w, tile) + assert torch.equal(y_t, full_y) and torch.equal(r_t, full_r), tile + fy, fr = full_y.view(torch.int32), full_r.view(torch.int32) + y_buf = torch.empty(_N_MAX + 1, N_HC, device="cuda") + r_buf = torch.empty(_N_MAX + 1, device="cuda") + mismatch = torch.zeros(_N_MAX + 1, dtype=torch.int64, device="cuda") + guard = torch.zeros(_N_MAX + 1, dtype=torch.int64, device="cuda") + for m in range(1, _N_MAX + 1): + y_buf[: m + 1].fill_(float("nan")) + r_buf[: m + 1].fill_(float("nan")) + op(x[:m], w, y_buf[:m], r_buf[:m], m, N_HC, K_HC, _target_tile(m)) + mismatch[m] = (y_buf[:m].view(torch.int32) != fy[:m]).sum() + ( + r_buf[:m].view(torch.int32) != fr[:m] + ).sum() + guard[m] = (~y_buf[m].isnan()).sum() + (~r_buf[m].isnan()).long() + bad = mismatch.nonzero().flatten().tolist() + over = guard.nonzero().flatten().tolist() + assert not bad, ( + f"{len(bad)} row counts differ from the {_N_MAX}-row result's rows, first {bad[:10]}" + ) + assert not over, f"{len(over)} row counts wrote the guard row, first {over[:10]}" + + +def test_row_sweep_harness_sees_a_one_row_shift() -> None: + """Control for the sweep: its bitwise per-row check rejects a read one row off, at every band.""" + x, w = _make(_N_MAX + 1, K_HC, seed=4243, flavor="gauss") + full_y, full_r = _call(x[:_N_MAX], w, 24) + for m in (1, 2, 127, 128, 129, 2048, 2049, 8192, 16384, _N_MAX): + y, r = _call(x[1 : m + 1], w, _target_tile(m)) + n_bad = int((y.view(torch.int32) != full_y.view(torch.int32)[:m]).sum().item()) + n_bad_r = int((r.view(torch.int32) != full_r.view(torch.int32)[:m]).sum().item()) + print( + f"[control row shift] M={m}: {n_bad} / {m * N_HC} y and {n_bad_r} / {m} r elements differ" + ) + assert n_bad > m * N_HC // 2 and n_bad_r > m // 2 + + +def _kernels_launched(fn) -> list: + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + return [e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA] + + +_PROBE_M = tuple( + sorted( + set(_CELL_BAND_M) + | set(_CAPTURE_M) + | {2, 3, 7, 130, 2050, 4096, 8191, 8193, 12288, 16383, 16385, 24576, _N_MAX - 1} + ) +) + + +def test_one_compiled_kernel_per_tile_band() -> None: + """Probe: within a band, every M launches the same single kernel; the three bands launch three. + + The op is compiled ahead of time (no JIT). This shows that M reaches the + launch only as the grid size, and the target's tile_n is the only kernel + choice. + """ + x, w = _make(_N_MAX, K_HC, seed=4244, flavor="gauss") + y = torch.empty(_N_MAX, N_HC, device="cuda") + r = torch.empty(_N_MAX, device="cuda") + per_band = {} + for m in _PROBE_M: + names = _kernels_launched( + lambda m=m: op(x[:m], w, y[:m], r[:m], m, N_HC, K_HC, _target_tile(m)) + ) + assert len(names) == 1, (m, names) + per_band.setdefault(_target_tile(m), set()).add(names[0]) + print(f"{len(_PROBE_M)} row counts probed") + for tile, names in sorted(per_band.items()): + print(f"tile_n={tile}: {sorted(names)}") + assert len(names) == 1, (tile, names) + assert f"mhcGemmSqrsumFmaKernel<{tile}>" in next(iter(names)), names + assert len(per_band) == 3 + + +@pytest.mark.parametrize("m", _CAPTURE_M) +def test_capture_union_replay_is_bitwise_eager(m: int) -> None: + """Captured at M (the target's tile) after an eager call at M, replayed on new + data: bitwise equal to an eager call on that data. + + The first replay, on the capture-time data, must differ from the second, on + new data, so a replay that reused stale inputs would fail. + """ + tile = _target_tile(m) + x_static, w = _make(m, K_HC, seed=6000 + m, flavor="hc_like") + y = torch.full((m, N_HC), float("nan"), device="cuda") + r = torch.full((m,), float("nan"), device="cuda") + op(x_static, w, y, r, m, N_HC, K_HC, tile, 0) # the eager call at this shape + torch.cuda.synchronize() + y_eager0, r_eager0 = y.clone(), r.clone() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(x_static, w, y, r, m, N_HC, K_HC, tile, 0) + y.fill_(float("nan")) + r.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(y, y_eager0) and torch.equal(r, r_eager0), m + y_old = y.clone() + x_new, _ = _make(m, K_HC, seed=7000 + m, flavor="hc_like") + x_static.copy_(x_new) + y.fill_(float("nan")) + r.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + y_eager, r_eager = _call(x_new, w, tile) + assert torch.equal(y, y_eager) and torch.equal(r, r_eager), m + assert not torch.equal(y, y_old), "the replay did not read the new stream" + _check(y, r, x_new, w) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/gemm/test_modeling_v2_mxfp8_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/gemm/test_modeling_v2_mxfp8_mxfp8_gemm.py new file mode 100644 index 000000000000..8e751797af3f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/gemm/test_modeling_v2_mxfp8_mxfp8_gemm.py @@ -0,0 +1,912 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the mxfp8_mxfp8_gemm entry (MXFP8 x MXFP8 dense GEMM, flashinfer.mm_mxfp8). + +The entry owns what is certified: `CELLS` is the list (every DeepSeek-V4.1 +dense FP8 linear at decode through prefill rows on each certified backend, the +8192-row prefill chunk, every kind of scale padding, fp16 output and a caller +`out`), `reference` (the exact product of the dequantized operands in true +fp32) and `compare` (the fp32-accumulation bound) are the gate. This file +builds the operands byte-first, with checkpoint-style 32x32 weight scales +expanded to the per-row form the op consumes, and keeps the checks a cell +cannot carry: the bit-exact regime, ignored padding bytes, fused projections, +the wrong-layout controls, determinism, CUDA-graph replay, the op's own +rejections and the guard, and the target's row-count sweeps, captures, +workspace stability and row-addressing controls at the core and drafter +linears. +""" + +import zlib + +import pytest +import torch +from torch.profiler import ProfilerActivity, profile, record_function + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.gemm.mxfp8_mxfp8_gemm import ( # noqa: E501 + mxfp8_mxfp8_gemm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "mxfp8_mxfp8_gemm requires a CUDA device" + +DEV = torch.device("cuda") +BLOCK = 32 # MXFP8: one UE8M0 scale per 32 contiguous elements along K +BACKENDS = ("cudnn", "cute-dsl", "cutlass") + +# The DeepSeek-V4.1-Flash dense FP8 linears, (N, K) in nn.Linear [out, in] order +# per rank: wq_a||wkv, wq_b (attention TP4 and attention DP), wo_b (attention +# TP4 and attention DP), shared w1||w3 and w2, indexer wq_b (also wq_b under +# attention TP8), engram wkv, wq_b and wo_b under attention TP2, wo_b under +# attention TP8, DSpark main_proj. wq_b is [64 / tp * 512, 1280] and wo_b +# [5120, 8 / tp * 1024] per rank (tp 1 under attention DP). +TARGET_SHAPES = ( + (1792, 5120), + (8192, 1280), + (32768, 1280), + (5120, 2048), + (5120, 8192), + (4608, 5120), + (5120, 2304), + (4096, 1280), + (25600, 6144), + (16384, 1280), + (5120, 4096), + (5120, 1024), + (5120, 15360), +) + + +def _pad_up(x: int, m: int) -> int: + return (x + m - 1) // m * m + + +def _swizzle(sf_2d: torch.Tensor, pad_fill: int = 0) -> torch.Tensor: + """[rows, cols] UE8M0 bytes -> the flat 128x4-swizzled buffer, in native torch. + + Offset of the scale of (row r, block c): + (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + + (r // 128) * 128 * pad_up(cols, 4) + Length pad_up(rows, 128) * pad_up(cols, 4); `pad_fill` goes to every offset + no real (r, c) addresses. + """ + rows, cols = sf_2d.shape + padded_cols = _pad_up(cols, 4) + r = torch.arange(rows, device=DEV).view(-1, 1) + c = torch.arange(cols, device=DEV).view(1, -1) + idx = ( + (c % 4) + + (c // 4) * (4 * 128) + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.full((_pad_up(rows, 128) * padded_cols,), pad_fill, dtype=torch.uint8, device=DEV) + out[idx.flatten()] = sf_2d.flatten() + return out + + +def _e4m3(rows: int, k: int, g: torch.Generator) -> torch.Tensor: + """Random e4m3 elements over the full finite code space (both NaN codes excluded).""" + b = torch.randint(0, 256, (rows, k), generator=g, device=DEV, dtype=torch.uint8) + b = torch.where((b & 0x7F) == 0x7F, b ^ 0x01, b) + return b.view(torch.float8_e4m3fn) + + +def _activation(m: int, k: int, seed: int, sf_lo: int = 116, sf_hi: int = 124): + """(e4m3 [M, K], linear UE8M0 [M, K/32]) built byte-first.""" + g = torch.Generator(device="cuda").manual_seed(seed) + q = _e4m3(m, k, g) + sf = torch.randint(sf_lo, sf_hi, (m, k // BLOCK), generator=g, device=DEV, dtype=torch.uint8) + return q, sf + + +def _checkpoint_weight(n: int, k: int, seed: int): + """(e4m3 [N, K], checkpoint 32x32 block scale [ceil(N/32), K/32], row-expanded [N, K/32]). + + The checkpoint stores one UE8M0 byte per 32x32 tile. Expanding it to the 1x32 + per-row vector this op consumes is `repeat_interleave(32, dim=0)[:N]`, which + is exact: every row of a tile carries that tile's scale. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + q = _e4m3(n, k, g) + blk = torch.randint( + 118, 126, (_pad_up(n, 32) // 32, k // BLOCK), generator=g, device=DEV, dtype=torch.uint8 + ) + return q, blk, blk.repeat_interleave(BLOCK, dim=0)[:n].contiguous() + + +def _values(q: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """Exact fp32 value of an MXFP8 operand from linear scale bytes, for the bit-exact regime.""" + return q.float() * torch.exp2(sf.float() - 127.0).repeat_interleave(BLOCK, dim=-1) + + +def _ref_and_bound(aq, asf, wq, wsf, out_dtype: torch.dtype = torch.bfloat16): + """The entry's `(ref, bound)` for operands given with linear `[rows, K/32]` scale bytes.""" + return op.reference(aq, _swizzle(asf), wq, _swizzle(wsf), out_dtype) + + +def _ratio(out: torch.Tensor, ref: torch.Tensor, bound: torch.Tensor) -> float: + """max |out - ref| / bound: the quantity the entry's gate holds at <= 1.""" + err = (out.float() - ref).abs() + return (err / bound.clamp_min(torch.finfo(torch.float32).tiny)).max().item() + + +def _check_against(out: torch.Tensor, ref: torch.Tensor, bound: torch.Tensor) -> float: + """The entry's gate on a precomputed (possibly row-sliced) reference; returns the ratio.""" + op.compare(out, (ref, bound)) + return _ratio(out, ref, bound) + + +def _check(out: torch.Tensor, aq, asf, wq, wsf) -> float: + """The entry's gate against its reference for these operands; returns max |err| / bound.""" + ref, bound = _ref_and_bound(aq, asf, wq, wsf, out.dtype) + return _check_against(out, ref, bound) + + +def _weight_rows(wq: torch.Tensor, rows: torch.Tensor) -> torch.Tensor: + """A row subset of an e4m3 weight (selected as bytes).""" + return wq.view(torch.uint8).index_select(0, rows).view(torch.float8_e4m3fn) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + n, k, m, backend = spec["N"], spec["K"], spec["M"], spec["backend"] + wq, _, wsf = _checkpoint_weight(n, k, seed=n + k) + aq, asf = _activation(m, k, zlib.crc32(cell.why.encode()) % 2**31, *spec["act_sf"]) + a_sw, w_sw = _swizzle(asf), _swizzle(wsf) + assert a_sw.numel() == _pad_up(m, 128) * _pad_up(k // BLOCK, 4) + buf = None + if spec["out"]: + buf = torch.full((m, n), float("nan"), dtype=spec["out_dtype"], device=DEV) + out = op(aq, a_sw, wq, w_sw, spec["out_dtype"], backend, buf) + assert out.shape == (m, n) and out.dtype == spec["out_dtype"] + if buf is not None: + assert out.data_ptr() == buf.data_ptr(), "a caller `out` is written in place and returned" + op.compare(out, op.reference(aq, a_sw, wq, w_sw, spec["out_dtype"], backend, buf)) + if backend == "cutlass": + # cute-dsl and cutlass are bitwise identical to each other at every cell. + assert torch.equal(out, op(aq, a_sw, wq, w_sw, spec["out_dtype"], "cute-dsl")), cell.why + + +def test_exact_regime_is_bit_exact() -> None: + """When the fp32 sum is exact, the bf16 output is exactly one rounding of it. + + Elements are +-{0.5, 1, 2} and scales 2^{-1, 0, 1}: every product is a + multiple of 2^-4 below 2^4, every partial sum over K <= 15360 fits in 22 + bits, so fp32 accumulation in any order is exact. Any misplaced scale, lost + block or wrong byte would change the bits. + """ + codes = torch.tensor( + [0x30, 0x38, 0x40, 0xB0, 0xB8, 0xC0], dtype=torch.uint8, device=DEV + ) # +-0.5, +-1, +-2 + for n, k, m in ((1792, 5120, 8), (5120, 15360, 33), (5120, 2304, 200), (5120, 576, 5)): + g = torch.Generator(device="cuda").manual_seed(n * 7 + k + m) + aq = codes[torch.randint(0, 6, (m, k), generator=g, device=DEV)].view(torch.float8_e4m3fn) + wq = codes[torch.randint(0, 6, (n, k), generator=g, device=DEV)].view(torch.float8_e4m3fn) + asf = torch.randint(126, 129, (m, k // BLOCK), generator=g, device=DEV, dtype=torch.uint8) + wsf = torch.randint( + 126, 129, (n // BLOCK, k // BLOCK), generator=g, device=DEV, dtype=torch.uint8 + ) + wsf = wsf.repeat_interleave(BLOCK, dim=0).contiguous() + exact = (_values(aq, asf).double() @ _values(wq, wsf).double().T).to(torch.bfloat16) + for backend in BACKENDS: + out = op(aq, _swizzle(asf), wq, _swizzle(wsf), backend=backend) + assert torch.equal(out, exact), (backend, n, k, m) + + +def test_scale_padding_bytes_are_ignored() -> None: + """Whatever fills the padded rows/columns of either scale buffer never reaches the result.""" + k, n, m = 576, 160, 7 # every kind of padding at once + wq, _, wsf = _checkpoint_weight(n, k, seed=8) + aq, asf = _activation(m, k, seed=8) + for backend in BACKENDS: + zero = op(aq, _swizzle(asf, 0x00), wq, _swizzle(wsf, 0x00), backend=backend) + # 0xFE is the largest finite E8M0 (2^127); 0xFF would be NaN. + loud = op(aq, _swizzle(asf, 0xFE), wq, _swizzle(wsf, 0xFE), backend=backend) + assert torch.equal(zero, loud), backend + _check(zero, aq, asf, wq, wsf) + + +def test_fused_projection_equals_separate_ones() -> None: + """Row-concatenating two weights (wq_a || wkv) and splitting the output is bitwise the two separate GEMMs.""" + k, m = 5120, 33 + wq_a, _, sf_a = _checkpoint_weight(1280, k, seed=10) + wkv, _, sf_kv = _checkpoint_weight(512, k, seed=11) + aq, asf = _activation(m, k, seed=12) + a_sw = _swizzle(asf) + for backend in BACKENDS: + fused = op( + aq, a_sw, torch.cat([wq_a, wkv]), _swizzle(torch.cat([sf_a, sf_kv])), backend=backend + ) + assert torch.equal(fused[:, :1280], op(aq, a_sw, wq_a, _swizzle(sf_a), backend=backend)), ( + backend + ) + assert torch.equal(fused[:, 1280:], op(aq, a_sw, wkv, _swizzle(sf_kv), backend=backend)), ( + backend + ) + + +def test_wrong_operand_layouts_fail() -> None: + """Deliberately wrong controls: each is accepted by the op and must land far outside the gate. + + On this model family a wrong layout produces plausible numbers rather than + an error, so the test proves the gate can see each one: + - weight scale fed row-major (unswizzled) at the right buffer size; + - activation scale fed row-major (unswizzled) at the right buffer size; + - the checkpoint's 32x32 tile scale expanded by tiling instead of + repeat_interleave (row r gets tile r mod N/32 instead of r // 32); + - every scale byte off by one (a wrong UE8M0 bias: the result doubles). + """ + n, k, m = 1792, 5120, 64 + wq, blk, wsf = _checkpoint_weight(n, k, seed=13) + aq, asf = _activation(m, k, seed=13) + ref, bound = _ref_and_bound(aq, asf, wq, wsf) + a_sw, w_sw = _swizzle(asf), _swizzle(wsf) + + def linear_in_swizzled_size(sf: torch.Tensor, size: int) -> torch.Tensor: + buf = torch.zeros(size, dtype=torch.uint8, device=DEV) + buf[: sf.numel()] = sf.flatten() + return buf + + wrong = { + "unswizzled weight scale": (a_sw, linear_in_swizzled_size(wsf, w_sw.numel())), + "unswizzled activation scale": (linear_in_swizzled_size(asf, a_sw.numel()), w_sw), + "tiled 32x32 expansion": (a_sw, _swizzle(blk.repeat(BLOCK, 1)[:n].contiguous())), + "scale bytes off by one": (a_sw, _swizzle(wsf + 1)), + } + for backend in BACKENDS: + assert _check_against(op(aq, a_sw, wq, w_sw, backend=backend), ref, bound) <= 1.0 + for name, (sa, sw) in wrong.items(): + out = op(aq, sa, wq, sw, backend=backend) + assert torch.isfinite(out).all(), (backend, name) # plausible numbers, not an error + with pytest.raises(AssertionError): + _check_against(out, ref, bound) + + +def test_deterministic_and_inputs_untouched() -> None: + """Repeating a call gives the same bits; no input is written.""" + n, k, m = 4608, 5120, 16 + wq, _, wsf = _checkpoint_weight(n, k, seed=14) + aq, asf = _activation(m, k, seed=14) + a_sw, w_sw = _swizzle(asf), _swizzle(wsf) + snap = [t.clone() for t in (aq.view(torch.uint8), a_sw, wq.view(torch.uint8), w_sw)] + for backend in BACKENDS: + first = op(aq, a_sw, wq, w_sw, backend=backend) + for _ in range(3): + assert torch.equal(first, op(aq, a_sw, wq, w_sw, backend=backend)), backend + for before, now in zip(snap, (aq.view(torch.uint8), a_sw, wq.view(torch.uint8), w_sw)): + assert torch.equal(before, now) + + +def test_cuda_graph_replay_matches_eager() -> None: + """After one eager call (JIT / cuDNN graph build), the call captures and replays bit-identically.""" + n, k = 8192, 1280 + wq, _, wsf = _checkpoint_weight(n, k, seed=15) + w_sw = _swizzle(wsf) + for m in (8, 256): + aq, asf = _activation(m, k, seed=15 + m) + a_sw = _swizzle(asf) + for backend in BACKENDS: + eager = op(aq, a_sw, wq, w_sw, backend=backend) + out = torch.empty_like(eager) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op(aq, a_sw, wq, w_sw, backend=backend, out=out) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(aq, a_sw, wq, w_sw, backend=backend, out=out) + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, eager), (backend, m) + + +def test_rejected_domains() -> None: + """Violations the op itself rejects loudly (each observed to raise on this machine).""" + k, n, m = 1024, 512, 8 + wq, _, wsf = _checkpoint_weight(n, k, seed=16) + aq, asf = _activation(m, k, seed=16) + a_sw, w_sw = _swizzle(asf), _swizzle(wsf) + for backend in BACKENDS: + with pytest.raises(ValueError): # no fp32 output on any backend + op(aq, a_sw, wq, w_sw, out_dtype=torch.float32, backend=backend) + with pytest.raises(ValueError): # activation scale not in the padded swizzled size + op(aq, asf.flatten(), wq, w_sw, backend=backend) + with pytest.raises(ValueError): # weight scale too short + op(aq, a_sw, wq, w_sw[:-4], backend=backend) + w96, _, sf96 = _checkpoint_weight(96, k, seed=17) + with pytest.raises(ValueError): # N < 128 + op(aq, a_sw, w96, _swizzle(sf96), backend=backend) + with pytest.raises(ValueError): # activation not e4m3 + op(_values(aq, asf).bfloat16(), a_sw, wq, w_sw, backend=backend) + + +def test_guard_rejects_the_silent_domains() -> None: + """The violations observed to be accepted and computed wrongly are rejected by `is_valid`.""" + k, n, m = 1024, 512, 8 + wq, _, wsf = _checkpoint_weight(n, k, seed=18) + aq, asf = _activation(m, k, seed=18) + a_sw, w_sw = _swizzle(asf), _swizzle(wsf) + wide_a = torch.cat([aq.view(torch.uint8)] * 2, 1).view(torch.float8_e4m3fn) + wide_w = torch.cat([wq.view(torch.uint8)] * 2, 1).view(torch.float8_e4m3fn) + wide_out = torch.empty(m, 2 * n, dtype=torch.bfloat16, device=DEV) + for backend in BACKENDS: + with pytest.raises(AssertionError, match="must be contiguous"): + with validating(op): + op(wide_a[:, :k], a_sw, wq, w_sw, backend=backend) + with pytest.raises(AssertionError, match="must be contiguous"): + with validating(op): + op(aq, a_sw, wide_w[:, :k], w_sw, backend=backend) + with pytest.raises(AssertionError, match="out must be contiguous"): + with validating(op): + op(aq, a_sw, wq, w_sw, backend=backend, out=wide_out[:, :n]) + for backend in ("trtllm", "auto", "b12x"): + with pytest.raises(AssertionError, match="backend must be one of"): + with validating(op): + op(aq, a_sw, wq, w_sw, backend=backend) + + +def test_zero_rows_raise() -> None: + """M == 0 raises on every backend (last: cute-dsl's rejection is a failed kernel launch).""" + k, n = 1024, 512 + wq, _, wsf = _checkpoint_weight(n, k, seed=19) + w_sw = _swizzle(wsf) + aq = torch.empty(0, k, dtype=torch.float8_e4m3fn, device=DEV) + a_sw = torch.empty(0, dtype=torch.uint8, device=DEV) + for backend in ("cutlass", "cudnn", "cute-dsl"): + with pytest.raises( + Exception + ): # RuntimeError (cutlass, cudnn); a CuTe DSL launch error (cute-dsl) + op(aq, a_sw, wq, w_sw, backend=backend) + torch.cuda.synchronize() + + +# ─── DeepSeek-V4.1 core target: every runtime row count ──────────────────── +# +# The core target's twelve dense FP8 linears (TARGET_SHAPES without the DSpark +# main_proj), over its members: attention TP 2 (pp2 stages), 4 or 8, or +# attention DP. The target picks the backend by row count: cute-dsl at M <= 32, +# cudnn above. M is the rank's token count this step: any of 1..max_num_tokens, +# at most 16384 per rank. In layers 21-39 the decoder bounded replay (on every +# member, attention TP or DP) runs decoder passes over the finishing requests' +# 128-row tails plus the step's generation rows; the target cuts them into +# passes of at most the step capacity, and 32768 (128 x 256) is the uncut row +# set. wq_a||wkv, wq_b and wo_b at every member, the shared w1||w3 and w2, and +# the indexer wq_b (index layers 24, 28, 32, 36) run there, so their N_MAX is +# 32768. The Engram wkv (layers 1 and 14) stays at 16384. +CORE_SHAPES = tuple(s for s in TARGET_SHAPES if s != (5120, 15360)) +N_MAX = { + (1792, 5120): 32768, + (8192, 1280): 32768, + (32768, 1280): 32768, + (5120, 2048): 32768, + (5120, 8192): 32768, + (4608, 5120): 32768, + (5120, 2304): 32768, + (4096, 1280): 32768, + (25600, 6144): 16384, + (16384, 1280): 32768, + (5120, 4096): 32768, + (5120, 1024): 32768, +} +DECODE_ROWS = tuple(range(1, 33)) +# FlashInfer's cudnn backend builds one cuDNN graph per (N, K, M bucket), at the +# bucket's top, and runs it at the real M. The buckets are powers of two to +# 256, steps of 256 to 2048, steps of 512 to 4096, then powers of two. This list +# covers every bucket from 33 to 32768. For each one it has the bucket's first +# row count (farthest below the size the plan was built for), one inside, and +# the top. +BUCKET_ROWS = ( + 33, 48, 61, 64, 65, 97, 128, 129, 211, 256, 257, 323, 509, 512, 513, 576, 767, 768, + 769, 862, 1000, 1024, 1025, 1153, 1280, 1281, 1423, 1536, 1537, 1709, 1792, 1793, + 1935, 2048, 2049, 2311, 2560, 2561, 2917, 3072, 3073, 3305, 3584, 3585, 3907, 4096, + 4097, 4617, 6523, 8179, 8191, 8192, 8193, 12289, 16383, 16384, 16385, 24577, 32767, + 32768, +) # fmt: skip +# Row counts above 32 of the core target's linears in its runtime-value dumps: +# the gate runs, and the engine legs at max_num_tokens 16384 and +# max_batch_size 256 (tep4, tp4, dep4; padded graphs to 256; L to 1M). +DUMP_ROWS = ( + 34, 48, 60, 96, 108, 127, 128, 129, 134, 144, 168, 192, 228, 240, 288, 330, 336, 384, + 432, 462, 480, 528, 546, 576, 597, 610, 624, 672, 720, 768, 780, 852, 1152, 1536, 1906, + 2033, 2038, 2048, 2049, 2778, 2784, 3870, 4208, 4245, 4388, 4394, 4591, 4868, 7835, + 7855, 7893, 7908, 8191, 8192, 8467, 9193, 9240, 9364, 9370, 9524, 10441, 10444, 10447, + 10451, 16298, 16308, 16316, 16320, 16322, 16327, 16328, 16329, 16333, 16340, 16378, + 16384, +) # fmt: skip +PREFILL_ROWS = tuple(sorted(set(BUCKET_ROWS) | set(DUMP_ROWS))) +# The captured batch sizes G, per rank: every G 1..32, plus TRT-LLM's padded +# default list for cuda_graph_config {enable_padding: true, max_batch_size: 256} +# (40..128 step 8, 192, 256); the padding-off list for 256 (1..32, 64, 128, 256) +# is a subset. A captured step is generation-only; each linear sees a multiple +# of G (no MXFP8 linear runs on dep4's gathered rows). +GSET = tuple(range(1, 33)) + (40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256) +# (N, K) and the multiples of G it is captured at: G (decode) and 6G (DSpark +# target verify, 1 + 5 drafts) for the trunk linears; 5G (the drafter's 5-token +# draft block) where the drafter has a linear of that shape; G and 5G for the +# drafter-only main_proj and fused wkv x3 (its bonus rows and its ring +# back-fill). Up to 1536 rows. +CAPTURE_CELLS = ( + ((1792, 5120), (1, 5, 6)), # wq_a || wkv: trunk, drafter + ((8192, 1280), (1, 6)), # wq_b, attention TP4 (16 heads per rank) + ((32768, 1280), (1, 5, 6)), # wq_b, attention DP (all 64 heads) and the drafter + ((5120, 2048), (1, 6)), # wo_b, attention TP4 + ((5120, 8192), (1, 5, 6)), # wo_b, attention DP and the drafter + ((4608, 5120), (1, 5, 6)), # shared w1 || w3: trunk, drafter + ((5120, 2304), (1, 5, 6)), # shared w2: trunk, drafter + ((4096, 1280), (1, 6)), # indexer wq_b; wq_b, attention TP8 (8 heads per rank) + ((25600, 6144), (1, 6)), # Engram wkv + ((16384, 1280), (1, 6)), # wq_b, attention TP2 (32 heads per rank) + ((5120, 4096), (1, 6)), # wo_b, attention TP2 + ((5120, 1024), (1, 6)), # wo_b, attention TP8 + ((5120, 15360), (1, 5)), # drafter main_proj + ((1536, 5120), (1, 5)), # drafter wkv x3, fused on the row axis +) + + +def _capture_rows(factors) -> tuple: + return tuple(sorted({f * g for g in GSET for f in factors})) + + +def _target_backend(m: int) -> str: + """The target's policy: cute-dsl at M <= 32, cudnn above.""" + return "cute-dsl" if m <= 32 else "cudnn" + + +@pytest.mark.parametrize("n,k", CORE_SHAPES) +def test_core_decode_every_row_count(n: int, k: int) -> None: + """M = 1..32 at each core linear, on cute-dsl (the target's backend there) and cudnn. + + Every M is held to the fp32-accumulation bound. cute-dsl is also held to a + measured row invariance: at every M each row's bits equal that row's bits + at M = 32. That holds across its three low-M kernels (N tile 8 / 16 / 32). + """ + wq, _, wsf = _checkpoint_weight(n, k, seed=n + 3 * k) + w_sw = _swizzle(wsf) + aq, asf = _activation(32, k, seed=k + 32) + ref, bound = _ref_and_bound(aq, asf, wq, wsf) + top = op(aq, _swizzle(asf), wq, w_sw, backend="cute-dsl") + for m in DECODE_ROWS: + a_sw = _swizzle(asf[:m]) + for backend in ("cute-dsl", "cudnn"): + out = op(aq[:m], a_sw, wq, w_sw, backend=backend) + assert out.shape == (m, n) and out.dtype == torch.bfloat16 + _check_against(out, ref[:m], bound[:m]) + if backend == "cute-dsl": + assert torch.equal(out, top[:m]), (n, k, m) + + +@pytest.mark.parametrize("n,k", CORE_SHAPES) +def test_core_prefill_every_cudnn_bucket(n: int, k: int) -> None: + """Every row count in PREFILL_ROWS up to the shape's N_MAX at each core linear, on all three backends. + + cudnn is the target's backend above 32 rows. It is checked against the + bound only, because its last bits may depend on the row count. cute-dsl + and cutlass must match each other bitwise, and each row's bits must equal + that row's bits at M = N_MAX (a measured row invariance). + """ + n_max = N_MAX[(n, k)] + wq, _, wsf = _checkpoint_weight(n, k, seed=2 * n + k) + w_sw = _swizzle(wsf) + aq, asf = _activation(n_max, k, seed=5 * k + 1) + ref, bound = _ref_and_bound(aq, asf, wq, wsf) + full = op(aq, _swizzle(asf), wq, w_sw, backend="cutlass") + _check_against(full, ref, bound) + worst = 0.0 + for m in (m for m in PREFILL_ROWS if m <= n_max): + a_sw = _swizzle(asf[:m]) + outs = {} + for backend in BACKENDS: + out = op(aq[:m], a_sw, wq, w_sw, backend=backend) + assert out.shape == (m, n) + ratio = _check_against(out, ref[:m], bound[:m]) + if backend == "cudnn": + worst = max(worst, ratio) + outs[backend] = out + assert torch.equal(outs["cute-dsl"], outs["cutlass"]), (n, k, m) + assert torch.equal(outs["cutlass"], full[:m]), (n, k, m) + del outs + print(f"[full width M 33..{n_max}] ({n}, {k}): worst cudnn |err| / bound = {worst:.4f}") + del ref, bound, full + torch.cuda.empty_cache() + + +def _cudnn_kernel_map(aq, asf, wq, w_sw, rows) -> dict: + """Profile one cudnn call per row count; return {M: the kernel names that call launched, in order}. + + Each call is isolated between two synchronizations inside its own + `record_function` range, so the device kernels whose start falls in that + range are exactly the ones the call launched (the range's own device-side + annotation and any PyTorch kernel are excluded by name). + """ + names = {} + for i in range(0, len(rows), 256): + chunk = rows[i : i + 256] + ops = [(m, _swizzle(asf[:m])) for m in chunk] + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for m, a_sw in ops: + with record_function(f"M={m}"): + op(aq[:m], a_sw, wq, w_sw, backend="cudnn") + torch.cuda.synchronize() + evs = prof.events() + ranges = [ + (int(e.name[2:]), e.time_range.start, e.time_range.end) + for e in evs + if e.name.startswith("M=") and e.device_type == torch.autograd.DeviceType.CPU + ] + kernels = sorted( + (e.time_range.start, e.name) + for e in evs + if e.device_type == torch.autograd.DeviceType.CUDA + and not e.name.startswith("M=") + and "at::native" not in e.name + ) + for m, lo, hi in ranges: + names[m] = tuple(nm for t, nm in kernels if lo <= t <= hi) + del ops, prof + torch.cuda.empty_cache() + return names + + +@pytest.mark.parametrize("n,k", CORE_SHAPES) +def test_core_cudnn_every_row_count_33_to_n_max(n: int, k: int) -> None: + """Every M from 33 to the shape's N_MAX (16384 or 32768) on cudnn, the target's backend there, at each core linear. + + Every row of every call is held to the bound on a column sample: one + column in each 64-column block, at a position that changes from block to + block, plus the last column. A row that is wrong, missing or misplaced + fails on every sampled column, as the row-addressing controls show. The + sample can miss an error confined to a few columns of one tile. Each output + is poisoned before it is freed. The next call's output often reuses the + block, and its rows sit at the same offsets, so a row the kernel skipped + would otherwise keep the previous call's correct bits. + + cuDNN picks its kernel by M (nvjet tile / cluster / split-K variants that + change inside FlashInfer's M buckets too), so the test then profiles every + call M 33..N_MAX, derives the kernel map (maximal runs of consecutive M + that launch the same kernel sequence), and checks the first, a middle and + the last M of every run at FULL width against the reference (every kernel + switch point at its edges). The full-width check at + the bucket edges and the dumped row counts on all three backends is + `test_core_prefill_every_cudnn_bucket`. + """ + n_max = N_MAX[(n, k)] + wq, _, wsf = _checkpoint_weight(n, k, seed=4 * n + k) + w_sw = _swizzle(wsf) + aq, asf = _activation(n_max, k, seed=7 * k + 3) + blocks = torch.arange(n // 64, device=DEV) + cols = torch.cat([blocks * 64 + (blocks * 37) % 64, torch.tensor([n - 1], device=DEV)]) + ref, bound = _ref_and_bound(aq, asf, _weight_rows(wq, cols), wsf.index_select(0, cols)) + bound = bound.clamp_min(torch.finfo(torch.float32).tiny) + worst = torch.zeros((), device=DEV) + for m in range(33, n_max + 1): + out = op(aq[:m], _swizzle(asf[:m]), wq, w_sw, backend="cudnn") + err = (out.index_select(1, cols).float() - ref[:m]).abs() + worst = torch.maximum(worst, (err / bound[:m]).max()) + out.fill_(float("nan")) + if m % 2048 == 0: + # Each call's output is larger than any cached block, so the allocator + # would otherwise keep a segment per size and fill the device. + del out, err + torch.cuda.empty_cache() + ratio = worst.item() + print( + f"[cudnn M 33..{n_max}] ({n}, {k}): max |err| / bound over {cols.numel()} sampled columns = {ratio:.4f}" + ) + assert ratio <= 1.0, ratio + del ref, bound + + rows = list(range(33, n_max + 1)) + names = _cudnn_kernel_map(aq, asf, wq, w_sw, rows) + runs = [] + for m in rows: + assert names.get(m), f"no kernel attributed to the cudnn call at M={m}" + if runs and runs[-1][2] == names[m]: + runs[-1][1] = m + else: + runs.append([m, m, names[m]]) + edges = sorted({m for lo, hi, _ in runs for m in (lo, (lo + hi) // 2, hi)}) + ref, bound = _ref_and_bound(aq, asf, wq, wsf) + edge_worst = 0.0 + for m in edges: + out = op(aq[:m], _swizzle(asf[:m]), wq, w_sw, backend="cudnn") + edge_worst = max(edge_worst, _check_against(out, ref[:m], bound[:m])) + del out + kinds = sorted({nm for _, _, nm in runs}) + print( + f"[cudnn kernel map ({n}, {k})] {len(runs)} runs of {len(kinds)} distinct kernel sequences over M 33..{n_max}; " + f"{len(edges)} run edges / middles at full width, worst |err| / bound {edge_worst:.4f}; " + f"switch points (first M of each run): {[lo for lo, _, _ in runs]}" + ) + for kind in kinds: + print(f" {' + '.join(nm[:120] for nm in kind)}") + del ref, bound + torch.cuda.empty_cache() + + +def _graph_sweep(n: int, k: int, rows, seed: int) -> dict: + """Capture one call per row count on the target's backend and replay it on new activations. + + It follows the engine's order: one eager call at the same M (its warmup), + then capture. The call is captured without `out`, as the target makes it, + so the output lives in the graph's pool. That output is poisoned before + every replay, then checked against the bound and bitwise against an eager + call on the same data. A stale replay is distinguishable: the warm-up's + output on the capture-time data (what a replay that did not rerun, or reran + on the old inputs, would hold) must differ from each replay's output on + most elements. Returns per-backend capture counts, the worst + |err| / bound and the smallest differing fraction. + """ + wq, _, wsf = _checkpoint_weight(n, k, seed=seed) + w_sw = _swizzle(wsf) + stats = {"cute-dsl": 0, "cudnn": 0, "worst": 0.0, "stale_min": 1.0} + for m in rows: + be = _target_backend(m) + aq, asf = _activation(m, k, seed=11 * m + k) + a_sw = _swizzle(asf) + stale = op(aq, a_sw, wq, w_sw, torch.bfloat16, be) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(aq, a_sw, wq, w_sw, torch.bfloat16, be) + for seed_ in (1, 2): + nq, nsf = _activation(m, k, seed=100000 * seed_ + 11 * m + k) + aq.view(torch.uint8).copy_(nq.view(torch.uint8)) + a_sw.copy_(_swizzle(nsf)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + stats["worst"] = max(stats["worst"], _check(out, nq, nsf, wq, wsf)) + assert torch.equal(out, op(nq, _swizzle(nsf), wq, w_sw, torch.bfloat16, be)), ( + n, + k, + m, + be, + ) + frac = (out != stale).float().mean().item() + stats["stale_min"] = min(stats["stale_min"], frac) + assert frac > 0.9, (n, k, m, be, frac) + stats[be] += 1 + del graph, out, stale + return stats + + +@pytest.mark.parametrize("n,k,factors", [(n, k, f) for (n, k), f in CAPTURE_CELLS]) +def test_graph_capture_at_every_captured_batch(n: int, k: int, factors) -> None: + """Each linear captured at every multiple of G it sees over GSET, on the target's backend (cute-dsl + at M <= 32, cudnn above), and replayed twice on new activations.""" + rows = _capture_rows(factors) + stats = _graph_sweep(n, k, rows, seed=n + k + 7) + print( + f"[capture ({n}, {k}) x G{factors}] {len(rows)} row counts 1..{rows[-1]}: {stats['cute-dsl']} on cute-dsl, " + f"{stats['cudnn']} on cudnn, {2 * len(rows)} replays bitwise = eager, worst |err| / bound " + f"{stats['worst']:.4f}, " + f"stale output differs on >= {stats['stale_min']:.4f} of elements" + ) + + +def _mm_mxfp8_workspace(): + """(bytes, address) of FlashInfer's cached `mm_mxfp8_workspace` buffer on this device, or None.""" + from flashinfer import utils as fi_utils + + for (name, _), buf in fi_utils._cache_buf.items(): + if name == "mm_mxfp8_workspace": + return buf.numel(), buf.data_ptr() + return None + + +def test_cudnn_graphs_survive_eager_calls_to_n_max() -> None: + """Graphs captured at a decode row count stay valid after eager cudnn calls at every bucket edge up to N_MAX. + + FlashInfer keeps one cached workspace buffer per device and grows it with + `resize_` (a new allocation at a new address) whenever a cudnn call asks + for more than it holds, so a graph captured before such a growth would + keep the freed address. The engine captures its decode graphs after its + warm-up and then runs eager prefill and decoder-pass steps of up to N_MAX + rows. Here one cudnn graph per core linear is captured at M = 96 after an + eager call; then every core linear runs eagerly at every BUCKET_ROWS value + up to its N_MAX (the first, an interior and the last M of every FlashInfer + bucket); the buffer's size and address must not change, and every graph + must then replay on new data bitwise equal to eager and inside the bound. + """ + graphs = [] + for n, k in CORE_SHAPES: + wq, _, wsf = _checkpoint_weight(n, k, seed=41 + n + k) + w_sw = _swizzle(wsf) + aq, asf = _activation(96, k, seed=42 + k) + a_sw = _swizzle(asf) + op(aq, a_sw, wq, w_sw, torch.bfloat16, "cudnn") + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(aq, a_sw, wq, w_sw, torch.bfloat16, "cudnn") + graphs.append((n, k, wq, wsf, w_sw, aq, a_sw, graph, out)) + before = _mm_mxfp8_workspace() + assert before is not None + calls = 0 + for n, k, wq, _, w_sw, *_ in graphs: + n_max = N_MAX[(n, k)] + big_q, big_sf = _activation(n_max, k, seed=43 + k) + for m in (m for m in BUCKET_ROWS if m <= n_max): + op(big_q[:m], _swizzle(big_sf[:m]), wq, w_sw, torch.bfloat16, "cudnn") + calls += 1 + assert _mm_mxfp8_workspace() == before, (n, k, m, before, _mm_mxfp8_workspace()) + del big_q, big_sf + torch.cuda.synchronize() + for n, k, wq, wsf, w_sw, aq, a_sw, graph, out in graphs: + nq, nsf = _activation(96, k, seed=44 + k) + aq.view(torch.uint8).copy_(nq.view(torch.uint8)) + a_sw.copy_(_swizzle(nsf)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + _check(out, nq, nsf, wq, wsf) + assert torch.equal(out, op(nq, _swizzle(nsf), wq, w_sw, torch.bfloat16, "cudnn")), (n, k) + print( + f"[workspace] {calls} eager cudnn calls at the bucket edges up to N_MAX after {len(graphs)} captures: " + f"buffer stayed at {before[0]} bytes, same address; every graph replayed bitwise = eager" + ) + + +def test_row_addressing_controls_at_runtime_row_counts() -> None: + """The bound catches a row-misaddressed activation at a decode and a prefill row count. + + Shared w1||w3 (4608 x 5120), at M = 7 on cute-dsl and M = 1935 on cudnn + (the target's backends there). The op accepts each wrong operand below, + and each result must land outside the bound: + - activation scale rows rolled by one (row m carries row m-1's scales); + - activation data rows rolled by one against their own scales; + - activation scale fed row-major at the swizzled size. + """ + n, k = 4608, 5120 + wq, _, wsf = _checkpoint_weight(n, k, seed=21) + w_sw = _swizzle(wsf) + for m, backend in ((7, "cute-dsl"), (1935, "cudnn")): + aq, asf = _activation(m, k, seed=22 + m) + a_sw = _swizzle(asf) + assert _check(op(aq, a_sw, wq, w_sw, backend=backend), aq, asf, wq, wsf) <= 1.0 + linear = torch.zeros_like(a_sw) + linear[: asf.numel()] = asf.flatten() + rolled_data = torch.roll(aq.view(torch.uint8), 1, 0).view(torch.float8_e4m3fn).contiguous() + wrong = { + "scale rows rolled by one": (aq, _swizzle(torch.roll(asf, 1, 0).contiguous())), + "data rows rolled by one": (rolled_data, a_sw), + "row-major activation scale": (aq, linear), + } + ref, bound = _ref_and_bound(aq, asf, wq, wsf) + for name, (qa, sa) in wrong.items(): + out = op(qa, sa, wq, w_sw, backend=backend) + assert torch.isfinite(out).all(), (backend, name) + ratio = ( + ((out.float() - ref).abs() / bound.clamp_min(torch.finfo(torch.float32).tiny)) + .max() + .item() + ) + print(f"[control M={m} {backend}] {name}: max |err| / bound = {ratio:.3g}") + assert ratio > 1.0, (backend, name) + + +# ─── DeepSeek-V4.1 DSpark drafter ────────────────────────────────────────── +# +# The drafter's dense MXFP8 linears (replicated on every rank) run inside the +# target's decode CUDA graphs, and eagerly on steps that are not captured. With +# G the drafter's batch (1..256) the row count is G (forward_batched's bonus +# rows through main_proj and the fused wkv) or 5G (the draft block through +# every linear, and the batched ring back-fill through main_proj and wkv), so +# 1..1280, on the target's backend policy. +# (5120, 15360) main_proj (stage 0), (1536, 5120) the three stages' wkv +# fused on the row axis (stage s is output columns [512 s, 512 (s + 1))) +# run eagerly at every M 1..1280 and captured at every M 1..160 (and at G and +# 5G over GSET, CAPTURE_CELLS). The other drafter linears share a core shape +# and run at M = 5G only: wq_a || wkv (1792, 5120), wq_b (32768, 1280), wo_b +# (5120, 8192), shared w1 || w3 (4608, 5120), shared w2 (5120, 2304); their +# eager rows are inside the core shapes' 1..N_MAX and their 5G captures are in +# CAPTURE_CELLS. +DSPARK_NEW_SHAPES = ((5120, 15360), (1536, 5120)) +DSPARK_ROWS = tuple(range(1, 1281)) +DSPARK_GRAPH_ROWS = tuple(range(1, 161)) + + +@pytest.mark.parametrize("n,k", DSPARK_NEW_SHAPES) +def test_dspark_every_row_count_1_to_1280(n: int, k: int) -> None: + """main_proj and the fused wkv x3 at every M 1..1280 on all three backends, within the bound. + + cute-dsl and cutlass are bitwise equal to each other at every M (the same + measured property as at the core shapes); cudnn is held to the bound only. + For the fused wkv x3, at every M, each stage's 512 output columns are + compared with that stage's own [512, 5120] GEMM on the same backend. On + cute-dsl and cutlass they are bitwise equal at every M. On cudnn they are + bitwise equal except at M 193..640, where the fused [1536, 5120] and the + separate [512, 5120] GEMMs launch different kernels (the fused one leaves + split-K at 193) and differ in the last bits of a small fraction of the + elements; there each separate GEMM is held to the bound, and the set of M + where they differ must be exactly 193..640 (the measured window). + """ + wq, _, wsf = _checkpoint_weight(n, k, seed=5 * n + k) + w_sw = _swizzle(wsf) + aq, asf = _activation(DSPARK_ROWS[-1], k, seed=3 * k + 5) + ref, bound = _ref_and_bound(aq, asf, wq, wsf) + stages = [] + if n == 1536: + stages = [ + ( + wq[512 * s : 512 * (s + 1)].contiguous(), + _swizzle(wsf[512 * s : 512 * (s + 1)].contiguous()), + ) + for s in range(3) + ] + worst = 0.0 + cudnn_differs = [] + frac_max = 0.0 + for m in DSPARK_ROWS: + a_sw = _swizzle(asf[:m]) + outs = {} + for backend in BACKENDS: + out = op(aq[:m], a_sw, wq, w_sw, backend=backend) + assert out.shape == (m, n) and out.dtype == torch.bfloat16 + worst = max(worst, _check_against(out, ref[:m], bound[:m])) + outs[backend] = out + assert torch.equal(outs["cute-dsl"], outs["cutlass"]), (n, k, m) + for backend in BACKENDS if stages else (): + differs = False + for s, (ws_, wsf_) in enumerate(stages): + one = op(aq[:m], a_sw, ws_, wsf_, backend=backend) + cols = slice(512 * s, 512 * (s + 1)) + if backend != "cudnn": + assert torch.equal(outs[backend][:, cols], one), (backend, m, s) + elif not torch.equal(outs[backend][:, cols], one): + differs = True + frac_max = max(frac_max, (outs[backend][:, cols] != one).float().mean().item()) + _check_against(one, ref[:m, cols], bound[:m, cols]) + if differs: + cudnn_differs.append(m) + print(f"[dspark ({n}, {k}) M 1..{DSPARK_ROWS[-1]}] max |err| / bound = {worst:.4f}") + if stages: + print( + f"[dspark wkv x3 stages] cute-dsl, cutlass: fused = separate at every M; cudnn: fused != separate at " + f"{len(cudnn_differs)} M ({cudnn_differs[:1]}..{cudnn_differs[-1:]}), at most {frac_max:.4f} of a " + "stage's elements" + ) + assert cudnn_differs == list(range(193, 641)), cudnn_differs + + +@pytest.mark.parametrize("n,k", DSPARK_NEW_SHAPES) +def test_dspark_graph_capture_every_row_count(n: int, k: int) -> None: + """main_proj and the fused wkv x3 captured at every M 1..160 on the target's backend.""" + _graph_sweep(n, k, DSPARK_GRAPH_ROWS, seed=n + k + 9) + + +def test_dspark_row_addressing_controls() -> None: + """At main_proj, a row-misaddressed activation lands outside the bound on both target backends + (M = 5 on cute-dsl, M = 160 on cudnn): scale rows rolled by one, a row-major scale buffer.""" + n, k = 5120, 15360 + wq, _, wsf = _checkpoint_weight(n, k, seed=31) + w_sw = _swizzle(wsf) + for m in (5, 160): + be = _target_backend(m) + aq, asf = _activation(m, k, seed=32 + m) + a_sw = _swizzle(asf) + ref, bound = _ref_and_bound(aq, asf, wq, wsf) + assert _check_against(op(aq, a_sw, wq, w_sw, backend=be), ref, bound) <= 1.0 + linear = torch.zeros_like(a_sw) + linear[: asf.numel()] = asf.flatten() + wrong = { + "scale rows rolled by one": _swizzle(torch.roll(asf, 1, 0).contiguous()), + "row-major activation scale": linear, + } + for name, sa in wrong.items(): + out = op(aq, sa, wq, w_sw, backend=be) + ratio = ( + ((out.float() - ref).abs() / bound.clamp_min(torch.finfo(torch.float32).tiny)) + .max() + .item() + ) + print(f"[dspark control M={m} {be}] {name}: max |err| / bound = {ratio:.3g}") + assert ratio > 1.0, (be, name) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/moe/test_modeling_v2_gate_forward.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/moe/test_modeling_v2_gate_forward.py new file mode 100644 index 000000000000..c8c0d89aab27 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/moe/test_modeling_v2_gate_forward.py @@ -0,0 +1,1131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gate_forward entry. + +Covers the two DeepSeek-V4.1 routing surfaces: the main MoE (384 experts, +top-6, route_scale 1.5) and the DSpark drafter MoE (128 experts, top-3), plus +the (256, 6) pair the op also dispatches. The entry owns what is certified: +`CELLS` is the list (token counts 0..262144 over three logit / bias +distributions, route_scale values, caller out buffers), `reference` (the +checkpoint's ``Gate.forward`` rebuilt from native torch ops) and `compare` +(bit-exact) are the gate: + + scores = sqrt(softplus(logits_fp32)) + ids = topk(scores + bias) + weights = scores[ids] / (sum + 1e-20) * route_scale + +This file keeps the checks a cell cannot carry: + +* the checkpoint's own torch.sum order, the op's edge cases (the softplus + threshold and a +inf logit, all-zero rows, ties, a dominant bias), the per-token bias as two calls, + determinism, CUDA graph replay, three deliberately wrong references that + must fail the gate, the guard and the op's loud rejections; +* the golden layer dumps: every layer of every golden case routed from the + golden ``ffn_in`` with the checkpoint's own gate weight and bias must + reproduce the golden ``topk_ids``/``topk_weights``; and the drafter's + ``mtp.{0,1,2}`` gates on proxy inputs (no golden dump covers the drafter). + This part skips, naming the directories, when the dumps or the checkpoint + are absent; +* the target's row axis (every prefix 1..8192 at (384, 6), 1..1280 at the + drafter's call), one kernel per surface, the dumped chunk sizes, and CUDA + graph capture at every captured batch. +""" + +import json +import os +import zlib + +import pytest +import torch +import torch.nn.functional as F +from torch.profiler import ProfilerActivity, profile +from utils.llm_data import llm_models_root + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.moe.gate_forward import ( # noqa: E501 + gate_forward as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "gate_forward requires a CUDA device" + +DEV = torch.device("cuda") +ROUTE_SCALE = 1.5 # routed_scaling_factor +MAIN = (384, 6) # n_routed_experts, num_experts_per_tok +DSPARK = (128, 3) # dspark_n_routed_experts, dspark_num_experts_per_tok + +# Reference activations dumped from the checkpoint's own inference code; no default +# location, so the cases that compare against them skip unless this is set. +_GOLDEN_DIR = os.environ.get("MODELING_V2_DSV41_GOLDEN_DIR", "") +# The checkpoint resolves under LLM_MODELS_ROOT, as the whole-model gates' does; +# MODELING_V2_DSV41_CKPT_DIR overrides it. A case that needs it skips when it is absent. +_MODELS_ROOT = llm_models_root() +_CKPT_DIR = os.environ.get("MODELING_V2_DSV41_CKPT_DIR") or ( + str(_MODELS_ROOT / "DeepSeek-V4.1-Flash-NVFP4") if _MODELS_ROOT else "" +) +_GOLDEN_CASES = ( + "capital_fr", + "arith_chat", + "repeat_ngram", + "long_retrieval", + "long_topk", + "boot_natalia", + "boot_capital", + "boot_passkey", +) +_N_LAYERS = 40 +_DSPARK_TARGET_LAYERS = (37, 38, 39) + + +# ─── reference ──────────────────────────────────────────────────────────────── + + +def _warp_sum(w: torch.Tensor) -> torch.Tensor: + """Sum of the topk weights in the kernel's order: a 32-lane xor butterfly + (offsets 16, 8, 4, 2, 1) with lane j holding slot j and the other lanes 0. + + For 6 addends this is ((w0 + w4) + w2) + ((w1 + w5) + w3); for 3 it is + (w0 + w2) + w1. Used here for the per-token-bias reference and the + controls; the certified reference is the entry's. + """ + lanes = [w[:, j] for j in range(w.shape[1])] + lanes += [torch.zeros_like(lanes[0])] * (32 - len(lanes)) + for offset in (16, 8, 4, 2, 1): + lanes = [lanes[j] + lanes[j ^ offset] for j in range(32)] + return lanes[0] + + +def _ref(logits: torch.Tensor, bias: torch.Tensor, topk: int, route_scale: float): + """The entry's reference: (weights fp32, ids int32).""" + return op.reference(logits, bias, topk, route_scale) + + +def _variant( + logits: torch.Tensor, + bias: torch.Tensor, + topk: int, + route_scale: float, + *, + score_fn=None, + weights_from_biased: bool = False, + apply_scale: bool = True, + torch_sum: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """The reference with one deliberate change: the wrong-reference controls, + and the checkpoint's own torch.sum normalization order.""" + x = logits.float() + scores = score_fn(x) if score_fn is not None else F.softplus(x).sqrt() + choice = scores + bias.float() + ids = torch.sort(choice, dim=-1, descending=True, stable=True).indices[:, :topk] + weights = (choice if weights_from_biased else scores).gather(1, ids) + total = weights.sum(-1) if torch_sum else _warp_sum(weights) + weights = weights / (total.unsqueeze(-1) + 1e-20) + if apply_scale: + weights = weights * route_scale + return weights, ids.to(torch.int32) + + +def _check_exact(weights, ids, ref_weights, ref_ids, what: str) -> None: + """The entry's bit-exact gate, with `what` naming the case on failure.""" + try: + op.compare((weights, ids), (ref_weights, ref_ids)) + except AssertionError as e: + raise AssertionError(f"{what}: {e}") from e + + +def _rand_case(num_tokens: int, num_experts: int, sigma: float, bias_center: float, seed: int): + g = torch.Generator(device=DEV).manual_seed(seed) + logits = torch.randn(num_tokens, num_experts, device=DEV, generator=g) * sigma + # The checkpoint's per-layer biases sit in a narrow band around a large + # offset (9.1..57.6 over the 40 layers, spread <= 2.2 inside one layer); + # the drafter's are exactly 0. + bias = bias_center + 0.3 * torch.randn(num_experts, device=DEV, generator=g) + return logits, bias + + +# ─── synthetic surface ──────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + e, k, t = spec["E"], spec["topk"], spec["T"] + logits, bias = _rand_case( + t, e, spec["sigma"], spec["center"], zlib.crc32(cell.why.encode()) % 2**31 + ) + logits_before, bias_before = logits.clone(), bias.clone() + out_w = out_i = None + if spec["out"]: + out_w = torch.full((t, k), -7.0, device=DEV) + out_i = torch.full((t, k), -7, dtype=torch.int32, device=DEV) + weights, ids = op(logits, bias, k, spec["route_scale"], out_w, out_i) + assert weights.shape == (t, k) and weights.dtype == torch.float32 + assert ids.shape == (t, k) and ids.dtype == torch.int32 + assert weights.is_contiguous() and ids.is_contiguous() + if spec["out"]: + assert weights is out_w and ids is out_i, "caller buffers are written and returned" + assert torch.equal(logits, logits_before) and torch.equal(bias, bias_before) + op.compare((weights, ids), op.reference(logits, bias, k, spec["route_scale"], out_w, out_i)) + if t == 0: + torch.cuda.synchronize() + # nothing left pending for the next launch to trip over + assert float(torch.ones(8, device=DEV).softmax(-1).sum()) == pytest.approx(1.0) + + +@pytest.mark.parametrize("out", [False, True], ids=["fresh", "caller_out"]) +@pytest.mark.parametrize("num_experts,topk", [MAIN, DSPARK]) +def test_zero_tokens_launch_nothing(num_experts, topk, out): + """T = 0 is answered without a launch. The kernel's grid is one block per four rows, and a + zero-block grid is rejected by CUDA as an invalid launch, surfacing at the next + synchronize; the entry returns the two empty outputs instead (the caller's own when given).""" + logits = torch.empty(0, num_experts, dtype=torch.float32, device=DEV) + bias = torch.zeros(num_experts, dtype=torch.float32, device=DEV) + out_w = torch.empty(0, topk, device=DEV) if out else None + out_i = torch.empty(0, topk, dtype=torch.int32, device=DEV) if out else None + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + weights, ids = op(logits, bias, topk, ROUTE_SCALE, out_w, out_i) + torch.cuda.synchronize() + launched = [ + e.name + for e in prof.events() + if e.device_type == torch.autograd.DeviceType.CUDA and e.self_device_time_total > 0 + ] + assert not launched, launched + assert weights.shape == (0, topk) and weights.dtype == torch.float32 + assert ids.shape == (0, topk) and ids.dtype == torch.int32 + if out: + assert weights is out_w and ids is out_i + # The context is still healthy: the next launch and synchronize succeed. + assert float(torch.ones(8, device=DEV).softmax(-1).sum()) == pytest.approx(1.0) + + +@pytest.mark.parametrize("num_experts,topk", [MAIN, DSPARK]) +def test_checkpoint_sum_order_matches(num_experts, topk): + """The checkpoint normalizes with torch.sum; on this build that is the + kernel's order too, so the op is bit-exact to the checkpoint's own code.""" + logits, bias = _rand_case(65536, num_experts, 3.0, 20.0, 11) + weights, ids = op(logits, bias, topk, ROUTE_SCALE) + _check_exact( + weights, ids, *_variant(logits, bias, topk, ROUTE_SCALE, torch_sum=True), "torch.sum order" + ) + + +@pytest.mark.parametrize("num_experts,topk", [MAIN, DSPARK]) +def test_wrong_references_fail(num_experts, topk): + """Three deliberately wrong references, each far outside the (exact) gate.""" + logits, bias = _rand_case(8192, num_experts, 3.0, 10.0, 13) + weights, ids = op(logits, bias, topk, ROUTE_SCALE) + controls = { + "weights from biased scores": _variant( + logits, bias, topk, ROUTE_SCALE, weights_from_biased=True + ), + "sigmoid instead of sqrt(softplus)": _variant( + logits, bias, topk, ROUTE_SCALE, score_fn=torch.sigmoid + ), + "missing x route_scale": _variant(logits, bias, topk, ROUTE_SCALE, apply_scale=False), + } + for name, (cw, ci) in controls.items(): + wrong_w = int((cw != weights).sum()) + wrong_rows = int((ci != ids).any(-1).sum()) + max_abs = float((cw - weights).abs().max()) + print( + f"[control E={num_experts} K={topk}] {name}: {wrong_w}/{weights.numel()} weights differ, " + f"max abs {max_abs:.4f}, {wrong_rows}/{ids.shape[0]} rows route differently" + ) + assert wrong_w > weights.numel() // 2 and max_abs > 1e-2, name + + +@pytest.mark.parametrize("num_experts,topk", [MAIN, DSPARK]) +def test_bias_selects_but_does_not_weight(num_experts, topk): + """A dominant bias forces an expert into the selection; its weight is still + its bare score's share.""" + logits, _ = _rand_case(64, num_experts, 1.0, 0.0, 17) + bias = torch.zeros(num_experts, device=DEV) + bias[num_experts - 1] = 100.0 + weights, ids = op(logits, bias, topk, ROUTE_SCALE) + assert torch.all(ids[:, 0] == num_experts - 1) + _check_exact(weights, ids, *_ref(logits, bias, topk, ROUTE_SCALE), "dominant bias") + scores = F.softplus(logits).sqrt() + torch.testing.assert_close( + weights[:, 0], + scores[:, -1] / scores.gather(1, ids.long()).sum(-1) * ROUTE_SCALE, + rtol=4 * torch.finfo(torch.float32).eps, # summation order of 6 fp32 addends + atol=0.0, + ) + + +@pytest.mark.parametrize("num_experts,topk", [MAIN, DSPARK]) +def test_softplus_threshold_matches_torch(num_experts, topk): + """The kernel applies torch's softplus threshold (x > 20 gives x, else + log1pf(expf(x))), so every finite logit is bit-exact to the reference: no + expf overflow at x ~ 88.72 any more. Only a +inf logit still makes its + score inf, which gives that expert inf / inf = NaN and the row's others 0. + + Each row puts the probed logit on expert 0, topk-1 partners at logit 0 and + the rest at -inf, so the kernel's score for x is visible through w0. + """ + xs = torch.linspace(-120.0, 120.0, 240001, device=DEV) + logits = torch.full((xs.numel(), num_experts), float("-inf"), device=DEV) + logits[:, 0] = xs + logits[:, 1:topk] = 0.0 + bias = torch.zeros(num_experts, device=DEV) + weights, ids = op(logits, bias, topk, 1.0) + _check_exact(weights, ids, *_ref(logits, bias, topk, 1.0), "x in [-120, 120]") + assert torch.isfinite(weights).all() + + inf_row = torch.full((1, num_experts), float("-inf"), device=DEV) + inf_row[0, 0] = float("inf") + inf_row[0, 1:topk] = 0.0 + weights, ids = op(inf_row, bias, topk, 1.0) + assert ids[0].tolist() == list(range(topk)), ids + assert torch.isnan(weights[0, 0]) and torch.all(weights[0, 1:] == 0.0), weights + print( + f"[softplus E={num_experts}] bit-exact over {xs.numel()} logits in [-120, 120]; " + "a +inf logit gives NaN for its expert and 0 for the others" + ) + + +@pytest.mark.parametrize("num_experts,topk", [MAIN, DSPARK]) +def test_all_zero_scores_give_zero_weights(num_experts, topk): + """The kernel divides by sum + 1e-20, as the checkpoint does: a row whose + selected scores are all exactly 0 (every logit <= ~-104, or -inf) gets + 0 / 1e-20 * route_scale = 0 weights, bit-exact to the reference. A NaN + logit is selected first and NaNs the row.""" + bias = torch.zeros(num_experts, device=DEV) + for val in [-200.0, float("-inf")]: + logits = torch.full((4, num_experts), val, device=DEV) + weights, ids = op(logits, bias, topk, ROUTE_SCALE) + _check_exact(weights, ids, *_ref(logits, bias, topk, ROUTE_SCALE), f"all logits {val}") + assert torch.all(weights == 0.0), val + logits, _ = _rand_case(4, num_experts, 1.0, 0.0, 19) + logits[:, 5] = float("nan") + weights, ids = op(logits, bias, topk, ROUTE_SCALE) + assert torch.all(ids[:, 0] == 5) and torch.isnan(weights).all() + + +@pytest.mark.parametrize("num_experts,topk", [MAIN, DSPARK]) +def test_ties_break_to_lower_expert(num_experts, topk): + bias = torch.zeros(num_experts, device=DEV) + """Equal selection scores come out in ascending expert order. (torch.topk + emitted [1, 0, 4, 5, 2, 3] / [2, 1, 0] for all-equal rows on this machine, + which is why _ref uses a stable sort.)""" + zeros = torch.zeros(3, num_experts, device=DEV) + weights, ids = op(zeros, bias, topk, ROUTE_SCALE) + assert ids.tolist() == [list(range(topk))] * 3 + _check_exact(weights, ids, *_ref(zeros, bias, topk, ROUTE_SCALE), "all-equal row") + logits = torch.full((1, num_experts), -5.0, device=DEV) + for e in [3, 70, 100, 127]: + logits[0, e] = 2.0 + weights, ids = op(logits, bias, topk, ROUTE_SCALE) + assert ids[0, : min(topk, 4)].tolist() == [3, 70, 100, 127][: min(topk, 4)] + _check_exact(weights, ids, *_ref(logits, bias, topk, ROUTE_SCALE), "4-way tie") + + +@pytest.mark.parametrize("num_experts,topk", [MAIN, DSPARK]) +def test_per_token_bias_as_two_calls_and_select(num_experts, topk): + """The op takes one [E] bias per call. A per-token choice between two + biases (bias for text, bias_vl for image-span tokens) is two calls plus a + per-token select, bit-exact to the per-token-bias reference.""" + logits, bias = _rand_case(2048, num_experts, 3.0, 10.0, 23) + bias_vl = bias + 4.0 * torch.randn( + num_experts, device=DEV, generator=torch.Generator(device=DEV).manual_seed(29) + ) + image = ( + torch.rand(2048, device=DEV, generator=torch.Generator(device=DEV).manual_seed(31)) < 0.5 + ) + w_t, i_t = op(logits, bias, topk, ROUTE_SCALE) + w_v, i_v = op(logits, bias_vl, topk, ROUTE_SCALE) + weights = torch.where(image[:, None], w_v, w_t) + ids = torch.where(image[:, None], i_v, i_t) + x = F.softplus(logits).sqrt() + choice = x + torch.where(image[:, None], bias_vl, bias) + ref_i = torch.sort(choice, dim=-1, descending=True, stable=True).indices[:, :topk] + ref_w = x.gather(1, ref_i) + ref_w = ref_w / (_warp_sum(ref_w).unsqueeze(-1) + 1e-20) * ROUTE_SCALE + _check_exact(weights, ids, ref_w, ref_i.to(torch.int32), "two calls + select") + assert int((i_t != i_v).any(-1).sum()) > 0 # the two biases do route differently + + +@pytest.mark.parametrize("num_experts,topk", [MAIN, DSPARK]) +def test_out_buffers_inputs_untouched_and_determinism(num_experts, topk): + logits, bias = _rand_case(8192, num_experts, 3.0, 10.0, 37) + logits_before, bias_before = logits.clone(), bias.clone() + out_w = torch.full((8192, topk), -7.0, device=DEV) + out_i = torch.full((8192, topk), -7, dtype=torch.int32, device=DEV) + weights, ids = op(logits, bias, topk, ROUTE_SCALE, out_w, out_i) + assert weights is out_w and ids is out_i + assert torch.equal(logits, logits_before) and torch.equal(bias, bias_before) + ref_w, ref_i = _ref(logits, bias, topk, ROUTE_SCALE) + _check_exact(weights, ids, ref_w, ref_i, "caller buffers") + again_w, again_i = op(logits, bias, topk, ROUTE_SCALE) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + alt_w, alt_i = op(logits, bias, topk, ROUTE_SCALE) + torch.cuda.current_stream().wait_stream(stream) + assert torch.equal(again_w, weights) and torch.equal(again_i, ids) + assert torch.equal(alt_w, weights) and torch.equal(alt_i, ids) + + +@pytest.mark.parametrize("num_experts,topk", [MAIN, DSPARK]) +def test_cuda_graph_replay(num_experts, topk): + num_tokens = 40 + static_logits, bias = _rand_case(num_tokens, num_experts, 3.0, 10.0, 41) + out_w = torch.empty(num_tokens, topk, device=DEV) + out_i = torch.empty(num_tokens, topk, dtype=torch.int32, device=DEV) + op(static_logits, bias, topk, ROUTE_SCALE, out_w, out_i) # warm-up outside capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(static_logits, bias, topk, ROUTE_SCALE, out_w, out_i) + for seed in (43, 47, 53): + fresh, _ = _rand_case(num_tokens, num_experts, 3.0, 10.0, seed) + static_logits.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + _check_exact( + out_w, out_i, *_ref(fresh, bias, topk, ROUTE_SCALE), f"graph replay seed {seed}" + ) + + +# ─── guards and loud rejections ─────────────────────────────────────────────── + + +@pytest.mark.parametrize("num_experts,topk", [MAIN, DSPARK]) +def test_guard_catches_silent_misroutes(num_experts, topk): + """Each guarded layout is accepted by the bare op and misroutes (probed + here on in-bounds reads through `raw_call`, outside validation); `is_valid` + rejects it before the call.""" + + def bare(x, b): + return op.raw_call(x, b, topk, ROUTE_SCALE) + + def guarded(*args): + with validating(op): + return op(*args) + + wide, bias = _rand_case(64, num_experts + 16, 3.0, 0.0, 59) + bias = bias[:num_experts].contiguous() + sliced = wide[:, :num_experts] + _, ids_bare = bare(sliced, bias) + _, ids_good = bare(sliced.contiguous(), bias) + assert not torch.equal(ids_bare, ids_good), "bare op no longer misroutes a column slice" + with pytest.raises(AssertionError, match="contiguous"): + guarded(sliced, bias, topk, ROUTE_SCALE) + + logits = sliced.contiguous() + strided_bias = torch.randn(2 * num_experts, device=DEV)[::2] + _, ids_bare = bare(logits, strided_bias) + _, ids_good = bare(logits, strided_bias.contiguous()) + assert not torch.equal(ids_bare, ids_good), "bare op no longer misroutes a strided bias" + with pytest.raises(AssertionError, match="contiguous"): + guarded(logits, strided_bias, topk, ROUTE_SCALE) + + cube = torch.randn(1, num_experts, num_experts, device=DEV) + _, ids_bare = bare(cube, bias) + assert ids_bare.shape[0] == 1 # the op routed one row of a 3-D input + with pytest.raises(AssertionError, match="2-D"): + guarded(cube, bias, topk, ROUTE_SCALE) + + with pytest.raises(AssertionError, match="bias must be"): + guarded(logits, bias[: num_experts - 32], topk, ROUTE_SCALE) + with pytest.raises(AssertionError, match="bias must be"): + guarded(logits, torch.empty(0, device=DEV), topk, ROUTE_SCALE) + narrow_w = torch.empty(64, topk - 1, device=DEV) + with pytest.raises(AssertionError, match="out_weights must be"): + guarded(logits, bias, topk, ROUTE_SCALE, narrow_w, None) + wide_i = torch.empty(64, topk + 2, dtype=torch.int32, device=DEV)[:, :topk] + with pytest.raises(AssertionError, match="out_indices must be"): + guarded(logits, bias, topk, ROUTE_SCALE, None, wide_i) + + +def test_op_rejections_are_loud(): + logits, bias = _rand_case(4, 384, 1.0, 0.0, 61) + for num_experts, topk in [(128, 6), (384, 3), (256, 3), (64, 3), (128, 2), (512, 6)]: + x, b = _rand_case(4, num_experts, 1.0, 0.0, 67) + with pytest.raises( + RuntimeError, + match=r"\(n_experts, topk\) must be one of \(256, 6\), \(384, 6\), \(128, 3\)", + ): + op(x, b, topk, ROUTE_SCALE) + with pytest.raises(RuntimeError, match="scores_in must be float32"): + op(logits.bfloat16(), bias, 6, ROUTE_SCALE) + with pytest.raises(RuntimeError, match="expected scalar type Float but found BFloat16"): + op(logits, bias.bfloat16(), 6, ROUTE_SCALE) + with pytest.raises(RuntimeError, match="out_weights must be float32"): + op(logits, bias, 6, ROUTE_SCALE, torch.empty(4, 6, dtype=torch.bfloat16, device=DEV), None) + with pytest.raises(RuntimeError, match="out_indices must be int32"): + op(logits, bias, 6, ROUTE_SCALE, None, torch.empty(4, 6, dtype=torch.int64, device=DEV)) + with pytest.raises(RuntimeError, match="All tensors must be CUDA tensors"): + op( + logits.cpu(), + bias, + 6, + ROUTE_SCALE, + torch.empty(4, 6, device=DEV), + torch.empty(4, 6, dtype=torch.int32, device=DEV), + ) + + +# ─── golden layer dumps and the checkpoint's gates ──────────────────────────── + + +def _golden_available() -> bool: + return all( + os.path.isfile(os.path.join(_GOLDEN_DIR, c, "layers", f"L{_N_LAYERS - 1:02d}.pt")) + for c in _GOLDEN_CASES + ) and os.path.isfile(os.path.join(_CKPT_DIR, "model.safetensors.index.json")) + + +def _skip_without_golden(): + if not _golden_available(): + pytest.skip(f"golden dumps ({_GOLDEN_DIR}) or checkpoint ({_CKPT_DIR}) not present") + + +def _ckpt_tensor(name: str) -> torch.Tensor: + from safetensors import safe_open + + with open(os.path.join(_CKPT_DIR, "model.safetensors.index.json")) as f: + shard = json.load(f)["weight_map"][name] + with safe_open(os.path.join(_CKPT_DIR, shard), "pt", device="cpu") as f: + return f.get_tensor(name) + + +def _golden_layer(case: str, layer: int) -> dict: + return torch.load( + os.path.join(_GOLDEN_DIR, case, "layers", f"L{layer:02d}.pt"), + map_location="cpu", + mmap=True, + weights_only=True, + ) + + +def _router_logits(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """linear(x.float(), W.float()) as the checkpoint computes it (IEEE fp32, + TF32 off, like the golden reference).""" + prev = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = False + try: + return F.linear(x.to(DEV).float(), weight.to(DEV).float()) + finally: + torch.backends.cuda.matmul.allow_tf32 = prev + + +_GATES: dict = {} + + +def _gate(prefix: str): + if prefix not in _GATES: + _GATES[prefix] = tuple( + _ckpt_tensor(f"{prefix}.ffn.gate.{n}").to(DEV) for n in ("weight", "bias", "bias_vl") + ) + return _GATES[prefix] + + +@pytest.mark.parametrize("case", _GOLDEN_CASES) +def test_golden_main_routing(case): + """Every layer of the case, routed from the golden ffn_in with the + checkpoint's gate weight and bias: bit-exact to _ref, and equal to the + golden topk_ids/topk_weights up to the order of experts whose fp32 + selection scores tie exactly; the three wrong references are not.""" + _skip_without_golden() + lo, hi, above20 = float("inf"), float("-inf"), 0 + bad = { + "weights from biased scores": 0, + "sigmoid instead of sqrt(softplus)": 0, + "missing x route_scale": 0, + } + rows = reordered = 0 + for layer in range(_N_LAYERS): + weight, bias, _ = _gate(f"layers.{layer}") + assert weight.shape == (384, 5120) and bias.dtype == torch.float32 + d = _golden_layer(case, layer) + logits = _router_logits(d["ffn_in"], weight) + lo, hi = min(lo, float(logits.min())), max(hi, float(logits.max())) + above20 += int((logits > 20).sum()) + weights, ids = op(logits, bias, 6, ROUTE_SCALE) + _check_exact( + weights, ids, *_ref(logits, bias, 6, ROUTE_SCALE), f"{case} L{layer:02d} vs _ref" + ) + + g_ids = d["topk_ids"].to(DEV).to(torch.int32) + g_w = d["topk_weights"].to(DEV) + rows += ids.shape[0] + # The golden used torch.topk, which orders exactly tied selection + # scores differently from the kernel (ascending expert id). The large + # per-layer bias (up to 57.6, fp32 ulp 3.8e-6) makes such ties real. + # Every differing row must hold the same expert set, and the selection + # scores slot by slot must be the same values. + same_order = (ids == g_ids).all(-1) + reordered += int((~same_order).sum()) + assert torch.equal(ids.sort(-1).values, g_ids.sort(-1).values), ( + f"{case} L{layer:02d}: expert sets differ" + ) + choice = F.softplus(logits).sqrt() + bias + assert torch.equal(choice.gather(1, ids.long()), choice.gather(1, g_ids.long())), ( + f"{case} L{layer:02d}: a reordered row is not an exact tie" + ) + assert torch.equal(weights[same_order], g_w[same_order]), ( + f"{case} L{layer:02d}: weights differ" + ) + # On a reordered row the kernel adds the same six fp32 scores in another + # slot order: the sum can land one ulp away (relative <= eps), and the + # division carries that and adds its own rounding, so a normalized + # weight may move by up to two ulps (measured max relative 1.94e-7). + k_sorted = weights.gather(1, ids.long().argsort(-1)) + g_sorted = g_w.gather(1, g_ids.long().argsort(-1)) + torch.testing.assert_close( + k_sorted, g_sorted, rtol=2 * torch.finfo(torch.float32).eps, atol=0.0 + ) + + for name, kw in [ + ("weights from biased scores", dict(weights_from_biased=True)), + ("sigmoid instead of sqrt(softplus)", dict(score_fn=torch.sigmoid)), + ("missing x route_scale", dict(apply_scale=False)), + ]: + cw, ci = _variant(logits, bias, 6, ROUTE_SCALE, **kw) + bad[name] += int(((cw != g_w).any(-1) | (ci != g_ids).any(-1)).sum()) + print( + f"[golden {case}] 40 layers x {rows // _N_LAYERS} tokens: bit-exact to _ref; vs golden " + f"{rows - reordered}/{rows} rows identical, {reordered} reordered exact ties (same set); router logits " + f"[{lo:.3f}, {hi:.3f}], {above20} above 20; wrong-reference rows: {bad}" + ) + for name, n in bad.items(): + assert n == rows, f"control '{name}' matched the golden on {rows - n} of {rows} rows" + + +def test_golden_dspark_gates_on_proxy_inputs(): + """The drafter's mtp.{0,1,2} gates (128 experts, top-3). No golden dump + covers the drafter, so the input is a proxy: the golden ffn_in of its + target layers 37/38/39 with that layer's ffn_norm weight divided out and + the drafter's ffn_norm weight multiplied in. Checked bit-exact against + _ref, and the drafter's biases are recorded.""" + _skip_without_golden() + for s in range(3): + weight, bias, bias_vl = _gate(f"mtp.{s}") + assert weight.shape == (128, 5120) + assert torch.all(bias == 0.0), f"mtp.{s}.ffn.gate.bias is not all zero" + norm = _ckpt_tensor(f"mtp.{s}.ffn_norm.weight").to(DEV).float() + lo, hi, rows = float("inf"), float("-inf"), 0 + for layer in _DSPARK_TARGET_LAYERS: + layer_norm = _ckpt_tensor(f"layers.{layer}.ffn_norm.weight").to(DEV).float() + for case in _GOLDEN_CASES: + x = _golden_layer(case, layer)["ffn_in"].to(DEV).float() / layer_norm * norm + logits = _router_logits(x, weight) + lo, hi = min(lo, float(logits.min())), max(hi, float(logits.max())) + for b in (bias, bias_vl): + weights, ids = op(logits, b, 3, ROUTE_SCALE) + _check_exact( + weights, ids, *_ref(logits, b, 3, ROUTE_SCALE), f"mtp.{s} {case} L{layer}" + ) + rows += logits.shape[0] + print( + f"[dspark mtp.{s}] {rows} proxy rows bit-exact (bias and bias_vl); proxy logits [{lo:.3f}, {hi:.3f}]; " + f"bias all zero; bias_vl [{float(bias_vl.min()):.4f}, {float(bias_vl.max()):.4f}]" + ) + + +def test_golden_bias_vl_two_calls(): + """Checkpoint bias/bias_vl of layers 0 and 39 on golden ffn_in with a + synthetic image span: two calls + select equals the per-token-bias + reference (the checkpoint's torch.where(image_mask, bias_vl, bias)).""" + _skip_without_golden() + for layer in (0, 19, 39): + weight, bias, bias_vl = _gate(f"layers.{layer}") + logits = _router_logits(_golden_layer("long_topk", layer)["ffn_in"], weight) + num_tokens = logits.shape[0] + image = torch.zeros(num_tokens, dtype=torch.bool, device=DEV) + image[100:400] = True + w_t, i_t = op(logits, bias, 6, ROUTE_SCALE) + w_v, i_v = op(logits, bias_vl, 6, ROUTE_SCALE) + weights = torch.where(image[:, None], w_v, w_t) + ids = torch.where(image[:, None], i_v, i_t) + x = F.softplus(logits).sqrt() + choice = x + torch.where(image[:, None], bias_vl, bias) + ref_i = torch.sort(choice, dim=-1, descending=True, stable=True).indices[:, :6] + ref_w = x.gather(1, ref_i) + ref_w = ref_w / (_warp_sum(ref_w).unsqueeze(-1) + 1e-20) * ROUTE_SCALE + _check_exact(weights, ids, ref_w, ref_i.to(torch.int32), f"bias_vl L{layer}") + print( + f"[bias_vl L{layer}] rows routed differently by bias_vl: {int((i_t != i_v).any(-1).sum())}/{num_tokens}" + ) + + +# ─── the core target's row axis 1..8192 and its CUDA-graph batches ──────────── +# The DeepSeek-V4.1 target calls gate_forward(logits fp32 [n, 384], bias fp32 +# [384], 6, 1.5) once per MoE chunk (`_route` from `_moe`), n <= 8192 (the +# chunk cap of `_moe_chunks`, at every max_num_tokens; the attention-DP members +# dep2 / dep4 / dep8 route the all-gathered chunk). Every member (tep2/4/8, +# dep2/4/8, tp2/4, with or without two pipeline stages) routes the full 384 +# experts on every rank, so the call's shape is the same on all of them. The `bias_vl` call is the same call with the +# other bias, made only on eager steps. The op's host launcher picks its kernel by +# (E, topk, is_hash) only; T only sizes the grid (one warp per row, 4 rows per +# 128-thread block). So the row axis is checked as: a run at N_max against the +# reference, every prefix 1..N_max bitwise equal to the N_max run, and a +# one-kernel probe, plus capture at every batch the engine captures. + +_CORE_NMAX = 8192 +# Every chunk size in the core runtime-value dump that no earlier cell named. +_CORE_DUMPED = (12, 16, 20, 24, 28, 32, 48, 56, 192, 1018, 1935, 3552, 3560) +# Every chunk size the batch-256 / long-context engine dumps show (tep4 graph and +# eager, dep4 and tp4 graph, L 1M at batch 8, L 128K at batch 32), captured and eager. +_CORE_DUMPED_B256 = ( + 1, + 2, + 4, + 5, + 6, + 7, + 8, + 12, + 16, + 18, + 24, + 28, + 29, + 32, + 34, + 48, + 60, + 64, + 96, + 108, + 128, + 134, + 144, + 168, + 192, + 228, + 240, + 288, + 330, + 336, + 384, + 432, + 462, + 480, + 528, + 546, + 576, + 610, + 624, + 672, + 720, + 768, + 780, + 804, + 844, + 852, + 960, + 1152, + 1172, + 1178, + 1344, + 1536, + 1728, + 1920, + 2112, + 2252, + 2255, + 2259, + 2304, + 2388, + 2496, + 2688, + 2778, + 2784, + 2880, + 3072, + 3870, + 4208, + 4245, + 4388, + 4394, + 4591, + 4608, + 4868, + 5328, + 6144, + 8116, + 8124, + 8128, + 8130, + 8137, + 8141, + 8148, + 8168, + 8192, +) +# CUDA-graph batches per rank: every G 1..32 plus TRT-LLM's padded list at +# max_batch_size 256 (`_generate_cuda_graph_batch_sizes`); the padding-off list +# at 256 (1..32, 64, 128, 256) is a subset. +_GSET = tuple(range(1, 33)) + (40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256) +# Capture union, per rank, of the rows the router sees: G (decode) and 6G +# (DSpark target-verify, 1 + 5 drafts) on a tep / tp rank, and after the +# attention-DP all-gather 2G and 12G (dep2), 4G and 24G (dep4), 8G and 48G +# (dep8). dep8's 48G at G 192 and 256 (9216, 12288) exceeds the 8192-row chunk +# cap and is two calls in one graph (`_CORE_CAPTURE_SPLIT`); every other size is +# one call. +_CORE_BATCHES = _GSET +_CORE_MULTIPLES = (1, 6, 2, 12, 4, 24, 8, 48) +_CORE_CAPTURE = tuple( + sorted({m * b for m in _CORE_MULTIPLES for b in _CORE_BATCHES if m * b <= _CORE_NMAX}) +) +_CORE_CAPTURE_SPLIT = tuple( + sorted({m * b for m in _CORE_MULTIPLES for b in _CORE_BATCHES if m * b > _CORE_NMAX}) +) +_POISON_W = 0x7FC0DEAD # a quiet-NaN bit pattern the op never produces +_POISON_I = -7 + + +def _bits(w: torch.Tensor) -> torch.Tensor: + return w.view(torch.int32) + + +@pytest.mark.parametrize("sigma,center,seed", [(3.0, 10.0, 101), (12.0, 57.0, 103)]) +def test_core_rows_1_to_8192_prefix_bitwise(sigma, center, seed): + """The row axis on the core surface (384, 6), route_scale 1.5. + + (a) the 8192-row call is bit-exact to _ref; (b) for EVERY r in 1..8192 the + call on the first r rows, written into the first r rows of a poisoned + 8192-row buffer pair, equals the first r rows of the 8192-row result bit for + bit, and rows r.. keep the poison (no write past T); (c) a one-row-shift + control is rejected at every r. Counts accumulate on device; one sync. + """ + E, K = MAIN + logits, bias = _rand_case(_CORE_NMAX, E, sigma, center, seed) + full_w, full_i = op(logits, bias, K, ROUTE_SCALE) + _check_exact(full_w, full_i, *_ref(logits, bias, K, ROUTE_SCALE), f"N_max sigma={sigma}") + full_wb = _bits(full_w) + shift_wb = full_wb.roll(-1, 0) # row r holds row r+1's weights + out_w = torch.empty(_CORE_NMAX, K, device=DEV) + out_i = torch.empty(_CORE_NMAX, K, dtype=torch.int32, device=DEV) + out_wb = _bits(out_w) + bad = torch.zeros(_CORE_NMAX + 1, dtype=torch.int64, device=DEV) + shifted_rows = torch.zeros(_CORE_NMAX + 1, dtype=torch.int64, device=DEV) + for r in range(1, _CORE_NMAX + 1): + out_wb.fill_(_POISON_W) + out_i.fill_(_POISON_I) + w, i = op(logits[:r], bias, K, ROUTE_SCALE, out_w[:r], out_i[:r]) + assert w.data_ptr() == out_w.data_ptr() and i.data_ptr() == out_i.data_ptr() + bad[r] = ( + (out_wb[:r] != full_wb[:r]).sum() + + (out_i[:r] != full_i[:r]).sum() + + (out_wb[r:] != _POISON_W).sum() + + (out_i[r:] != _POISON_I).sum() + ) + shifted_rows[r] = (out_wb[:r] != shift_wb[:r]).any(-1).sum() + torch.cuda.synchronize() + failing = torch.nonzero(bad).flatten().tolist() + assert not failing, ( + f"prefix/poison mismatches at r = {failing[:20]} (counts {bad[failing[:20]].tolist()})" + ) + r_all = torch.arange(_CORE_NMAX + 1, device=DEV) + # A reference shifted by one row differs on every row of every prefix. + assert torch.equal(shifted_rows[1:], r_all[1:]), ( + "one-row-shift control was not rejected on every row" + ) + print( + f"[rows sigma={sigma} bias~{center}] r=1..{_CORE_NMAX}: every prefix bitwise equal to the " + f"N_max run, tail poison intact; shift control rejected on {int(shifted_rows.sum())} rows" + ) + + +def test_core_one_kernel_for_every_row_count(): + """The op is compiled ahead of time (no per-shape JIT); this + pins that one and the same kernel serves every row count, including the + dumped ones, with no second launch.""" + E, K = MAIN + names = set() + for t in (1, 2, 3, 4, 5, 7, 8, 1018, 1536, 1935, 3560, 6144, _CORE_NMAX): + logits, bias = _rand_case(t, E, 3.0, 10.0, 200 + t) + op(logits, bias, K, ROUTE_SCALE) + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + op(logits, bias, K, ROUTE_SCALE) + torch.cuda.synchronize() + launched = [ + e.name + for e in prof.events() + if e.device_type == torch.autograd.DeviceType.CUDA and e.self_device_time_total > 0 + ] + assert len(launched) == 1, (t, launched) + names.add(launched[0]) + assert len(names) == 1, names + (name,) = names + assert "gate_forward_kernel" in name and "384" in name, name + print(f"[kernel] one kernel for every T: {name}") + + +def test_core_dumped_chunk_sizes_fresh_inputs(): + """Every dumped chunk size, on fresh inputs, bit-exact to _ref (the prefix + test already covers them; this also checks them standalone, allocating + outputs as the target does).""" + E, K = MAIN + for t in sorted(set(_CORE_DUMPED) | set(_CORE_DUMPED_B256)): + for sigma, center, seed in [(3.0, 10.0, 7), (12.0, 57.0, 9)]: + logits, bias = _rand_case(t, E, sigma, center, seed * 1000 + t) + weights, ids = op(logits, bias, K, ROUTE_SCALE) + _check_exact( + weights, ids, *_ref(logits, bias, K, ROUTE_SCALE), f"dumped T={t} sigma={sigma}" + ) + + +def test_core_cuda_graph_capture_union(): + """Capture at the core call: one graph per T in the capture union (167 + sizes, 1..6144: G and 6G on a tep / tp rank and the gathered 2G / 12G of dep2, + 4G / 24G of dep4 and 8G / 48G of dep8 up to the 8192-row chunk cap, G over the + 46 batches of `_GSET`), each warmed eagerly first, captured into one shared pool + largest first as the engine does, with the target's allocate-in-capture + call (no out buffers). Three rounds of interleaved replays, each on new + logits AND a new bias written in place: every replay is bitwise equal to an + eager call on the same inputs and to _ref, and the previous round's + reference is rejected (the replay read the new data).""" + E, K = MAIN + pool = torch.cuda.graph_pool_handle() + graphs = {} + for t in sorted(_CORE_CAPTURE, reverse=True): + logits, bias = _rand_case(t, E, 3.0, 10.0, 300 + t) + op(logits, bias, K, ROUTE_SCALE) # eager warm-up at this T + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + out = op(logits, bias, K, ROUTE_SCALE) + graphs[t] = (graph, logits, bias, out) + prev = {} + for rnd, (sigma, center) in enumerate([(1.0, 0.0), (4.0, 12.0), (12.0, 57.0)]): + for t in _CORE_CAPTURE: + graph, logits, bias, (gw, gi) = graphs[t] + fresh_l, fresh_b = _rand_case(t, E, sigma, center, 1000 * rnd + t) + logits.copy_(fresh_l) + bias.copy_(fresh_b) + graph.replay() + torch.cuda.synchronize() + eager_w, eager_i = op(fresh_l, fresh_b, K, ROUTE_SCALE) + assert torch.equal(_bits(gw), _bits(eager_w)) and torch.equal(gi, eager_i), ( + f"replay != eager at T={t} round {rnd}" + ) + _check_exact( + gw, gi, *_ref(fresh_l, fresh_b, K, ROUTE_SCALE), f"graph T={t} round {rnd}" + ) + if t in prev: + stale_w, _ = prev[t] + assert (_bits(stale_w) != _bits(gw)).any(-1).all(), ( + f"stale-input control not rejected at T={t}" + ) + prev[t] = _ref(fresh_l, fresh_b, K, ROUTE_SCALE) + assert len(graphs) == 167 and max(graphs) == 6144, (len(graphs), max(graphs)) + print( + f"[graphs] {len(graphs)} captured T {list(_CORE_CAPTURE)}: 3 rounds, replay == eager == _ref bitwise" + ) + + +def test_core_cuda_graph_capture_dep8_two_chunks(): + """Capture where one captured step routes two chunks: dep8's gathered 48G + at G 192 and 256 (9216 and 12288 rows), cut as the target cuts them + (`_moe_chunks`: 8192 + 1024 and 8192 + 4096). One graph per total, holding one + call per chunk on its rows of one static logits buffer, each call allocating + its outputs in the graph pool as the target's does, warmed eagerly chunk by + chunk first, captured into one shared pool. A replay launches one + `gate_forward_kernel<384, 6, false>` per chunk. Three rounds of replays, each + on new logits AND a new bias written in place: the chunks' outputs, + concatenated, are bitwise equal to eager calls on the same chunks and to _ref + over the whole batch, and the previous round's reference is rejected on every + row.""" + E, K = MAIN + assert _CORE_CAPTURE_SPLIT == (9216, 12288), _CORE_CAPTURE_SPLIT + pool = torch.cuda.graph_pool_handle() + graphs = {} + for t in sorted(_CORE_CAPTURE_SPLIT, reverse=True): + chunks = [(a, min(t, a + _CORE_NMAX)) for a in range(0, t, _CORE_NMAX)] + logits, bias = _rand_case(t, E, 3.0, 10.0, 400 + t) + for a, b in chunks: + op(logits[a:b], bias, K, ROUTE_SCALE) # eager warm-up at each chunk + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + outs = [op(logits[a:b], bias, K, ROUTE_SCALE) for a, b in chunks] + graphs[t] = (graph, logits, bias, outs, chunks) + for t, (graph, _, _, _, chunks) in graphs.items(): + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + graph.replay() + torch.cuda.synchronize() + launched = [ + e.name + for e in prof.events() + if e.device_type == torch.autograd.DeviceType.CUDA and e.self_device_time_total > 0 + ] + assert len(launched) == len(chunks) == 2, (t, launched) + assert all("gate_forward_kernel" in n and "384" in n for n in launched), launched + prev = {} + for rnd, (sigma, center) in enumerate([(1.0, 0.0), (4.0, 12.0), (12.0, 57.0)]): + for t in _CORE_CAPTURE_SPLIT: + graph, logits, bias, outs, chunks = graphs[t] + fresh_l, fresh_b = _rand_case(t, E, sigma, center, 2000 * (rnd + 1) + t) + logits.copy_(fresh_l) + bias.copy_(fresh_b) + graph.replay() + torch.cuda.synchronize() + gw = torch.cat([w for w, _ in outs]) + gi = torch.cat([i for _, i in outs]) + eager = [op(fresh_l[a:b], fresh_b, K, ROUTE_SCALE) for a, b in chunks] + assert torch.equal(_bits(gw), _bits(torch.cat([w for w, _ in eager]))), ( + f"replay != eager at T={t}" + ) + assert torch.equal(gi, torch.cat([i for _, i in eager])), f"replay != eager at T={t}" + _check_exact( + gw, + gi, + *_ref(fresh_l, fresh_b, K, ROUTE_SCALE), + f"two-chunk graph T={t} round {rnd}", + ) + if t in prev: + stale_w, _ = prev[t] + assert (_bits(stale_w) != _bits(gw)).any(-1).all(), ( + f"stale-input control not rejected at T={t}" + ) + prev[t] = _ref(fresh_l, fresh_b, K, ROUTE_SCALE) + print( + f"[two-chunk graphs] T {list(_CORE_CAPTURE_SPLIT)} as " + f"{ {t: [b - a for a, b in g[4]] for t, g in graphs.items()} }: two launches per replay, " + f"3 rounds, replay == eager == _ref bitwise" + ) + + +# ─── the DSpark drafter's call: (128, 3), zero bias, route_scale 1.5 ────────── +# The drafter routes every stage's draft block through +# gate_forward(cublas_mm(x.float(), W_gate^T) -> fp32 [5G, 128], zeros [128] fp32, 3, 1.5) +# with G requests x 5 draft queries (`DSparkDraftModel._moe` in the unified +# target), G 1..256 (T 5..1280) eagerly, and inside its CUDA graph at every G of +# `_GSET`. mtp.{0,1,2}.ffn.gate.bias is all zeros in the checkpoint. The +# drafter's proxy logits span [-9.8, 9.6]; sigma 3 covers that and beyond, +# sigma 1 the narrow end. N_max for this call is 5 x 256 = 1280. + +_DS_NMAX = 1280 +_DS_CAPTURE = tuple(5 * g for g in _GSET) # 5G over the 46 graph batches, 5..1280 + + +def _ds_logits(num_tokens: int, sigma: float, seed: int) -> torch.Tensor: + g = torch.Generator(device=DEV).manual_seed(seed) + return torch.randn(num_tokens, DSPARK[0], device=DEV, generator=g) * sigma + + +@pytest.mark.parametrize("sigma,seed", [(3.0, 501), (1.0, 503)]) +def test_dspark_rows_1_to_1280_prefix_bitwise(sigma, seed): + """The row axis at the drafter's call: (a) the 1280-row call is bit-exact to + _ref; (b) for EVERY r in 1..1280 the call on the first r rows, written into + the first r rows of a poisoned 1280-row buffer pair, equals the first r rows + of the 1280-row result bit for bit and leaves rows r.. poisoned; (c) a + one-row-shift control is rejected at every r.""" + E, K = DSPARK + logits = _ds_logits(_DS_NMAX, sigma, seed) + bias = torch.zeros(E, device=DEV) + full_w, full_i = op(logits, bias, K, ROUTE_SCALE) + _check_exact(full_w, full_i, *_ref(logits, bias, K, ROUTE_SCALE), f"DSpark N_max sigma={sigma}") + full_wb = _bits(full_w) + shift_wb = full_wb.roll(-1, 0) + out_w = torch.empty(_DS_NMAX, K, device=DEV) + out_i = torch.empty(_DS_NMAX, K, dtype=torch.int32, device=DEV) + out_wb = _bits(out_w) + bad = torch.zeros(_DS_NMAX + 1, dtype=torch.int64, device=DEV) + shifted_rows = torch.zeros(_DS_NMAX + 1, dtype=torch.int64, device=DEV) + for r in range(1, _DS_NMAX + 1): + out_wb.fill_(_POISON_W) + out_i.fill_(_POISON_I) + w, i = op(logits[:r], bias, K, ROUTE_SCALE, out_w[:r], out_i[:r]) + assert w.data_ptr() == out_w.data_ptr() and i.data_ptr() == out_i.data_ptr() + bad[r] = ( + (out_wb[:r] != full_wb[:r]).sum() + + (out_i[:r] != full_i[:r]).sum() + + (out_wb[r:] != _POISON_W).sum() + + (out_i[r:] != _POISON_I).sum() + ) + shifted_rows[r] = (out_wb[:r] != shift_wb[:r]).any(-1).sum() + torch.cuda.synchronize() + failing = torch.nonzero(bad).flatten().tolist() + assert not failing, ( + f"prefix/poison mismatches at r = {failing[:20]} (counts {bad[failing[:20]].tolist()})" + ) + r_all = torch.arange(_DS_NMAX + 1, device=DEV) + assert torch.equal(shifted_rows[1:], r_all[1:]), ( + "one-row-shift control was not rejected on every row" + ) + print( + f"[DSpark rows sigma={sigma}] r=1..{_DS_NMAX}: every prefix bitwise equal to the N_max run, " + f"tail poison intact; shift control rejected on {int(shifted_rows.sum())} rows" + ) + + +def test_dspark_one_kernel_for_every_row_count(): + """At (128, 3): one launch of one kernel at every drafter T.""" + E, K = DSPARK + bias = torch.zeros(E, device=DEV) + names = set() + for t in (1, 2, 3, 4) + _DS_CAPTURE: + logits = _ds_logits(t, 3.0, 600 + t) + op(logits, bias, K, ROUTE_SCALE) + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + op(logits, bias, K, ROUTE_SCALE) + torch.cuda.synchronize() + launched = [ + e.name + for e in prof.events() + if e.device_type == torch.autograd.DeviceType.CUDA and e.self_device_time_total > 0 + ] + assert len(launched) == 1, (t, launched) + names.add(launched[0]) + assert len(names) == 1, names + (name,) = names + assert "gate_forward_kernel" in name and "128" in name, name + print(f"[DSpark kernel] one kernel for every T: {name}") + + +def test_dspark_cuda_graph_capture_5g(): + """Capture at the drafter's call: one graph per T = 5G, G over the 46 + batches of `_GSET` (46 sizes, 5..1280), each warmed eagerly first, captured into one shared pool largest + first with the drafter's allocate-in-capture call (no out buffers, the + persistent zero bias). Three rounds of interleaved replays on new logits + written in place: every replay is bitwise equal to an eager call on the + same inputs and to _ref, and the previous round's reference is rejected on + every row (the replay read the new logits).""" + E, K = DSPARK + bias = torch.zeros(E, device=DEV) + pool = torch.cuda.graph_pool_handle() + graphs = {} + for t in sorted(_DS_CAPTURE, reverse=True): + logits = _ds_logits(t, 3.0, 700 + t) + op(logits, bias, K, ROUTE_SCALE) # eager warm-up at this T + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + out = op(logits, bias, K, ROUTE_SCALE) + graphs[t] = (graph, logits, out) + prev = {} + for rnd, sigma in enumerate((3.0, 1.0, 4.0)): + for t in _DS_CAPTURE: + graph, logits, (gw, gi) = graphs[t] + fresh = _ds_logits(t, sigma, 10000 * (rnd + 1) + t) + logits.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + eager_w, eager_i = op(fresh, bias, K, ROUTE_SCALE) + assert torch.equal(_bits(gw), _bits(eager_w)) and torch.equal(gi, eager_i), ( + f"replay != eager at T={t} round {rnd}" + ) + _check_exact( + gw, gi, *_ref(fresh, bias, K, ROUTE_SCALE), f"DSpark graph T={t} round {rnd}" + ) + if t in prev: + stale_w, _ = prev[t] + assert (_bits(stale_w) != _bits(gw)).any(-1).all(), ( + f"stale-input control not rejected at T={t}" + ) + prev[t] = _ref(fresh, bias, K, ROUTE_SCALE) + assert len(graphs) == 46 and max(graphs) == 1280, (len(graphs), max(graphs)) + print( + f"[DSpark graphs] {len(graphs)} captured T {list(_DS_CAPTURE)}: 3 rounds, replay == eager == _ref bitwise" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/norm/test_modeling_v2_engram_gate.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/norm/test_modeling_v2_engram_gate.py new file mode 100644 index 000000000000..5964aaaf9299 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/norm/test_modeling_v2_engram_gate.py @@ -0,0 +1,1079 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the engram_gate entry of the DeepSeek-V4.1 target. + +What is certified: one Triton launch computes the DeepSeek-V4.1 Engram gate per +(token, hc copy) and the gated residual update ``out = x + gate * value`` -- +``Engram.forward`` of the checkpoint's ``inference/model.py`` after the ``wkv`` +projection, with ``token_mask`` False shutting the gate. + +The entry owns what is certified: ``CELLS`` is the list, ``reference`` and +``compare`` are the gate. This file turns a cell's spec into tensors and keeps +the checks a cell cannot carry: the prefix sweeps over every T up to 16384, the +CUDA-graph captures, in-place and strided layouts, the layout guards run against +the unguarded launch, and the golden layers. + +Two sources of truth: + +* synthetic inputs, against the entry's ``reference`` (torch ops in the + reference's own order: per-copy ``mean`` over ``D``, ``rsqrt``, the product + sum, the signed sqrt with clamp, ``sigmoid``, ``masked_fill``, one cast); +* the golden layer dumps of layers 1 and 14: the key and value are rebuilt from + the checkpoint (the golden hash ids -> the rows of the Engram table read + through a memory map -> FP8 1x32 activation quantize -> IEEE fp32 GEMM with + the dequantized ``wkv``), the rebuilt value must equal the golden + ``engram_value`` bit for bit, and the kernel on the golden ``h_in`` with the + real ``q_weight`` / ``k_weight`` must reproduce the golden ``engram_gate`` + and ``h_after_engram``. This part skips, naming the directories, when the + golden dumps or the checkpoint are absent. + +Gates. The kernel differs from the torch reference only in the order of its +three fp32 sums over ``D`` (the pointwise tail uses the CUDA math library +functions torch's float kernels call), so ``op.compare`` holds the gate at the +default fp32 tolerance and the output at its dtype's default tolerance, and +``_check`` adds the output bit for bit against ``x + gate * value`` formed in +torch from the kernel's own gate. Where only exact arithmetic is involved (rows +with one or two nonzero terms), the gate is compared bit for bit. Deliberately +wrong variants -- the RMS taken jointly over the 4 copies, no signed sqrt, the +gate not zeroed on masked tokens, no clamp -- go through the same gates and +must fail them; how far outside each lands is printed. +""" + +import functools +import json +import os +from contextlib import contextmanager + +import numpy as np +import pytest +import torch +from utils.llm_data import llm_models_root + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.norm.engram_gate import ( # noqa: E501 + _NUM_WARPS, + _engram_gate_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.norm.engram_gate import ( # noqa: E501 + engram_gate as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "engram_gate requires a CUDA device" + +HC = 4 # hc_mult +DIM = 5120 # hidden_size +EPS = 1e-20 # norm_eps +CLAMP = 1e-6 # Engram.clamp_value + +# Reference activations dumped from the checkpoint's own inference code; no default +# location, so the cases that compare against them skip unless this is set. +_GOLDEN_DIR = os.environ.get("MODELING_V2_DSV41_GOLDEN_DIR", "") +# The checkpoint resolves under LLM_MODELS_ROOT, as the whole-model gates' does; +# MODELING_V2_DSV41_CKPT_DIR overrides it. A case that needs it skips when it is absent. +_MODELS_ROOT = llm_models_root() +_CKPT_DIR = os.environ.get("MODELING_V2_DSV41_CKPT_DIR") or ( + str(_MODELS_ROOT / "DeepSeek-V4.1-Flash-NVFP4") if _MODELS_ROOT else "" +) +_GOLDEN_CASES = ("capital_fr", "repeat_ngram", "long_retrieval") +_ENGRAM_LAYERS = (1, 14) + +DEV = torch.device("cuda") + + +def test_device_and_launch_config(): + cap = torch.cuda.get_device_capability() + print( + f"\n device {torch.cuda.get_device_name()} sm_{cap[0]}{cap[1]}; launch: grid T*HC, " + f"BLOCK next_pow2(D) = {1 << (DIM - 1).bit_length()}, num_warps {_NUM_WARPS}, no autotune" + ) + + +# ─── reference: the entry's own; controls: deliberately wrong variants of it ─── + + +def _ref(x, key, value, q_weight, k_weight, token_mask=None, variant=None): + """Returns (out in x.dtype, gate fp32 [T, HC]). With ``variant`` None this is the + entry's ``reference``; a ``variant`` names a deliberately wrong reference used as a + control (``mask_ignored`` is the entry's reference without the mask).""" + if variant in (None, "mask_ignored"): + mask = None if variant == "mask_ignored" else token_mask + return op.reference(x, key, value, q_weight, k_weight, EPS, CLAMP, mask) + key = key.float() + weight = q_weight.float() * k_weight.float() + h = x.float() + if variant == "joint_rms": + rstd = torch.rsqrt(h.square().mean((-2, -1), keepdim=True) + EPS) * torch.rsqrt( + key.square().mean((-2, -1), keepdim=True) + EPS + ) + rstd = rstd.squeeze(-1) + else: + rstd = torch.rsqrt(h.square().mean(-1) + EPS) * torch.rsqrt(key.square().mean(-1) + EPS) + dot = (h * weight * key).sum(-1) * rstd * DIM**-0.5 + if variant == "no_signed_sqrt": + gate = torch.sigmoid(dot) + elif variant == "no_clamp": + gate = torch.sigmoid(torch.copysign(dot.abs().sqrt(), dot)) + else: + gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(CLAMP).sqrt(), dot)) + if token_mask is not None: + gate = gate.masked_fill(~token_mask.bool().unsqueeze(-1), 0) + out = (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype) + return out, gate + + +def _ulps(a, b): + """Elementwise distance in units of the last place (bf16 / fp16 / fp32), via + the order-preserving integer image of the bit pattern.""" + assert a.dtype == b.dtype + itype = {torch.bfloat16: torch.int16, torch.float16: torch.int16, torch.float32: torch.int32}[ + a.dtype + ] + + def ordered(t): + i = t.contiguous().view(itype).long() + top = torch.iinfo(itype).min + return torch.where(i < 0, top - i, i) + + return (ordered(a) - ordered(b)).abs() + + +def _check(out, gate, ref_out, ref_gate, x, value): + """The certification gates. + + 1. and 3. ``op.compare``: the gate at the default fp32 tolerance (the kernel's + three sums over D run in a different order than torch's, nothing else + differs) and the output against the reference at the default tolerance of + its dtype. + 2. The output bit for bit against ``x + gate * value`` formed in fp32 from the + kernel's own gate and rounded once, in torch: the update itself admits no + tolerance (no FMA contraction, one rounding). + """ + own = (x.float() + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(out.dtype) + assert torch.equal(out, own), f"out != x + gate * value on {(out != own).sum().item()} elements" + op.compare((out, gate), (ref_out, ref_gate)) + + +def _stats(out, gate, ref_out, ref_gate): + exact = (out.reshape(-1) == ref_out.reshape(-1)).float().mean().item() + d = (out.float() - ref_out.float()).abs() + return ( + f"gate max|diff| {(gate - ref_gate).abs().max().item():.3g} (bit-exact " + f"{(gate == ref_gate).float().mean().item():.4f}), out bit-exact {exact:.6f}, max|diff| {d.max().item():.3g}, " + f"elements > 1 ulp {(_ulps(out, ref_out) > 1).float().mean().item():.2e}" + ) + + +def _assert_controls_fail(out, gate, x, key, value, qw, kw, mask, variants, tag): + for v in variants: + w_out, w_gate = _ref(x, key, value, qw, kw, mask, variant=v) + with pytest.raises(AssertionError): + _check(out, gate, w_out, w_gate, x, value) + with pytest.raises( + AssertionError + ): # the output alone, at bf16's default tolerance, also sees it + torch.testing.assert_close(out, w_out) + print(f" control {tag} {v}: rejected; {_stats(out, gate, w_out, w_gate)}") + + +# ─── synthetic inputs shaped like the checkpoint's ─── + + +def _inputs(t_len, seed, x_dtype=torch.bfloat16, kv_dtype=torch.bfloat16, w_dtype=torch.bfloat16): + """x [T, HC, D]; key/value as the two split views of one fused wkv output + [T, (HC + 1) * D]; q/k weights [HC, D]. The copies get different scales (as + the real stream's copies do), so a joint RMS is visibly wrong; each key is + noise plus a small random multiple of its x row, so the gates spread over (0, 1).""" + g = torch.Generator(device=DEV).manual_seed(seed) + x_scale = torch.tensor([0.25, 1.0, 3.0, 10.0], device=DEV).view(1, HC, 1) + k_scale = torch.tensor([2.0, 0.5, 1.0, 0.1], device=DEV).view(1, HC, 1) + xn = torch.randn(t_len, HC, DIM, device=DEV, generator=g) + x = (xn * x_scale * torch.exp(torch.randn(t_len, 1, 1, device=DEV, generator=g))).to(x_dtype) + rho = (torch.rand(t_len, HC, 1, device=DEV, generator=g) - 0.5) * 0.1 + keyf = (torch.randn(t_len, HC, DIM, device=DEV, generator=g) + rho * xn) * k_scale + valf = torch.randn(t_len, DIM, device=DEV, generator=g) + kv = torch.cat([keyf.reshape(t_len, HC * DIM), valf], 1).to(kv_dtype) + key = kv[:, : HC * DIM].view(t_len, HC, DIM) + value = kv[:, HC * DIM :] + qw = (1.0 + 0.3 * torch.randn(HC, DIM, device=DEV, generator=g)).to(w_dtype) + kw = (1.0 + 0.3 * torch.randn(HC, DIM, device=DEV, generator=g)).to(w_dtype) + return x, key, value, qw, kw + + +def _mask(t_len, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + m = torch.rand(t_len, device=DEV, generator=g) > 0.2 + if t_len > 1: + m[0] = False + m[-1] = True + return m + + +def _call_cell_seeds(spec): + """The input and mask seeds the configuration was measured with.""" + if spec["mask_stride"] != 1: + return 11, 12 + if (spec["x_dtype"], spec["kv_dtype"], spec["w_dtype"]) != (torch.bfloat16,) * 3: + return 7, 8 + return spec["T"], spec["T"] + 1 + + +def _drive_call_cell(spec): + """Plain call form: a new ``out``, ``gate_out`` given, key and value the split views of + one fused wkv buffer, the dtypes and mask layout of the cell. Masked rows: gate exactly 0 + and the output rows bit-identical to x; the mask may be bool or integer (nonzero = keep) + and strided.""" + t_len, stride = spec["T"], spec["mask_stride"] + seed, mask_seed = _call_cell_seeds(spec) + x, key, value, qw, kw = _inputs(t_len, seed, spec["x_dtype"], spec["kv_dtype"], spec["w_dtype"]) + keep = _mask(t_len, seed=mask_seed) if spec["masked"] else None + mask = None + if keep is not None: + buf = torch.zeros(t_len, stride, dtype=spec["mask_dtype"], device=DEV) + buf[:, stride // 2] = keep.to(spec["mask_dtype"]) + mask = buf[:, stride // 2] + assert mask.stride() == (stride,) + gate = torch.full((t_len, HC), float("nan"), device=DEV) + out = op(x, key, value, qw, kw, EPS, CLAMP, token_mask=mask, out=None, gate_out=gate) + assert out.dtype == x.dtype and out.shape == x.shape and out.is_contiguous() + ref_out, ref_gate = op.reference(x, key, value, qw, kw, EPS, CLAMP, mask) + _check(out, gate, ref_out, ref_gate, x, value) + print( + f"\n T={t_len} mask={None if mask is None else (mask.dtype, mask.stride())} " + f"x {x.dtype} kv {key.dtype} w {qw.dtype}: {_stats(out, gate, ref_out, ref_gate)}; " + f"gate range [{ref_gate.min().item():.4f}, {ref_gate.max().item():.4f}]" + ) + if keep is not None: + assert torch.equal(gate[~keep], torch.zeros_like(gate[~keep])) + assert torch.equal(out[~keep], x[~keep]) + assert (gate[keep] > 0).all() + if stride == 1: + variants = ["joint_rms", "no_signed_sqrt"] + ( + ["mask_ignored"] if keep is not None and not bool(keep.all()) else [] + ) + _assert_controls_fail(out, gate, x, key, value, qw, kw, mask, variants, f"T={t_len}") + + +def _drive_edge_cell(): + """Rows built so every sum has one nonzero term (exact in any order): x = e0, + key = c*e0 + e1 (fp32), so dot = c * sqrt(D) * (1 + c^2)^-1/2 exactly as torch + computes it. |dot| sweeps across the 1e-6 clamp on both signs, plus zero rows + (dot = 0 -> sigmoid(+1e-3)) and |dot| up to 1e2. The gate must equal the torch + reference bit for bit; the no-clamp and no-signed-sqrt controls must not.""" + targets = [0.0, 1e-12, 1e-9, 1e-7, 5e-7, 1e-6, 2e-6, 1e-5, 1e-4, 1e-2, 0.1, 1.0, 4.0, 100.0] + cs = [s * v / DIM**0.5 for v in targets for s in (1.0, -1.0)] + n = len(cs) + 3 + x = torch.zeros(n, HC, DIM, dtype=torch.bfloat16, device=DEV) + key = torch.zeros(n, HC, DIM, dtype=torch.float32, device=DEV) + x[:, :, 0] = 1.0 + for i, c in enumerate(cs): + key[i, :, 0] = c + key[i, :, 1] = 1.0 + x[len(cs)] = 0.0 # x row zero: rstd_x = 1e10, dot = 0 + key[len(cs), :, 1] = 1.0 + key[len(cs) + 1] = 0.0 # key row zero + x[len(cs) + 2] = 0.0 # both zero + value = torch.linspace(-3, 3, n * DIM, device=DEV).view(n, DIM) + qw = torch.ones(HC, DIM, dtype=torch.bfloat16, device=DEV) + kw = torch.ones(HC, DIM, dtype=torch.bfloat16, device=DEV) + gate = torch.empty(n, HC, dtype=torch.float32, device=DEV) + out = op(x, key, value, qw, kw, EPS, CLAMP, gate_out=gate) + ref_out, ref_gate = op.reference(x, key, value, qw, kw, EPS, CLAMP) + assert torch.equal(gate, ref_gate), (gate[:, 0] - ref_gate[:, 0]).abs().max().item() + assert torch.equal(out, ref_out) + _check(out, gate, ref_out, ref_gate, x, value) + lo = torch.sigmoid(torch.tensor([CLAMP], device=DEV).sqrt()).item() + print( + f"\n {n} rows: gate bit-exact incl. dot = +-{targets[1]}..{targets[-1]}; " + f"zero rows gate {gate[-3:, 0].tolist()} " + f"(sigmoid(sqrt(fp32(1e-6))) = {lo})" + ) + assert torch.equal(gate[-3:], torch.full_like(gate[-3:], lo)) + for v in ("no_clamp", "no_signed_sqrt"): + w_out, w_gate = _ref(x, key, value, qw, kw, variant=v) + with pytest.raises(AssertionError): + _check(out, gate, w_out, w_gate, x, value) + print(f" control {v}: rejected; gate max|diff| {(gate - w_gate).abs().max().item():.3g}") + + +def _drive_target_masked_cell(kind, t_len): + """The multimodal target's masked call at T = 16384 with one mask: in place, split wkv + views, a contiguous bool keep. (a) It equals the same call with ``gate_out`` (for the + gate) bit for bit, and passes the certification gates against the reference; (b) masked + rows: gate exactly 0 and the output rows equal x bit for bit; (c) keep rows equal the + UNMASKED call's rows bit for bit (the mask perturbs no live row); an all-True mask + tensor equals ``token_mask=None`` bitwise. Controls against the reference built with a + wrong mask -- ignored, shifted by one row, inverted -- are rejected by the same gates + wherever the wrong mask differs from the right one.""" + x, kv, key3, value, qw, kw = _mm_state(t_len, seed=9100 + _MM_MASKS.index(kind)) + keep = _mm_mask(kind, t_len) + assert keep.dtype == torch.bool and keep.is_contiguous() and keep.shape == (t_len,) + out = _mm_call(x, key3, value, qw, kw, keep) + gate = torch.full((t_len, HC), float("nan"), device=DEV) + with_gate = x.clone() + op(with_gate, key3, value, qw, kw, EPS, CLAMP, keep, with_gate, gate_out=gate) + assert torch.equal(with_gate, out) + ref_out, ref_gate = op.reference(x, key3, value, qw, kw, EPS, CLAMP, keep) + _check(out, gate, ref_out, ref_gate, x, value) + dead = ~keep + assert torch.equal(gate[dead], torch.zeros_like(gate[dead])) + assert torch.equal(out[dead], x[dead]) + plain = _mm_call(x, key3, value, qw, kw, None) + assert torch.equal(out[keep], plain[keep]), ( + f"{int((out[keep] != plain[keep]).sum())} live elements differ" + ) + if bool(keep.all()): + assert torch.equal(out, plain) + n_dead = int(dead.sum()) + print( + f"\n T={t_len} {kind}: {n_dead} masked rows; {_stats(out, gate, ref_out, ref_gate)}; masked rows == x, " + f"live rows == unmasked call, bit for bit" + ) + wrong = {"mask_ignored": None, "mask_shifted": torch.roll(keep, 1), "mask_inverted": ~keep} + for name, w_mask in wrong.items(): + w_keep = torch.ones_like(keep) if w_mask is None else w_mask + n_diff = int((w_keep != keep).sum()) + if n_diff == 0: # a constant mask is its own shift; all-True ignored is itself + print(f" control {kind} {name}: identical to the right mask, nothing to reject") + continue + w_out, w_gate = op.reference(x, key3, value, qw, kw, EPS, CLAMP, w_mask) + with pytest.raises(AssertionError): + _check(out, gate, w_out, w_gate, x, value) + with pytest.raises(AssertionError): + torch.testing.assert_close(out, w_out) + print( + f" control {kind} {name} ({n_diff} rows differ): rejected; {_stats(out, gate, w_out, w_gate)}" + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Every cell through ``op`` (guarded by ``is_valid``), ``op.reference`` and ``op.compare``, + with the structural checks and wrong-variant controls its configuration carries.""" + certifying(op) + spec = cell.spec + assert (spec["HC"], spec["D"], spec["eps"], spec["clamp_value"]) == (HC, DIM, EPS, CLAMP) + if spec["kind"] == "call": + _drive_call_cell(spec) + elif spec["kind"] == "edges": + _drive_edge_cell() + else: + assert spec["kind"] == "target_masked", spec["kind"] + _drive_target_masked_cell(spec["mask"], spec["T"]) + + +def test_in_place_and_strided_views(): + """out = x (in place), x / out as views into wider buffers, key / value as split + views of one fused buffer, gate_out strided: all bit-identical to contiguous + copies, and nothing outside the views is written.""" + t_len = 96 + x, key, value, qw, kw = _inputs(t_len, seed=21) + mask = _mask(t_len, seed=22) + base = op( + x.contiguous(), key.contiguous(), value.contiguous(), qw, kw, EPS, CLAMP, token_mask=mask + ) + + xi = x.clone() + res = op(xi, key, value, qw, kw, EPS, CLAMP, token_mask=mask, out=xi) + assert res.data_ptr() == xi.data_ptr() and torch.equal(xi, base) + + wide = torch.randn(t_len, HC + 2, DIM + 64, device=DEV).to(torch.bfloat16) + wide_before = wide.clone() + xv = wide[:, 1 : HC + 1, 32 : 32 + DIM] + xv.copy_(x) + wide_before[:, 1 : HC + 1, 32 : 32 + DIM] = x + ob = torch.full((t_len, HC, DIM + 16), 7.0, dtype=torch.bfloat16, device=DEV) + ov = ob[:, :, 8 : 8 + DIM] + gb = torch.full((HC, 2 * t_len), 7.0, device=DEV) + gv = gb[:, ::2].t() # [T, HC], strides (2, 2T) + op(xv, key, value, qw, kw, EPS, CLAMP, token_mask=mask, out=ov, gate_out=gv) + assert torch.equal(ov, base) + assert torch.equal(wide, wide_before) # x view is read only + assert (ob[:, :, :8] == 7).all() and (ob[:, :, 8 + DIM :] == 7).all() + assert (gb[:, 1::2] == 7).all() + _, ref_gate = _ref(x, key, value, qw, kw, mask) + torch.testing.assert_close(gv, ref_gate) + # inputs are not modified by an out-of-place call + snap = [t.clone() for t in (x, key, value, qw, kw)] + op(x, key, value, qw, kw, EPS, CLAMP) + assert all(torch.equal(a, b) for a, b in zip(snap, (x, key, value, qw, kw))) + print( + "\n in-place, strided x/out/gate_out, split key/value views: bit-identical, no stray writes" + ) + + +def test_batch_invariance(): + """Each (token, copy) row is computed by one program with a fixed block, so a + token's result does not depend on how many tokens share the call: rows of a + 1024-token call equal the same rows computed alone (T=1) and in a 37-token call.""" + x, key, value, qw, kw = _inputs(1024, seed=31) + mask = _mask(1024, seed=32) + full = op(x, key, value, qw, kw, EPS, CLAMP, token_mask=mask) + for t in (0, 1, 500, 1023): + one = op( + x[t : t + 1], + key[t : t + 1], + value[t : t + 1], + qw, + kw, + EPS, + CLAMP, + token_mask=mask[t : t + 1], + ) + assert torch.equal(one, full[t : t + 1]) + part = op( + x[300:337], key[300:337], value[300:337], qw, kw, EPS, CLAMP, token_mask=mask[300:337] + ) + assert torch.equal(part, full[300:337]) + again = op(x, key, value, qw, kw, EPS, CLAMP, token_mask=mask) + assert torch.equal(again, full) + print("\n rows of T=1024 == the same rows at T=1 and T=37; repeat call bit-identical") + + +@pytest.mark.parametrize("t_len", [1, 5, 160]) +def test_cuda_graph_capture_replay(t_len): + """Captured once, replayed on new data written into the same buffers (in place, + with a mask): bit-identical to an eager call on that data.""" + x, key, value, qw, kw = _inputs(t_len, seed=41 + t_len) + mask = _mask(t_len, seed=42) + sx, skey, sval, smask = x.clone(), key.clone(), value.clone(), mask.clone() + sgate = torch.empty(t_len, HC, device=DEV) + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + op( + sx.clone(), skey, sval, qw, kw, EPS, CLAMP, token_mask=smask, gate_out=sgate + ) # warm-up / compile + torch.cuda.current_stream().wait_stream(s) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(sx, skey, sval, qw, kw, EPS, CLAMP, token_mask=smask, out=sx, gate_out=sgate) + for seed in (1, 2): + x2, key2, value2, _, _ = _inputs(t_len, seed=100 + seed) + mask2 = _mask(t_len, seed=200 + seed) + sx.copy_(x2) + skey.copy_(key2) + sval.copy_(value2) + smask.copy_(mask2) + graph.replay() + g2 = torch.empty(t_len, HC, device=DEV) + eager = op(x2, key2, value2, qw, kw, EPS, CLAMP, token_mask=mask2, gate_out=g2) + assert torch.equal(sx, eager) and torch.equal(sgate, g2) + print(f"\n T={t_len}: graph replay bit-identical to eager on 2 fresh inputs") + + +# ─── every token count 1..16384 is the same per-row function ─── +# +# The core target calls this op on one rank's T tokens, T = 1..max_num_tokens, and +# max_num_tokens reaches 16384 per rank (every member gates only its own tokens, so no +# gathered count reaches it; Engram layers 1 and 14 are before the decoder bounded-replay +# pass). The grid is T * HC programs, one per (token, copy) row, and T is not a kernel +# argument at all: it only sizes the grid. So every prefix of a 16384-token call must +# return exactly the first r rows of the 16384-token result, which is held to the reference. +_N_MAX = 16384 + + +def _jit_variants(fn) -> int: + """Compiled variants Triton holds for ``fn`` on every device (Triton >= 3.2 keeps them + in ``device_caches[device][0]``, older releases in ``cache[device]``).""" + dc = getattr(fn, "device_caches", None) + if dc: + return sum(len(v[0]) for v in dc.values()) + return sum(len(v) for v in fn.cache.values()) + + +_POISON = 0x7FA5 # a bf16 NaN bit pattern; the kernel writes finite values for finite inputs +_GUARD = 64 # token rows past the prefix that must still hold the poison after each call + + +def test_every_token_count_is_a_prefix_of_16384(): + """16384 tokens in the core's layout (key and value the split views of one fused + [T, 25600] wkv buffer, no mask). The 16384-token call, run with ``gate_out``, passes + the certification gates against the reference, and the same call without + ``gate_out`` (the core's form) is bit-identical to it. Then every prefix r = 1..16384 + is called into a token-row-prefix view of one poisoned output buffer: its r rows + must equal the 16384-token rows bit for bit and the next 64 token rows keep the + poison (mismatches summed on the device, one synchronization). The kernel's + compiled variants are counted: none may appear after the first prefix call.""" + n_max = _N_MAX + x, key, value, qw, kw = _inputs(n_max, seed=8192) + gate = torch.empty(n_max, HC, dtype=torch.float32, device=DEV) + with_gate = op(x, key, value, qw, kw, EPS, CLAMP, gate_out=gate) + ref_out, ref_gate = _ref(x, key, value, qw, kw) + _check(with_gate, gate, ref_out, ref_gate, x, value) + full = op(x, key, value, qw, kw, EPS, CLAMP) + assert torch.equal(full, with_gate) + print(f"\n T={n_max}: {_stats(full, gate, ref_out, ref_gate)}") + full_bits = full.view(torch.int16) + buf = torch.empty(n_max + _GUARD, HC, DIM, dtype=torch.bfloat16, device=DEV) + bits = buf.view(torch.int16) + bits.fill_(_POISON) + bad = torch.zeros((), dtype=torch.int64, device=DEV) + over = torch.zeros((), dtype=torch.int64, device=DEV) + before = _jit_variants(_engram_gate_kernel) + settled = None + for r in range(1, n_max + 1): + res = op(x[:r], key[:r], value[:r], qw, kw, EPS, CLAMP, None, buf[:r]) + assert res.data_ptr() == buf.data_ptr() + bad += (bits[:r] != full_bits[:r]).sum() + over += (bits[r : r + _GUARD] != _POISON).sum() + if r == 1: + settled = _jit_variants(_engram_gate_kernel) + torch.cuda.synchronize() + after = _jit_variants(_engram_gate_kernel) + print( + f" r = 1..{n_max}: {bad.item()} elements differ from the {n_max}-token rows, {over.item()} guard " + f"elements written; kernel variants {before} before, {settled} after r = 1, {after} after r = {n_max}" + ) + assert bad.item() == 0 and over.item() == 0 + assert after == settled and settled - before <= 1, (before, settled, after) + + +def test_row_sweep_harness_sees_a_one_token_shift(): + """Control for the sweep above: the same bitwise per-row check rejects a call that + reads its inputs one token row off, at every r tried, including r = 1.""" + x, key, value, qw, kw = _inputs(65, seed=65) + full = op(x[:64], key[:64], value[:64], qw, kw, EPS, CLAMP) + for r in (1, 2, 7, 63, 64): + shifted = op(x[1 : r + 1], key[1 : r + 1], value[1 : r + 1], qw, kw, EPS, CLAMP) + n_bad = int((shifted.view(torch.int16) != full.view(torch.int16)[:r]).sum().item()) + print(f" [control row shift] r={r}: {n_bad} / {r * HC * DIM} elements differ") + assert n_bad > r * HC * DIM // 2 + + +# ─── the multimodal target's masked call: every T 1..16384, eager ─── +# +# The multimodal target calls, on one rank's T = 1..16384 tokens of a step that holds a +# context chunk of an image request, +# engram_gate(h, key3, value, q_weight, k_weight, 1e-20, 1e-6, keep, h) +# in place on the bf16 stream h [T, 4, 5120], key3 / value the split views of ONE fused +# [T, 25600] wkv buffer, and keep a contiguous bool [T] that is False on image-span rows +# (gate 0) and True on every other row (text, delimiters, generation rows of a mixed step). +# Such steps are never captured: the graph decision is all-rank and captured steps are +# generation-only, which carry no image row, so the masked form runs eagerly only. T only +# sizes the grid and the mask is read at row t alone, so the prefix rule applies: the 16384-row +# masked call against the reference, and every prefix r (with the mask's prefix) bitwise +# equal to the first r rows of that call. + +_MM_T = _N_MAX + + +def _spans(t_len, spans): + """keep [T] bool: False on each [a, b) of ``spans`` (image-span rows).""" + m = torch.ones(t_len, dtype=torch.bool, device=DEV) + for a, b in spans: + m[a:b] = False + return m + + +def _random_spans(t_len, g, n_spans, max_len): + """``n_spans`` image spans of 1..max_len rows at random starts (may overlap / touch).""" + out = [] + for _ in range(n_spans): + ln = int(torch.randint(1, max_len + 1, (1,), generator=g)) + a = int(torch.randint(0, max(1, t_len - ln + 1), (1,), generator=g)) + out.append((a, min(t_len, a + ln))) + return out + + +def _mm_mask(kind, t_len=_MM_T): + """The named keep masks (False = image-span row). ``rand`` are seeded random mixes of + span counts, span lengths and Bernoulli densities.""" + if ( + kind == "none_masked" + ): # an image request's chunk that holds no span row: all-True tensor, not None + return torch.ones(t_len, dtype=torch.bool, device=DEV) + if kind == "all_masked": # the whole step is image-span rows + return torch.zeros(t_len, dtype=torch.bool, device=DEV) + if kind == "span_at_row0": + return _spans(t_len, [(0, 1517)]) + if kind == "span_at_end": + return _spans(t_len, [(t_len - 2611, t_len)]) + if kind == "one_span_mid": + return _spans(t_len, [(1000, 1000 + 2305)]) + if kind == "multi_span": # several images, delimiters / text rows between them, mixed-step tail + return _spans( + t_len, + [ + (0, 3), + (5, 1524), + (1530, 1531), + (2000, 4611), + (4613, 7046), + (8190, 8194), + (9000, 13311), + (t_len - 2, t_len), + ], + ) + if kind == "alternating": + return torch.arange(t_len, device=DEV) % 2 == 0 + g = torch.Generator().manual_seed(4200 + int(kind[4:])) + k = int(kind[4:]) + if k % 3 == 0: # Bernoulli rows at a random density + p = float(torch.rand((), generator=g)) + return (torch.rand(t_len, generator=g) > p).to(DEV) + n_spans = int(torch.randint(1, 17, (1,), generator=g)) + return _spans( + t_len, _random_spans(t_len, g, n_spans, 1 + int(torch.randint(1, 3000, (1,), generator=g))) + ) + + +_MM_MASKS = [ + "none_masked", + "all_masked", + "span_at_row0", + "span_at_end", + "one_span_mid", + "multi_span", + "alternating", +] + [f"rand{k}" for k in range(9)] + + +def _mm_state(t_len, seed): + """The target's buffers: the stream x [T, HC, D], ONE fused wkv buffer [T, 25600] and its + split views key3 [T, HC, D] (row stride 25600) and value [T, D] (row stride 25600).""" + x, key0, value0, qw, kw = _inputs(t_len, seed=seed) + kv = torch.empty(t_len, (HC + 1) * DIM, dtype=torch.bfloat16, device=DEV) + kv[:, : HC * DIM].copy_(key0.reshape(t_len, HC * DIM)) + kv[:, HC * DIM :].copy_(value0) + key3 = kv[:, : HC * DIM].view(t_len, HC, DIM) + value = kv[:, HC * DIM :] + assert key3.stride() == ((HC + 1) * DIM, DIM, 1) and value.stride() == ((HC + 1) * DIM, 1) + return x, kv, key3, value, qw, kw + + +def _mm_call(x, key3, value, qw, kw, keep): + """The target's call on a copy of ``x``: in place, keep as token_mask, no gate_out.""" + h = x.clone() + res = op(h, key3, value, qw, kw, EPS, CLAMP, keep, h) + assert res.data_ptr() == h.data_ptr() + return h + + +@pytest.mark.parametrize("kind", ["multi_span", "rand3"]) +def test_multimodal_masked_every_token_count_is_a_prefix_of_16384(kind): + """The prefix rule in the target's masked form. The 16384-row call (in place, split views of + one fused wkv buffer, contiguous bool keep) passes the reference gates. Then every + r = 1..16384 is run in the same form on one poisoned stream buffer: rows [0, r) get x's + rows, the call runs in place on buf[:r] with key3[:r], value[:r], keep[:r]; its rows + must equal the 16384-row result bit for bit and the next 64 rows keep the poison + (mismatches summed on the device, one synchronization). The masked variant is compiled + by the 16384-row call in exactly this form, so the sweep must compile nothing.""" + x, kv, key3, value, qw, kw = _mm_state(_MM_T, seed=9300 + len(kind)) + keep = _mm_mask(kind) + buf = torch.empty(_MM_T + _GUARD, HC, DIM, dtype=torch.bfloat16, device=DEV) + bits = buf.view(torch.int16) + bits.fill_(_POISON) + buf[:_MM_T].copy_(x) + op(buf[:_MM_T], key3, value, qw, kw, EPS, CLAMP, keep, buf[:_MM_T]) # the 16384-row call + full = buf[:_MM_T].clone() + assert (bits[_MM_T:] == _POISON).all() + gate = torch.empty(_MM_T, HC, dtype=torch.float32, device=DEV) + with_gate = x.clone() + op(with_gate, key3, value, qw, kw, EPS, CLAMP, keep, with_gate, gate_out=gate) + assert torch.equal(with_gate, full) + ref_out, ref_gate = _ref(x, key3, value, qw, kw, keep) + _check(full, gate, ref_out, ref_gate, x, value) + assert torch.equal(full[~keep], x[~keep]) + print( + f"\n T={_MM_T} {kind} ({int((~keep).sum())} masked rows): {_stats(full, gate, ref_out, ref_gate)}" + ) + full_bits = full.view(torch.int16) + bits.fill_(_POISON) + bad = torch.zeros((), dtype=torch.int64, device=DEV) + over = torch.zeros((), dtype=torch.int64, device=DEV) + before = _jit_variants(_engram_gate_kernel) + for r in range(1, _MM_T + 1): + h = buf[:r] + h.copy_(x[:r]) + res = op(h, key3[:r], value[:r], qw, kw, EPS, CLAMP, keep[:r], h) + assert res.data_ptr() == buf.data_ptr() + bad += (bits[:r] != full_bits[:r]).sum() + over += (bits[r : r + _GUARD] != _POISON).sum() + torch.cuda.synchronize() + after = _jit_variants(_engram_gate_kernel) + print( + f" masked r = 1..{_MM_T}: {bad.item()} elements differ from the {_MM_T}-row result, {over.item()} guard " + f"elements written; kernel variants {before} before the sweep, {after} after" + ) + assert bad.item() == 0 and over.item() == 0 + assert after == before, (before, after) + + +def test_multimodal_masked_sweep_harness_sees_a_shifted_mask(): + """Control for the masked sweep: the same bitwise per-row check rejects a prefix call whose + mask is read one row off (keep[1:r+1]) at every r tried where the shifted mask differs.""" + x, kv, key3, value, qw, kw = _mm_state(_MM_T, seed=9300 + len("rand3")) + keep = _mm_mask("rand3") + full = _mm_call(x, key3, value, qw, kw, keep) + for r in (2, 7, 64, 1517, 8191, 16383): + n_rows = int((keep[1 : r + 1] != keep[:r]).sum()) + shifted = _mm_call(x[:r], key3[:r], value[:r], qw, kw, keep[1 : r + 1].contiguous()) + n_bad = int((shifted.view(torch.int16) != full.view(torch.int16)[:r]).sum().item()) + print( + f" [control shifted mask] r={r}: mask differs on {n_rows} rows, {n_bad} / {r * HC * DIM} elements differ" + ) + assert n_rows > 0 and n_bad > 0 + + +# ─── capture at every graph batch the gates capture: the core's exact call ─── + +# The generation-only CUDA-graph batches an engine captures per rank, Gset: every batch +# 1..32 plus TRT-LLM's padded default list for cuda_graph_config {enable_padding: true, +# max_batch_size: 256} (1, 2, 4, 8, 16..128 step 8, 192, 256; the padding-off list at 256, +# 1..32, 64, 128, 256, is a subset). A decode step gates G tokens and a DSpark target-verify +# step 6 G (1 + 5 drafts), so T = G and 6 G: 82 counts up to 1536. Every member (tep4, dep4, +# tp4) gates only its own tokens, so no gathered row count reaches this op. +_GRAPH_BATCHES = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_CAPTURE_T = sorted(set(_GRAPH_BATCHES) | {6 * b for b in _GRAPH_BATCHES}) +assert len(_GRAPH_BATCHES) == 46 and len(_CAPTURE_T) == 82 and _CAPTURE_T[-1] == 1536 + + +@pytest.mark.parametrize("t_len", _CAPTURE_T) +def test_cuda_graph_at_every_capture_batch(t_len): + """The core target's call: key and value are the two split views of ONE static fused + ``[T, 25600]`` wkv buffer, no mask, no ``gate_out``, in place on the stream. Warmed + up eagerly at the same T, captured once, then replayed on 3 fresh inputs written in + place: each replay equals the same eager call on the same data bit for bit, and + passes the certification gates against the reference (the gate taken from an eager + call with ``gate_out``, whose ``out`` must equal the call without it bit for bit). + The joint-RMS and no-signed-sqrt controls are run against the first replay.""" + kv_n = (HC + 1) * DIM + x0, key0, value0, qw, kw = _inputs(t_len, seed=700 + t_len) + sx = x0.clone() + skv = torch.empty(t_len, kv_n, dtype=torch.bfloat16, device=DEV) + skv[:, : HC * DIM].copy_(key0.reshape(t_len, HC * DIM)) + skv[:, HC * DIM :].copy_(value0) + skey = skv[:, : HC * DIM].view(t_len, HC, DIM) # reshape(split(kv)[0], [T, HC, D]) in the core + sval = skv[:, HC * DIM :] # split(kv)[1]: a [T, D] view with row stride 25600 + assert skey.data_ptr() == skv.data_ptr() and sval.stride(0) == kv_n + out0 = x0.clone() + op(out0, skey, sval, qw, kw, EPS, CLAMP, None, out0) # the capture-time data's result + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + warm = sx.clone() + op(warm, skey, sval, qw, kw, EPS, CLAMP, None, warm) # engine warm-up at this T + torch.cuda.current_stream().wait_stream(s) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(sx, skey, sval, qw, kw, EPS, CLAMP, None, sx) + for seed in (1, 2, 3): + x2, key2, value2, _, _ = _inputs(t_len, seed=800 + 10 * t_len + seed) + sx.copy_(x2) + skv[:, : HC * DIM].copy_(key2.reshape(t_len, HC * DIM)) + skv[:, HC * DIM :].copy_(value2) + graph.replay() + torch.cuda.synchronize() + eager = x2.clone() + op(eager, skey, sval, qw, kw, EPS, CLAMP, None, eager) # the same call form, eagerly + assert torch.equal(sx, eager), f"replay != eager on {(sx != eager).sum().item()} elements" + gate = torch.empty(t_len, HC, dtype=torch.float32, device=DEV) + with_gate = x2.clone() + op(with_gate, skey, sval, qw, kw, EPS, CLAMP, None, with_gate, gate_out=gate) + assert torch.equal(with_gate, eager) + ref_out, ref_gate = _ref(x2, key2, value2, qw, kw) + _check(sx, gate, ref_out, ref_gate, x2, value2) + stale = (sx != out0).float().mean().item() + assert stale > 0.5, "a replay on new data returned the capture-time data's result" + if seed == 1: + print( + f"\n T={t_len} graph: replay == eager bit for bit; {_stats(sx, gate, ref_out, ref_gate)}; " + f"differs from the capture-time data's output on {stale:.4f} of elements" + ) + _assert_controls_fail( + sx, + gate, + x2, + key2, + value2, + qw, + kw, + None, + ["joint_rms", "no_signed_sqrt"], + f"graph T={t_len}", + ) + + +def test_empty_batch(): + x, key, value, qw, kw = _inputs(4, seed=51) + out = op( + x[:0], + key[:0], + value[:0], + qw, + kw, + EPS, + CLAMP, + token_mask=torch.ones(0, dtype=torch.bool, device=DEV), + gate_out=torch.empty(0, HC, device=DEV), + ) + assert out.shape == (0, HC, DIM) and out.dtype == torch.bfloat16 + + +# ─── is_valid: each rejected layout, run through the unguarded launch first ─── + + +def _interleaved(t): + """A view with ``t``'s values and stride 2 along the last dim (other, deterministic data between).""" + return torch.stack([t, (t.float() * -1.7 + 0.3).to(t.dtype)], -1).flatten(-2)[..., ::2] + + +def _call(args, **kwargs): + return op(args["x"], args["key"], args["value"], args["qw"], args["kw"], EPS, CLAMP, **kwargs) + + +def test_is_valid_rejects_layouts_the_kernel_misreads(): + """Every rejected layout below is run through ``op`` outside ``validating`` first (the + launch alone, no guard), and it returns without an error: a stride-2 view is read as if + contiguous (the result is wrong against the reference on the tensor as given), and a + view with fewer rows / columns / copies than ``x`` implies is read or written past its + end (the result equals the intended call's on the parent buffer the view was cut from). + Inside ``validating`` the entry's ``is_valid`` refuses each one.""" + t_len = 16 + x, key, value, qw, kw = _inputs(t_len, seed=61) + mask = _mask(t_len, seed=62) + base = dict(x=x, key=key, value=value, qw=qw, kw=kw) + with validating(op): + intended = op(x, key, value, qw, kw, EPS, CLAMP, token_mask=mask) # the intended tensors + strided = { + "x inner stride 2": ("x", "x must be"), + "key inner stride 2": ("key", "key must be"), + "value inner stride 2": ("value", "value must be"), + "q_weight inner stride 2": ("qw", "q_weight must be"), + } + for name, (arg, match) in strided.items(): + args = dict(base) + args[arg] = _interleaved(base[arg]) + got = _call(args, token_mask=mask) + torch.cuda.synchronize() + worst = _ulps(got, intended).max().item() # the view holds the intended values + assert worst > 1, name + with pytest.raises(AssertionError, match=match): + with validating(op): + _call(args, token_mask=mask) + print( + f" rejected: {name} (unguarded: no error, bit-exact vs the view's values " + f"{(got.view(-1) == intended.view(-1)).float().mean().item():.4f}, max ulps {worst})" + ) + short = { + "key with HC-1 copies": ("key", key[:, : HC - 1], "key must be"), + "value with D/2 columns": ("value", value[:, : DIM // 2], "value must be"), + "q_weight with HC-1 rows": ("qw", qw[: HC - 1], "q_weight must be"), + } + for name, (arg, view, match) in short.items(): + args = dict(base) + args[arg] = view + got = _call(args, token_mask=mask) + torch.cuda.synchronize() + assert torch.equal(got, intended), name # it used the parent's data beyond the view + with pytest.raises(AssertionError, match=match): + with validating(op): + _call(args, token_mask=mask) + print(f" rejected: {name} (unguarded: no error, read the parent buffer past the view)") + got = op(x, key, value, qw, kw, EPS, CLAMP, token_mask=mask[: t_len // 2]) + torch.cuda.synchronize() + assert torch.equal(got, intended) + with pytest.raises(AssertionError, match="token_mask must be"): + with validating(op): + op(x, key, value, qw, kw, EPS, CLAMP, token_mask=mask[: t_len // 2]) + print( + f" rejected: token_mask with T/2 entries (unguarded read {t_len - t_len // 2} entries past it)" + ) + ob = torch.full((t_len * 2, HC, DIM), 7.0, dtype=torch.bfloat16, device=DEV) + op(x, key, value, qw, kw, EPS, CLAMP, token_mask=mask, out=ob[: t_len // 2]) + torch.cuda.synchronize() + assert torch.equal(ob[:t_len], intended) and (ob[t_len:] == 7).all() + with pytest.raises(AssertionError, match="out must be"): + with validating(op): + op(x, key, value, qw, kw, EPS, CLAMP, token_mask=mask, out=ob[: t_len // 2]) + print(f" rejected: out with T/2 rows (unguarded wrote {t_len - t_len // 2} rows past it)") + gb = torch.full((t_len * 2, HC), 7.0, device=DEV) + op(x, key, value, qw, kw, EPS, CLAMP, gate_out=gb[: t_len // 2]) + torch.cuda.synchronize() + assert (gb[t_len // 2 : t_len] != 7).all() and (gb[t_len:] == 7).all() + with pytest.raises(AssertionError, match="gate_out must be"): + with validating(op): + op(x, key, value, qw, kw, EPS, CLAMP, gate_out=gb[: t_len // 2]) + print(" rejected: gate_out with T/2 rows (unguarded wrote past it)") + # CPU tensors are refused by Triton's own launcher, loudly + with pytest.raises(ValueError, match="cannot be accessed from Triton") as err: + op(x.cpu(), key, value, qw, kw, EPS, CLAMP) + print(f" CPU x: ValueError: {err.value}") + + +# ─── golden layer dumps (layers 1 and 14), key/value rebuilt from the checkpoint ─── + + +def _golden_available() -> bool: + return all( + os.path.isfile(os.path.join(_GOLDEN_DIR, c, "layers", "L14.pt")) for c in _GOLDEN_CASES + ) and os.path.isfile(os.path.join(_CKPT_DIR, "model.safetensors.index.json")) + + +def _skip_without_golden(): + if not _golden_available(): + pytest.skip(f"golden dumps ({_GOLDEN_DIR}) or checkpoint ({_CKPT_DIR}) not present") + + +@functools.lru_cache(maxsize=None) +def _st_entry(name: str): + """(file, absolute byte offset, dtype string, shape) of one safetensors tensor.""" + with open(os.path.join(_CKPT_DIR, "model.safetensors.index.json")) as f: + fname = os.path.join(_CKPT_DIR, json.load(f)["weight_map"][name]) + with open(fname, "rb") as f: + n = int.from_bytes(f.read(8), "little") + hdr = json.loads(f.read(n)) + m = hdr[name] + return fname, 8 + n + m["data_offsets"][0], m["dtype"], tuple(m["shape"]) + + +def _st_rows(name: str, rows: np.ndarray) -> torch.Tensor: + """Rows of a 2-D tensor as raw bytes, through a memory map (only the touched pages are read).""" + fname, off, _, shape = _st_entry(name) + table = np.memmap( + fname, dtype=np.uint8, mode="r", offset=off, shape=(shape[0], _row_bytes(name)) + ) + return torch.from_numpy(np.ascontiguousarray(table[rows])) + + +def _row_bytes(name): + _, _, dt, shape = _st_entry(name) + return shape[1] * {"F8_E4M3": 1, "F8_E8M0": 1, "BF16": 2}[dt] + + +def _st_tensor(name: str) -> torch.Tensor: + fname, off, dt, shape = _st_entry(name) + rows = np.arange(shape[0]) + raw = _st_rows(name, rows) + return raw.view( + {"F8_E4M3": torch.float8_e4m3fn, "F8_E8M0": torch.uint8, "BF16": torch.bfloat16}[dt] + ).reshape(shape) + + +@contextmanager +def _ieee_fp32_matmul(): + """The golden reference runs its fp32 GEMMs with TF32 off.""" + prev_allow = torch.backends.cuda.matmul.allow_tf32 + prev_prec = getattr(torch.backends.cuda.matmul, "fp32_precision", None) + torch.backends.cuda.matmul.allow_tf32 = False + if prev_prec is not None: + torch.backends.cuda.matmul.fp32_precision = "ieee" + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = prev_allow + if prev_prec is not None: + torch.backends.cuda.matmul.fp32_precision = prev_prec + + +def _e8m0(u8): + """UE8M0 byte b -> 2^(b-127) from the fp32 exponent bits (b=0 -> 2^-127).""" + b = u8.to(torch.int32) + bits = torch.where(b == 0, torch.full_like(b, 0x00400000), b << 23) + return bits.view(torch.float32) + + +def _fp8_act_qdq(x): + """The reference's FP8 1x32 activation quantize-dequantize: amax floored at 1e-4, + scale 2^ceil(log2(amax / 448)) from the fp32 bits, E4M3 of the clamped quotient.""" + xf = x.float().unflatten(-1, (-1, 32)) + v = xf.abs().amax(-1).clamp_min(1e-4) * (1.0 / 448.0) + bits = v.contiguous().view(torch.int32) + e = ((bits >> 23) & 0xFF) - 127 + ((bits & 0x7FFFFF) != 0).to(torch.int32) + sc = ((e + 127) << 23).view(torch.float32) + q = (xf / sc.unsqueeze(-1)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + return (q.float() * sc.unsqueeze(-1)).flatten(-2) + + +@functools.lru_cache(maxsize=None) +def _layer_weights(layer: int): + """(wkv dequantized to fp32 on the GPU, q_weight, k_weight) of one Engram layer.""" + p = f"layers.{layer}.engram." + ww = _st_tensor(p + "wkv.weight").to(DEV) + ws = _st_tensor(p + "wkv.scale").to(DEV) + n, k = ww.shape + sc = _e8m0(ws).repeat_interleave(32, 0)[:n].repeat_interleave(32, 1)[:, :k] + wkv = ww.float() * sc + return wkv, _st_tensor(p + "q_weight").to(DEV), _st_tensor(p + "k_weight").to(DEV) + + +def _golden_layer(case: str, layer: int): + """(h_in, key, value, golden gate, golden h_after_engram) with key/value rebuilt + from the checkpoint and value checked bit-exact against the golden dump.""" + li = _ENGRAM_LAYERS.index(layer) + hids = torch.load(os.path.join(_GOLDEN_DIR, case, "engram_hash.pt"), weights_only=True)[ + "hash_ids" + ][:, li, :] + t_len, lookups = hids.shape + uniq, inv = torch.unique(hids.reshape(-1), return_inverse=True) + p = f"layers.{layer}.engram.embed." + w = _st_rows(p + "weight", uniq.numpy()).view(torch.float8_e4m3fn).to(DEV) + s = _e8m0(_st_rows(p + "scale", uniq.numpy()).to(DEV)) + rows = (w.float().unflatten(-1, (-1, 32)) * s.unsqueeze(-1)).flatten(-2).to(torch.bfloat16) + gathered = rows[inv.to(DEV)].view(t_len, lookups * rows.shape[-1]) + wkv, _, _ = _layer_weights(layer) + with _ieee_fp32_matmul(): + kv = torch.matmul(_fp8_act_qdq(gathered), wkv.t()).to(torch.bfloat16) + key = kv[:, : HC * DIM].view(t_len, HC, DIM) + value = kv[:, HC * DIM :] + g = torch.load( + os.path.join(_GOLDEN_DIR, case, "layers", f"L{layer:02d}.pt"), + map_location="cpu", + mmap=True, + weights_only=True, + ) + gold_value = g["engram_value"].to(DEV) + assert torch.equal(value, gold_value), f"{case} L{layer}: rebuilt value != golden engram_value" + return g["h_in"].to(DEV), key, value, g["engram_gate"].to(DEV), g["h_after_engram"].to(DEV) + + +@pytest.mark.parametrize("layer", _ENGRAM_LAYERS) +@pytest.mark.parametrize("case", _GOLDEN_CASES) +def test_golden_layer(case, layer): + _skip_without_golden() + h_in, key, value, gold_gate, gold_out = _golden_layer(case, layer) + _, qw, kw = _layer_weights(layer) + t_len = h_in.shape[0] + gate = torch.empty(t_len, HC, device=DEV) + out = op(h_in, key, value, qw, kw, EPS, CLAMP, gate_out=gate) + ref_out, ref_gate = _ref(h_in, key, value, qw, kw) + # the torch reference rebuilt here is the golden computation (same ops, same inputs) + same = torch.equal(ref_gate, gold_gate) and torch.equal(ref_out, gold_out) + torch.testing.assert_close(ref_gate, gold_gate, rtol=0, atol=0) + torch.testing.assert_close(ref_out, gold_out, rtol=0, atol=0) + _check(out, gate, gold_out, gold_gate, h_in, value) + rms = h_in.float().square().mean(-1).sqrt() + print( + f"\n {case} L{layer:02d} T={t_len}: value bit-exact vs golden, " + f"reference rebuilt here bit-exact vs golden {same}; " + f"kernel vs golden {_stats(out, gate, gold_out, gold_gate)}; " + f"gate range [{gold_gate.min().item():.4f}, {gold_gate.max().item():.4f}]; per-copy RMS of h_in (token 1) " + f"{[round(v, 3) for v in rms[min(1, t_len - 1)].tolist()]}" + ) + _assert_controls_fail( + out, + gate, + h_in, + key, + value, + qw, + kw, + None, + ["joint_rms", "no_signed_sqrt"], + f"golden {case} L{layer}", + ) + # masking on real data: every third token shut, those rows must be h_in exactly + mask = torch.arange(t_len, device=DEV) % 3 != 1 + gm = torch.empty(t_len, HC, device=DEV) + om = op(h_in, key, value, qw, kw, EPS, CLAMP, token_mask=mask, gate_out=gm) + assert torch.equal(om[~mask], h_in[~mask]) and torch.equal(om[mask], out[mask]) + assert torch.equal(gm[~mask], torch.zeros_like(gm[~mask])) + with pytest.raises(AssertionError): + _check( + om, gm, gold_out, gold_gate, h_in, value + ) # the unmasked golden = the mask-ignored control + print( + f" control golden {case} L{layer} mask_ignored: rejected; {_stats(om, gm, gold_out, gold_gate)}" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/norm/test_modeling_v2_mhc_big_fuse.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/norm/test_modeling_v2_mhc_big_fuse.py new file mode 100644 index 000000000000..01b7d45e9d49 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/norm/test_modeling_v2_mhc_big_fuse.py @@ -0,0 +1,524 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the mhc_big_fuse catalog entry. + +The entry owns what is certified: `CELLS` is the list, `reference` and `compare` are the gate. +This file turns a cell's spec into tensors and keeps the checks a cell cannot carry: bit-identity +across launch shapes and between the full and post/comb-only modes, the `sinkhorn_repeat <= 0` +equivalence, the wrong-reference controls, M = 0, the op's own `num_splits` rejection, the +contiguity guard, CUDA-graph replay, and the DeepSeek-V4.1 call at every row count 1..32768 +(prefix sweep and its control, one compiled kernel per call, the CUDA-graph capture union). +""" + +import zlib + +import pytest +import torch +from torch.profiler import ProfilerActivity, profile + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.norm.mhc_big_fuse import ( # noqa: E501 + layer_input_error_ratio, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.norm.mhc_big_fuse import ( # noqa: E501 + mhc_big_fuse as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "mhc_big_fuse requires a CUDA device" + +DEV = "cuda" +HC = 4 +HIDDEN = 5120 +K = HC * HIDDEN +SENTINEL = 7.0 +# DeepSeek-V4.1: rms_norm_eps, hc_eps (pre and Sinkhorn), post = 2 * sigmoid, 20 Sinkhorn passes. +V41 = dict( + rms_eps=1e-20, + hc_pre_eps=1e-6, + hc_sinkhorn_eps=1e-6, + hc_post_mult_value=2.0, + sinkhorn_repeat=20, +) +# The fp32 gate: torch.testing.assert_close defaults for float32. +FP32_ATOL = 1e-5 + + +def _inputs(M, seed, num_splits=1, hidden=HIDDEN, zero_rows=0): + """Stream, split-K mixes and square sums, built the way the hc mixes GEMM produces them.""" + g = torch.Generator(device="cuda").manual_seed(seed) + k = HC * hidden + # Per-token stream magnitudes spread over four decades, as the real stream's are. + mag = torch.exp(torch.empty(M, 1, 1, device="cuda").uniform_(-3.0, 7.0, generator=g)) + residual = (torch.randn(M, HC, hidden, device="cuda", generator=g) * mag).to(torch.bfloat16) + if zero_rows: + residual[M - zero_rows :] = 0 # CUDA-graph padding rows + fn = torch.randn(24, k, device="cuda", generator=g) * 0.05 + xf = residual.reshape(M, k).float() + step = k // num_splits + parts = [slice(s * step, (s + 1) * step) for s in range(num_splits)] + y_acc = torch.stack([xf[:, p] @ fn[:, p].t() for p in parts]).contiguous() + r_acc = torch.stack([xf[:, p].square().sum(-1) for p in parts]).contiguous() + if num_splits == 1: + y_acc, r_acc = y_acc[0].contiguous(), r_acc[0].contiguous() + scale = torch.tensor([0.1, 0.05, 0.25], device="cuda") + base = torch.empty(24, device="cuda").uniform_(-5.0, 3.0, generator=g) + # Strongly negative comb diagonal (the checkpoint's hc bases reach -43.1): near-permutation + # rows, where Sinkhorn converges slowly enough for an iteration-count error to be visible. + base[8::5] = -40.0 + return y_acc, r_acc, residual, scale, base + + +def _extreme_inputs(M, seed): + """Logits of magnitude ~1e3-1e4 at unit RMS (mixes == y), a zero base.""" + g = torch.Generator(device=DEV).manual_seed(seed) + y = torch.randn(M, 24, device=DEV, generator=g) * 2e4 + r = torch.full((M,), float(K), device=DEV) + res = torch.randn(M, HC, HIDDEN, device=DEV, generator=g).to(torch.bfloat16) + scale = torch.tensor([0.1, 0.05, 0.25], device=DEV) + base = torch.zeros(24, device=DEV) + return y, r, res, scale, base + + +def _call_args(y_acc, r_acc, residual, scale, base, M, hidden, num_splits, block_size, pad=3, **kw): + """The op's arguments in order, with sentinel-filled outputs `pad` rows longer than M.""" + c = dict(V41) + c.update(kw) + post = torch.full((M + pad, HC), SENTINEL, device=DEV) + comb = torch.full((M + pad, HC, HC), SENTINEL, device=DEV) + li = torch.full((M + pad, max(hidden, 1)), SENTINEL, dtype=torch.bfloat16, device=DEV) + if hidden == 0: + li = torch.empty(M + pad, 0, dtype=torch.bfloat16, device=DEV) + return ( + y_acc, + r_acc, + residual, + scale, + base, + post, + comb, + li, + M, + K, + hidden, + c["rms_eps"], + c["hc_pre_eps"], + c["hc_sinkhorn_eps"], + c["hc_post_mult_value"], + c["sinkhorn_repeat"], + num_splits, + block_size, + ) + + +def _assert_untouched(out, M): + """Rows at and beyond M keep their sentinel.""" + post, comb, li = out + assert torch.all(post[M:] == SENTINEL) and torch.all(comb[M:] == SENTINEL) + assert torch.all(li[M:] == SENTINEL) + + +def _run(y_acc, r_acc, residual, scale, base, M, hidden, num_splits, block_size, pad=3, **kw): + """Call the op on sentinel-filled outputs with `pad` extra rows; check those stay untouched.""" + args = _call_args( + y_acc, r_acc, residual, scale, base, M, hidden, num_splits, block_size, pad, **kw + ) + out = op(*args) + torch.cuda.synchronize() + _assert_untouched(out, M) + return tuple(t[:M] for t in out) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + s = cell.spec + M, hidden = s["M"], s["hidden_size"] + seed = zlib.crc32(cell.why.encode()) + if s["kind"] == "extreme": + y, r, res, scale, base = _extreme_inputs(M, seed) + else: + y, r, res, scale, base = _inputs( + M, seed, num_splits=s["num_splits"], zero_rows=s["zero_rows"] + ) + if hidden == 0: + res = torch.empty(M, HC, 0, dtype=torch.bfloat16, device=DEV) + args = _call_args( + y, + r, + res, + scale, + base, + M, + hidden, + s["num_splits"], + s["block_size"], + sinkhorn_repeat=s["sinkhorn_repeat"], + hc_post_mult_value=s["hc_post_mult_value"], + ) + out = op(*args) + torch.cuda.synchronize() + assert all(o is a for o, a in zip(out, args[5:8])), ( + "the op writes in place and returns its outputs" + ) + _assert_untouched(out, M) + assert all(torch.isfinite(t[:M]).all() for t in out) + op.compare(out, op.reference(*args)) + + +@pytest.mark.parametrize("block_size", [0, 128, 256, 512]) +@pytest.mark.parametrize("M", [1, 64, 65, 300]) +def test_block_sizes_give_identical_bits(M, block_size): + """Every launch shape (incl. the 512-thread bulk-copy variant at M <= 64) gives identical bits.""" + y, r, res, scale, base = _inputs(M, seed=100 + M) + post, comb, li = _run(y, r, res, scale, base, M, HIDDEN, 1, block_size) + p0, c0, l0 = _run(y, r, res, scale, base, M, HIDDEN, 1, 0) + assert torch.equal(post, p0) and torch.equal(comb, c0) and torch.equal(li, l0) + + +@pytest.mark.parametrize("block_size", [0, 128, 256]) +@pytest.mark.parametrize("M", [1, 6, 4096]) +def test_post_comb_only_mode(M, block_size): + """hidden_size=0: post/comb bit-identical to the full call; residual unread, layer_input unwritten.""" + y, r, res, scale, base = _inputs(M, seed=200 + M) + p_full, c_full, _ = _run(y, r, res, scale, base, M, HIDDEN, 1, block_size) + empty_res = torch.empty(M, HC, 0, dtype=torch.bfloat16, device="cuda") + post, comb, li = _run(y, r, empty_res, scale, base, M, 0, 1, block_size) + assert li.shape == (M, 0) + assert torch.equal(post, p_full) and torch.equal(comb, c_full) + # A real stream and a real layer_input buffer: still not written when hidden_size=0. + keep = torch.full((M, HIDDEN), SENTINEL, dtype=torch.bfloat16, device="cuda") + post2 = torch.empty(M, HC, device="cuda") + comb2 = torch.empty(M, HC, HC, device="cuda") + op(y, r, res, scale, base, post2, comb2, keep, M, K, 0, *V41.values(), 1, block_size) + torch.cuda.synchronize() + assert torch.all(keep == SENTINEL) + assert torch.equal(post2, p_full) and torch.equal(comb2, c_full) + + +def test_sinkhorn_repeat_below_one_runs_one_pass(): + """sinkhorn_repeat <= 0 is accepted and computes exactly what sinkhorn_repeat=1 does.""" + M = 16 + y, r, res, scale, base = _inputs(M, seed=450) + one = _run(y, r, res, scale, base, M, HIDDEN, 1, 0, sinkhorn_repeat=1) + for repeat in (0, -3): + got = _run(y, r, res, scale, base, M, HIDDEN, 1, 0, sinkhorn_repeat=repeat) + assert all(torch.equal(a, b) for a, b in zip(got, one)) + + +def test_controls_fail(): + """Deliberately wrong references must land far outside the gate, or the gate certifies nothing. + + (a) 19 Sinkhorn passes instead of 20; (b) comb transposed; (c) layer_input taken to be the + stream collapsed with a LAGGED pre (the previous sublayer's mixes), which is what DeepSeek-V4.1 + actually feeds its sublayer -- the op's layer_input is the fresh-pre collapse instead. + Each wrong reference is put through the entry's own `compare`. + """ + M = 256 + y, r, res, scale, base = _inputs(M, seed=600) + args = _call_args(y, r, res, scale, base, M, HIDDEN, 1, 0, pad=0) + out = op(*args) + _, comb, li = out + good = op.reference(*args) + op.compare(out, good) + + wrong19 = op.reference(*args[:15], 19, *args[16:]) + margin19 = (comb - wrong19[1]).abs().max().item() / FP32_ATOL + with pytest.raises(AssertionError): + op.compare(out, wrong19) + + wrong_t = (good[0], good[1].transpose(-1, -2), good[2], good[3]) + margin_t = (comb - wrong_t[1]).abs().max().item() / FP32_ATOL + with pytest.raises(AssertionError): + op.compare(out, wrong_t) + + # The lagged collapse: the previous sublayer's mixes, this sublayer's stream. Its post/comb are + # handed back unchanged so that only the layer_input gate is under test. + y_prev, r_prev, _, _, _ = _inputs(M, seed=601) + lag = op.reference(y_prev, r_prev, *args[2:]) + with pytest.raises(AssertionError, match="layer_input outside the gate"): + op.compare((lag[0], lag[1], li), lag) + lag_ratio = layer_input_error_ratio(li, lag[2], lag[3]) + lag_rel = ((li.double() - lag[2]).norm() / lag[2].norm()).item() + + print( + f"control margins: 19-pass comb max|diff| = {margin19:.1f} x atol, " + f"transposed comb = {margin_t:.1f} x atol, lagged-pre layer_input max err/bound = " + f"{lag_ratio:.3e} (rel L2 {lag_rel:.3e})" + ) + assert margin19 > 10.0 and margin_t > 10.0 and lag_ratio > 10.0 + + +def test_m_zero_is_a_no_op(): + y, r, res, scale, base = _inputs(4, seed=700) + post, comb, li = _run(y, r, res, scale, base, 0, HIDDEN, 1, 0, pad=4) + assert post.numel() == 0 and comb.numel() == 0 and li.numel() == 0 + + +@pytest.mark.parametrize("num_splits", [0, 3, 32]) +def test_rejects_num_splits_outside_the_set(num_splits): + M = 8 + y, r, res, scale, base = _inputs(M, seed=800) + with pytest.raises(RuntimeError, match=r"only num_splits .* supported"): + _run(y, r, res, scale, base, M, HIDDEN, num_splits, 0) + + +@pytest.mark.parametrize( + "which", ["y_acc", "r_acc", "residual", "post_mix", "comb_mix", "layer_input"] +) +def test_is_valid_rejects_strided_views(which): + """The guard refuses a strided view of each of the six data tensors.""" + M = 8 + y, r, res, scale, base = _inputs(M, seed=900) + a = dict( + y_acc=y, + r_acc=r, + residual=res, + post_mix=torch.empty(M, HC, device="cuda"), + comb_mix=torch.empty(M, HC, HC, device="cuda"), + layer_input=torch.empty(M, HIDDEN, dtype=torch.bfloat16, device="cuda"), + ) + strided = dict( + y_acc=torch.zeros(M, 32, device="cuda")[:, :24], + r_acc=torch.zeros(M, 2, device="cuda")[:, 0], + residual=res.transpose(0, 1).contiguous().transpose(0, 1), + post_mix=torch.empty(HC, M, device="cuda").t(), + comb_mix=torch.empty(M, HC, HC, device="cuda").transpose(1, 2), + layer_input=torch.empty(M, 2 * HIDDEN, dtype=torch.bfloat16, device="cuda")[:, :HIDDEN], + ) + a[which] = strided[which] + with pytest.raises(AssertionError, match="must be contiguous"): + with validating(op): + op( + a["y_acc"], + a["r_acc"], + a["residual"], + scale, + base, + a["post_mix"], + a["comb_mix"], + a["layer_input"], + M, + K, + HIDDEN, + *V41.values(), + 1, + 0, + ) + + +@pytest.mark.parametrize("hidden", [HIDDEN, 0]) +def test_cuda_graph_capture(hidden): + """Captured once, replayed on new inputs copied into the static buffers.""" + M = 6 + y, r, res, scale, base = _inputs(M, seed=1000) + post = torch.empty(M, HC, device="cuda") + comb = torch.empty(M, HC, HC, device="cuda") + li = torch.empty(M, hidden, dtype=torch.bfloat16, device="cuda") + res_arg = res if hidden else torch.empty(M, HC, 0, dtype=torch.bfloat16, device="cuda") + args = (y, r, res_arg, scale, base, post, comb, li, M, K, hidden, *V41.values(), 1, 0) + + def call(): + op(*args) + + call() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + y2, r2, res2, _, _ = _inputs(M, seed=1001) + y.copy_(y2) + r.copy_(r2) + res.copy_(res2) + post.fill_(SENTINEL) + comb.fill_(SENTINEL) + graph.replay() + torch.cuda.synchronize() + op.compare((post, comb, li), op.reference(*args)) + + +# ─── The target's call at every row count 1..32768, and the capture union ── +# +# DeepSeek-V4.1's core target calls this op twice per layer in post/comb-only +# mode (hidden_size = 0, zero-size residual and layer_input) with +# block_size = 128, num_splits = 1 and M = the rows of the pass: the tokens the +# rank holds this step (any integer up to the engine's max_num_tokens, 16384 at +# most in the envelope) or, in layers 21-39 of a step that runs the decoder +# bounded-replay pass, that pass's rows (up to 32768). The DSpark drafter makes +# the same call at N = 5G rows. With block_size fixed by the caller, M only +# sizes the grid (one CTA per token, CTA t reads row t of y_acc / r_acc and +# writes row t of post_mix / comb_mix). The tests below check that on the GPU +# at every M from 1 to 32768. + +_N_MAX = 32768 +_TARGET_BS = 128 +# The row counts of the target's call that the entry's cells certify (its runtime-value dumps +# plus the tops of the step and decoder-pass ranges). +_TARGET_M = tuple(sorted({c.spec["M"] for c in op.CELLS if c.why.startswith("target call")})) +# The capture union per rank. G ranges over every CUDA-graph batch size an +# envelope engine captures: every G 1..32, plus TRT-LLM's padded default list +# for cuda_graph_config {enable_padding: true, max_batch_size: 256} +# (40..128 step 8, 192, 256); the padding-off list at 256 (1..32, 64, 128, 256) +# is a subset. A captured step is generation-only, so the trunk runs at G (a +# decode step) and 6G (the DSpark target-verify step, 1 + 5 draft tokens per +# request), and the drafter at 5G. +_GSET = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_CAPTURE_M = tuple(sorted({k * g for g in _GSET for k in (1, 5, 6)})) +assert len(_GSET) == 46 and len(_CAPTURE_M) == 114 and max(_CAPTURE_M) == 1536 + + +def _target_args(y, r, scale, base, post, comb, m): + """The core target's exact call: hidden_size 0, zero-size residual / layer_input, block_size 128.""" + no_res = torch.empty(m, HC, 0, dtype=torch.bfloat16, device="cuda") + no_li = torch.empty(m, 0, dtype=torch.bfloat16, device="cuda") + return (y, r, no_res, scale, base, post, comb, no_li, m, K, 0, *V41.values(), 1, _TARGET_BS) + + +def _target_call(y, r, scale, base, post, comb, m): + return op(*_target_args(y, r, scale, base, post, comb, m)) + + +def test_target_call_every_row_count_is_the_same_per_row_function() -> None: + """Every prefix M = 1..32768 of the target's call returns exactly the first M + rows of the 32768-row result, which is held to the fp32 reference. + + Outputs are NaN-filled before every call and the row after the last is a + guard row that must stay NaN: a skipped row or a write past row M - 1 is a + mismatch. + """ + y, r, res, scale, base = _inputs(_N_MAX, seed=4242, zero_rows=1) + del res + full_post = torch.full((_N_MAX, HC), float("nan"), device="cuda") + full_comb = torch.full((_N_MAX, HC, HC), float("nan"), device="cuda") + full_args = _target_args(y, r, scale, base, full_post, full_comb, _N_MAX) + op.compare(op(*full_args), op.reference(*full_args)) + fp, fc = full_post.view(torch.int32), full_comb.view(torch.int32) + post_buf = torch.empty(_N_MAX + 1, HC, device="cuda") + comb_buf = torch.empty(_N_MAX + 1, HC, HC, device="cuda") + mismatch = torch.zeros(_N_MAX + 1, dtype=torch.int64, device="cuda") + guard = torch.zeros(_N_MAX + 1, dtype=torch.int64, device="cuda") + for m in range(1, _N_MAX + 1): + post_buf[: m + 1].fill_(float("nan")) + comb_buf[: m + 1].fill_(float("nan")) + _target_call(y[:m], r[:m], scale, base, post_buf[:m], comb_buf[:m], m) + mismatch[m] = (post_buf[:m].view(torch.int32) != fp[:m]).sum() + ( + comb_buf[:m].view(torch.int32) != fc[:m] + ).sum() + guard[m] = (~post_buf[m].isnan()).sum() + (~comb_buf[m].isnan()).sum() + bad = mismatch.nonzero().flatten().tolist() + over = guard.nonzero().flatten().tolist() + assert not bad, ( + f"{len(bad)} row counts differ from the {_N_MAX}-row result's rows, first {bad[:10]}" + ) + assert not over, f"{len(over)} row counts wrote the guard row, first {over[:10]}" + + +def test_target_call_row_sweep_harness_sees_a_one_row_shift() -> None: + """Control for the sweep: its bitwise per-row check rejects a read one row off.""" + n = _N_MAX + y, r, _, scale, base = _inputs(n + 1, seed=4243) + full_post = torch.empty(n, HC, device="cuda") + full_comb = torch.empty(n, HC, HC, device="cuda") + _target_call(y[:n], r[:n], scale, base, full_post, full_comb, n) + for m in (1, 2, 16, 17, 2048, 2049, 16384, n): + post = torch.empty(m, HC, device="cuda") + comb = torch.empty(m, HC, HC, device="cuda") + _target_call( + y[1 : m + 1].contiguous(), r[1 : m + 1].contiguous(), scale, base, post, comb, m + ) + n_bad = int((comb.view(torch.int32) != full_comb.view(torch.int32)[:m]).sum().item()) + n_bad_p = int((post.view(torch.int32) != full_post.view(torch.int32)[:m]).sum().item()) + print( + f"[control row shift] M={m}: {n_bad} / {m * HC * HC} comb and {n_bad_p} / {m * HC} post elements differ" + ) + assert n_bad > m * HC * HC // 2 and n_bad_p > m * HC // 2 + + +_PROBE_M = tuple( + sorted( + set(_TARGET_M) + | set(_CAPTURE_M) + | { + 9, + 15, + 17, + 64, + 65, + 129, + 4096, + 8191, + 8192, + 8193, + 12288, + 16383, + 16385, + 24576, + 32767, + _N_MAX, + } + ) +) + + +def test_target_call_one_compiled_kernel_at_every_row_count() -> None: + """Probe: the target's call launches the same single kernel at every M (compiled ahead of time, no JIT).""" + y, r, _, scale, base = _inputs(_N_MAX, seed=4244) + post = torch.empty(_N_MAX, HC, device="cuda") + comb = torch.empty(_N_MAX, HC, HC, device="cuda") + names = set() + for m in _PROBE_M: + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _target_call(y[:m], r[:m], scale, base, post[:m], comb[:m], m) + torch.cuda.synchronize() + launched = [ + e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA + ] + assert len(launched) == 1, (m, launched) + names.add(launched[0]) + print(f"{len(_PROBE_M)} row counts probed; kernels launched: {sorted(names)}") + assert len(names) == 1 and "mhcBigFuseKernel<1, 128, false, false, false>" in next( + iter(names) + ), names + + +@pytest.mark.parametrize("M", _CAPTURE_M) +def test_target_call_capture_union_replay_is_bitwise_eager(M: int) -> None: + """The target's call captured at M after an eager call at M, replayed on new + data: bitwise equal to an eager call on that data. + + The replay on the capture-time data must equal the eager call before the + capture and differ from the replay on new data, so a replay that reused + stale inputs would fail. Rows past the first carry one all-zero padding row, + as a padded CUDA-graph batch does. + """ + y, r, _, scale, base = _inputs(M, seed=6000 + M, zero_rows=1 if M > 1 else 0) + post = torch.full((M, HC), SENTINEL, device="cuda") + comb = torch.full((M, HC, HC), SENTINEL, device="cuda") + _target_call(y, r, scale, base, post, comb, M) # the eager call at this shape + torch.cuda.synchronize() + post0, comb0 = post.clone(), comb.clone() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _target_call(y, r, scale, base, post, comb, M) + post.fill_(SENTINEL) + comb.fill_(SENTINEL) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(post, post0) and torch.equal(comb, comb0), M + comb_old = comb.clone() + y2, r2, _, _, _ = _inputs(M, seed=7000 + M, zero_rows=1 if M > 1 else 0) + y.copy_(y2) + r.copy_(r2) + post.fill_(SENTINEL) + comb.fill_(SENTINEL) + graph.replay() + torch.cuda.synchronize() + post_e = torch.full((M, HC), SENTINEL, device="cuda") + comb_e = torch.full((M, HC, HC), SENTINEL, device="cuda") + _target_call(y2, r2, scale, base, post_e, comb_e, M) + torch.cuda.synchronize() + assert torch.equal(post, post_e) and torch.equal(comb, comb_e), M + assert not torch.equal(comb, comb_old), "the replay did not read the new mixes" + args = _target_args(y2, r2, scale, base, post, comb, M) + op.compare((post, comb, args[7]), op.reference(*args)) # the replay's output diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/norm/test_modeling_v2_mhc_hc_head_apply.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/norm/test_modeling_v2_mhc_hc_head_apply.py new file mode 100644 index 000000000000..a40010cd7673 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/norm/test_modeling_v2_mhc_hc_head_apply.py @@ -0,0 +1,681 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the mhc_hc_head_apply catalog entry. + +What is certified is the cross-stream use: ``mixes`` and ``sqrsum`` are the +statistics of one stream (the *producer*) and ``x`` is a different stream (the +*consumer*). That is DeepSeek-V4.1's lagged ``hc_pre``: the attention input of +layer L is collapsed with the ``pre`` computed on the stream before layer L-1's +FFN, the FFN input with the ``pre`` computed on the stream before this layer's +attention, and the final-norm input with layer 39's FFN ``pre``. + +Two sources of truth: + +* synthetic producer/consumer pairs drawn independently, so nothing of ``x`` + can hide inside ``pre`` -- the entry's ``reference`` and ``compare`` are the + gate, and its ``CELLS`` the certified configurations; +* the golden layer dumps of the checkpoint (``golden/``, a pure-torch port of + the vendor reference forward), with the real ``hc_*`` parameters read from + the checkpoint. These cases need two external directories and skip, naming + them, when either is absent. + +This file also keeps the checks a cell cannot carry: the deliberately wrong +variants (fresh instead of lagged mixes, the consuming sublayer's parameters, +the consumer's RMS statistic, no ``+eps``, copies rolled) rejected by the same +gate, the parameter-slicing and in-place bounds, CUDA-graph replay, the guard +and the op's dtype rejections, and the DeepSeek-V4.1 call at every row count +1..32768 (prefix sweep and its control, one compiled kernel per call, the +CUDA-graph capture union). +""" + +import json +import os +import zlib + +import pytest +import torch +from torch.profiler import ProfilerActivity, profile +from utils.llm_data import llm_models_root + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.norm.mhc_hc_head_apply import ( # noqa: E501 + collapse_error_ratio, + pre_weights, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.norm.mhc_hc_head_apply import ( # noqa: E501 + mhc_hc_head_apply as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "mhc_hc_head_apply requires a CUDA device" + +DEV = "cuda" +MULT = 4 +HIDDEN = 5120 +K = MULT * HIDDEN +NORM_EPS = 1e-20 # DeepSeek-V4.1 rms_norm_eps, used by hc_mixes +HC_EPS = 1e-6 # DeepSeek-V4.1 hc_eps +# The op is not bit-identical to torch: it is built with --use_fast_math, so +# sigmoid's exp and reciprocal and the rsqrt are the approximate hardware +# versions, and its 4-term fp32 sum is a fused-multiply-add chain. pre then +# differs from torch's by about one fp32 ulp, which moves a bf16 rounding +# decision on a small fraction of elements. Measured on sm_103: between 0.99988 +# and 1.0 of elements bit-identical, synthetic and golden cases alike. The hard +# gate below sits under that and far above every wrong variant the tests reject, +# the closest of which kept 0.40 of elements bit-identical. (Dropping +eps at +# the production hc_eps kept 0.998: under this gate, but too near it to call +# that discrimination -- see test_eps_is_added_after_sigmoid.) +MIN_BITEXACT = 0.999 + +# Reference activations dumped from the checkpoint's own inference code; no default +# location, so the cases that compare against them skip unless this is set. +_GOLDEN_DIR = os.environ.get("MODELING_V2_DSV41_GOLDEN_DIR", "") +# The checkpoint resolves under LLM_MODELS_ROOT, as the whole-model gates' does; +# MODELING_V2_DSV41_CKPT_DIR overrides it. A case that needs it skips when it is absent. +_MODELS_ROOT = llm_models_root() +_CKPT_DIR = os.environ.get("MODELING_V2_DSV41_CKPT_DIR") or ( + str(_MODELS_ROOT / "DeepSeek-V4.1-Flash-NVFP4") if _MODELS_ROOT else "" +) +_ENGRAM_LAYERS = (1, 14) + + +def _producer_stats(stream: torch.Tensor, fn: torch.Tensor): + """What the hc mixes GEMM hands the collapse: raw mixes (no rsqrt applied) + for the first `MULT` rows of `fn`, as a contiguous [T, MULT] copy, and the + sum of squares of the flattened stream, both fp32.""" + flat = stream.reshape(stream.shape[0], -1).float() + with true_fp32_matmul(): + y = flat @ fn.float().t() + return y[:, :MULT].contiguous(), flat.square().sum(-1) + + +def _pre(mixes, sqrsum, scale, base, eps=HC_EPS): + return pre_weights(mixes, sqrsum, scale, base, MULT, K, NORM_EPS, eps) + + +def _ref_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + xf = x.float() + return (weight * (xf * torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps))).to(x.dtype) + + +def _args(mixes, sqrsum, x, scale, base, *, k=K, norm_eps=NORM_EPS, eps=HC_EPS): + """The op's arguments in order, collapsing every row of `x` into a NaN-filled `out`.""" + m = x.shape[0] + out = torch.full((m, HIDDEN), float("nan"), dtype=torch.bfloat16, device=DEV) + return (mixes, sqrsum, x, out, scale, base, m, MULT, HIDDEN, k, norm_eps, eps) + + +def _run(mixes, sqrsum, x, scale, base, **kw): + return op(*_args(mixes, sqrsum, x, scale, base, **kw)) + + +def _bitexact(a: torch.Tensor, b: torch.Tensor) -> float: + return (a == b).float().mean().item() + + +def _row_rel(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + a, b = a.float(), b.float() + return (a - b).norm(dim=-1) / b.norm(dim=-1).clamp_min(1e-30) + + +def _assert_matches(out: torch.Tensor, ref: torch.Tensor) -> float: + """The golden comparisons: default bf16 assert_close plus the bit-exact floor.""" + torch.testing.assert_close(out, ref) + frac = _bitexact(out, ref) + assert frac >= MIN_BITEXACT, f"bit-exact fraction {frac} < {MIN_BITEXACT}" + return frac + + +def _assert_rejected(tag: str, out: torch.Tensor, ref: tuple) -> None: + """A deliberately wrong pairing of output and reference must fail the entry's gate.""" + rel = _row_rel(out, ref[0]) + print( + f" control {tag}: max|d|={(out.float() - ref[0].float()).abs().max().item():.4g} " + f"row-rel median={rel.median().item():.4g} max={rel.max().item():.4g} " + f"bit-exact={_bitexact(out, ref[0]):.4f}" + ) + with pytest.raises(AssertionError): + op.compare(out, ref) + + +def _hc_params(seed: int, saturated: bool = False): + """Checkpoint-shaped parameters: fn [24, K], scale [3], base [24] (fp32).""" + g = torch.Generator(device="cuda").manual_seed(seed) + fn = torch.randn(24, K, device="cuda", generator=g) * (1.0 / K**0.5) + scale = torch.rand(3, device="cuda", generator=g) * 1.5 + 0.25 + base = torch.randn(24, device="cuda", generator=g) + if saturated: + # sigmoid pinned near 0 and 1 on two copies: pre ~ eps and ~ 1 + eps + base[0], base[3] = -40.0, 40.0 + return fn, scale, base + + +def _stream(m: int, seed: int, row_scale: float) -> torch.Tensor: + """bf16 [m, MULT, HIDDEN] with per-row magnitudes spread over two decades; + for m >= 3 row 1 is all zeros (a padded token: sqrsum 0) and row 2 is an + attention-sink-like row with a ~2e6 stream norm.""" + g = torch.Generator(device="cuda").manual_seed(seed) + s = torch.randn(m, MULT, HIDDEN, device="cuda", generator=g) + s *= row_scale * torch.logspace(-1, 1, m, device="cuda").view(m, 1, 1) + if m >= 3: + s[1] = 0.0 + s[2] *= 2.0e6 / s[2].norm() + return s.to(torch.bfloat16) + + +def _lagged_inputs(m: int, seed: int, saturated: bool = False): + fn, scale, base = _hc_params(seed=seed, saturated=saturated) + producer = _stream(m, seed=seed + 1, row_scale=3.0) + consumer = _stream(m, seed=seed + 2, row_scale=1.0) + mixes, sqrsum = _producer_stats(producer, fn) + return mixes, sqrsum, consumer, scale, base + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + s = cell.spec + mixes, sqrsum, x, scale, base = _lagged_inputs( + s["M"], seed=zlib.crc32(cell.why.encode()) % 2**31, saturated=s["saturated"] + ) + if s["saturated"]: + pre = _pre(mixes, sqrsum, scale, base) + assert pre[:, 0].max().item() < 2 * HC_EPS and pre[:, 3].min().item() > 1.0 - 1e-6 + args = _args(mixes, sqrsum, x, scale, base, eps=s["eps"]) + out = op(*args) + assert out is args[3], "the op writes out in place and returns it" + op.compare(out, op.reference(*args)) + + +# ─── synthetic: independent producer and consumer streams ────────────────── + + +@pytest.mark.parametrize("m", [5, 6, 32, 1000, 8192]) +def test_lagged_collapse_rejects_fresh_wiring(m: int) -> None: + """The gate the op passes rejects the V4 (fresh) wirings of the same collapse.""" + mixes, sqrsum, consumer, scale, base = _lagged_inputs(m, seed=10 * m) + args = _args(mixes, sqrsum, consumer, scale, base) + out = op(*args) + op.compare(out, op.reference(*args)) + fn, _, _ = _hc_params(seed=10 * m) + fm, fs = _producer_stats(consumer, fn) + # fresh: the consumer's own mixes and statistic (V4 semantics) + _assert_rejected("fresh", out, op.reference(fm, fs, *args[2:])) + # the producer's mixes normalized by the consumer's RMS statistic + _assert_rejected("rstd-from-x", out, op.reference(mixes, fs, *args[2:])) + + +def test_eps_is_added_after_sigmoid() -> None: + m, big_eps = 64, 0.25 + mixes, sqrsum, consumer, scale, base = _lagged_inputs(m, seed=3) + args = _args(mixes, sqrsum, consumer, scale, base, eps=big_eps) + out = op(*args) + op.compare(out, op.reference(*args)) + _assert_rejected("no-eps", out, op.reference(*args[:11], 0.0)) + # At the production hc_eps the same omission is below bf16 resolution, so no + # output-level gate can see it; record how close it lands. + args6 = _args(mixes, sqrsum, consumer, scale, base) + out6 = op(*args6) + ref6_noeps = op.reference(*args6[:11], 0.0)[0] + print( + f" at eps={HC_EPS}: no-eps variant max|d|=" + f"{(out6.float() - ref6_noeps.float()).abs().max().item():.4g} " + f"bit-exact={_bitexact(out6, ref6_noeps):.6f}" + ) + + +def test_reads_scale0_and_first_mult_bases_only() -> None: + m = 32 + fn, scale, base = _hc_params(seed=5) + producer = _stream(m, seed=15, row_scale=3.0) + consumer = _stream(m, seed=25, row_scale=1.0) + mixes, sqrsum = _producer_stats(producer, fn) + full = _run(mixes, sqrsum, consumer, scale, base) + trimmed = _run(mixes, sqrsum, consumer, scale[:1].contiguous(), base[:MULT].contiguous()) + assert torch.equal(full, trimmed) + + +def test_in_place_bounds() -> None: + m, rows = 40, 48 + fn, scale, base = _hc_params(seed=9) + producer = _stream(rows, seed=19, row_scale=3.0) + consumer = _stream(rows, seed=29, row_scale=1.0) + mixes, sqrsum = _producer_stats(producer, fn) + snap = [t.clone() for t in (mixes, sqrsum, consumer, scale, base)] + out = torch.full((rows, HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + args = (mixes, sqrsum, consumer, out, scale, base, m, MULT, HIDDEN, K, NORM_EPS, HC_EPS) + op(*args) + op.compare(out, op.reference(*args)) + assert bool((out[m:] == 7.0).all()), "rows >= M of out were written" + for before, after in zip(snap, (mixes, sqrsum, consumer, scale, base)): + assert torch.equal(before, after), "an input was modified" + # M = 0 launches an empty grid: out untouched, and later work in the process is unaffected + op(mixes, sqrsum, consumer, out, scale, base, 0, MULT, HIDDEN, K, NORM_EPS, HC_EPS) + torch.cuda.synchronize() + assert bool((out[m:] == 7.0).all()) + + +def test_cuda_graph_replay_reads_current_inputs() -> None: + m = 8 + fn, scale, base = _hc_params(seed=11) + mixes = torch.empty(m, MULT, device="cuda") + sqrsum = torch.empty(m, device="cuda") + x = torch.empty(m, MULT, HIDDEN, dtype=torch.bfloat16, device="cuda") + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda") + + def fill(step: int) -> None: + pm, ps = _producer_stats(_stream(m, seed=100 + step, row_scale=3.0), fn) + mixes.copy_(pm) + sqrsum.copy_(ps) + x.copy_(_stream(m, seed=200 + step, row_scale=1.0)) + + fill(0) + args = (mixes, sqrsum, x, out, scale, base, m, MULT, HIDDEN, K, NORM_EPS, HC_EPS) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(*args) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(*args) + for step in range(1, 4): + fill(step) + graph.replay() + torch.cuda.synchronize() + op.compare(out, op.reference(*args)) + + +def test_is_valid_rejects_layouts_the_op_miscomputes_silently() -> None: + m = 16 + fn, scale, base = _hc_params(seed=13) + producer = _stream(m, seed=33, row_scale=3.0) + consumer = _stream(m, seed=43, row_scale=1.0) + flat = producer.reshape(m, -1).float() + with true_fp32_matmul(): + y24 = flat @ fn.t() # the GEMM's full [M, 24] output + sqrsum = flat.square().sum(-1) + good_args = _args(y24[:, :MULT].contiguous(), sqrsum, consumer, scale, base) + good_ref = op.reference(*good_args) + op.compare(op(*good_args), good_ref) + + # Evidence for the guard: outside validation the op takes the natural slice y24[:, :4] + # (row stride 24) without complaint and computes something else. + raw = op(*_args(y24[:, :MULT], sqrsum, consumer, scale, base)) + _assert_rejected("op-on-strided-mixes", raw, good_ref) + + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda") + wide = torch.empty(m, 2 * HIDDEN, dtype=torch.bfloat16, device="cuda") + sq2 = torch.stack([sqrsum, sqrsum], dim=1) + x_t = consumer.transpose(0, 1).contiguous().transpose(0, 1) # same values, non-contiguous + bad_calls = { + "strided mixes": (y24[:, :MULT], sqrsum, consumer, out, m), + "[M,24] mixes": (y24, sqrsum, consumer, out, m), + "strided sqrsum": (y24[:, :MULT].contiguous(), sq2[:, 0], consumer, out, m), + "non-contiguous x": (y24[:, :MULT].contiguous(), sqrsum, x_t, out, m), + "strided out": (y24[:, :MULT].contiguous(), sqrsum, consumer, wide[:, :HIDDEN], m), + "M > out rows": (y24[:, :MULT].contiguous(), sqrsum, consumer, out[: m - 2], m), + } + for name, (mx, sq, xx, oo, mm) in bad_calls.items(): + with pytest.raises(AssertionError, match="must be a contiguous"): + with validating(op): + op(mx, sq, xx, oo, scale, base, mm, MULT, HIDDEN, K, NORM_EPS, HC_EPS) + print(f" is_valid rejected: {name}") + + +def test_wrong_dtype_raises_before_launch() -> None: + m = 4 + fn, scale, base = _hc_params(seed=17) + mixes, sqrsum = _producer_stats(_stream(m, seed=37, row_scale=3.0), fn) + x = _stream(m, seed=47, row_scale=1.0) + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda") + cases = { + "mixes bf16": ( + dict(mixes=mixes.bfloat16()), + "expected scalar type Float but found BFloat16", + ), + "sqrsum bf16": ( + dict(sqrsum=sqrsum.bfloat16()), + "expected scalar type Float but found BFloat16", + ), + "scale bf16": ( + dict(scale=scale.bfloat16()), + "expected scalar type Float but found BFloat16", + ), + "x fp32": (dict(x=x.float()), "expected scalar type BFloat16 but found Float"), + "out fp32": (dict(out=out.float()), "expected scalar type BFloat16 but found Float"), + } + for name, (override, text) in cases.items(): + a = dict(mixes=mixes, sqrsum=sqrsum, x=x, out=out, scale=scale, base=base) + a.update(override) + with pytest.raises(RuntimeError, match=text): + op( + a["mixes"], + a["sqrsum"], + a["x"], + a["out"], + a["scale"], + a["base"], + m, + MULT, + HIDDEN, + K, + NORM_EPS, + HC_EPS, + ) + print(f" {name}: raised {text!r}") + torch.cuda.synchronize() + + +# ─── golden: the checkpoint's real streams and hc parameters ─────────────── + +_SITES = { + # attention collapse of layer 1: pre from the stream before layer 0's FFN; x after layer 1's Engram + "attn@1": ("attn", 1), + "attn@20": ("attn", 20), + # FFN collapse of layer 14: pre from the stream entering layer 14's attention (after its Engram) + "ffn@14": ("ffn", 14), + "ffn@39": ("ffn", 39), + # final-norm input: pre from the stream before layer 39's FFN; x = the last layer's output + "final": ("final", 39), +} +_GOLDEN_CASES = [(c, s) for c in ("capital_fr", "long_retrieval") for s in _SITES] + [ + ("boot_passkey", "final") +] + + +class _Ckpt: + def __init__(self, root: str): + from safetensors import safe_open + + self._root = root + self._open = safe_open + with open(os.path.join(root, "model.safetensors.index.json")) as f: + self._map = json.load(f)["weight_map"] + + def get(self, name: str) -> torch.Tensor: + with self._open( + os.path.join(self._root, self._map[name]), framework="pt", device="cpu" + ) as f: + return f.get_tensor(name).cuda() + + def hc(self, layer: int, which: str): + p = f"layers.{layer}.hc_{which}_" + return self.get(p + "fn"), self.get(p + "scale"), self.get(p + "base") + + +def _golden_layer(case: str, layer: int) -> dict: + return torch.load( + os.path.join(_GOLDEN_DIR, case, "layers", f"L{layer:02d}.pt"), + map_location="cpu", + mmap=True, + weights_only=True, + ) + + +def _block_input(d: dict, layer: int) -> torch.Tensor: + return d["h_after_engram"] if layer in _ENGRAM_LAYERS else d["h_in"] + + +@pytest.mark.parametrize("case,site", _GOLDEN_CASES, ids=[f"{c}-{s}" for c, s in _GOLDEN_CASES]) +def test_golden_lagged_collapse(case: str, site: str) -> None: + if not os.path.isdir(os.path.join(_GOLDEN_DIR, case)) or not os.path.isfile( + os.path.join(_CKPT_DIR, "model.safetensors.index.json") + ): + pytest.skip(f"golden dumps ({_GOLDEN_DIR}) or checkpoint ({_CKPT_DIR}) not present") + ck = _Ckpt(_CKPT_DIR) + kind, layer = _SITES[site] + cur = _golden_layer(case, layer) + if kind == "attn": + prev = _golden_layer(case, layer - 1) + producer, x = prev["h_mid"], _block_input(cur, layer) + golden_pre = cur["pre_mix_in"] + lag_params, own_params = ck.hc(layer - 1, "ffn"), ck.hc(layer, "attn") + golden_after = cur["attn_in"] + norm_w = ck.get(f"layers.{layer}.attn_norm.weight") + elif kind == "ffn": + producer, x = _block_input(cur, layer), cur["h_mid"] + golden_pre = cur["attn_hc_pre"] + lag_params, own_params = ck.hc(layer, "attn"), ck.hc(layer, "ffn") + golden_after = cur["ffn_in"] + norm_w = ck.get(f"layers.{layer}.ffn_norm.weight") + else: + fin = torch.load( + os.path.join(_GOLDEN_DIR, case, "final.pt"), + map_location="cpu", + mmap=True, + weights_only=True, + ) + producer, x = cur["h_mid"], fin["h_last_out"] + golden_pre = cur["pre_mix_out"] + # there are no hc_head parameters: the collapse reuses layer 39's FFN pre + lag_params, own_params = ck.hc(layer, "ffn"), ck.hc(layer, "attn") + golden_after = fin["h_final"] + norm_w = None + producer, x, golden_pre, golden_after = ( + t.cuda() for t in (producer, x, golden_pre, golden_after) + ) + fn, scale, base = lag_params + assert fn.dtype == scale.dtype == base.dtype == torch.float32 + assert tuple(fn.shape) == (24, K) and tuple(scale.shape) == (3,) and tuple(base.shape) == (24,) + + mixes, sqrsum = _producer_stats(producer, fn) + # the recipe reproduces the pre the reference carried to this site + torch.testing.assert_close(_pre(mixes, sqrsum, scale, base), golden_pre) + + args = _args(mixes, sqrsum, x, scale, base) + out = op(*args) + ref = op.reference(*args) + op.compare(out, ref) + expected = ref[0] # the torch collapse with the recipe's pre, checked equal to the golden pre + frac = _assert_matches(out, expected) + rel = _row_rel(out, expected) + msg = ( + f"\n{case} {site} T={x.shape[0]}: bit-exact vs hc_pre(x, pre) {frac:.6f}, " + f"row-rel max {rel.max().item():.3g}" + ) + if norm_w is None: + frac_g = _assert_matches(out, golden_after) # golden h_final itself + msg += f"; bit-exact vs golden h_final {frac_g:.6f}" + else: + normed = _ref_rmsnorm(out, norm_w, NORM_EPS) + torch.testing.assert_close(normed, golden_after) + msg += ( + f"; rmsnorm(out) vs golden {'attn_in' if kind == 'attn' else 'ffn_in'} " + f"bit-exact {_bitexact(normed, golden_after):.6f}" + ) + print(msg) + + # controls: each wires one thing the way a V4 port would, and must be rejected + fm, fs = _producer_stats(x, own_params[0] if kind != "final" else fn) + fresh_scale, fresh_base = (own_params[1], own_params[2]) if kind != "final" else (scale, base) + _assert_rejected("fresh-mixes(op)", _run(fm, fs, x, fresh_scale, fresh_base), ref) + _assert_rejected( + "consumer-params(op)", _run(mixes, sqrsum, x, own_params[1], own_params[2]), ref + ) + + +# ─── The target's call at every row count 1..32768, and the capture union ── +# +# DeepSeek-V4.1's core target collapses the stream with this op twice per layer +# and once before the final norm, at M = the rows of the pass: the tokens the +# rank holds this step (any integer up to the engine's max_num_tokens, 16384 at +# most in the envelope) or, in layers 21-39 and the final collapse of a step +# that runs the decoder bounded-replay pass, that pass's rows (up to 32768). +# The DSpark drafter makes the same call at N = 5G rows. The launch is one CTA +# of 256 threads per token (CTA t reads row t of mixes, sqrsum and x and writes +# row t of out), so M only sizes the grid. The tests below check that on the +# GPU at every M from 1 to 32768. + +_N_MAX = 32768 +# The row counts of the target's call that the entry's cells certify (its runtime-value dumps +# plus the tops of the step and decoder-pass ranges). +_TARGET_M = tuple(sorted({c.spec["M"] for c in op.CELLS if c.why.startswith("target call")})) +# The capture union per rank. G ranges over every CUDA-graph batch size an +# envelope engine captures: every G 1..32, plus TRT-LLM's padded default list +# for cuda_graph_config {enable_padding: true, max_batch_size: 256} +# (40..128 step 8, 192, 256); the padding-off list at 256 (1..32, 64, 128, 256) +# is a subset. A captured step is generation-only, so the trunk runs at G (a +# decode step) and 6G (the DSpark target-verify step, 1 + 5 draft tokens per +# request), and the drafter at 5G. +_GSET = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_CAPTURE_M = tuple(sorted({k * g for g in _GSET for k in (1, 5, 6)})) +assert len(_GSET) == 46 and len(_CAPTURE_M) == 114 and max(_CAPTURE_M) == 1536 + + +def _target_args(mixes, sqrsum, x, out, scale, base, m): + return (mixes, sqrsum, x, out, scale, base, m, MULT, HIDDEN, K, NORM_EPS, HC_EPS) + + +def _target_call(mixes, sqrsum, x, out, scale, base, m): + return op(*_target_args(mixes, sqrsum, x, out, scale, base, m)) + + +def test_target_call_every_row_count_is_the_same_per_row_function() -> None: + """Every prefix M = 1..32768 of the target's call returns exactly the first M + rows of the 32768-row result, which is held to the reference. + + The output buffer is NaN-filled before every call and the row after the + last is a guard row that must stay NaN: a skipped row or a write past row + M - 1 is a mismatch. + """ + mixes, sqrsum, x, scale, base = _lagged_inputs(_N_MAX, seed=4242) + full = torch.full((_N_MAX, HIDDEN), float("nan"), dtype=torch.bfloat16, device="cuda") + full_args = _target_args(mixes, sqrsum, x, full, scale, base, _N_MAX) + op.compare(op(*full_args), op.reference(*full_args)) + fb = full.view(torch.int16) + buf = torch.empty(_N_MAX + 1, HIDDEN, dtype=torch.bfloat16, device="cuda") + mismatch = torch.zeros(_N_MAX + 1, dtype=torch.int64, device="cuda") + guard = torch.zeros(_N_MAX + 1, dtype=torch.int64, device="cuda") + differs = torch.empty( + _N_MAX, HIDDEN, dtype=torch.bool, device="cuda" + ) # one buffer; .any() keeps the reduction free of a widened copy + for m in range(1, _N_MAX + 1): + buf[: m + 1].fill_(float("nan")) + _target_call(mixes[:m], sqrsum[:m], x[:m], buf[:m], scale, base, m) + mismatch[m] = torch.ne(buf[:m].view(torch.int16), fb[:m], out=differs[:m]).any() + guard[m] = (~buf[m].isnan()).sum() + bad = mismatch.nonzero().flatten().tolist() + over = guard.nonzero().flatten().tolist() + assert not bad, ( + f"{len(bad)} row counts differ from the {_N_MAX}-row result's rows, first {bad[:10]}" + ) + assert not over, f"{len(over)} row counts wrote the guard row, first {over[:10]}" + + +def test_target_call_row_sweep_harness_sees_a_one_row_shift() -> None: + """Control for the sweep: its bitwise per-row check rejects a read one row off.""" + n = _N_MAX + mixes, sqrsum, x, scale, base = _lagged_inputs(n + 1, seed=4243) + full = torch.empty(n, HIDDEN, dtype=torch.bfloat16, device="cuda") + _target_call(mixes[:n], sqrsum[:n], x[:n], full, scale, base, n) + for m in (1, 2, 7, 1024, 2048, 16384, n): + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda") + _target_call(mixes[1 : m + 1], sqrsum[1 : m + 1], x[1 : m + 1], out, scale, base, m) + n_bad = int((out.view(torch.int16) != full.view(torch.int16)[:m]).sum().item()) + print(f"[control row shift] M={m}: {n_bad} / {m * HIDDEN} elements differ") + assert n_bad > m * HIDDEN // 2 + + +def test_target_call_term_gate_rejects_wrong_variants() -> None: + """Discrimination of the entry's gate: the op passes it, and each deliberately + wrong collapse lands far outside it. + + (a) fresh: the consumer's own mixes and statistic; (b) the producer's mixes + with the consumer's RMS statistic; (c) another sublayer's scale / base; + (d) pre applied to the wrong copies (rolled by one). + """ + m = 2048 + fn, scale, base = _hc_params(seed=4245) + producer = _stream(m, seed=4246, row_scale=3.0) + x = _stream(m, seed=4247, row_scale=1.0) + mixes, sqrsum = _producer_stats(producer, fn) + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda") + args = _target_args(mixes, sqrsum, x, out, scale, base, m) + op(*args) + ref = op.reference(*args) + op.compare(out, ref) + fm, fs = _producer_stats(x, fn) + _, scale_o, base_o = _hc_params(seed=4248) + rest = args[6:] + wrong = { + "fresh": op.reference(fm, fs, x, out, scale, base, *rest), + "rstd-from-x": op.reference(mixes, fs, x, out, scale, base, *rest), + "other-sublayer-params": op.reference(mixes, sqrsum, x, out, scale_o, base_o, *rest), + # copy j weighted with pre[j - 1]: the mixes and bases rolled by one copy + "copies-rolled": op.reference( + mixes.roll(1, dims=1).contiguous(), sqrsum, x, out, scale, base[:MULT].roll(1), *rest + ), + } + print(f"\ncorrect: max err / bound {collapse_error_ratio(out, ref[1], ref[2]):.3f}") + for name, ref_w in wrong.items(): + r_w = collapse_error_ratio(out, ref_w[1], ref_w[2]) + print(f" control {name}: max err / bound {r_w:.4g}") + assert r_w > 100.0, (name, r_w) + + +_PROBE_M = tuple( + sorted( + set(_TARGET_M) + | set(_CAPTURE_M) + | {9, 64, 65, 129, 4096, 8191, 8192, 8193, 12288, 16383, 16385, 24576, 32767, _N_MAX} + ) +) + + +def test_target_call_one_compiled_kernel_at_every_row_count() -> None: + """Probe: the target's call launches the same single kernel at every M (compiled ahead of time, no JIT).""" + mixes, sqrsum, x, scale, base = _lagged_inputs(_N_MAX, seed=4244) + out = torch.empty(_N_MAX, HIDDEN, dtype=torch.bfloat16, device="cuda") + names = set() + for m in _PROBE_M: + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _target_call(mixes[:m], sqrsum[:m], x[:m], out[:m], scale, base, m) + torch.cuda.synchronize() + launched = [ + e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA + ] + assert len(launched) == 1, (m, launched) + names.add(launched[0]) + print(f"{len(_PROBE_M)} row counts probed; kernels launched: {sorted(names)}") + assert len(names) == 1 and "mhcHcHeadApplyKernel" in next(iter(names)), names + + +@pytest.mark.parametrize("m", _CAPTURE_M) +def test_target_call_capture_union_replay_is_bitwise_eager(m: int) -> None: + """The target's call captured at M after an eager call at M, replayed on new + data: bitwise equal to an eager call on that data. + + The replay on the capture-time data must equal the eager call before the + capture and differ from the replay on new data, so a replay that reused + stale inputs would fail. + """ + mixes, sqrsum, x, scale, base = _lagged_inputs(m, seed=6000 + m) + out = torch.full((m, HIDDEN), float("nan"), dtype=torch.bfloat16, device="cuda") + _target_call(mixes, sqrsum, x, out, scale, base, m) # the eager call at this shape + torch.cuda.synchronize() + out0 = out.clone() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _target_call(mixes, sqrsum, x, out, scale, base, m) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, out0), m + out_old = out.clone() + mixes2, sqrsum2, x2, _, _ = _lagged_inputs(m, seed=7000 + m) + mixes.copy_(mixes2) + sqrsum.copy_(sqrsum2) + x.copy_(x2) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + eager = torch.full((m, HIDDEN), float("nan"), dtype=torch.bfloat16, device="cuda") + _target_call(mixes2, sqrsum2, x2, eager, scale, base, m) + torch.cuda.synchronize() + assert torch.equal(out, eager), m + assert not torch.equal(out, out_old), "the replay did not read the new inputs" + args = _target_args(mixes2, sqrsum2, x2, out, scale, base, m) + op.compare(out, op.reference(*args)) # the replay's output diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/norm/test_modeling_v2_mhc_post_mapping.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/norm/test_modeling_v2_mhc_post_mapping.py new file mode 100644 index 000000000000..e65eaffad4d2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/norm/test_modeling_v2_mhc_post_mapping.py @@ -0,0 +1,349 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the mhc_post_mapping catalog entry. + +The entry owns what is certified: `CELLS` is the list, `reference` and `compare` are the gate. +This file turns a cell's spec into tensors and keeps the checks a cell cannot carry: the +transposed-comb and rolled-post controls, row independence under non-finite values, T = 0, +CUDA-graph replay, the guard and the op's dtype rejections, and the DeepSeek-V4.1 aliased call +at every row count 1..32768 (prefix sweep and its control, one compiled kernel per call, the +CUDA-graph capture union). +""" + +import zlib + +import pytest +import torch +from torch.profiler import ProfilerActivity, profile + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.norm.mhc_post_mapping import ( # noqa: E501 + mhc_post_mapping as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "mhc_post_mapping requires a CUDA device" + +HC = 4 +# The checkpoint's width (DeepSeek-V4.1-Flash hidden_size). +H_CKPT = 5120 + + +def _sinkhorn(t: int, gen: torch.Generator) -> torch.Tensor: + """A doubly-stochastic [t, 4, 4] fp32 comb, built the way the checkpoint builds it.""" + c = torch.softmax(torch.randn(t, HC, HC, device="cuda", generator=gen), -1) + 1e-6 + c = c / (c.sum(-2, keepdim=True) + 1e-6) + for _ in range(19): + c = c / (c.sum(-1, keepdim=True) + 1e-6) + c = c / (c.sum(-2, keepdim=True) + 1e-6) + return c.contiguous() + + +def _inputs(t: int, h: int, seed: int): + gen = torch.Generator(device="cuda").manual_seed(seed) + residual = torch.randn(t, HC, h, device="cuda", generator=gen) * 3 + if t > 1: + # an attention-sink row: the checkpoint's BOS stream reaches |x| ~ 3.5e5 + residual[0] *= 1e5 + residual = residual.bfloat16() + x = (torch.randn(t, h, device="cuda", generator=gen) * 2).bfloat16() + post = 2 * torch.sigmoid(torch.randn(t, HC, device="cuda", generator=gen)) + return residual, x, post.contiguous(), _sinkhorn(t, gen) + + +def _check(out, residual, x, post, comb): + """The entry's gate against its own reference.""" + op.compare(out, op.reference(residual, x, post, comb)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + s = cell.spec + t, h = s["T"], s["H"] + residual, x, post, comb = _inputs(t, h, seed=zlib.crc32(cell.why.encode()) % 2**31) + if s["flat"]: + post, comb = post.view(t, HC, 1), comb.view(t, HC * HC) + buf = None + if s["out"] == "alias": + residual = residual.clone() + out = residual + elif s["out"] == "view": + buf = torch.full((t + 2, HC, h), 7.0, device="cuda", dtype=torch.bfloat16) + out = buf[:t] + else: + out = None + args = (residual, x, post, comb, out) + ref = op.reference(*args) # before the call: an aliased call overwrites residual + ret = op(*args) + assert ret.shape == (t, HC, h) and ret.dtype == torch.bfloat16 + if out is not None: + assert ret.data_ptr() == out.data_ptr(), "the op writes the given out and returns it" + if buf is not None: + assert bool((buf[t:] == 7.0).all()), "rows past T of the parent buffer were written" + op.compare(ret, ref) + + +def test_control_transposed_comb_is_detected() -> None: + # The deliberately wrong control: comb applied untransposed (out[j] mixes + # comb[j, k] instead of comb[k, j]). A doubly-stochastic comb makes the two + # orientations numerically close, so this is the error a plausible-looking + # result would hide; the gate must reject it. + residual, x, post, comb = _inputs(64, H_CKPT, seed=7) + out = op(residual, x, post, comb) + _check(out, residual, x, post, comb) + wrong = comb.transpose(1, 2).contiguous() + with pytest.raises(AssertionError): + op.compare(out, op.reference(residual, x, post, wrong)) + # a wrong post (x scaled by the next copy's coefficient) is caught too + with pytest.raises(AssertionError): + op.compare(out, op.reference(residual, x, post.roll(1, dims=1), comb)) + + +def test_rows_are_independent() -> None: + # CUDA-graph padding rows may carry garbage; a non-finite value must stay in its row + t = 64 + residual, x, post, comb = _inputs(t, H_CKPT, seed=21) + clean = op(residual, x, post, comb) + post[10, 2] = float("nan") + comb[20, 1, 3] = float("inf") + x[30, 100] = float("nan") + residual[40, 0, 7] = float("inf") + dirty = op(residual, x, post, comb) + keep = torch.ones(t, dtype=torch.bool, device="cuda") + keep[[10, 20, 30, 40]] = False + assert torch.equal(dirty[keep], clean[keep]) + assert not bool(torch.isfinite(dirty[10, 2].float()).any()) + assert not bool(torch.isfinite(dirty[20, 3].float()).any()) + + +def test_zero_tokens() -> None: + residual = torch.empty(0, HC, H_CKPT, device="cuda", dtype=torch.bfloat16) + x = torch.empty(0, H_CKPT, device="cuda", dtype=torch.bfloat16) + post = torch.empty(0, HC, device="cuda") + comb = torch.empty(0, HC, HC, device="cuda") + out = op(residual, x, post, comb) + assert out.shape == (0, HC, H_CKPT) + torch.cuda.synchronize() + # nothing left pending: the next call runs and is exact + r, xx, p, c = _inputs(4, H_CKPT, seed=4) + _check(op(r, xx, p, c), r, xx, p, c) + torch.cuda.synchronize() + + +def test_cuda_graph_replay() -> None: + t = 32 + residual, x, post, comb = _inputs(t, H_CKPT, seed=31) + out = torch.empty_like(residual) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(residual, x, post, comb, out=out) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(residual, x, post, comb, out=out) + r2, x2, p2, c2 = _inputs(t, H_CKPT, seed=32) + residual.copy_(r2) + x.copy_(x2) + post.copy_(p2) + comb.copy_(c2) + out.zero_() + graph.replay() + torch.cuda.synchronize() + _check(out, r2, x2, p2, c2) + + +def test_is_valid_rejects_silently_wrong_layouts() -> None: + # Each of these returns wrong values (or writes out of bounds) with no error + # when handed to the raw op; the guard refuses them before launch. + t = 8 + residual, x, post, comb = _inputs(t, H_CKPT, seed=41) + wide = torch.zeros(t, 2 * H_CKPT, device="cuda", dtype=torch.bfloat16) + mixes = torch.zeros(t, 24, device="cuda") + bad = [ + dict(x=wide[:, :H_CKPT]), # column slice of a wider buffer + dict(post_mix=mixes[:, 4:8]), # column slice of a [T, 24] mixes buffer + dict(comb_mix=comb.transpose(1, 2)), # transpose view + dict(residual=residual[:, :2].contiguous(), x=x), # hc = 2 + dict(x=x[: t // 2].contiguous()), # too few rows + dict(post_mix=post[:, :2].contiguous()), + dict(comb_mix=comb[:, :, :2].contiguous()), + dict(out=torch.empty(t // 2, HC, H_CKPT, device="cuda", dtype=torch.bfloat16)), + dict(out=torch.empty(t, HC, 2 * H_CKPT, device="cuda", dtype=torch.bfloat16)[..., :H_CKPT]), + ] + for override in bad: + args = dict(residual=residual, x=x, post_mix=post, comb_mix=comb, out=None) + args.update(override) + with pytest.raises(AssertionError, match="must"): + with validating(op): + op(**args) + + +@pytest.mark.parametrize( + "field,dtype", + [ + ("x", torch.float32), + ("x", torch.float16), + ("residual", torch.float32), + ("post_mix", torch.bfloat16), + ("comb_mix", torch.bfloat16), + ("out", torch.float32), + ], +) +def test_wrong_dtype_raises(field: str, dtype: torch.dtype) -> None: + residual, x, post, comb = _inputs(4, H_CKPT, seed=51) + args = dict( + residual=residual, x=x, post_mix=post, comb_mix=comb, out=torch.empty_like(residual) + ) + args[field] = args[field].to(dtype) + with pytest.raises(RuntimeError, match="expected scalar type"): + op(**args) + + +# ─── The target's aliased call at every row count 1..32768, and the capture union ── +# +# DeepSeek-V4.1's core target updates the stream in place twice per layer, +# `mhc_post_mapping(h, sublayer_out, post, comb, out=h)`, at T = the rows of +# the pass: the tokens the rank holds this step (any integer up to the engine's +# max_num_tokens, 16384 at most in the envelope) or, in layers 21-39 of a step +# that runs the decoder bounded-replay pass, that pass's rows (up to 32768). +# The DSpark drafter makes the same call at N = 5G rows. The launch is one CTA +# of 256 threads per token (CTA t reads row t of every input and writes row t +# of out), so T only sizes the grid. The tests below check that on the GPU at +# every T from 1 to 32768, in the aliased form. + +_N_MAX = 32768 +# The row counts of the target's call that the entry's cells certify (its runtime-value dumps +# plus the tops of the step and decoder-pass ranges). +_TARGET_T = tuple(sorted({c.spec["T"] for c in op.CELLS if c.why.startswith("target call")})) +# The capture union per rank. G ranges over every CUDA-graph batch size an +# envelope engine captures: every G 1..32, plus TRT-LLM's padded default list +# for cuda_graph_config {enable_padding: true, max_batch_size: 256} +# (40..128 step 8, 192, 256); the padding-off list at 256 (1..32, 64, 128, 256) +# is a subset. A captured step is generation-only, so the trunk runs at G (a +# decode step) and 6G (the DSpark target-verify step, 1 + 5 draft tokens per +# request), and the drafter at 5G. +_GSET = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_CAPTURE_T = tuple(sorted({k * g for g in _GSET for k in (1, 5, 6)})) +assert len(_GSET) == 46 and len(_CAPTURE_T) == 114 and max(_CAPTURE_T) == 1536 + + +def _aliased(stream, x, post, comb): + """The core target's exact call: out is the stream itself.""" + ret = op(stream, x, post, comb, out=stream) + assert ret.data_ptr() == stream.data_ptr() + return ret + + +def test_target_call_every_row_count_is_the_same_per_row_function() -> None: + """Every prefix T = 1..32768 of the aliased call returns exactly the first T + rows of the 32768-row result, which matches the kernel-order reference bit + for bit. + + Each call runs on the first T rows of a buffer whose row T is a NaN guard + row that must stay NaN. A row the kernel skipped would keep its old stream + value, which differs from the updated one. + """ + residual, x, post, comb = _inputs(_N_MAX, H_CKPT, seed=4242) + full = residual.clone() + full_ref = op.reference(full, x, post, comb, full) + op.compare(_aliased(full, x, post, comb), full_ref) + fb = full.view(torch.int16) + buf = torch.empty(_N_MAX + 1, HC, H_CKPT, dtype=torch.bfloat16, device="cuda") + mismatch = torch.zeros(_N_MAX + 1, dtype=torch.int64, device="cuda") + guard = torch.zeros(_N_MAX + 1, dtype=torch.int64, device="cuda") + differs = torch.empty( + _N_MAX, HC, H_CKPT, dtype=torch.bool, device="cuda" + ) # one buffer; .any() keeps the reduction free of a widened copy + for t in range(1, _N_MAX + 1): + buf[:t].copy_(residual[:t]) + buf[t].fill_(float("nan")) + _aliased(buf[:t], x[:t], post[:t], comb[:t]) + mismatch[t] = torch.ne(buf[:t].view(torch.int16), fb[:t], out=differs[:t]).any() + guard[t] = (~buf[t].isnan()).sum() + bad = mismatch.nonzero().flatten().tolist() + over = guard.nonzero().flatten().tolist() + assert not bad, ( + f"{len(bad)} row counts differ from the {_N_MAX}-row result's rows, first {bad[:10]}" + ) + assert not over, f"{len(over)} row counts wrote the guard row, first {over[:10]}" + + +def test_target_call_row_sweep_harness_sees_a_one_row_shift() -> None: + """Control for the sweep: its bitwise per-row check rejects a read one row off.""" + n = _N_MAX + residual, x, post, comb = _inputs(n + 1, H_CKPT, seed=4243) + full = residual[:n].clone() + _aliased(full, x[:n], post[:n], comb[:n]) + for t in (1, 2, 7, 1024, 2048, 16384, n): + s = residual[1 : t + 1].clone() + _aliased(s, x[1 : t + 1], post[1 : t + 1], comb[1 : t + 1]) + n_bad = int((s.view(torch.int16) != full.view(torch.int16)[:t]).sum().item()) + print(f"[control row shift] T={t}: {n_bad} / {t * HC * H_CKPT} elements differ") + assert n_bad > t * HC * H_CKPT // 2 + + +_PROBE_T = tuple( + sorted( + set(_TARGET_T) + | set(_CAPTURE_T) + | {9, 64, 65, 129, 4096, 8191, 8192, 8193, 12288, 16383, 16385, 24576, 32767, _N_MAX} + ) +) + + +def test_target_call_one_compiled_kernel_at_every_row_count() -> None: + """Probe: the aliased call launches the same single kernel at every T (compiled ahead of time, no JIT).""" + residual, x, post, comb = _inputs(_N_MAX, H_CKPT, seed=4244) + names = set() + for t in _PROBE_T: + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _aliased(residual[:t], x[:t], post[:t], comb[:t]) + torch.cuda.synchronize() + launched = [ + e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA + ] + assert len(launched) == 1, (t, launched) + names.add(launched[0]) + print(f"{len(_PROBE_T)} row counts probed; kernels launched: {sorted(names)}") + assert len(names) == 1 and "mhcPostMappingKernel" in next(iter(names)), names + + +@pytest.mark.parametrize("num_tokens", _CAPTURE_T) +def test_target_call_capture_union_replay_is_bitwise_eager(num_tokens: int) -> None: + """The aliased call captured at T after an eager aliased call at T, replayed + on new data: bitwise equal to an eager call on that data. + + The replay on the capture-time data must differ from the replay on new + data, so a replay that reused stale inputs would fail. + """ + t = num_tokens + residual, x, post, comb = _inputs(t, H_CKPT, seed=6000 + t) + stream = residual.clone() + ref = op.reference(residual, x, post, comb) + _aliased(stream, x, post, comb) # the eager call at this shape + torch.cuda.synchronize() + op.compare(stream, ref) + stream.copy_(residual) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _aliased(stream, x, post, comb) + graph.replay() + torch.cuda.synchronize() + op.compare(stream, ref) + old = stream.clone() + r2, x2, p2, c2 = _inputs(t, H_CKPT, seed=7000 + t) + stream.copy_(r2) + x.copy_(x2) + post.copy_(p2) + comb.copy_(c2) + graph.replay() + torch.cuda.synchronize() + eager = r2.clone() + _aliased(eager, x2, p2, c2) + torch.cuda.synchronize() + assert torch.equal(stream, eager), t + assert not torch.equal(stream, old), "the replay did not read the new inputs" + op.compare(stream, op.reference(r2, x2, p2, c2)) # the replay's output diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_e4m3_quantize_pow2_dynamic.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_e4m3_quantize_pow2_dynamic.py new file mode 100644 index 000000000000..c1f50a1174a2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_e4m3_quantize_pow2_dynamic.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the e4m3_quantize_pow2_dynamic catalog entry. + +The entry owns what is certified: `CELLS` is the list, `reference` and +`compare` are the gate. This file turns a cell's spec into tensors and keeps +the checks a cell cannot carry: CUDA-graph replay on new data, the round-trip +error bound, the wrong-variant controls and the op's own rejections. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.quantization.e4m3_quantize_pow2_dynamic import ( # noqa: E501 + e4m3_quantize_pow2_dynamic as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "e4m3_quantize_pow2_dynamic requires a CUDA device" + +DEV = "cuda" +E4M3_MAX = 448.0 +FLOOR = 1e-4 + + +def _query_like(T, H, seed, spread=6): + """bf16 [T, H, 512] with per-(token, head) magnitudes over 2^-spread .. 2^spread (the golden + queries' |q| reaches 18..88 per layer) and a few exact zeros.""" + g = torch.Generator(device=DEV).manual_seed(seed) + mag = torch.randint(-spread, spread + 1, (T, H, 1), device=DEV, generator=g).float() + x = torch.randn(T, H, 512, device=DEV, generator=g) * torch.exp2(mag) + x[:, :, 7] = 0.0 + return x.to(torch.bfloat16) + + +def _build(spec, seed): + if spec["kind"] == "query": + x = _query_like(spec["T"], spec["H"], seed).to(spec["dtype"]) + else: + amax = spec["amax"] + x = torch.zeros(4, 16, 512, dtype=spec["dtype"], device=DEV) + x[1, 3, 5] = amax + x[2, 0, 0] = -amax / 3 + out = torch.empty(x.shape, dtype=torch.float8_e4m3fn, device=DEV) + scale_out = torch.full((1,), -1.0, dtype=torch.float32, device=DEV) + return x, out, scale_out, spec["scale_mult"] + + +def _scale_of(x): + """The power-of-two scale `s` the op is specified to pick for `x`.""" + _, ref_scale = op.reference(x, x, torch.zeros(1, device=DEV), 1.0) + return ref_scale[0].item() + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + x, out, scale_out, scale_mult = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + result = op(x, out, scale_out, scale_mult) + assert result[0] is out and result[1] is scale_out, "the op writes in place and returns both" + op.compare(result, op.reference(x, out, scale_out, scale_mult)) + + s = _scale_of(x) + if cell.spec["kind"] == "query": + # No value saturates: the largest |x| / s is in [224, 448]. + top = x.float().abs().max().item() / s + assert 224.0 <= top <= 448.0, top + elif cell.spec["amax"] == 449.0: + assert s == 2.0 + elif cell.spec["amax"] <= FLOOR: + assert s == 2.0**-22 # 2^ceil(log2(1e-4 / 448)) + + +def test_cuda_graph_capture_replay_follows_the_data(): + """The scale is recomputed on the device at every replay (no host sync, no stale amax).""" + x = _query_like(96, 16, seed=11) + out = torch.empty(x.shape, dtype=torch.float8_e4m3fn, device=DEV) + so = torch.empty(1, dtype=torch.float32, device=DEV) + op(x, out, so, 0.5) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(x, out, so, 0.5) + for seed, spread in ((12, 2), (13, 9), (14, 0)): + x.copy_(_query_like(96, 16, seed=seed, spread=spread)) + graph.replay() + op.compare((out, so), op.reference(x, out, so, 0.5)) + + +def test_round_trip_error_is_e4m3_rounding(): + """out * s reproduces x to E4M3's 3-bit mantissa: relative error <= 2^-4 in the normal range.""" + x = _query_like(512, 16, seed=21) + out = torch.empty(x.shape, dtype=torch.float8_e4m3fn, device=DEV) + so = torch.empty(1, dtype=torch.float32, device=DEV) + op(x, out, so, 1.0) + s = so.item() + back = out.float() * s + xf = x.float() + normal = xf.abs() >= 2.0**-6 * s + rel = ((back - xf).abs() / xf.abs().clamp_min(1e-30))[normal] + assert rel.max().item() <= 2.0**-4 + 1e-7 + + +def test_wrong_variants_are_detected(): + """Controls: a non-power-of-two scale (amax / 448), and a scale rounded down instead of up.""" + x = _query_like(64, 16, seed=31) + out = torch.empty(x.shape, dtype=torch.float8_e4m3fn, device=DEV) + so = torch.empty(1, dtype=torch.float32, device=DEV) + op(x, out, so, 1.0) + amax = x.float().abs().max().item() + for name, s in ( + ("amax/448", amax / 448.0), + ("floor pow2", 2.0 ** torch.floor(torch.log2(torch.tensor(amax / 448.0))).item()), + ): + q = (x.float() / s).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn).view(torch.uint8) + frac = (q != out.view(torch.uint8)).float().mean().item() + assert frac > 0.3, f"control '{name}' differs on only {frac:.3f}" + + +def test_rejects_bad_inputs(): + x = _query_like(4, 16, seed=41) + so = torch.empty(1, dtype=torch.float32, device=DEV) + with pytest.raises(ValueError, match="float8_e4m3fn"): + op(x, torch.empty_like(x), so, 1.0) + with pytest.raises(ValueError, match="float32"): + op( + x, + torch.empty(x.shape, dtype=torch.float8_e4m3fn, device=DEV), + torch.empty(1, dtype=torch.bfloat16, device=DEV), + 1.0, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_mxfp8_dequantize.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_mxfp8_dequantize.py new file mode 100644 index 000000000000..abe679d3c72e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_mxfp8_dequantize.py @@ -0,0 +1,2434 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the mxfp8_dequantize catalog entry. + +The entry owns what is certified: `CELLS` is the list (value domain, widths, the checkpoint +recipe's rows, paged gathers, the 528 B layouts, strided slots and out, the shadow pool, scatter, +zero rows), `reference` and `compare` are the bitwise gate. This file turns a cell's spec into +tensors and keeps the checks a cell cannot carry: NaN encodings, CUDA-graph replay, the op's own +rejections, the deliberately-wrong decodes, and the native pools of the live DeepSeek-V4.1 cache +manager (far-end slots, pool-size invariance, every row count, per-query window reads, the +fixed-size window buffer, CUDA-graph capture at every graph batch, renumbered reads over random +multi-request steps). + +Every comparison is bitwise (bf16 compared as int16): each output element is one e4m3 value times +a power of two, so there is no tolerance to choose. The entry's reference decodes with the torch +float8_e4m3fn -> fp64 cast and torch.ldexp(code, byte - 127) in fp64, then fp64 -> fp32 (exact: at +most 4 significant bits, exponent >= -136) -> bf16 (torch's round-to-nearest-even). The kernel +instead builds the scale from the IEEE exponent bits, so the two derivations are independent. + +The fake-quantized operands (codes + UE8M0 bytes of window-KV-like rows) are also produced here in +native torch, with the checkpoint reference's recipe (amax floored at 1e-4, scale +2^ceil(log2(amax / 448)), exponent from torch.frexp), so "dequantize(native) == the bf16 +container's value" is checked without calling any other catalog entry. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.quantization.mxfp8_dequantize import ( # noqa: E501 + mxfp8_dequantize as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "mxfp8_dequantize requires a CUDA device" + +BLOCK = 32 +E4M3_MAX = 448.0 +AMAX_FLOOR = 1e-4 +W_KV = 512 # DeepSeek-V4.1 window KV head_dim (448 nope + 64 RoPE, all quantized) +NB_KV = W_KV // BLOCK +E4M3_NAN_CODES = (0x7F, 0xFF) +DEV = "cuda" +# The synthetic pools are sized far above any batch (320 pages x 128 = 40960 slots); the cache +# manager's own pool sizes are certified on live managers further down. +POOL_PAGES = 320 + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device=DEV).manual_seed(seed) + + +def _ref(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """The entry's reference decode of every slot of a [N, W] pool, as a dense bf16 [N, W].""" + return op.reference(codes, scales, _sentinel((codes.shape[0], codes.shape[-1]))) + + +def _control_decode( + codes: torch.Tensor, + scales: torch.Tensor, + *, + bias: int = 127, + code_dtype=torch.float8_e4m3fn, + scale_index: torch.Tensor | None = None, +) -> torch.Tensor: + """A deliberately-wrong decode of uint8 codes [N, W] and UE8M0 bytes [N, W/32] -> bf16 [N, W]. + + `bias`, `code_dtype` and `scale_index` (which byte serves 32-block b) select the control. + """ + n, w = codes.shape + v = codes.contiguous().view(code_dtype).float().double().view(n, w // BLOCK, BLOCK) + e = (scales.long() - bias).double() + if scale_index is not None: + e = e[:, scale_index] + return torch.ldexp(v, e.unsqueeze(-1)).view(n, w).float().to(torch.bfloat16) + + +def _random_codes(n: int, w: int, g: torch.Generator) -> torch.Tensor: + """Uniform e4m3 code bytes, the two NaN codes replaced by +-0 (they are tested separately).""" + c = torch.randint(0, 256, (n, w), device=DEV, dtype=torch.int32, generator=g).to(torch.uint8) + for nan in E4M3_NAN_CODES: + c[c == nan] = nan & 0x80 + return c + + +def _random_bytes(n: int, nb: int, g: torch.Generator, lo: int = 0, hi: int = 254) -> torch.Tensor: + return torch.randint(lo, hi + 1, (n, nb), device=DEV, dtype=torch.int32, generator=g).to( + torch.uint8 + ) + + +def _fake_quant_native(x: torch.Tensor): + """Checkpoint-recipe MXFP8 of bf16 [T, W]: (bf16 container values, uint8 codes, uint8 UE8M0 bytes).""" + t, w = x.shape + xf = x.float().view(t, w // BLOCK, BLOCK) + a = xf.abs().amax(-1).clamp_min(AMAX_FLOOR) + m, e = torch.frexp( + a.double() / E4M3_MAX + ) # exact ceil(log2(.)): mantissa 0.5 is an exact power of two + e = torch.where(m == 0.5, e - 1, e).to(torch.int32) + s = torch.exp2(e.double()).float().unsqueeze(-1) + q = (xf / s).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) + y = (q.float() * s).view(t, w).to(torch.bfloat16) + return y, q.view(t, w).view(torch.uint8), (e + 127).to(torch.uint8) + + +def _window_kv_like(t: int, w: int = W_KV, seed: int = 0) -> torch.Tensor: + """bf16 [t, w] rows with block magnitudes 2^-30 .. 2^20, all-zero blocks, amax = 448 * 2^k and sub-1e-4 + blocks (scale byte 105, the recipe's floor), so every scale path of the recipe is present.""" + g = _gen(seed) + nb = w // BLOCK + mag = torch.randint(-30, 21, (t, nb, 1), device=DEV, generator=g).float() + x = torch.randn(t, nb, BLOCK, device=DEV, generator=g) * torch.exp2(mag) + for r in range(0, t, 4): + k = r % 7 - 3 + b = (r // 4) % nb + x[r, b] = 0.0 + x[r, (b + 1) % nb, 0] = 448.0 * 2.0**k + x[r, (b + 2) % nb] = torch.randn(BLOCK, device=DEV, generator=g) * 1e-6 + return x.view(t, w).to(torch.bfloat16) + + +def _n_bad(actual: torch.Tensor, expected: torch.Tensor) -> int: + """Elements whose bf16 bits differ, counted in chunks (a bool reduction materializes an int64 copy).""" + a = actual.reshape(-1).view(torch.int16) + e = expected.reshape(-1).view(torch.int16) + n = torch.zeros((), dtype=torch.int64, device=a.device) + for i in range(0, a.numel(), 1 << 26): + n += (a[i : i + (1 << 26)] != e[i : i + (1 << 26)]).sum() + return int(n) + + +def _assert_bitwise(actual: torch.Tensor, expected: torch.Tensor) -> None: + n_bad = _n_bad(actual, expected) + assert n_bad == 0, f"{n_bad} of {actual.numel()} elements differ bitwise" + + +def _sentinel(shape, device=DEV) -> torch.Tensor: + """bf16 filled with a NaN bit pattern the op never writes (0x7FAB).""" + return torch.full(shape, 0x7FAB, dtype=torch.int16, device=device).view(torch.bfloat16) + + +def _slots( + rows: int, num_slots: int, g: torch.Generator, dtype=torch.int32, invalid_every: int = 0 +) -> torch.Tensor: + """`rows` distinct random slots of a `num_slots` pool; every `invalid_every`-th entry is invalid + (alternately -1 and num_slots, i.e. negative and one past the end).""" + s = torch.randperm(num_slots, device=DEV, generator=g)[:rows].to(torch.int64) + if invalid_every: + bad = torch.arange(0, rows, invalid_every, device=DEV) + s[bad] = torch.where(bad % (2 * invalid_every) == 0, -1, num_slots) + return s.to(dtype) + + +def _pool(pages: int, page_size: int, w: int, g: torch.Generator, lo: int = 0, hi: int = 254): + codes = _random_codes(pages * page_size, w, g) + scales = _random_bytes(pages * page_size, w // BLOCK, g, lo, hi) + if page_size == 1: + return codes, scales + return codes.view(pages, page_size, w), scales.view(pages, page_size, w // BLOCK) + + +def _gather_ref(codes, scales, src: torch.Tensor) -> torch.Tensor: + """The entry's reference rows for `src` (invalid -> zeros) from a [slots, W] / [pages, ps, W] pool.""" + return op.reference(codes, scales, _sentinel((src.shape[0], codes.shape[-1])), src) + + +# ─── certified cells ────────────────────────────────────────────────────────── + + +def _build(spec, seed: int): + """A cell's spec -> (codes, scales, out, src_slots, dst_slots, extra). + + `extra` carries what the structural checks need: `pool` (the tensors whose bytes must not + change), and per kind the bf16 `container` or the `wide` parent of a strided out. + """ + g = _gen(seed) + kind = spec["kind"] + src = dst = None + extra = {} + if kind == "identity": + n, w = spec["rows"], spec["width"] + codes, scales = _pool(n, 1, w, g) + if spec["pin_extremes"]: + # bytes 0 / 1 / 254: fp32-subnormal scale, bf16-subnormal products, overflow + scales[:3] = torch.tensor([0, 1, 254], dtype=torch.uint8, device=DEV).view(3, 1) + codes[:3, :256] = torch.arange(256, device=DEV, dtype=torch.int32).to(torch.uint8) + for nan in E4M3_NAN_CODES: + codes[:3][codes[:3] == nan] = 0 + out = _sentinel((n, w)) + elif kind == "recipe": + t = spec["rows"] + y, codes, scales = _fake_quant_native(_window_kv_like(t, seed=seed)) + extra["container"] = y + out = _sentinel((t, W_KV)) + elif kind == "gather": + ps, rows = spec["page_size"], spec["rows"] + pages = POOL_PAGES * 128 // ps + codes, scales = _pool(pages, ps, W_KV, g, lo=100, hi=160) + src = _slots(rows, pages * ps, g, spec["slot_dtype"], invalid_every=16 if rows >= 16 else 0) + out = _sentinel((rows, W_KV)) + elif kind == "interleaved": + # The native SWA row: 512 code bytes then 16 UE8M0 bytes, 528 B per slot. + n_slots = POOL_PAGES * 128 + buf = torch.randint(0, 256, (n_slots, 528), device=DEV, dtype=torch.int32, generator=g).to( + torch.uint8 + ) + buf[:, :W_KV] = _random_codes(n_slots, W_KV, g) + buf[:, W_KV:] = _random_bytes(n_slots, NB_KV, g) + if spec["paged"]: + buf = buf.view(POOL_PAGES, 128, 528) + codes = buf[..., :W_KV].view(torch.float8_e4m3fn) + scales = buf[..., W_KV:].view(spec["scale_dtype"]) + src = _slots(4096, n_slots, g, invalid_every=64) + out = _sentinel((4096, W_KV)) + extra["pool"] = (buf,) + elif kind == "footer": + # One uint8 page = page_size code rows of 512 B, then page_size scale rows of 16 B. + ps, pages = spec["page_size"], 24 + buf = torch.randint( + 0, 256, (pages, ps * 528), device=DEV, dtype=torch.int32, generator=g + ).to(torch.uint8) + codes = buf[:, : ps * W_KV].view(pages, ps, W_KV) + scales = buf[:, ps * W_KV :].view(pages, ps, NB_KV) + for nan in E4M3_NAN_CODES: + codes[codes == nan] = 0 + scales[scales == 255] = 127 + src = _slots(2000, pages * ps, g, torch.int64, invalid_every=50) + out = _sentinel((2000, W_KV)) + extra["pool"] = (buf,) + elif kind == "strided": + codes, scales = _pool(POOL_PAGES, 128, W_KV, g, lo=100, hi=160) + src = _slots(2 * 777, POOL_PAGES * 128, g)[::2] + wide = _sentinel((777, 2 * W_KV)) + out = wide[:, :W_KV] + extra["wide"] = wide + elif kind == "shadow": + codes, scales = _pool(POOL_PAGES, 128, W_KV, g, lo=100, hi=160) + src = dst = _slots(4 * 128, POOL_PAGES * 128, g) + src[::37] = -1 + out = _sentinel((POOL_PAGES, 128, W_KV)) + elif kind == "scatter": + n_src, n_dst, rows = 4096, 2048, 600 + codes, scales = _pool(n_src, 1, W_KV, g, lo=100, hi=160) + src = _slots(rows, n_src, g, torch.int64) + dst = _slots(rows, n_dst, g, torch.int32) + src[0:600:5] = -1 # zero rows + dst[1:600:7] = -1 # skipped + dst[2:600:11] = n_dst # skipped + out = _sentinel((n_dst, W_KV)) + elif kind == "dst_only": + rows = 600 + codes, scales = _pool(rows, 1, W_KV, g, lo=100, hi=160) + dst = _slots(rows, 16 * 128, g, torch.int32) + out = _sentinel((16, 128, W_KV)) + else: + assert kind == "zero_rows", kind + codes, scales = _pool(64, 1, W_KV, g) + src = torch.empty(0, dtype=torch.int32, device=DEV) + out = _sentinel((4, W_KV)) + extra.setdefault("pool", (codes, scales)) + return codes, scales, out, src, dst, extra + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + codes, scales, out, src, dst, extra = _build( + cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + ) + before = out.clone() + pool_before = [t.view(torch.uint8).clone() for t in extra["pool"]] + result = op(codes, scales, out, src, dst) + assert result is out, "the op writes out in place and returns it" + op.compare(out, op.reference(codes, scales, before, src, dst)) + for t, b in zip(extra["pool"], pool_before): + assert torch.equal(t.view(torch.uint8), b), "the source pool must not be written" + if "container" in extra: + _assert_bitwise(out, extra["container"]) + if cell.spec["rows"] >= 32: + assert int((scales == 105).sum()) > 0, "the 1e-4 floor block (byte 105) must be present" + if "wide" in extra: + _assert_bitwise(extra["wide"][:, W_KV:], _sentinel((777, W_KV))) + + +# ─── special encodings, graphs, rejections ──────────────────────────────────── + + +def test_nan_encodings() -> None: + """UE8M0 byte 255 is NaN (whole block NaN); e4m3 codes 0x7F / 0xFF are NaN (that element only).""" + g = _gen(255) + codes, scales = _pool(8, 1, W_KV, g, lo=110, hi=140) + scales[1, 3] = 255 + codes[2, 7] = 0x7F + codes[2, 100] = 0xFF + out = _sentinel((8, W_KV)) + op(codes, scales, out) + nan = torch.isnan(out.float()) + expect_nan = torch.zeros_like(nan) + expect_nan[1, 96:128] = True + expect_nan[2, 7] = True + expect_nan[2, 100] = True + assert torch.equal(nan, expect_nan) + ref = _ref(codes, scales) + assert torch.equal(torch.isnan(ref.float()), expect_nan) + _assert_bitwise(out[~expect_nan], ref[~expect_nan]) + + +def test_cuda_graph_replay() -> None: + """Capture after one eager warm-up; replay after the pool and the slot list change in place.""" + g = _gen(17) + codes, scales = _pool(POOL_PAGES, 128, W_KV, g, lo=100, hi=160) + src = _slots(1024, POOL_PAGES * 128, g, invalid_every=32) + out = _sentinel((1024, W_KV)) + op(codes, scales, out, src) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(codes, scales, out, src) + for step in range(2): + new_codes, new_scales = _pool(POOL_PAGES, 128, W_KV, g, lo=90 + step, hi=170) + codes.copy_(new_codes) + scales.copy_(new_scales) + src.copy_(_slots(1024, POOL_PAGES * 128, g, invalid_every=32 + step)) + out.copy_(_sentinel((1024, W_KV))) + graph.replay() + torch.cuda.synchronize() + op.compare(out, _gather_ref(codes, scales, src)) + + +@pytest.mark.parametrize( + "case, match", + [ + ("codes_cpu", "codes must be a CUDA tensor"), + ("codes_bf16", "codes must be float8_e4m3fn or uint8"), + ("width_48", "positive multiple of 32"), + ("codes_col_stride", "unit stride along its last dim"), + ("codes_4d", r"codes must be \[slots, C\] or \[pages, page_size, C\]"), + ("scales_int32", "scales must be uint8 or float8_e8m0fnu"), + ("scales_shape", "scales must be"), + ("scales_cpu", "scales must be on"), + ("out_fp32", "out must be bfloat16"), + ("out_width", "out's last dim must be 512"), + ("out_cpu", "out must be on"), + ("src_float", "src_slots must be int32 or int64"), + ("src_2d", "src_slots must be 1-D"), + ("dst_cpu", "dst_slots must be on"), + ("len_mismatch", "same length"), + ("out_too_small", "without dst_slots out needs >= 16 slots"), + ], +) +def test_rejections(case: str, match: str) -> None: + g = _gen(1) + codes, scales = _pool(16, 1, W_KV, g) + out = _sentinel((16, W_KV)) + src = torch.arange(16, device=DEV, dtype=torch.int32) + args = dict(codes=codes, scales=scales, out=out, src_slots=None, dst_slots=None) + if case == "codes_cpu": + args["codes"] = codes.cpu() + elif case == "codes_bf16": + args["codes"] = codes.to(torch.bfloat16) + elif case == "width_48": + args["codes"], args["scales"] = codes[:, :48], scales[:, :1] + elif case == "codes_col_stride": + args["codes"] = torch.empty(16, 2 * W_KV, dtype=torch.uint8, device=DEV)[:, ::2] + elif case == "codes_4d": + args["codes"], args["scales"] = codes.view(2, 2, 4, W_KV), scales.view(2, 2, 4, NB_KV) + elif case == "scales_int32": + args["scales"] = scales.to(torch.int32) + elif case == "scales_shape": + args["scales"] = scales[:, :8] + elif case == "scales_cpu": + args["scales"] = scales.cpu() + elif case == "out_fp32": + args["out"] = out.float() + elif case == "out_width": + args["out"] = out[:, :256] + elif case == "out_cpu": + args["out"] = out.cpu() + elif case == "src_float": + args["src_slots"] = src.float() + elif case == "src_2d": + args["src_slots"] = src.view(4, 4) + elif case == "dst_cpu": + args["dst_slots"] = src.cpu() + elif case == "len_mismatch": + args["src_slots"], args["dst_slots"] = src, src[:8] + elif case == "out_too_small": + args["out"] = out[:8] + before = out.clone() + with pytest.raises(ValueError, match=match): + op(**args) + _assert_bitwise(out, before) + + +# ─── deliberately-wrong controls ────────────────────────────────────────────── + + +def test_wrong_recipes_are_rejected_by_the_gate() -> None: + """The bitwise gate that passes the op fails each wrong decode of the same native rows.""" + t = 4096 + y, codes, scales = _fake_quant_native(_window_kv_like(t, seed=11)) + g = _gen(12) + n_slots = POOL_PAGES * 128 + pool_codes = torch.zeros(n_slots, W_KV, dtype=torch.uint8, device=DEV) + pool_scales = torch.full((n_slots, NB_KV), 127, dtype=torch.uint8, device=DEV) + src = _slots(t, n_slots, g) + pool_codes[src.long()] = codes + pool_scales[src.long()] = scales + # a distinct neighbour row in every slot next to a used one, so an off-by-one slot is visible + nb = (src.long() + 1) % n_slots + free = ~torch.isin(nb, src.long()) + y2, c2, s2 = _fake_quant_native(_window_kv_like(t, seed=13)) + pool_codes[nb[free]] = c2[free] + pool_scales[nb[free]] = s2[free] + + out = _sentinel((t, W_KV)) + paged_codes = pool_codes.view(POOL_PAGES, 128, W_KV) + paged_scales = pool_scales.view(POOL_PAGES, 128, NB_KV) + op(paged_codes, paged_scales, out, src) + _assert_bitwise(out, y) + op.compare(out, _gather_ref(paged_codes, paged_scales, src)) + + wrong = { + "exponent bias 126 (scale one step high)": _control_decode(codes, scales, bias=126), + "exponent bias 128 (scale one step low)": _control_decode(codes, scales, bias=128), + "e5m2 decode of the code bytes": _control_decode( + codes, scales, code_dtype=torch.float8_e5m2 + ), + "one scale byte per 64 codes (block b reads byte b // 2)": _control_decode( + codes, scales, scale_index=torch.arange(NB_KV, device=DEV) // 2 + ), + "slot off by one": _ref( + pool_codes[(src.long() + 1) % n_slots], pool_scales[(src.long() + 1) % n_slots] + ), + } + report = {name: _n_bad(out, w) for name, w in wrong.items()} + print("wrong-recipe elements differing from the op (of", t * W_KV, "):", report) + for name, n_bad in report.items(): + assert n_bad > 0, f"control '{name}' is indistinguishable from the op's result" + + +# ─── native storage on the live DeepSeek-V4.1 cache manager ────────────────── +# +# With TRTLLM_DSV41_KV_STORAGE=native the V4.1 DeepseekV4CacheManager keeps the +# SWA window KV natively: one uint8 pool [B, 128, 528] for all 40 layers, 512 +# E4M3 codes then 16 UE8M0 bytes per row, and get_v41_native_kv_views(L, SWA) +# returns (codes [B, 128, 512], scales [B, 128, 16]), uint8 views with strides +# (67584, 528, 1). The native-storage target writes it with mxfp8_fake_quantize's +# native form and reads the rows a step's attention references with this op, +# either compact (src_slots = the referenced pool rows, the attention's window +# indices renumbered to rows of out) or into a bf16 shadow pool of the native +# pool's geometry (src_slots = dst_slots, the indices unchanged). +# B comes from free memory. The tables are the manager's native SWA pages at +# 0.5 x this GB300's total memory for every (max_num_tokens, max_batch_size, +# DSpark max_draft_len): the native gates reach 732,280 (accuracy, (8192, 32, +# 5)); over max_batch_size {8, 32} the largest is 854,440 ((8192, 8, 5)), and +# over every max_batch_size 1..32 it is B_max = 941,800 +# ((8192, 1, 0): 120,550,400 slots; the pool 63.6 GB, a shadow of its geometry +# 123.4 GB). The pool does not grow monotonically as max_batch_size falls. +# +# Round trips here write the pool with torch.ops.trtllm.mxfp8_fake_quantize +# (the writer this op pairs with) and hold the reads to that writer's own bf16 +# container output of the same rows and to this file's native-torch recipe. + +import os # noqa: E402 +import random # noqa: E402 +from contextlib import contextmanager # noqa: E402 + +import tensorrt_llm._torch.custom_ops.mxfp8_dequantize as _dq_op # noqa: E402 (the Triton kernel, variant probe) +import tensorrt_llm._torch.custom_ops.mxfp8_fake_quantize # noqa: E402,F401 -- registers the writer (round-trip inputs) + +_STORAGE_ENV = "TRTLLM_DSV41_KV_STORAGE" +PAGE = 128 +NATIVE_ROW = 528 +NATIVE_PAGE_BYTES = PAGE * NATIVE_ROW # 67584 +BYTE_SENTINEL = 0xFF # the writer never writes it for a finite row (E4M3 -NaN / UE8M0 NaN) +SENTINEL16 = 0x7FAB # bf16 NaN bit pattern of _sentinel(), as int16 +_BMAX_DEVICE_TOTAL = 297_020_948_480 +_NATIVE_SWA_PAGES = { + (2048, 8, 0): 399_640, + (2048, 8, 5): 549_080, + (2048, 32, 0): 366_400, + (2048, 32, 5): 569_440, + (8192, 8, 0): 776_000, + (8192, 8, 5): 854_440, + (8192, 32, 0): 585_840, + (8192, 32, 5): 732_280, +} +MBS_AXIS = tuple(range(1, 33)) # every max_batch_size the engine configuration permits (1..32) +MBS_LISTED = (1, 2, 3, 4, 8, 16, 32) +# Native SWA pages at 0.5 x device total per (max_num_tokens, DSpark max_draft_len), for max_batch_size 1..32 +# (measured on this GB300; re-measured by test_native_swa_pools_from_the_cache_manager). +_SWA_PAGES_BY_MBS = { + (2048, 0): ( + 549_600, + 549_600, + 549_600, + 549_600, + 549_600, + 488_040, + 439_360, + 399_640, + 399_640, + 549_080, + 506_880, + 506_880, + 506_880, + 506_880, + 506_880, + 439_360, + 439_360, + 439_360, + 439_360, + 412_080, + 387_720, + 387_720, + 462_720, + 462_720, + 439_360, + 439_360, + 439_360, + 418_520, + 418_520, + 382_280, + 366_400, + 366_400, + ), + (2048, 5): ( + 549_600, + 549_600, + 549_600, + 549_600, + 732_280, + 599_240, + 599_240, + 549_080, + 549_080, + 549_080, + 506_880, + 628_040, + 628_040, + 628_040, + 549_080, + 549_080, + 549_080, + 646_400, + 646_400, + 610_160, + 578_400, + 578_400, + 549_080, + 628_040, + 599_240, + 599_240, + 573_440, + 573_440, + 573_440, + 591_800, + 569_440, + 569_440, + ), + (8192, 0): ( + 941_800, + 879_240, + 879_240, + 879_240, + 879_240, + 824_160, + 824_160, + 776_000, + 776_000, + 732_280, + 732_280, + 809_760, + 809_760, + 809_760, + 732_280, + 732_280, + 732_280, + 732_280, + 732_280, + 699_040, + 668_760, + 640_960, + 615_120, + 676_200, + 650_880, + 650_880, + 650_880, + 650_880, + 606_200, + 606_200, + 585_840, + 585_840, + ), + (8192, 5): ( + 879_240, + 879_240, + 879_240, + 879_240, + 879_240, + 905_080, + 905_080, + 854_440, + 854_440, + 809_760, + 809_760, + 879_240, + 879_240, + 879_240, + 798_840, + 798_840, + 798_840, + 798_840, + 859_880, + 824_160, + 790_880, + 732_280, + 732_280, + 732_280, + 757_600, + 757_600, + 757_600, + 757_600, + 708_960, + 708_960, + 732_280, + 732_280, + ), +} +_NATIVE_SWA_PAGES_AXIS = { + (t, s, d): pages[s - 1] for (t, d), pages in _SWA_PAGES_BY_MBS.items() for s in MBS_AXIS +} +assert all(_NATIVE_SWA_PAGES_AXIS[k] == v for k, v in _NATIVE_SWA_PAGES.items()) +NATIVE_SWA_BMAX = 941_800 # over every max_batch_size 1..32 +NATIVE_SWA_BMAX_KEY = (8192, 1, 0) # the configuration that reaches it +NATIVE_SWA_BMAX_8_32 = 854_440 # the B_max over max_batch_size {8, 32} only +NATIVE_SWA_GATE_MAX = 732_280 +# The native SWA ceiling of the engine envelope: the measured maximum over its scanned tuples at 0.5 x the +# device's memory, PP2 stage managers included (stage alone), max_num_tokens {2048, 4096, 8192, 16384} x +# max_batch_size {1..8, 12, 16, 24, 32, 48, 64, 96, 128, 160, 192, 224, 256} x DSpark {0, 5} x max_seq_len {16K, +# 32K, 64K, 128K, 256K, 1M}, unsplit and per PP2 stage. It is stage 1's (layers 20-39) at NATIVE_SWA_CEILING_KEY +# (the SWA pool does not follow max_seq_len). Quotas above 0.5 x and unsampled max_batch_size values are not +# covered. +NATIVE_SWA_CEILING = 1_648_320 +NATIVE_SWA_CEILING_KEY = (16384, 1, 0) +NATIVE_SWA_CEILING_STAGE, NATIVE_SWA_CEILING_LAYER = 1, 20 +NATIVE_SWA_UNSPLIT = 1_098_720 # the pp_size 1 ceiling, at (16384, 1, 0) +# The shadow form (a bf16 pool of the native pool's geometry) runs on a native pool of this many pages: a shadow of +# the stage ceiling is 216 GB and does not fit one GB300 beside its 111 GB source; the unsplit ceiling's does. +SWA_SHADOW_PAGES = NATIVE_SWA_UNSPLIT +# Other pool sizes the invariance test runs: an engine dump's (16384, 256, 5) pool, envelope managers' pools, the +# unsplit ceiling, stage 0's largest and stage 1's largest in a coupled PP2 engine. +_SWA_EXTRA_B = (233_840, 617_600, 1_051_040, 1_098_720, 878_760, 1_098_220) +# Distinct SWA rows one step can reference: T + 127 per request (each request's +# window reaches 127 rows before its first new token), T <= 8192, <= 32 requests. +MAX_STEP_ROWS = 8192 + 127 * 32 # 12256 +SWEEP_ROWS = 16384 + + +def _jit_variants(fn) -> int: + """Number of compiled variants Triton holds for a @triton.jit function, over all devices.""" + dc = getattr(fn, "device_caches", None) + if dc is not None: + return sum(len(v[0]) if isinstance(v, tuple) else len(v) for v in dc.values()) + return sum(len(v) for v in fn.cache.values()) + + +def _write(x, codes, scales, slots): + torch.ops.trtllm.mxfp8_fake_quantize(x, codes, slots, scales) + + +def _container(x): + """The writer's own bf16-container output of x (identity form, a fixed 8192-row out per call size).""" + out = torch.empty(max(8192, x.shape[0]), W_KV, dtype=torch.bfloat16, device=DEV) + torch.ops.trtllm.mxfp8_fake_quantize(x, out, None, None) + return out[: x.shape[0]] + + +@contextmanager +def _native_storage(): + saved = os.environ.get(_STORAGE_ENV) + os.environ[_STORAGE_ENV] = "native" + try: + yield + finally: + if saved is None: + os.environ.pop(_STORAGE_ENV, None) + else: + os.environ[_STORAGE_ENV] = saved + + +def _v41_native_manager( + max_num_tokens: int, max_batch_size: int, max_draft_len: int, max_seq_len=None, pp_stage=None +): + """A V4.1 manager with native KV storage at the gate settings and 0.5 x the device's total memory. + max_seq_len defaults to the gates' 16384, as the engine passes it: + 14 with DSpark.""" + from types import SimpleNamespace + + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, + KvCacheConfig, + ) + from tensorrt_llm.mapping import Mapping + + total = torch.cuda.mem_get_info()[1] + assert total == _BMAX_DEVICE_TOTAL, ( + f"B_max constants are for a {_BMAX_DEVICE_TOTAL}-byte GB300, got {total}" + ) + text = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=128, + head_dim=512, + index_head_dim=128, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, + ) + config = SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **{k: v for k, v in text.items() if k != "model_type"}, + ) + with _native_storage(): + mgr = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + max_gpu_total_bytes=total // 2, + host_cache_size=0, + event_buffer_max_size=0, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=40, + num_kv_heads=1, + head_dim=512, + tokens_per_block=128, + max_seq_len=max_seq_len or (16398 if max_draft_len else 16384), + max_batch_size=max_batch_size, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + if pp_stage is None + else Mapping(world_size=2, rank=pp_stage, tp_size=1, pp_size=2, gpus_per_node=4), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=max_num_tokens, + sparse_attn_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ), + pretrained_config=config, + spec_config=DSparkDecodingConfig(max_draft_len=max_draft_len, block_size=max_draft_len) + if max_draft_len + else None, + max_cuda_graph_batch_size=max_batch_size, + disable_overlap_scheduler=False, + ) + assert mgr.v41_layout.kv_storage == "native" + return mgr + + +def _native_swa(mgr, layer: int = 0): + """(the manager's SWA pool [B, 128, 528] uint8 as `layer` sees it, codes view, scales view), geometry + asserted.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + pool = mgr.get_buffers(layer, AT.SWA) + codes, scales = mgr.get_v41_native_kv_views(layer, AT.SWA) + b = pool.shape[0] + assert pool.dtype == torch.uint8 and tuple(pool.shape) == (b, PAGE, NATIVE_ROW), pool.shape + assert pool.stride() == (NATIVE_PAGE_BYTES, NATIVE_ROW, 1), pool.stride() + assert tuple(codes.shape) == (b, PAGE, W_KV) and codes.stride() == pool.stride() + assert tuple(scales.shape) == (b, PAGE, NB_KV) and scales.stride() == pool.stride() + assert codes.data_ptr() == pool.data_ptr() and scales.data_ptr() == pool.data_ptr() + W_KV + return pool, codes, scales + + +def _ceiling_swa_manager(): + """The PP2 stage-1 manager whose native SWA pool is the envelope's ceiling (NATIVE_SWA_CEILING), built as the + stage probe builds it (Mapping(world_size 2, rank 1, tp_size 1, pp_size 2)); (manager, pool, codes, scales), + the page count asserted.""" + mgr = _v41_native_manager(*NATIVE_SWA_CEILING_KEY, pp_stage=NATIVE_SWA_CEILING_STAGE) + try: + pool, codes, scales = _native_swa(mgr, NATIVE_SWA_CEILING_LAYER) + assert pool.shape[0] == NATIVE_SWA_CEILING, pool.shape + except BaseException: + mgr.shutdown() + raise + return mgr, pool, codes, scales + + +def _raw_native_pool(pages: int): + """(uint8 [pages + 1, 128, 528] with a guard page, codes view, scales view), every byte the sentinel.""" + buf = torch.full((pages + 1, PAGE, NATIVE_ROW), BYTE_SENTINEL, dtype=torch.uint8, device=DEV) + return buf, buf[:pages, :, :W_KV], buf[:pages, :, W_KV:] + + +def _native_hazard_slots(num_slots: int) -> list[int]: + """Slots of a 528-byte-row pool on both sides of every byte offset 2^31..2^36, both ends and the last page.""" + cands = {0, 1, PAGE - 1, PAGE} + for limit in (2**31, 2**32, 2**33, 2**34, 2**35, 2**36): + s = limit // NATIVE_ROW + cands |= {s - 1, s, s + 1} + cands |= {num_slots - PAGE - 1, num_slots - PAGE, num_slots - 2, num_slots - 1} + return sorted(c for c in cands if 0 <= c < num_slots) + + +def _shadow_hazard_slots(num_slots: int) -> list[int]: + """Slots of a 1024-byte-row bf16 shadow on both sides of every byte offset 2^31..2^37.""" + cands = set() + for limit in (2**31, 2**32, 2**33, 2**34, 2**35, 2**36, 2**37): + s = limit // (2 * W_KV) + cands |= {s - 1, s} + return sorted(c for c in cands if 0 <= c < num_slots) + + +def _count_not(buf: torch.Tensor, value: int) -> torch.Tensor: + """Device count of elements of the contiguous `buf` that are not `value`, in 2^30-element chunks.""" + whole = buf.view(-1) + step = ( + 1 << 26 + ) # a bool reduction materializes an int64 copy (8 bytes per element): keep chunks small + n = torch.zeros((), dtype=torch.int64, device=DEV) + for i in range(0, whole.numel(), step): + n += (whole[i : i + step] != value).sum() + return n + + +def _shadow(slots: int) -> torch.Tensor: + """A bf16 shadow [slots // 128 + 1, 128, 512] (the native pool's geometry plus a guard page), all sentinel.""" + return _sentinel((slots // PAGE + 1, PAGE, W_KV)) + + +def _check_shadow( + shadow: torch.Tensor, slots: torch.Tensor, expected: torch.Tensor, num_slots: int +) -> None: + """Named (valid) slots hold `expected` bitwise; restored to the sentinel, nothing else in `shadow` changed.""" + flat = shadow.view(-1, W_KV) + s64 = slots.to(torch.int64) + ok = (s64 >= 0) & (s64 < num_slots) + got = flat[s64[ok]] + n_bad = _n_bad(got, expected[ok]) + assert n_bad == 0, f"{n_bad} of {got.numel()} shadow elements in the named slots differ bitwise" + flat.view(torch.int16)[s64[ok]] = SENTINEL16 + n_out = int(_count_not(shadow.view(torch.int16), SENTINEL16)) + assert n_out == 0, f"{n_out} shadow elements outside the named slots were written" + + +def _spread_slots(num_slots: int, n: int, g: torch.Generator) -> torch.Tensor: + """n distinct slots: the last page in order, every hazard slot, the rest spread uniformly over the pool.""" + last_page = torch.arange(num_slots - PAGE, num_slots, device=DEV) + fixed = torch.unique( + torch.cat([last_page, torch.tensor(_native_hazard_slots(num_slots), device=DEV)]) + ) + rest = torch.randint(0, num_slots, (3 * n,), device=DEV, generator=g).unique() + rest = rest[~torch.isin(rest, fixed)] + rest = rest[torch.randperm(rest.numel(), device=DEV, generator=g)][: n - fixed.numel()] + s = torch.cat([last_page, fixed[~torch.isin(fixed, last_page)], rest]).to(torch.int64) + assert s.numel() == n and torch.unique(s).numel() == n + return s + + +def test_native_swa_pools_from_the_cache_manager() -> None: + """Every native configuration's SWA pool from the manager, over the whole max_batch_size axis: for + (max_num_tokens, max_batch_size, DSpark max_draft_len) in {2048, 8192} x 1..32 x {0, 5} (128 managers), + [B, 128, 528] uint8 for all 40 layers, the views its two halves; the pages equal the table. B_max over + the axis is 941,800 at (8192, 1, 0); over max_batch_size {8, 32} it is 854,440; 732,280 at the gates.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + got = {} + for key in _NATIVE_SWA_PAGES_AXIS: + mgr = _v41_native_manager(*key) + try: + pool, _, _ = _native_swa(mgr) + for layer in range(40): + assert mgr.get_buffers(layer, AT.SWA).data_ptr() == pool.data_ptr(), layer + got[key] = pool.shape[0] + finally: + mgr.shutdown() + del mgr + torch.cuda.empty_cache() + over = {k: v for k, v in sorted(got.items()) if v > NATIVE_SWA_BMAX_8_32} + print( + f"\n[native SWA pages at 0.5 x device total, 128 configurations] {got}\n" + f"[above the max_batch_size {{8, 32}} bound {NATIVE_SWA_BMAX_8_32}] {over}" + ) + assert got == _NATIVE_SWA_PAGES_AXIS + assert max(got.values()) == NATIVE_SWA_BMAX and got[NATIVE_SWA_BMAX_KEY] == NATIVE_SWA_BMAX + assert max(got[k] for k in _NATIVE_SWA_PAGES) == NATIVE_SWA_BMAX_8_32 + assert max(got[k] for k in ((2048, 8, 5), (8192, 32, 5))) == NATIVE_SWA_GATE_MAX + + +@pytest.mark.parametrize("rows", ["hazard_int32", "hazard_int64", "spread_12256"]) +def test_native_round_trip_far_end_on_the_manager(rows: str) -> None: + """On the live PP2 stage-1 manager's ceiling native SWA pool (NATIVE_SWA_CEILING_KEY; 1,648,320 pages, + 111.4 GB): the writer's native form puts rows at the far end (both sides of every 2^31..2^36 byte offset of + the pool, the first / last page, the last slot, and the slots on every 2^31..2^37 byte offset of a shadow of + its geometry) or 12,256 rows spread over the whole pool; this op reads them back compact (src_slots, -1 and + past-the-end entries among them give zero rows) into a sentinel out, bitwise equal to the writer's own + bf16-container output of the same rows (and to the recipe), and the native pool is left as the writer left + it. + + The shadow form needs a bf16 shadow of the source pool's geometry, which for the stage ceiling (216 GB) does + not fit on one GB300 beside its source. So after the manager is shut down, the shadow form runs on a + separate native pool of SWA_SHADOW_PAGES (the unsplit ceiling, [1098720, 128, 528], the manager's strides, + a guard page) with its own far-end rows: src_slots = dst_slots into a 144.0 GB shadow, the named slots + equal to the container and nothing else in the shadow or its guard page written. + """ + + def far_rows(num_slots, g): + if rows.startswith("hazard"): + return torch.tensor( + sorted(set(_native_hazard_slots(num_slots)) | set(_shadow_hazard_slots(num_slots))), + dtype=torch.int64, + device=DEV, + ) + return _spread_slots(num_slots, MAX_STEP_ROWS, g) + + sdt = torch.int32 if rows != "hazard_int64" else torch.int64 + mgr, pool, codes, scales = _ceiling_swa_manager() + try: + num_slots = pool.shape[0] * PAGE + pool.fill_(BYTE_SENTINEL) + g = _gen(8540 + len(rows)) + valid = far_rows(num_slots, g) + x = _window_kv_like(valid.numel(), seed=8540) + _write(x, codes, scales, valid) + y = _container(x) + _assert_bitwise(y, _fake_quant_native(x)[0]) + src = torch.cat([valid, torch.tensor([-1, num_slots], device=DEV)]) + expected = torch.cat([y, torch.zeros(2, W_KV, dtype=torch.bfloat16, device=DEV)]) + out = _sentinel((src.numel(), W_KV)) + op(codes, scales, out, src.to(sdt)) + _assert_bitwise(out, expected) + written = pool.view(-1, NATIVE_ROW)[valid] + assert torch.equal(written[:, :W_KV], _fake_quant_native(x)[1]) and torch.equal( + written[:, W_KV:], _fake_quant_native(x)[2] + ) + pool.view(-1, NATIVE_ROW)[valid] = BYTE_SENTINEL + assert int(_count_not(pool, BYTE_SENTINEL)) == 0, "the read changed the native pool" + pool_shape, pool_stride, n_live = tuple(pool.shape), pool.stride(), valid.numel() + finally: + mgr.shutdown() + del mgr, pool, codes, scales + torch.cuda.synchronize() + torch.cuda.empty_cache() + buf = shadow = None + try: + buf, codes, scales = _raw_native_pool(SWA_SHADOW_PAGES) + assert buf[:-1].stride() == pool_stride + ns = SWA_SHADOW_PAGES * PAGE + g = _gen(8547 + len(rows)) + valid = far_rows(ns, g) + x = _window_kv_like(valid.numel(), seed=8547) + _write(x, codes, scales, valid) + y = _container(x) + src = torch.cat([valid, torch.tensor([-1, ns], device=DEV)]) + expected = torch.cat([y, torch.zeros(2, W_KV, dtype=torch.bfloat16, device=DEV)]) + shadow = _shadow(ns) + op(codes, scales, shadow[:-1], src.to(sdt), src.to(sdt)) + torch.cuda.synchronize() + _check_shadow(shadow, src, expected, ns) + print( + f"\n[{rows}] {n_live} rows at the far end of the stage manager's {pool_shape} native pool: compact read " + f"== the writer's bf16 container; {valid.numel()} far-end rows of a ({SWA_SHADOW_PAGES}, 128, 528) pool, " + f"shadow {tuple(shadow.shape)} read == the container, nothing else written" + ) + finally: + buf = codes = scales = shadow = None + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +def test_native_harness_sees_a_wrapped_offset() -> None: + """Control: a read from where a mod-2^32 byte offset would put the row, and a shadow write to where a + mod-2^32 offset would put it, both fail the checks above.""" + buf, codes, scales = _raw_native_pool(NATIVE_SWA_CEILING) + num_slots = NATIVE_SWA_CEILING * PAGE + true = torch.tensor( + [s for s in _native_hazard_slots(num_slots) if s * NATIVE_ROW >= 2**32], + dtype=torch.int64, + device=DEV, + ) + wrapped = (true * NATIVE_ROW % 2**32) // NATIVE_ROW + assert true.numel() >= 9 and not torch.isin(wrapped, true).any() + x = _window_kv_like(2 * true.numel(), seed=8541) + _write(x, codes, scales, torch.cat([true, wrapped])) + y = _container(x)[: true.numel()] + out = _sentinel((true.numel(), W_KV)) + op(codes, scales, out, wrapped) + n_bad = _n_bad(out, y) + print(f"\n[wrapped read] {n_bad} of {out.numel()} elements differ from the true rows") + assert n_bad > out.numel() // 2 + del buf, codes, scales + torch.cuda.empty_cache() + num_slots = SWA_SHADOW_PAGES * PAGE # the shadow form's largest geometry + shadow = _shadow(num_slots) + src = torch.arange(true.numel(), device=DEV) + small, small_s = _pool(1, true.numel(), W_KV, _gen(8542), lo=100, hi=160) + ref = _ref(small.view(-1, W_KV), small_s.view(-1, NB_KV)) + s_true = torch.tensor( + [s for s in _shadow_hazard_slots(num_slots) if s * 2 * W_KV >= 2**32], + dtype=torch.int64, + device=DEV, + ) + assert s_true.numel() >= 8 and s_true.numel() <= true.numel() + s_wrapped = (s_true * 2 * W_KV % 2**32) // (2 * W_KV) + op(small.view(-1, W_KV), small_s.view(-1, NB_KV), shadow[:-1], src[: s_true.numel()], s_wrapped) + with pytest.raises(AssertionError, match="shadow elements in the named slots differ"): + _check_shadow(shadow, s_true, ref[: s_true.numel()], num_slots) + shadow.view(-1, W_KV).view(torch.int16)[s_true] = SENTINEL16 + with pytest.raises(AssertionError, match="outside the named slots were written"): + _check_shadow(shadow, s_true[:0], ref[:0], num_slots) + del shadow + torch.cuda.empty_cache() + + +@pytest.mark.parametrize("t", [32, MAX_STEP_ROWS]) +def test_native_pool_size_invariance(t: int) -> None: + """The same rows at the same slots read through codes[:B] / scales[:B] of the live PP2 stage-1 manager's + ceiling pool (NATIVE_SWA_CEILING_KEY, 1,648,320 pages), compact: the same bits at every B and equal to the + writer's container. The shadow form (into shadow[:B]) runs, after the manager is shut down, on a separate + native pool of SWA_SHADOW_PAGES (the largest whose shadow fits beside it), with the same rows at the same + slots, at every B up to SWA_SHADOW_PAGES: the same bits, nothing else written. Then from a separate native + pool of exactly the smallest size. + + B = the smallest pool holding the slots, the page counts whose byte count first passes 2^31..2^36, a mid + size, the gates' largest (732,280), the max_batch_size {8, 32} bound (854,440), the gates' axis B_max + (941,800), _SWA_EXTRA_B and the ceiling. + """ + mgr, pool, codes, scales = _ceiling_swa_manager() + try: + pool.fill_(BYTE_SENTINEL) + small = -(-t // PAGE) + g = _gen(8543 + t) + last = small * PAGE - 1 + perm = torch.randperm(small * PAGE, device=DEV, generator=g) + slots = torch.cat([torch.tensor([last], device=DEV), perm[perm != last]])[:t] + slots = slots[torch.randperm(t, device=DEV, generator=g)].to(torch.int64) + x = _window_kv_like(t, seed=8543 + t) + _write(x, codes, scales, slots) + y = _container(x) + src = slots.clone() + src[::16] = -1 + expected = torch.where((src >= 0)[:, None], y, torch.zeros_like(y)) + sizes = sorted( + { + small, + *( + lim // NATIVE_PAGE_BYTES + 1 + for lim in (2**31, 2**32, 2**33, 2**34, 2**35, 2**36) + ), + NATIVE_SWA_CEILING // 2 + 1, + NATIVE_SWA_GATE_MAX, + NATIVE_SWA_BMAX_8_32, + NATIVE_SWA_BMAX, + *_SWA_EXTRA_B, + NATIVE_SWA_CEILING, + } + ) + first = None + for b in sizes: + out = _sentinel((src.numel(), W_KV)) + op(codes[:b], scales[:b], out, src) + _assert_bitwise(out, expected) + first = out if first is None else first + assert torch.equal(out.view(torch.int16), first.view(torch.int16)), f"B={b} differs" + exact = pool.view(-1, NATIVE_ROW)[: small * PAGE].clone().view(small, PAGE, NATIVE_ROW) + finally: + mgr.shutdown() + del mgr, pool, codes, scales + torch.cuda.synchronize() + torch.cuda.empty_cache() + _buf = shadow = None + try: + _buf, codes, scales = _raw_native_pool(SWA_SHADOW_PAGES) + _write(x, codes, scales, slots) + shadow = _shadow(SWA_SHADOW_PAGES * PAGE) + for b in (v for v in sizes if v <= SWA_SHADOW_PAGES): + op(codes[:b], scales[:b], shadow[:b], src, src) + torch.cuda.synchronize() + _check_shadow(shadow, src, expected, b * PAGE) + finally: + _buf = codes = scales = shadow = None + torch.cuda.synchronize() + torch.cuda.empty_cache() + out = _sentinel((src.numel(), W_KV)) + op(exact[..., :W_KV], exact[..., W_KV:], out, src) + _assert_bitwise(out, first) + print( + f"\n[native t={t}] B = {sizes} (compact: views of the stage manager's ceiling pool; shadow[:B] for B <= " + f"{SWA_SHADOW_PAGES}: a pool of that geometry) and a separate [{small}, 128, 528] pool: " + f"{int((src >= 0).sum())} " + f"rows bitwise identical and == the writer's container" + ) + + +@pytest.mark.parametrize("form", ["compact", "shadow"]) +def test_native_every_row_count_1_to_16384(form: str) -> None: + """Over the step's referenced rows, every r in 1..16384 (>= the 12,256 distinct SWA rows a + step can reference) reads exactly the first r rows of the 16384-row call. + + The pool: a native [160, 128, 528] (20,480 slots) filled by the writer at every slot. src: 16,384 + distinct slots, 1/16 of them invalid (-1 or one past the end). compact: out is a fixed [16384, 512]; + before each call its first r rows are reset to the sentinel, after it they must equal the first r + rows of the full call (invalid -> zero). shadow: src = dst into a [160, 128, 512] sentinel shadow; + after each call the named valid slots hold the full call's rows, the number of non-sentinel elements + is 512 per valid row, and the named slots are reset. One compiled variant per form over the sweep. + """ + pages = 160 + num_slots = pages * PAGE + g = _gen(8544) + buf = torch.empty(pages, PAGE, NATIVE_ROW, dtype=torch.uint8, device=DEV) + codes, scales = buf[..., :W_KV], buf[..., W_KV:] + x = _window_kv_like(num_slots, seed=8544) + _write(x, codes, scales, torch.arange(num_slots, device=DEV)) + y = _container(x) + src = torch.randperm(num_slots, device=DEV, generator=g)[:SWEEP_ROWS].to(torch.int64) + bad = torch.arange(0, SWEEP_ROWS, 16, device=DEV) + src[bad] = torch.where(bad % 32 == 0, -1, num_slots) + src[1] = -1 + ok = (src >= 0) & (src < num_slots) + full = torch.where( + ok[:, None], + y[src.clamp(0, num_slots - 1)], + torch.zeros(1, W_KV, dtype=torch.bfloat16, device=DEV), + ) + kernel = _dq_op._mxfp8_dequantize_kernel + mism = torch.zeros((), dtype=torch.int64, device=DEV) + cnt_err = torch.zeros((), dtype=torch.int64, device=DEV) + if form == "compact": + out = _sentinel((SWEEP_ROWS, W_KV)) + op(codes, scales, out, src) + _assert_bitwise(out, full) + n1 = _jit_variants(kernel) + o16, f16 = out.view(torch.int16), full.view(torch.int16) + for r in range(1, SWEEP_ROWS + 1): + o16[:r] = SENTINEL16 + op(codes, scales, out, src[:r]) + mism += (o16[:r] != f16[:r]).sum() + n2 = _jit_variants(kernel) + else: + shadow = _sentinel((pages, PAGE, W_KV)) + s16 = shadow.view(-1, W_KV).view(torch.int16) + op(codes, scales, shadow, src, src) + _check_shadow(shadow, src, full, num_slots) + n1 = _jit_variants(kernel) + ok_pos = torch.nonzero(ok).squeeze(1) + vs = src[ok_pos] + f16 = full[ok_pos].view(torch.int16) + cnt = torch.cumsum(ok.to(torch.int64), 0).tolist() + for r in range(1, SWEEP_ROWS + 1): + op(codes, scales, shadow, src[:r], src[:r]) + k = cnt[r - 1] + mism += (s16[vs[:k]] != f16[:k]).sum() + cnt_err += ((s16 != SENTINEL16).sum() - k * W_KV).abs() + s16[vs[:k]] = SENTINEL16 + n2 = _jit_variants(kernel) + print( + f"\n[{form} rows 1..{SWEEP_ROWS}] mismatches {int(mism)}, count error {int(cnt_err)}; variants {n1} -> {n2}" + ) + assert int(mism) == 0 and int(cnt_err) == 0 + assert n2 == n1, "a row count compiled a new variant" + + +def test_native_per_query_window_rows_up_to_1048576() -> None: + """The per-query window form: src (and dst) = a step's [T, 128] window table flattened, -1 and + duplicate entries included, up to 8192 x 128 = 1,048,576 rows. + + One 1,048,576-row call against the writer's container (compact: invalid -> zero; shadow: every + duplicate entry names the same source row, so its writes are identical and the slot holds that row). + Then prefixes r at 128 T - 1, 128 T, 128 T + 1 for T in {128, 256, 512, 1024, 2048, 4096, 8191, 8192} + and 32 random r in (16384, 2^20]: each compact call equals the first r rows of the full call. One + compiled variant. + """ + pages = 160 + num_slots = pages * PAGE + g = _gen(8545) + buf = torch.empty(pages, PAGE, NATIVE_ROW, dtype=torch.uint8, device=DEV) + codes, scales = buf[..., :W_KV], buf[..., W_KV:] + x = _window_kv_like(num_slots, seed=8545) + _write(x, codes, scales, torch.arange(num_slots, device=DEV)) + y = _container(x) + t = 8192 + pos = torch.sort(torch.randint(0, num_slots - 256, (t,), device=DEV, generator=g)).values + win = pos[:, None] - torch.arange(127, -1, -1, device=DEV)[None, :] + win = torch.where(win >= 0, win, -1) + win[::7, :40] = -1 # windows cut by the start of a request + src = win.reshape(-1).to(torch.int32) + ok = src >= 0 + full_ref = torch.where( + ok[:, None], + y[src.long().clamp_min(0)], + torch.zeros(1, W_KV, dtype=torch.bfloat16, device=DEV), + ) + out = _sentinel((src.numel(), W_KV)) + kernel = _dq_op._mxfp8_dequantize_kernel + op(codes, scales, out, src) + _assert_bitwise(out, full_ref) + shadow = _sentinel((pages + 1, PAGE, W_KV)) + op(codes, scales, shadow[:-1], src, src) + torch.cuda.synchronize() + uniq = torch.unique(src[ok].long()) + _check_shadow(shadow, uniq, y[uniq], num_slots) + n1 = _jit_variants(kernel) + rs = sorted( + { + r + for tt in (128, 256, 512, 1024, 2048, 4096, 8191, 8192) + for r in (128 * tt - 1, 128 * tt, 128 * tt + 1) + if r <= src.numel() + } + | {int(v) for v in torch.randint(16385, 2**20 + 1, (32,), generator=g, device=DEV)} + ) + o16 = out.view(torch.int16) + mism = 0 + for r in rs: + o16[:r] = SENTINEL16 + op(codes, scales, out, src[:r]) + mism += _n_bad(out[:r], full_ref[:r]) + n2 = _jit_variants(kernel) + print( + f"\n[per-query window rows] {src.numel()} rows ({int(ok.sum())} valid, {uniq.numel()} distinct): compact " + f"and shadow exact; {len(rs)} prefixes, mismatches {int(mism)}; variants {n1} -> {n2}" + ) + assert int(mism) == 0 and n2 == n1 + + +# Captured batch sizes: every G in 1..32. At max_batch_size s the engine captures every G in 1..s with +# padding off (padding on captures a subset of those plus s), so over the max_batch_size axis 1..32 the +# union is every G in 1..32; the gates' capture set {1..8, 16, 24, 32} is part of it. +_GRAPH_BS = MBS_AXIS +# Rows a captured decode / DSpark-verify step reads: its requests' windows as distinct rows (128 per +# request, 133 with 6-token verify), or per query (128 per query: bs and 6 x bs queries). +_GRAPH_ROWS = tuple( + sorted( + {128 * b for b in _GRAPH_BS} | {133 * b for b in _GRAPH_BS} | {768 * b for b in _GRAPH_BS} + ) +) + + +@pytest.mark.parametrize("form", ["compact", "shadow"]) +def test_native_cuda_graph_every_graph_batch_on_the_bmax_pool(form: str) -> None: + """Per captured row count: an eager warm-up, the capture, then 2 replays after the writer has put new + rows at new far-end slots of a native pool (hazard slots, the last pages): the stage ceiling's geometry + ([1648320, 128, 528]) compact, SWA_SHADOW_PAGES ([1098720, 128, 528]) into a shadow of that geometry. Every + replay reads exactly the writer's container (shadow: nothing else written) and equals the same call run + eagerly.""" + pages = NATIVE_SWA_CEILING if form == "compact" else SWA_SHADOW_PAGES + buf, codes, scales = _raw_native_pool(pages) + num_slots = pages * PAGE + shadow = _shadow(num_slots) if form == "shadow" else None + g = _gen(8546) + hazard = torch.tensor(_native_hazard_slots(num_slots), device=DEV) + near = num_slots - 1 - torch.arange(256 * PAGE, device=DEV) + cand = torch.unique(torch.cat([hazard, near])) + cand = cand[cand != num_slots - 1] + table = torch.full((max(_GRAPH_ROWS),), -1, dtype=torch.int64, device=DEV) + + def pick(n): + """n slots: the pool's last slot, then distinct hazard / last-256-page slots; entry 2 is -1.""" + s = torch.cat( + [ + torch.tensor([num_slots - 1], device=DEV), + cand[torch.randperm(cand.numel(), device=DEV, generator=g)][: n - 1], + ] + ) + if n > 2: + s[2] = -1 + return s + + for n in _GRAPH_ROWS: + src = table[:n] + s0 = pick(n) + src.copy_(s0) + x = _fast_rows(n, seed=8546 + n) + _write(x, codes, scales, s0) + out = _sentinel((n, W_KV)) + if form == "compact": + op(codes, scales, out, src) + else: + op(codes, scales, shadow[:-1], src, src) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + if form == "compact": + op(codes, scales, out, src) + else: + op(codes, scales, shadow[:-1], src, src) + if shadow is not None: + shadow.view(torch.int16).fill_(SENTINEL16) + for rep in (1, 2): + s = pick(n) + x = _fast_rows(n, seed=8546 + 1000 * rep + n) + _write(x, codes, scales, s) + y = _container(x) + expected = torch.where((s >= 0)[:, None], y, torch.zeros_like(y)) + src.copy_(s) + out.view(torch.int16).fill_(SENTINEL16) + graph.replay() + torch.cuda.synchronize() + if form == "compact": + replayed = out.clone() + _assert_bitwise(replayed, expected) + out.view(torch.int16).fill_(SENTINEL16) + op(codes, scales, out, src) + _assert_bitwise(out, replayed) + else: + live = s[s >= 0] + replayed = shadow.view(-1, W_KV)[live].clone() + _check_shadow(shadow, s, expected, num_slots) + op(codes, scales, shadow[:-1], src, src) + torch.cuda.synchronize() + _assert_bitwise(shadow.view(-1, W_KV)[live], replayed) + _check_shadow(shadow, s, expected, num_slots) + del graph + print( + f"\n[{form} graphs on the [{pages}, 128, 528] native pool] rows {list(_GRAPH_ROWS)}: every " + f"replay == eager == the writer's container" + ) + del buf, codes, scales, shadow + torch.cuda.empty_cache() + + +# ─── reads over random multi-request steps on the live manager ─── + + +def _random_step( + rng: random.Random, max_tokens: int = 8192, max_requests: int = 32, max_len: int = 16384 +): + """([(cached, new)] contexts first (random order) then generation requests, num_contexts, gen_len).""" + k = rng.randint(1, max_requests) + gen_len = rng.choice([1, 6]) + g = rng.randint(0, k) if rng.random() < 0.7 else 0 + c = k - g + budget = max_tokens - g * gen_len + t_ctx = rng.randint(c, budget) if c else 0 + cuts = sorted(rng.sample(range(1, t_ctx), c - 1)) if c > 1 else [] + news = [b - a for a, b in zip([0] + cuts, cuts + [t_ctx])] if c else [] + reqs = [(0 if rng.random() < 0.4 else rng.randint(0, max_len - n), n) for n in news] + rng.shuffle(reqs) + reqs += [(rng.randint(0, max_len - gen_len), gen_len) for _ in range(g)] + return reqs, c, gen_len + + +_EDGE_STEPS = [ + ([(0, 8192)], 1, 1), + ([(0, 1)], 1, 1), + ([(16383, 1)], 1, 1), + ([(0, 256)] * 32, 32, 1), + ([(16384 - 256, 256)] * 32, 32, 1), + ([(16383, 1)] * 32, 0, 1), + ([(16378, 6)] * 32, 0, 6), + ([(513, 7000)] + [(16000, 6)] * 31, 1, 6), + ([(0, 127), (0, 128), (0, 129), (127, 1), (128, 1), (8000, 192)], 6, 1), +] +_NUM_RANDOM_STEPS = 200 + + +def _run_step(mgr, reqs, num_contexts: int, gen_len: int, uid: int, max_num_tokens: int = 8192): + """Allocate the step's requests on the manager and build its DeepSeek-V4.1 attention metadata as the + engine builds it (max_num_requests = the manager's max_batch_size, the engine's max_num_tokens; each + request's first token unique over the test, so block reuse shares no page across steps).""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + DeepseekV4TrtllmAttentionMetadata, + ) + from tensorrt_llm._torch.metadata import KVCacheParams + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from tensorrt_llm.bindings import SamplingConfig + from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig + from tensorrt_llm.mapping import Mapping + + requests = [] + for rid, (c, n) in enumerate(reqs): + r = LlmRequest( + request_id=rid, + max_new_tokens=64, + input_tokens=[uid + rid] + list(range(1, c + n)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + assert mgr.prepare_context(r) and mgr.resize_context(r, r.context_chunk_size), (rid, c, n) + requests.append(r) + md = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor([n for _, n in reqs], dtype=torch.int), + request_ids=[r.py_request_id for r in requests], + max_num_requests=mgr.max_batch_size, + num_contexts=num_contexts, + prompt_lens=[c + n for c, n in reqs], + max_num_tokens=max_num_tokens, + kv_cache_manager=mgr, + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=[c for c, _ in reqs] + ), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + sparse_attention_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ), + ) + if num_contexts < len(reqs) and gen_len > 1: + md.update_spec_dec_param( + batch_size=len(reqs), + is_spec_decoding_enabled=False, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, + max_draft_len=gen_len - 1, + max_total_draft_tokens=gen_len - 1, + ) + md.prepare() + return requests, md + + +def test_native_reads_over_random_multi_request_steps() -> None: + """Random multi-request steps on the live accuracy-gate manager (8192, 32, DSpark 5; native SWA 732,280 pages). + + Per step (9 edge steps + 200 seeded random partitions of T <= 8192 into k <= 32 requests, contexts + with random cached prefixes, decode or DSpark-verify generation requests) and per layer 0, 20, 39: + the attention's window rows are taken from the step's metadata (req_idx_per_token, + swa_local_indices_cuda, the layer's SWA sliding block table: page * 128 + position % 128), the + writer's native form puts the step's new rows at md.v41_swa_slot_mapping[local(L)][:T] and the + earlier rows the windows reach at theirs, and this op reads the step's referenced rows + + - compact over the distinct referenced rows, and per query ([T * 128], -1 -> zero rows); + - into a bf16 shadow of the pool's geometry (96 GB): src = dst = the distinct rows, and again + src = dst = the per-query [T * 128] table (duplicates and -1 included); + + every read equal to the writer's container of the row, nothing else in the shadow written; after + each step the native pool outside the step's rows is still the sentinel. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + mgr = _v41_native_manager(8192, 32, 5) + rng = random.Random(8547) + steps = list(_EDGE_STEPS) + [_random_step(rng) for _ in range(_NUM_RANDOM_STEPS)] + kernel = _dq_op._mxfp8_dequantize_kernel + shadow = None + try: + pool, _, _ = _native_swa(mgr) + num_slots = pool.shape[0] * PAGE + assert pool.shape[0] == NATIVE_SWA_GATE_MAX + pool.fill_(BYTE_SENTINEL) + flat = pool.view(-1, NATIVE_ROW) + shadow = _shadow(num_slots) + s16 = shadow.view(-1, W_KV).view(torch.int16) + stats = dict(steps=0, max_t=0, max_k=0, max_distinct=0, max_per_query=0, rows=0) + n_variants = None + for i, (reqs, nctx, gen_len) in enumerate(steps): + requests, md = _run_step(mgr, reqs, nctx, gen_len, uid=500_000 + 32 * i) + try: + t = sum(n for _, n in reqs) + req = md.req_idx_per_token[:t].long() + loc = md.swa_local_indices_cuda[:t].long() + for layer in (0, 20, 39): + local = mgr.layer_offsets[layer] + codes, scales = mgr.get_v41_native_kv_views(layer, AT.SWA) + bt = md.sliding_block_tables[local, AT.SWA.value] + page = torch.where(loc >= 0, bt[req[:, None], loc.clamp_min(0) // PAGE], -1) + win = torch.where((loc >= 0) & (page >= 0), page * PAGE + loc % PAGE, -1) + assert int(((loc >= 0) & (page < 0)).sum()) == 0, ( + "a window position without a page" + ) + new = md.v41_swa_slot_mapping[local][:t] + uniq = torch.unique(win[win >= 0]) + old = uniq[~torch.isin(uniq, new)] + assert torch.isin(new, uniq).all() and uniq.numel() <= MAX_STEP_ROWS + x_new = _fast_rows(t, seed=600_000 + 8 * i + layer) + x_old = _fast_rows(max(old.numel(), 1), seed=700_000 + 8 * i + layer)[ + : old.numel() + ] + _write(x_old, codes, scales, old) + _write(x_new, codes, scales, new) # the layer's forward write + # The writer's container of every referenced row, by slot. + slots_all = torch.cat([new.long(), old]) + y_all = ( + torch.cat([_container(x_new), _container(x_old)]) + if old.numel() + else _container(x_new) + ) + order = torch.argsort(slots_all) + slots_sorted, y_sorted = slots_all[order], y_all[order] + + def lookup(s): + idx = torch.searchsorted(slots_sorted, s.clamp_min(0)) + return torch.where( + (s >= 0)[:, None], + y_sorted[idx.clamp(max=slots_sorted.numel() - 1)], + torch.zeros(1, W_KV, dtype=torch.bfloat16, device=DEV), + ) + + out = _sentinel((uniq.numel(), W_KV)) + op(codes, scales, out, uniq) + _assert_bitwise(out, lookup(uniq)) + per_query = win.reshape(-1) + out_q = _sentinel((per_query.numel(), W_KV)) + op(codes, scales, out_q, per_query) + _assert_bitwise(out_q, lookup(per_query)) + op(codes, scales, shadow[:-1], uniq, uniq) + got = shadow.view(-1, W_KV)[uniq] + _assert_bitwise(got, lookup(uniq)) + s16[uniq] = SENTINEL16 + op(codes, scales, shadow[:-1], per_query, per_query) + _assert_bitwise(shadow.view(-1, W_KV)[uniq], lookup(uniq)) + s16[uniq] = SENTINEL16 + if n_variants is None: + n_variants = _jit_variants(kernel) + flat[slots_all] = BYTE_SENTINEL + stats["max_distinct"] = max(stats["max_distinct"], uniq.numel()) + stats["max_per_query"] = max(stats["max_per_query"], per_query.numel()) + stats["rows"] += uniq.numel() + n_out = int(_count_not(shadow.view(torch.int16), SENTINEL16)) + assert n_out == 0, ( + f"step {i}: {n_out} shadow elements outside the referenced rows were written" + ) + n_out = int(_count_not(pool, BYTE_SENTINEL)) + assert n_out == 0, ( + f"step {i}: {n_out} native bytes outside the step's rows were written" + ) + stats["steps"] += 1 + stats["max_t"] = max(stats["max_t"], t) + stats["max_k"] = max(stats["max_k"], len(reqs)) + finally: + for r_ in requests: + mgr.free_resources(r_) + print( + f"\n[native reads over multi-request steps, SWA {tuple(pool.shape)}, shadow {tuple(shadow.shape)}] " + f"{stats}; variants {n_variants} -> {_jit_variants(kernel)}" + ) + assert stats["steps"] == len(_EDGE_STEPS) + _NUM_RANDOM_STEPS and stats["max_t"] == 8192 + assert stats["max_k"] == 32 + finally: + mgr.shutdown() + del mgr + shadow = None + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +def _fast_rows(t: int, seed: int) -> torch.Tensor: + """bf16 [t, 512] window-KV-like rows, vectorized: one magnitude 2^-30..2^20 per 32-block, one zero block.""" + g = _gen(seed) + mag = torch.randint(-30, 21, (t, NB_KV, 1), device=DEV, generator=g).float() + x = torch.randn(t, NB_KV, BLOCK, device=DEV, generator=g) * torch.exp2(mag) + pick = torch.randint(0, NB_KV, (t, 1, 1), device=DEV, generator=g) + x = torch.where(torch.arange(NB_KV, device=DEV).view(1, NB_KV, 1) == pick, 0.0, x) + return x.view(t, W_KV).to(torch.bfloat16) + + +def test_native_variant_map() -> None: + """Which calls compile a new Triton variant. Rows are grid axis 0 and never a kernel argument, but the + out's slot count is one, and Triton specializes integers (== 1, % 16 == 0, other): a compact out sized + to the step's rows compiles one variant per class (and none for a second member of a class); a + fixed-size compact out or a shadow of the pool's geometry compiles one variant whatever the rows.""" + kernel = _dq_op._mxfp8_dequantize_kernel + buf = torch.empty(64, PAGE, NATIVE_ROW, dtype=torch.uint8, device=DEV) + codes, scales = buf[..., :W_KV], buf[..., W_KV:] + _write(_fast_rows(64 * PAGE, seed=1), codes, scales, torch.arange(64 * PAGE, device=DEV)) + fixed = _sentinel((8192, W_KV)) + shadow = _sentinel((64, PAGE, W_KV)) + added = {} + for form in ("fixed", "shadow"): + for n in (8192, 1, 2, 16, 17, 4096, 8191): + src = torch.randperm(64 * PAGE, device=DEV)[:n] + n0 = _jit_variants(kernel) + if form == "fixed": + op(codes, scales, fixed, src) + else: + op(codes, scales, shadow, src, src) + added[(form, n)] = _jit_variants(kernel) - n0 + classes = {"1": (1, 1), "16k": (16, 32, 48, 8192), "other": (2, 17, 33, 8191)} + for name, ns in classes.items(): + for i, n in enumerate(ns): + src = torch.randperm(64 * PAGE, device=DEV)[:n] + n0 = _jit_variants(kernel) + op(codes, scales, _sentinel((n, W_KV)), src) + added[(f"exact[{name}]#{i}", n)] = _jit_variants(kernel) - n0 + print(f"\n[variant map] new variants per call: {added}") + assert all( + v == 0 for (form, n), v in added.items() if form in ("fixed", "shadow") and n != 8192 + ) + for name, ns in classes.items(): + assert all(added[(f"exact[{name}]#{i}", n)] == 0 for i, n in enumerate(ns) if i), name + + +# ─── the native-storage target's attention read: a fixed-size out, a renumbered index table ─── +# +# The native-storage target (the unified sm_103 DeepSeek-V4.1 modeling with native KV storage) reads a layer's +# window KV in one call per layer and step: +# +# mxfp8_dequantize(codes, scales, swa_c, src_slots) +# +# - (codes, scales) = mgr.get_v41_native_kv_views(L, SWA), the live manager's native views; +# - swa_c is ONE persistent bf16 buffer [max_num_tokens * 128, 512], allocated once per manager from the +# metadata (md.swa_local_indices_cuda is [max_num_tokens, 128]) and zero-filled, never reset; +# - src_slots is the window half of the attention's int32 index table, flattened to [T * 128]: the table is +# [T, 128] at a ratio-0 layer (a view) and [T, 128 + 512] at a compressed layer (the window half split +# off, so the flatten is a copy). Entries are pool slots (page * 128 + position % 128) or -1. +# +# Query t's column c lands at row t * 128 + c (a -1 column gives a zero row), and the attention is handed +# the index table with its window half replaced by that fixed numbering, arange(max_num_tokens * 128) +# .view(max_num_tokens, 128)[:T], and its compressed half unchanged. Rows >= T * 128 keep whatever an +# earlier call left and are never named. The attention on bf16 storage instead reads the bf16-container +# pool at the original index; the tests below hold the two reads bitwise equal. + +_KVQ_ROWS = { + 2048: 262_144, + 8192: 1_048_576, +} # swa_c rows at each max_num_tokens of the gates' envelope +# swa_c rows at every max_num_tokens of the engine envelope (max_num_tokens * 128, measured below). +_ENVELOPE_ROWS = {2048: 262_144, 4096: 524_288, 8192: 1_048_576, 16384: 2_097_152} +# Every graph batch of the engine envelope: every G 1..32 and TRT-LLM's padded list to max_batch_size 256. +_ENVELOPE_GSET = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +assert len(_ENVELOPE_GSET) == 46 +# Row counts up to this buffer size are run one by one; the 2,097,152-row buffer's call takes over 1 ms +# (measured in the test), so its sweep is the sampled set allowed for a call of over 1 ms (_kvq_row_counts). +_ENVELOPE_EVERY_ROW_MAX = 1_048_576 +_KVQ_LAYERS = { + 0: 0, + 2: 2, + 20: 1, + 39: 1, +} # layer -> compress ratio (0: a [T, 128] table, else [T, 640]) +_KVQ_LOGICAL_PAGES = ( + 129 # logical pages of a request at max_seq_len 16398 (DSpark): positions <= 16397 +) +S64 = 0x7FAB7FAB7FAB7FAB # four SENTINEL16 as one int64 + + +def _kvq_random_step( + rng: random.Random, max_tokens: int = 8192, max_requests: int = 32, seq_len: int = 16384 +): + """_random_step over the engine's envelope with DSpark: contexts end at <= seq_len tokens, decode + requests (next_n 1) at <= seq_len, DSpark-verify requests (next_n 6) at <= seq_len + 14 (positions to + 16397 at the gates' 16384).""" + k = rng.randint(1, max_requests) + gen_len = rng.choice([1, 6]) + g = rng.randint(0, k) if rng.random() < 0.7 else 0 + c = k - g + budget = max_tokens - g * gen_len + t_ctx = rng.randint(c, budget) if c else 0 + cuts = sorted(rng.sample(range(1, t_ctx), c - 1)) if c > 1 else [] + news = [b - a for a, b in zip([0] + cuts, cuts + [t_ctx])] if c else [] + reqs = [(0 if rng.random() < 0.4 else rng.randint(0, seq_len - n), n) for n in news] + rng.shuffle(reqs) + end = seq_len + 14 if gen_len == 6 else seq_len + reqs += [(rng.randint(0, end - gen_len), gen_len) for _ in range(g)] + return reqs, c, gen_len + + +_KVQ_EDGE_STEPS = list(_EDGE_STEPS) + [ + ([(16392, 6)] * 32, 0, 6), # DSpark verify at the engine's max_seq_len 16398 + ([(8384, 8000)] + [(16392, 6)] * 31, 1, 6), + ([(0, 8192 - 31 * 6)] + [(16392 - 97 * i, 6) for i in range(31)], 1, 6), +] + + +def _kvq_page_candidates(num_pages: int, n: int, g: torch.Generator) -> torch.Tensor: + """n distinct pool pages: every page holding a hazard slot (both sides of the 2^31..2^35 byte lines) in + random order, then the pool's last pages and random pages alternately.""" + hz = torch.tensor( + sorted({s // PAGE for s in _native_hazard_slots(num_pages * PAGE)}), device=DEV + ) + hz = hz[torch.randperm(hz.numel(), device=DEV, generator=g)] + last = num_pages - 1 - torch.arange(n, device=DEV) + last = last[~torch.isin(last, hz)] + rnd = torch.randint(0, num_pages, (4 * n + 64,), device=DEV, generator=g).unique() + rnd = rnd[~torch.isin(rnd, torch.cat([hz, last]))] + rnd = rnd[torch.randperm(rnd.numel(), device=DEV, generator=g)][: last.numel()] + tail = torch.stack([last[: rnd.numel()], rnd], 1).reshape(-1) + cand = torch.cat([hz, tail, last[rnd.numel() :]])[:n] + assert cand.numel() == n and torch.unique(cand).numel() == n + return cand.to(torch.int64) + + +def _kvq_window(reqs, num_pages: int, g: torch.Generator) -> torch.Tensor: + """[T, 128] int64 window table of a step whose request b holds positions c_b .. c_b + n_b - 1: the query + at position p names positions p - 127 .. p (negative -> -1). Each (request, logical page) a window + touches gets its own pool page from _kvq_page_candidates (so the hazard pages are always among the rows + read); position q of request b is at slot page(b, q // 128) * 128 + q % 128.""" + req_of = torch.cat( + [torch.full((n,), b, dtype=torch.int64) for b, (_, n) in enumerate(reqs)] + ).to(DEV) + pos = torch.cat([torch.arange(c, c + n) for c, n in reqs]).to(DEV) + wpos = pos[:, None] - (PAGE - 1) + torch.arange(PAGE, device=DEV)[None] + q = wpos.clamp_min(0) + key = req_of[:, None] * _KVQ_LOGICAL_PAGES + q // PAGE + touched = torch.unique(key[wpos >= 0]) + page_of = torch.full((len(reqs) * _KVQ_LOGICAL_PAGES,), -1, dtype=torch.int64, device=DEV) + page_of[touched] = _kvq_page_candidates(num_pages, touched.numel(), g) + return torch.where(wpos >= 0, page_of[key] * PAGE + q % PAGE, -1) + + +def _kvq_partition(t: int, rng: random.Random, k_max: int = 32): + """t query tokens as contexts of up to k_max requests ending at <= 16384: request 0 from position 0 + (its first windows reach before the sequence: -1 columns), the others at random cached prefixes.""" + k = min(k_max, t, rng.randint(1, k_max)) + cuts = sorted(rng.sample(range(1, t), k - 1)) if k > 1 else [] + news = [b - a for a, b in zip([0] + cuts, cuts + [t])] + return [(0 if b == 0 else rng.randint(0, 16384 - n), n) for b, n in enumerate(news)] + + +def _kvq_expected( + win_flat: torch.Tensor, slots_sorted: torch.Tensor, y_sorted: torch.Tensor +) -> torch.Tensor: + """The bf16-container row of each entry of win_flat (slots sorted with their container rows; -1 -> +0.0).""" + idx = torch.searchsorted(slots_sorted, win_flat.clamp_min(0)).clamp( + max=slots_sorted.numel() - 1 + ) + out = y_sorted[idx] + out[win_flat < 0] = 0.0 + return out + + +def test_kvq_read_buffer_rows_from_the_metadata() -> None: + """The window read buffer the target allocates, [md.swa_local_indices_cuda.shape[0] * 128, 512], over the + whole max_batch_size axis: in every configuration (max_num_tokens, max_batch_size, DSpark max_draft_len) + of {2048, 8192} x 1..32 x {0, 5} (128 managers; max_seq_len 16384 without DSpark, 16398 with it), + metadata built as the engine builds it on the native manager of that configuration, and the CUDA-graph + metadata the engine derives from it (create_cuda_graph_metadata at G = 1 and G = max_batch_size), give + swa_local_indices_cuda [max_num_tokens, 128] int32. So the buffer is max_num_tokens * 128 rows -- + 262,144 at 2048 and 1,048,576 at 8192 -- whatever the max_batch_size or DSpark setting. The engine's + max_seq_len and the compressed tables' [max_batch_size, 132] are recorded alongside.""" + got, seen = {}, set() + for key in _NATIVE_SWA_PAGES_AXIS: + max_t, seqs, draft = key + mgr = _v41_native_manager(*key) + try: + assert mgr.max_seq_len == (16398 if draft else 16384) and mgr.max_batch_size == seqs + requests, md = _run_step(mgr, [(0, 100)], 1, 1, uid=950_000, max_num_tokens=max_t) + try: + rows = set() + for G in sorted({1, seqs}): + gmd = md.create_cuda_graph_metadata(G, False, draft) + for m in (md, gmd): + sw = m.swa_local_indices_cuda + assert tuple(sw.shape) == (max_t, PAGE) and sw.dtype == torch.int32, ( + key, + G, + sw.shape, + ) + rows.add(sw.shape[0] * PAGE) + seen.add( + tuple(m.compress_block_tables[1].shape) + (mgr.max_blocks_per_seq,) + ) + assert tuple(m.compress_block_tables[1].shape) == ( + seqs, + mgr.max_blocks_per_seq, + ) + assert len(rows) == 1 + got[key] = rows.pop() + finally: + for r_ in requests: + mgr.free_resources(r_) + finally: + mgr.shutdown() + del mgr + torch.cuda.empty_cache() + per_t = {t: sorted({v for k, v in got.items() if k[0] == t}) for t in _KVQ_ROWS} + print( + f"\n[native-storage window read buffer rows over 128 configurations] rows per max_num_tokens {per_t}; " + f"compressed tables (max_batch_size, cols, max_blocks_per_seq) seen {sorted(seen)}" + ) + assert got == {k: _KVQ_ROWS[k[0]] for k in _NATIVE_SWA_PAGES_AXIS} + assert {c for _, c, _ in seen} == {132} + + +def _kvq_row_counts(n_rows: int) -> list[int]: + """Every row count 1..n_rows up to _ENVELOPE_EVERY_ROW_MAX; above it (a call of over 1 ms) every + r in 1..16384 and every r = {0, 1, -1} mod 128 (the rows of one query's window) up to n_rows, which holds + every row count the target's call makes (128 T) and every dumped one.""" + if n_rows <= _ENVELOPE_EVERY_ROW_MAX: + return list(range(1, n_rows + 1)) + rs = set(range(1, 16385)) + for t in range(1, n_rows // PAGE + 1): + rs |= {t * PAGE - 1, t * PAGE, t * PAGE + 1} + return sorted(r for r in rs if r <= n_rows) + + +@pytest.mark.parametrize("n_rows", sorted(_ENVELOPE_ROWS.values())) +def test_kvq_fixed_out_every_row_count(n_rows: int) -> None: + """Every row count in the target's form, on the live accuracy-gate manager's native pool (8192, 32, 5): + out is the fixed [n_rows, 512] read buffer (every max_num_tokens 2048, 4096, 8192, 16384 of the + envelope), src_slots an int32 window table of n_rows / 128 queries (up to 32 requests, -1 columns, rows + repeated by neighbouring queries, the pool's hazard and last pages) flattened. The n_rows call equals the + writer's bf16 container; its time (median of 11 CUDA-event timings) is printed. Then every r of + _kvq_row_counts(n_rows): out[:r] reset to the sentinel, the call with src_slots[:r], out[:r] bitwise the + first r rows of the n_rows call and row r still the sentinel (a write past r would persist until its own + row's check); a guard past the buffer untouched; no compiled variant added over the sweep.""" + mgr = _v41_native_manager(8192, 32, 5) + try: + pool, codes, scales = _native_swa(mgr) + g = _gen(15002 + n_rows) + rng = random.Random(15002 + n_rows) + reqs = _kvq_partition(n_rows // PAGE, rng) + win = _kvq_window(reqs, pool.shape[0], g) + uniq = torch.unique(win[win >= 0]) + x = _fast_rows(uniq.numel(), seed=15003 + n_rows) + _write(x, codes, scales, uniq) + src = win.reshape(-1).to(torch.int32) + big = _sentinel((n_rows + PAGE, W_KV)) + out = big[:n_rows] + op(codes, scales, out, src) + _assert_bitwise(out, _kvq_expected(win.reshape(-1), uniq, _container(x))) + full = out.clone() + times = [] + for _ in range(11): + ev0, ev1 = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + ev0.record() + op(codes, scales, out, src) + ev1.record() + torch.cuda.synchronize() + times.append(ev0.elapsed_time(ev1)) + assert torch.equal(out.view(torch.int16), full.view(torch.int16)) + assert n_rows <= _ENVELOPE_EVERY_ROW_MAX or sorted(times)[5] > 1.0, ( + f"the sampled buffer's call takes {times} ms" + ) + kernel = _dq_op._mxfp8_dequantize_kernel + n1 = _jit_variants(kernel) + o64, f64 = out.view(torch.int64), full.view(torch.int64) + # One preallocated bool buffer and a logical reduction: a bool .sum() would materialize an int64 + # copy of every compare (1 GB at the largest r) and churn the allocator for a million sizes. + neq = torch.empty((n_rows, W_KV // 4), dtype=torch.bool, device=DEV) + mism = torch.zeros((), dtype=torch.int64, device=DEV) + past = torch.zeros((), dtype=torch.int64, device=DEV) + o64.fill_(S64) + counts = _kvq_row_counts(n_rows) + for r in counts: + o64[:r].fill_(S64) + op(codes, scales, out, src[:r]) + torch.ne(o64[:r], f64[:r], out=neq[:r]) + mism += neq[:r].any() + if r < n_rows: + past += (o64[r] != S64).any() + n2 = _jit_variants(kernel) + guard = int(_count_not(big[n_rows:].view(torch.int16), SENTINEL16)) + print( + f"\n[fixed out [{n_rows}, 512], {len(reqs)} requests, {uniq.numel()} distinct rows, " + f"{int((src < 0).sum())} -1 entries; the {n_rows}-row call {sorted(times)[5]:.3f} ms] " + f"{len(counts)} row counts ({'every r' if n_rows <= _ENVELOPE_EVERY_ROW_MAX else 'sampled'} in " + f"1..{n_rows}): with a prefix mismatch {int(mism)}, with a write past r {int(past)}, guard {guard}; " + f"variants {n1} -> {n2}" + ) + assert int(mism) == 0 and int(past) == 0 and guard == 0 + assert n2 == n1, "a row count compiled a new variant" + finally: + mgr.shutdown() + del mgr + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +def _kvq_window_read_steps(mgr, steps, max_num_tokens: int, uid0: int, seed0: int): + """Run the target's window read over `steps` on the live manager `mgr` (see + test_kvq_window_read_renumbered_over_random_steps for what each step checks). Returns + (pool shape, buffer shape, stats, controls).""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + container = swa_c = None + try: + pool, _, _ = _native_swa(mgr) + container = torch.empty((pool.shape[0], PAGE, W_KV), dtype=torch.bfloat16, device=DEV) + cflat = container.view(-1, W_KV) + g = _gen(seed0) + stats = dict( + steps=0, + layer_steps=0, + max_t=0, + max_k=0, + max_pos=0, + max_valid_cols=0, + rows=0, + zero_rows=0, + max_distinct=0, + ) + ctrl = {"renumbered off by one": [0, 0], "stale buffer row": [0, 0]} + for i, (reqs, nctx, gen_len) in enumerate(steps): + requests, md = _run_step( + mgr, reqs, nctx, gen_len, uid=uid0 + 64 * i, max_num_tokens=max_num_tokens + ) + try: + if swa_c is None: + max_t = md.swa_local_indices_cuda.shape[0] + assert max_t == max_num_tokens + swa_c = torch.zeros((max_t * PAGE, W_KV), dtype=torch.bfloat16, device=DEV) + win_idx = torch.arange(max_t * PAGE, dtype=torch.int32, device=DEV).view( + max_t, PAGE + ) + t = sum(n for _, n in reqs) + assert len(reqs) <= mgr.max_batch_size and t <= max_num_tokens + req = md.req_idx_per_token[:t].long() + loc = md.swa_local_indices_cuda[:t].long() + for layer, ratio in _KVQ_LAYERS.items(): + local = mgr.layer_offsets[layer] + codes, scales = mgr.get_v41_native_kv_views(layer, AT.SWA) + bt = md.sliding_block_tables[local, AT.SWA.value] + page = torch.where(loc >= 0, bt[req[:, None], loc.clamp_min(0) // PAGE], -1) + win = torch.where((loc >= 0) & (page >= 0), page * PAGE + loc % PAGE, -1) + assert int(((loc >= 0) & (page < 0)).sum()) == 0, ( + "a window position without a page" + ) + new = md.v41_swa_slot_mapping[local][:t].long() + uniq = torch.unique(win[win >= 0]) + old = uniq[~torch.isin(uniq, new)] + assert torch.isin(new, uniq).all() + x_new = _fast_rows(t, seed=seed0 + 2_000_000 + 8 * i + layer) + x_old = _fast_rows(max(old.numel(), 1), seed=seed0 + 3_000_000 + 8 * i + layer)[ + : old.numel() + ] + for x, s in ((x_old, old), (x_new, new)): + _write(x, codes, scales, s) + torch.ops.trtllm.mxfp8_fake_quantize(x, container, s, None) + if ratio: + gidx = torch.randint( + -1, 1 << 20, (t, PAGE + 512), dtype=torch.int32, device=DEV, generator=g + ) + gidx[:, :PAGE] = win + win_v, rest = gidx.split([PAGE, 512], 1) + src = win_v.reshape(t * PAGE) + else: + gidx, rest = win.to(torch.int32), None + src = gidx.reshape(t * PAGE) + before = swa_c.clone() + op(codes, scales, swa_c, src) + ren = win_idx[:t] + table = ren if rest is None else torch.cat([ren, rest], 1) + if rest is not None: + assert torch.equal(table[:, PAGE:], gidx[:, PAGE:]) + named = table[:, :PAGE].reshape(-1).long() + wf = win.reshape(-1) + ref = cflat[wf.clamp_min(0)] + ref[wf < 0] = 0.0 + n_bad = _n_bad(swa_c[named], ref) + assert n_bad == 0, ( + f"step {i} layer {layer}: {n_bad} elements differ from the bf16 path" + ) + n_tail = _n_bad(swa_c[t * PAGE :], before[t * PAGE :]) + assert n_tail == 0, ( + f"step {i} layer {layer}: {n_tail} elements past row {t * PAGE} changed" + ) + valid = bool((wf >= 0).any()) + for name, rows_of in ( + ( + "renumbered off by one", + lambda: swa_c[(named + 1).clamp(max=swa_c.shape[0] - 1)], + ), + ("stale buffer row", lambda before=before: before[named]), + ): + c = _n_bad(rows_of(), ref) + ctrl[name][0] += int(c > 0 or not valid) + ctrl[name][1] += c + del before, ref + stats["layer_steps"] += 1 + stats["rows"] += t * PAGE + stats["zero_rows"] += int((wf < 0).sum()) + stats["max_valid_cols"] = max( + stats["max_valid_cols"], int((win >= 0).sum(1).max()) + ) + stats["max_distinct"] = max(stats["max_distinct"], uniq.numel()) + stats["steps"] += 1 + stats["max_t"] = max(stats["max_t"], t) + stats["max_k"] = max(stats["max_k"], len(reqs)) + stats["max_pos"] = max(stats["max_pos"], max(c + n - 1 for c, n in reqs)) + finally: + for r_ in requests: + mgr.free_resources(r_) + torch.cuda.empty_cache() # per-step gathers vary in size; keep the cache from crowding the pools + return tuple(pool.shape), tuple(swa_c.shape), stats, ctrl + finally: + container = swa_c = None + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +def test_kvq_window_read_renumbered_over_random_steps() -> None: + """Random multi-request steps in the target's exact call, on the live accuracy-gate manager (8192, 32, DSpark 5). + + Per step (12 edge steps, among them DSpark verify at positions up to 16397, + 200 seeded random + partitions of T <= 8192 into k <= 32 requests: contexts with random cached prefixes, decode (next_n 1) + or DSpark-verify (next_n 6) requests) and per layer 0 (ratio 0), 2 (ratio 2), 20 and 39 (ratio 1): + + - the window half of the attention's index table from the step's metadata (req_idx_per_token, + swa_local_indices_cuda, the layer's SWA sliding block table: page * 128 + position % 128, or -1); + - the forward's window writes, twice from the same bf16 rows at the same slots: the writer's native form + into the manager's pool, and its bf16-container form into a bf16 pool of the pool's geometry (96 GB); + - the target's read: its [T, 128] or [T, 640] int32 table, window half flattened, into ONE persistent + zero-filled [md.swa_local_indices_cuda.shape[0] * 128, 512] buffer shared by every layer and step; + - the renumbered table: window half = arange(max_t * 128).view(max_t, 128)[:T], compressed half + unchanged. + + For every (t, c): the buffer row the renumbered table names is bitwise the bf16-container pool row the + original table names (+0.0 where it is -1), and the buffer past T * 128 is exactly as before the call. + Controls on the same state: the renumbered table off by one row, and the buffer as it was before this + call (a stale row), each differ on every layer-step with a valid window row. + """ + mgr = _v41_native_manager(8192, 32, 5) + rng = random.Random(15004) + steps = list(_KVQ_EDGE_STEPS) + [_kvq_random_step(rng) for _ in range(_NUM_RANDOM_STEPS)] + try: + assert _native_swa(mgr)[0].shape[0] == NATIVE_SWA_GATE_MAX + pool_shape, buf_shape, stats, ctrl = _kvq_window_read_steps( + mgr, steps, 8192, uid0=1_000_000, seed0=15004 + ) + print( + f"\n[native-storage window read over multi-request steps, SWA {pool_shape}, buffer " + f"{buf_shape}] {stats}; controls [layer-steps detected, elements]: {ctrl}" + ) + assert stats["steps"] == len(_KVQ_EDGE_STEPS) + _NUM_RANDOM_STEPS and stats["max_t"] == 8192 + assert ( + stats["max_k"] == 32 and stats["max_pos"] == 16397 and stats["max_valid_cols"] == PAGE + ) + for name, (hit, _) in ctrl.items(): + assert hit == stats["layer_steps"], ( + f"control '{name}' missed {stats['layer_steps'] - hit} layer-steps" + ) + finally: + mgr.shutdown() + del mgr + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +# The max_batch_size axis: the window read on a manager built at each max_batch_size. +_KVQ_MBS_CONFIGS = [(8192, s, 5) for s in MBS_AXIS] + [ + (2048, 1, 5) +] # (2048, 1, 5): the native-storage witness +_KVQ_MBS_RANDOM = {s: (_NUM_RANDOM_STEPS if s in MBS_LISTED else 20) for s in MBS_AXIS} + + +def _kvq_axis_edge_steps(max_t: int, seqs: int): + """Edge steps that fit a (max_num_tokens, max_batch_size) configuration: the file's edge steps with at + most `seqs` requests and `max_t` tokens, plus a full step of max_t tokens from position 0 (-1 window + columns) and one ending at position 16383, `seqs` decode requests at 16383, `seqs` DSpark-verify + requests at positions 16392..16397, and (seqs > 1) a context filling the step beside seqs - 1 verifies.""" + cands = [ + st for st in _KVQ_EDGE_STEPS if len(st[0]) <= seqs and sum(n for _, n in st[0]) <= max_t + ] + cands += [ + ([(0, max_t)], 1, 1), + ([(16384 - max_t, max_t)], 1, 1), + ([(16383, 1)] * seqs, 0, 1), + ([(16392, 6)] * seqs, 0, 6), + ] + if seqs > 1: + cands.append( + ([(0, max_t - 6 * (seqs - 1))] + [(16392 - 97 * i, 6) for i in range(seqs - 1)], 1, 6) + ) + steps = [] + for st in cands: + if st not in steps: + steps.append(st) + return steps + + +@pytest.mark.parametrize( + "cfg", _KVQ_MBS_CONFIGS, ids=[f"{t}x{s}x{d}" for t, s, d in _KVQ_MBS_CONFIGS] +) +def test_kvq_window_read_renumbered_at_every_max_batch_size(cfg) -> None: + """The max_batch_size axis of the target's window read (random multi-request steps at every + max_batch_size): the live native manager built at (max_num_tokens, max_batch_size, DSpark 5) for every + max_batch_size 1..32 at 8192, and the native witness's (2048, 1, 5). Per configuration: the edge steps that fit it + (_kvq_axis_edge_steps: a full max_num_tokens step, max_batch_size requests at the DSpark tail) and + seeded random partitions of T <= max_num_tokens into k <= max_batch_size requests (200 at max_batch_size + 1, 2, 3, 4, 8, 16, 32 and at the witness, 20 at every other), each step exactly as in + test_kvq_window_read_renumbered_over_random_steps: layers 0 / 2 / 20 / 39, the buffer + [md.swa_local_indices_cuda.shape[0] * 128, 512] (asserted = max_num_tokens * 128), every renumbered row + bitwise the bf16-container row of that manager's pool geometry (+0.0 for -1), rows past 128 T + unchanged, and both controls seen on every layer-step.""" + max_t, seqs, draft = cfg + n_random = 200 if cfg == (2048, 1, 5) else _KVQ_MBS_RANDOM[seqs] + rng = random.Random(15_100 + 100 * seqs + max_t) + edges = _kvq_axis_edge_steps(max_t, seqs) + steps = edges + [ + _kvq_random_step(rng, max_tokens=max_t, max_requests=seqs) for _ in range(n_random) + ] + mgr = _v41_native_manager(*cfg) + try: + pages = _native_swa(mgr)[0].shape[0] + assert pages == _NATIVE_SWA_PAGES_AXIS[cfg] and pages <= NATIVE_SWA_BMAX + pool_shape, buf_shape, stats, ctrl = _kvq_window_read_steps( + mgr, steps, max_t, uid0=2_000_000 + 100_000 * seqs, seed0=15_200 + 1_000 * seqs + max_t + ) + print( + f"\n[native-storage window read, max_batch_size {seqs}, max_num_tokens {max_t}, SWA {pool_shape}, buffer " + f"{buf_shape}] {len(edges)} edge + {n_random} random steps: {stats}; controls: {ctrl}" + ) + assert buf_shape == (max_t * PAGE, W_KV) + assert stats["steps"] == len(steps) and stats["max_t"] == max_t and stats["max_k"] == seqs + assert stats["max_pos"] == 16397 and stats["max_valid_cols"] == PAGE + for name, (hit, _) in ctrl.items(): + assert hit == stats["layer_steps"], ( + f"control '{name}' missed {stats['layer_steps'] - hit} layer-steps" + ) + finally: + mgr.shutdown() + del mgr + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +@pytest.mark.parametrize("n_rows", sorted(_ENVELOPE_ROWS.values())) +def test_kvq_window_read_cuda_graph_every_graph_batch_on_the_bmax_pool(n_rows: int) -> None: + """The target's read under CUDA-graph capture, on the live PP2 stage-1 manager's native pool + (NATIVE_SWA_CEILING_KEY; 1,648,320 pages, the envelope's measured maximum), into the fixed [n_rows, 512] + buffer of every + max_num_tokens 2048, 4096, 8192, 16384 (zero-filled, a sentinel guard past it). + + First the buffer's largest step, eager: n_rows / 128 queries whose windows lie on the pool's hazard + and last pages, equal to the writer's container. That call is the only warm-up. Then per batch in every + G of _ENVELOPE_GSET (1..32 and 40..128 step 8, 192, 256: the capture sets of every max_batch_size up to 256, + padding on or off; the graph depends only on T) x next_n {1, 6} (decode, DSpark verify; T = G x next_n + queries, 1..1536, at positions up to 16383 / 16397) and per table form + ([T, 128] view; [T, 640] with the window half split and flattened inside the graph): capture, then 2 + replays after the writer puts new rows at new far-end slots and the table changes in place. Each + replay: rows < T * 128 equal the container (+0.0 for -1), rows past them unchanged, and the same call + run eagerly gives the same bits. No compiled variant is added after the warm-up, and the guard stays + the sentinel. + """ + mgr, pool, codes, scales = _ceiling_swa_manager() + try: + g = _gen(15005 + n_rows) + rng = random.Random(15005 + n_rows) + max_t = n_rows // PAGE + big = torch.zeros((n_rows + PAGE, W_KV), dtype=torch.bfloat16, device=DEV) + big[n_rows:].view(torch.int16).fill_(SENTINEL16) + swa_c = big[:n_rows] + tab640 = torch.full((max_t, PAGE + 512), -1, dtype=torch.int32, device=DEV) + tab128 = torch.full((max_t, PAGE), -1, dtype=torch.int32, device=DEV) + kernel = _dq_op._mxfp8_dequantize_kernel + + def fresh(reqs): + win = _kvq_window(reqs, pool.shape[0], g) + uniq = torch.unique(win[win >= 0]) + x = _fast_rows(uniq.numel(), seed=rng.randint(0, 2**30)) + _write(x, codes, scales, uniq) + return win, _kvq_expected(win.reshape(-1), uniq, _container(x)) + + win, expected = fresh(_kvq_partition(max_t, rng)) + tab640[:, :PAGE] = win + op(codes, scales, swa_c, tab640.split([PAGE, 512], 1)[0].reshape(n_rows)) + _assert_bitwise(swa_c, expected) + n1 = _jit_variants(kernel) + graphs = 0 + for bs in _ENVELOPE_GSET: + for next_n in (1, 6): + t = bs * next_n + end = 16398 if next_n == 6 else 16384 + for form in ("ratio0", "compressed"): + tab = tab128 if form == "ratio0" else tab640 + win, _ = fresh([(rng.randint(0, end - next_n), next_n) for _ in range(bs)]) + tab[:t, :PAGE] = win + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + src = ( + tab[:t].reshape(t * PAGE) + if form == "ratio0" + else tab[:t].split([PAGE, 512], 1)[0].reshape(t * PAGE) + ) + op(codes, scales, swa_c, src) + graphs += 1 + for rep in (1, 2): + kv = [rng.randint(0, end - next_n) for _ in range(bs)] + if rep == 1: + kv[0] = end - next_n # the last position the engine reaches + win, expected = fresh([(c, next_n) for c in kv]) + tab[:t, :PAGE] = win + before = swa_c.clone() + graph.replay() + torch.cuda.synchronize() + _assert_bitwise(swa_c[: t * PAGE], expected) + _assert_bitwise(swa_c[t * PAGE :], before[t * PAGE :]) + replayed = swa_c[: t * PAGE].clone() + swa_c.copy_(before) + eager_src = ( + tab[:t].reshape(t * PAGE) + if form == "ratio0" + else tab[:t].split([PAGE, 512], 1)[0].reshape(t * PAGE) + ) + op(codes, scales, swa_c, eager_src) + _assert_bitwise(swa_c[: t * PAGE], replayed) + del before + del graph + n2 = _jit_variants(kernel) + guard = int(_count_not(big[n_rows:].view(torch.int16), SENTINEL16)) + print( + f"\n[native-storage window read graphs, pool {tuple(pool.shape)}, buffer [{n_rows}, 512]] eager {max_t} " + f"queries on far-end pages == the container; {graphs} graphs (bs {list(_ENVELOPE_GSET)} x next_n 1/6 x 2 " + f"table forms): every replay == eager == the container, rows past T * 128 unchanged; variants after " + f"the first eager call {n1} -> {n2}; guard {guard}" + ) + assert n2 == n1 and guard == 0 and graphs == 2 * 2 * len(_ENVELOPE_GSET) + finally: + mgr.shutdown() + del mgr + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +# ─── the engine envelope: max_num_tokens to 16,384, max_batch_size to 256, max_seq_len to 1,048,576 ─── +# +# The target's call is the one above, unchanged: mxfp8_dequantize(codes, scales, swa_c, win.reshape(T * 128)), +# with swa_c [md.swa_local_indices_cuda.shape[0] * 128, 512]. The envelope moves T (up to max_num_tokens +# 16,384: 2,097,152 rows; under capture T = G x next_n for every G of _ENVELOPE_GSET: 128 * T up to 196,608 rows) and +# the buffer (every max_num_tokens 2048, 4096, 8192, 16384). The call has no width that follows max_seq_len. + +# Native SWA pages at 0.5 x this GB300's total memory per (max_num_tokens, max_batch_size, DSpark +# max_draft_len, max_seq_len as L), measured. +_ENVELOPE_SWA_PAGES = { + (2048, 1, 0, 16384): 549_600, + (2048, 1, 5, 16384): 549_600, + (2048, 32, 0, 16384): 366_400, + (2048, 32, 5, 16384): 569_440, + (2048, 33, 0, 16384): 366_400, + (2048, 33, 5, 16384): 569_440, + (2048, 64, 0, 16384): 366_400, + (2048, 64, 5, 16384): 561_000, + (2048, 128, 0, 16384): 342_560, + (2048, 128, 5, 16384): 543_120, + (2048, 256, 0, 16384): 326_680, + (2048, 256, 5, 16384): 552_560, + (4096, 1, 0, 16384): 659_320, + (4096, 1, 5, 16384): 659_320, + (4096, 32, 0, 16384): 506_880, + (4096, 32, 5, 16384): 628_040, + (4096, 33, 0, 16384): 506_880, + (4096, 33, 5, 16384): 606_200, + (4096, 64, 0, 16384): 399_640, + (4096, 64, 5, 16384): 582_840, + (4096, 128, 0, 16384): 361_440, + (4096, 128, 5, 16384): 573_440, + (4096, 256, 0, 16384): 349_520, + (4096, 256, 5, 16384): 559_040, + (8192, 1, 0, 16384): 941_800, + (8192, 1, 5, 16384): 879_240, + (8192, 32, 0, 16384): 585_840, + (8192, 32, 5, 16384): 732_280, + (8192, 33, 0, 16384): 585_840, + (8192, 33, 5, 16384): 732_280, + (8192, 64, 0, 16384): 503_400, + (8192, 64, 5, 16384): 663_280, + (8192, 128, 0, 16384): 423_480, + (8192, 128, 5, 16384): 601_240, + (8192, 256, 0, 16384): 370_840, + (8192, 256, 5, 16384): 583_360, + (16384, 1, 0, 16384): 1_098_720, + (16384, 1, 5, 16384): 1_051_040, + (16384, 32, 0, 16384): 751_640, + (16384, 32, 5, 16384): 857_400, + (16384, 33, 0, 16384): 751_640, + (16384, 33, 5, 16384): 837_040, + (16384, 64, 0, 16384): 628_040, + (16384, 64, 5, 16384): 744_200, + (16384, 128, 0, 16384): 506_880, + (16384, 128, 5, 16384): 661_280, + (16384, 256, 0, 16384): 420_520, + (16384, 256, 5, 16384): 617_600, + (16384, 256, 0, 20480): 358_440, + (16384, 256, 5, 20480): 536_200, + (16384, 256, 0, 32768): 249_720, + (16384, 256, 5, 32768): 384_760, + (16384, 256, 0, 65536): 137_520, + (16384, 256, 5, 65536): 221_400, + (16384, 256, 0, 131072): 72_480, + (16384, 256, 5, 131072): 119_120, + (16384, 256, 0, 262144): 37_200, + (16384, 256, 5, 262144): 62_040, + (16384, 256, 0, 1048576): 10_920, + (16384, 256, 5, 1048576): 15_880, +} + + +def test_envelope_window_buffer_rows_from_the_metadata() -> None: + """The window read buffer over the engine envelope: in 60 native managers -- max_num_tokens 2048, + 4096, 8192, 16384 x max_batch_size 1, 32, 33, 64, 128, 256 x DSpark 0 / 5 at max_seq_len 16384, and + (16384, 256) x DSpark 0 / 5 at every max_seq_len 20480 .. 1,048,576 (+ 14 with DSpark, as the engine + passes it) -- metadata built as the engine builds it and its CUDA-graph copies at G = 1 and G = + max_batch_size have swa_local_indices_cuda [max_num_tokens, 128] int32. So swa_c is max_num_tokens * 128 + rows (262,144 / 524,288 / 1,048,576 / 2,097,152) whatever max_batch_size, DSpark or max_seq_len. The SWA + pool pages at 0.5 x the device's memory are pinned in _ENVELOPE_SWA_PAGES.""" + got, pages = {}, {} + for key in _ENVELOPE_SWA_PAGES: + max_t, seqs, draft, L = key + lm = L + 14 if draft else L + mgr = _v41_native_manager(max_t, seqs, draft, max_seq_len=lm) + try: + assert mgr.max_seq_len == lm and mgr.max_batch_size == seqs + pages[key] = _native_swa(mgr)[0].shape[0] + requests, md = _run_step(mgr, [(0, 100)], 1, 1, uid=980_000, max_num_tokens=max_t) + try: + rows = set() + for G in sorted({1, seqs}): + gmd = md.create_cuda_graph_metadata(G, False, draft) + for m in (md, gmd): + sw = m.swa_local_indices_cuda + assert tuple(sw.shape) == (max_t, PAGE) and sw.dtype == torch.int32, ( + key, + G, + sw.shape, + ) + rows.add(sw.shape[0] * PAGE) + assert len(rows) == 1 + got[key] = rows.pop() + finally: + for r_ in requests: + mgr.free_resources(r_) + finally: + mgr.shutdown() + del mgr + torch.cuda.empty_cache() + per_t = {t: sorted({v for k, v in got.items() if k[0] == t}) for t in _ENVELOPE_ROWS} + print( + f"\n[envelope window read buffer rows over {len(got)} configurations] rows per max_num_tokens {per_t}; " + f"SWA pages {pages}; largest {max(pages.values())}" + ) + assert got == {k: _ENVELOPE_ROWS[k[0]] for k in _ENVELOPE_SWA_PAGES} + assert pages == _ENVELOPE_SWA_PAGES + assert max(pages.values()) == NATIVE_SWA_UNSPLIT == pages[NATIVE_SWA_CEILING_KEY + (16384,)] + + +# KV tokens one step of the 256-request test may hold in total. The native manager at (16384, 256, DSpark +# 5, max_seq_len 32,782) held 22 requests of 32,768 tokens and refused the 23rd (about 720K tokens), so a +# step of 256 requests cannot put every request at the far end; steps keep their total under this bound. +_ENVELOPE_STEP_TOKENS = 560_000 + + +def _envelope_window_edge_steps(max_t: int, seqs: int, L: int): + """Edge steps of a (max_num_tokens, max_batch_size) configuration at max_seq_len L with DSpark, inside + _ENVELOPE_STEP_TOKENS: a full max_t step from position 0 (-1 window columns) and one ending at L; seqs decode + requests, 8 at L - 1 and the rest below position 1152; seqs DSpark verifies, 8 at positions up to L + 13 + and the rest below 1130; and a context filling the step beside seqs - 1 verifies (7 at the far end).""" + lm, rest = L + 14, seqs - 8 + steps = [ + ([(0, max_t)], 1, 1), + ([(L - max_t, max_t)], 1, 1), + ([(L - 1, 1)] * 8 + [(127 + (61 * i) % 1024, 1) for i in range(rest)], 0, 1), + ([(lm - 6, 6)] * 8 + [(100 + (37 * i) % 1024, 6) for i in range(rest)], 0, 6), + ( + [(0, max_t - 6 * (seqs - 1))] + + [(lm - 6, 6)] * 7 + + [(100 + (37 * i) % 1024, 6) for i in range(rest)], + 1, + 6, + ), + ] + for reqs, _, _ in steps: + assert sum(a + b for a, b in reqs) <= _ENVELOPE_STEP_TOKENS, reqs[:3] + return steps + + +def _envelope_fit(step, rng: random.Random): + """A random step with its total KV tokens capped at _ENVELOPE_STEP_TOKENS: its requests, in a random order, + keep their cached prefix while the rest can still fit; the others get a random prefix below 1024. The + request order (contexts first) and every request's new tokens are unchanged.""" + reqs, nctx, gen_len = step + if sum(a + n for a, n in reqs) <= _ENVELOPE_STEP_TOKENS: + return step + out, used = list(reqs), 0 + reserve = sum(n + 1024 for _, n in reqs) + order = list(range(len(reqs))) + rng.shuffle(order) + for i in order: + a, n = reqs[i] + reserve -= n + 1024 + if used + a + n + reserve > _ENVELOPE_STEP_TOKENS: + a = rng.randint(0, 1023) + out[i] = (a, n) + used += a + n + assert used <= _ENVELOPE_STEP_TOKENS + return out, nctx, gen_len + + +def test_envelope_window_read_renumbered_over_random_steps_256_requests() -> None: + """Random multi-request steps at the envelope's step axes, in the target's exact call, on the live native + manager of the engine configuration max_num_tokens 16384, max_batch_size 256, DSpark 5, max_seq_len 32768 + 14: + _envelope_window_edge_steps and 200 seeded random partitions of T <= 16,384 into k <= 256 requests (contexts + first with random cached prefixes, decode next_n 1 at positions < 32,768 or DSpark verify at < 32,782), + each step's KV tokens capped at _ENVELOPE_STEP_TOKENS (_envelope_fit). Each step exactly as + test_kvq_window_read_renumbered_over_random_steps: layers 0 / 2 / 20 / 39, the buffer + [md.swa_local_indices_cuda.shape[0] * 128, 512] (asserted 2,097,152 rows), every renumbered row bitwise + the bf16-container row the original index names (+0.0 for -1), rows past 128 T unchanged, and both + controls seen on every layer-step.""" + L, max_t, seqs = 32768, 16384, 256 + rng = random.Random(16_400 + seqs + L) + edges = _envelope_window_edge_steps(max_t, seqs, L) + steps = edges + [ + _envelope_fit(_kvq_random_step(rng, max_tokens=max_t, max_requests=seqs, seq_len=L), rng) + for _ in range(_NUM_RANDOM_STEPS) + ] + mgr = _v41_native_manager(max_t, seqs, 5, max_seq_len=L + 14) + try: + pages = _native_swa(mgr)[0].shape[0] + pool_shape, buf_shape, stats, ctrl = _kvq_window_read_steps( + mgr, steps, max_t, uid0=4_000_000, seed0=16_500 + ) + print( + f"\n[envelope window read, ({max_t}, {seqs}, 5) at max_seq_len {L + 14}, SWA {pool_shape}, buffer " + f"{buf_shape}] {len(edges)} edge + {_NUM_RANDOM_STEPS} random steps: {stats}; controls: {ctrl}" + ) + assert pages == pool_shape[0] == _ENVELOPE_SWA_PAGES[(max_t, seqs, 5, L)] and buf_shape == ( + max_t * PAGE, + W_KV, + ) + assert stats["steps"] == len(steps) and stats["max_t"] == max_t and stats["max_k"] == seqs + assert stats["max_pos"] == L + 13 and stats["max_valid_cols"] == PAGE + for name, (hit, _) in ctrl.items(): + assert hit == stats["layer_steps"], ( + f"control '{name}' missed {stats['layer_steps'] - hit} layer-steps" + ) + finally: + mgr.shutdown() + del mgr + torch.cuda.synchronize() + torch.cuda.empty_cache() diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_mxfp8_fake_quantize.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_mxfp8_fake_quantize.py new file mode 100644 index 000000000000..409da144d50e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_mxfp8_fake_quantize.py @@ -0,0 +1,2898 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the mxfp8_fake_quantize catalog entry. + +The entry owns what is certified: `CELLS` is the list (bf16 container row counts and widths, in +place, strided and 3-D `x`, paged scatter, strided pools, out-of-range slots, the native 528 B +layouts and identity form, zero rows), `reference` and `compare` are the bitwise gate. This file +turns a cell's spec into tensors and keeps the checks a cell cannot carry: non-finite row +isolation, CUDA-graph replay, the deliberately-wrong recipes, the op's own rejections, and the +engine-scale cases (the largest SWA pools, every row count, CUDA-graph capture at every graph +batch, the DSpark drafter's ring and block writes, native storage on the live DeepSeek-V4.1 cache +manager). + +Every comparison is bitwise (bf16 compared as int16, codes and scale bytes as uint8): the op +computes one exact value per element, so there is no tolerance to choose. The entry's reference +takes the scale exponent from torch.frexp of the float64 quotient amax / 448 rather than from the +IEEE-bit trick the kernel uses, so the two derivations are independent. (A float64 torch.log2 is +NOT usable for this: torch.ceil(torch.log2(amax.double() / 448)) gave -2 for amax = 56, i.e. for +56 / 448 = 2^-3, one too high at an exact power of two, amax = 7 * 2^k.) +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.quantization.mxfp8_fake_quantize import ( # noqa: E501 + mxfp8_fake_quantize as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "mxfp8_fake_quantize requires a CUDA device" + +BLOCK = 32 +E4M3_MAX = 448.0 +AMAX_FLOOR = 1e-4 +WINDOW_KV_WIDTH = 512 # DeepSeek-V4.1 window KV head_dim (448 nope + 64 RoPE, all quantized) + + +def _fake_quant(x: torch.Tensor): + """The entry's reference for the rows of a bf16 [T, W] x. + + Returns (bf16 dequantized [T, W], e4m3 codes as uint8 [T, W], UE8M0 bytes uint8 [T, W/32]). + """ + t, w = x.reshape(-1, x.shape[-1]).shape + y = op.reference(x, torch.empty(t, w, dtype=torch.bfloat16, device=x.device)) + codes, sbytes = op.reference( + x, + torch.empty(t, w, dtype=torch.uint8, device=x.device), + None, + torch.empty(t, w // BLOCK, dtype=torch.uint8, device=x.device), + ) + return y, codes, sbytes + + +def _scale_exponent( + amax: torch.Tensor, floor: float | None = AMAX_FLOOR, rounding=torch.ceil +) -> torch.Tensor: + """ceil (or floor) of log2(max(amax, floor) / 448), int32, exact. + + amax is fp32 and floored in fp32 (as the reference does). The float64 + quotient of a bf16 amax by 448 is a power of two exactly when amax = 7 * 2^k + and otherwise sits far from one, so frexp (mantissa in [0.5, 1)) classifies + it exactly: v = m * 2^E has ceil(log2 v) = E, or E - 1 when m == 0.5, and + floor(log2 v) = E - 1. + """ + a = amax.float() + if floor is not None: + a = a.clamp_min(floor) + m, e = torch.frexp(a.double() / E4M3_MAX) + if rounding is torch.floor: + return (e - 1).to(torch.int32) + return torch.where(m == 0.5, e - 1, e).to(torch.int32) + + +def _control_quant( + x: torch.Tensor, + *, + block: int = BLOCK, + floor: float | None = AMAX_FLOOR, + rounding=torch.ceil, + scale_floor_exp: int | None = None, +) -> torch.Tensor: + """A deliberately-wrong MXFP8 fake quantization of a bf16 [T, W] tensor -> bf16 [T, W]. + + `floor`, `rounding`, `block` and `scale_floor_exp` select the control. + """ + t, w = x.shape + xf = x.float().view(t, w // block, block) + e = _scale_exponent(xf.abs().amax(-1), floor, rounding) + if scale_floor_exp is not None: + e = e.clamp_min(scale_floor_exp) + s = torch.exp2(e.double()).float().unsqueeze(-1) + q = (xf / s).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) + return (q.float() * s).view(t, w).to(torch.bfloat16) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.view(torch.int16) if t.dtype == torch.bfloat16 else t.view(torch.uint8) + + +def _assert_bitwise(actual: torch.Tensor, expected: torch.Tensor) -> None: + a, e = _bits(actual), _bits(expected) + n_bad = (a != e).sum().item() + assert n_bad == 0, f"{n_bad} of {a.numel()} elements differ bitwise" + + +def _window_kv_like(t: int, w: int = WINDOW_KV_WIDTH, seed: int = 0) -> torch.Tensor: + """bf16 [t, w] rows whose 32-blocks cover every path of the recipe. + + Block magnitudes span 2^-30 .. 2^20 (so the 1e-4 floor, e4m3 subnormals and + BOS-like massive rows are all present), and a quarter of the rows carry + hand-placed blocks: all-zero, signed zeros, amax exactly 448 * 2^k (the + ceil boundary), amax one bf16 step above it, e4m3 round-to-even ties, and + sub-1e-4 blocks with a wide dynamic range (where the floor changes values). + """ + g = torch.Generator(device="cuda").manual_seed(seed) + nb = w // BLOCK + mag = torch.randint(-30, 21, (t, nb, 1), device="cuda", generator=g).float() + x = torch.randn(t, nb, BLOCK, device="cuda", generator=g) * torch.exp2(mag) + for r in range(0, t, 4): + k = r % 7 - 3 + b = (r // 4) % nb + x[r, b] = 0.0 + x[r, (b + 1) % nb] = torch.randn(BLOCK, device="cuda", generator=g) * 2.0**k + x[r, (b + 1) % nb, 0] = 448.0 * 2.0**k + x[r, (b + 2) % nb] = torch.randn(BLOCK, device="cuda", generator=g) * 2.0**k + x[r, (b + 2) % nb, 0] = 448.0 * 2.0**k * (1 + 2**-7) + # e4m3 ties at scale 2^k: 1.0625 (-> 1.0), 1.1875 (-> 1.25), 3.25*2^-9 subnormal (-> 3*2^-9 or 4*2^-9) + x[r, (b + 3) % nb] = 2.0**k * torch.tensor( + [448.0, 1.0625, -1.1875, 3.25 * 2**-9, -2.5 * 2**-9, 13.5 * 2**-9, 0.5 * 2**-9, -0.0] + * 4, + device="cuda", + ) + x[r, (b + 4) % nb] = torch.randn(BLOCK, device="cuda", generator=g) * 1e-6 + x[r, (b + 4) % nb, 1::2] *= 1e-3 + x[r, (b + 5) % nb] = torch.randn(BLOCK, device="cuda", generator=g) * 2.0**-10 + return x.view(t, w).to(torch.bfloat16) + + +# Pool widths are chosen plainly larger than every batch here: 300 pages of 128 slots (38400 +# slots) and 150 pages of 256 slots, the two tokens_per_block values the DeepSeek-V4 cache +# manager accepts, against batches of 1..8192 rows. The engine's own pool sizes are certified +# further down. +POOL_PAGES = {128: 300, 256: 150} + + +def _pool_slots(t: int, num_slots: int, seed: int, skips: int) -> torch.Tensor: + g = torch.Generator(device="cuda").manual_seed(seed) + slots = torch.randperm(num_slots, device="cuda", generator=g)[:t].to(torch.int64) + if skips: + slots[torch.randperm(t, device="cuda", generator=g)[:skips]] = -1 + return slots + + +# ─── certified cells ────────────────────────────────────────────────────────── + + +def _snapshot(t: torch.Tensor) -> torch.Tensor: + return t.clone() if t.dtype == torch.bfloat16 else t.view(torch.uint8).clone().view(t.dtype) + + +def _build(spec, seed: int): + """A cell's spec -> (x, out, slot_mapping, out_scale, extra). + + `extra` carries what the structural checks need: `untouched` (regions of a wider parent outside + `out` that must keep their bytes), `decodes_to_container` (decode the native slots in torch and + hold them to the bf16 container) and `floor_byte` (the 2^-22 floor block must be present). + """ + g = torch.Generator(device="cuda").manual_seed(seed) + kind, t = spec["kind"], spec.get("rows", 0) + w = WINDOW_KV_WIDTH + bf16 = dict(dtype=torch.bfloat16, device="cuda") + slots = out_scale = None + extra = {} + if kind == "container": + x = _window_kv_like(t, w=spec["width"], seed=seed) + out = torch.empty_like(x) + elif kind == "in_place": + x = _window_kv_like(t, seed=seed) + out = x + elif kind in ("x_column_slice", "x_3d"): + base = _window_kv_like(t, w=1024, seed=seed) + if kind == "x_column_slice": + x = base[:, 256:768] + assert x.stride() == (1024, 1) + else: + x = base[:, None, :512] + out = torch.empty(t, w, **bf16) + elif kind == "scatter": + ps = spec["page_size"] + pages = POOL_PAGES[ps] + x = _window_kv_like(t, seed=seed) + slots = _pool_slots(t, pages * ps, seed=seed, skips=t // 8).to(spec["slot_dtype"]) + out = torch.randn(pages, ps, w, device="cuda", generator=g).to(torch.bfloat16) + elif kind in ("pool_left_half", "pool_column_slice"): + pages, page = 40, 128 + x = _window_kv_like(t, seed=seed) + slots = _pool_slots(t, pages * page, seed=seed, skips=50) + if kind == "pool_left_half": + wide = torch.randn(pages, page, 2 * w, device="cuda", generator=g).to(torch.bfloat16) + out = wide[..., :w] + extra["untouched"] = [wide[..., w:]] + else: + slots = slots.to(torch.int32) + flat_wide = torch.randn(pages * page, 640, device="cuda", generator=g).to( + torch.bfloat16 + ) + out = flat_wide[:, 64:576] + extra["untouched"] = [flat_wide[:, :64], flat_wide[:, 576:]] + elif kind == "out_of_range": + x = _window_kv_like(t, seed=seed) + out = torch.randn(2, 128, w, device="cuda", generator=g).to(torch.bfloat16) + slots = torch.tensor( + [0, 256, 257, 10**9, -1, -(10**9), 255, 3], dtype=torch.int64, device="cuda" + ) + elif kind == "native_interleaved": + # [slots, 528] uint8 rows: 512 code bytes then 16 scale bytes. + x = _window_kv_like(t, seed=seed) + num_slots = 2 * t + 64 + slots = _pool_slots(t, num_slots, seed=seed, skips=t // 8) + rows = torch.randint( + 0, 256, (num_slots, 528), dtype=torch.uint8, device="cuda", generator=g + ) + out, out_scale = rows[:, :512], rows[:, 512:] + extra["decodes_to_container"] = True + elif kind == "native_footer": + # Page = page_size code rows (512 B) then page_size scale rows (16 B), carved from one uint8 pool. + ps, pages = spec["page_size"], 24 + x = _window_kv_like(t, seed=seed) + slots = _pool_slots(t, pages * ps, seed=seed, skips=100).to(torch.int32) + buf = torch.randint( + 0, 256, (pages, ps * 528), dtype=torch.uint8, device="cuda", generator=g + ) + out = buf[:, : ps * 512].view(pages, ps, 512).view(torch.float8_e4m3fn) + out_scale = buf[:, ps * 512 :].view(pages, ps, 16).view(torch.float8_e8m0fnu) + elif kind == "native_identity": + x = _window_kv_like(t, seed=seed) + out = torch.empty(t, 512, dtype=torch.float8_e4m3fn, device="cuda") + out_scale = torch.empty(t, 16, dtype=torch.uint8, device="cuda") + extra["decodes_to_container"] = True + extra["floor_byte"] = True + else: + assert kind == "zero_rows", kind + x = torch.empty(0, w, **bf16) + out = torch.randn(4, 128, w, device="cuda", generator=g).to(torch.bfloat16) + if spec["mapped"]: + slots = torch.empty(0, dtype=torch.int32, device="cuda") + return x, out, slots, out_scale, extra + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + x, out, slots, out_scale, extra = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + x_before = x.clone() + out_before = _snapshot(out) + scale_before = None if out_scale is None else _snapshot(out_scale) + untouched = [(u, u.clone()) for u in extra.get("untouched", [])] + result = op(x, out, slots, out_scale) + if out_scale is None: + assert result is out, "the bf16 form writes out in place and returns it" + else: + assert result[0] is out and result[1] is out_scale, ( + "the native form returns (out, out_scale)" + ) + op.compare(result, op.reference(x_before, out_before, slots, scale_before)) + if out is not x: + _assert_bitwise(x, x_before) + for u, before in untouched: + _assert_bitwise(u, before) + if extra.get("decodes_to_container"): + # native dequantization e4m3(code) * 2^(byte - 127) is bit-identical to the bf16 container + y = _fake_quant(x_before)[0] + s = slots.long() if slots is not None else torch.arange(x.shape[0], device="cuda") + live = (s >= 0) & (s < out.shape[0]) + codes = out.view(torch.uint8)[s[live]].view(torch.float8_e4m3fn).float().view(-1, 16, 32) + scale = torch.exp2(out_scale[s[live]].view(torch.uint8).float() - 127).unsqueeze(-1) + _assert_bitwise((codes * scale).view(-1, 512).to(torch.bfloat16), y[live]) + if extra.get("floor_byte"): + assert (out_scale == 105).any(), "the test data must contain floor-scaled (2^-22) blocks" + + +# ─── non-finite input, CUDA graph ───────────────────────────────────────────── + + +def test_non_finite_row_is_isolated() -> None: + """A NaN/inf row (e.g. CUDA-graph padding garbage) changes no other row.""" + x = _window_kv_like(64, seed=5) + x[7, 3] = float("nan") + x[9, 100] = float("inf") + x[11, :] = float("-inf") + out = torch.empty_like(x) + op(x, out) + keep = torch.ones(64, dtype=torch.bool, device="cuda") + keep[[7, 9, 11]] = False + _assert_bitwise(out[keep], _fake_quant(x[keep].contiguous())[0]) + + +def test_cuda_graph_capture_replay() -> None: + """Captured after a warm-up call (Triton compiles on first use), replayed on new rows and new slots.""" + t, pages, page = 64, 16, 128 + x = torch.empty(t, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda") + slots = torch.empty(t, dtype=torch.int32, device="cuda") + pool = torch.zeros(pages, page, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda") + x.copy_(_window_kv_like(t, seed=1)) + slots.copy_(_pool_slots(t, pages * page, seed=1, skips=4)) + op(x, pool, slots) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(x, pool, slots) + + for seed in (2, 3): + new_x = _window_kv_like(t, seed=seed) + new_slots = _pool_slots(t, pages * page, seed=seed, skips=4) + pool.zero_() + x.copy_(new_x) + slots.copy_(new_slots) + graph.replay() + torch.cuda.synchronize() + op.compare(pool, op.reference(new_x, torch.zeros_like(pool), new_slots)) + + +# ─── deliberately-wrong controls ────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "variant", + [ + pytest.param(dict(floor=None), id="no_1e-4_floor"), + pytest.param(dict(rounding=torch.floor), id="floor_not_ceil_scale"), + pytest.param(dict(block=64), id="block_64"), + pytest.param(dict(floor=None, scale_floor_exp=-13), id="scale_floored_at_2^-13"), + ], +) +def test_wrong_recipes_are_rejected(variant: dict) -> None: + """Each plausible wrong recipe differs from the op on this data, so the bitwise gate can see it. + + `scale_floored_at_2^-13` floors max(amax/448, 1e-4) instead of amax: a + different value set on every block with amax < 448 * 2^-14. + """ + x = _window_kv_like(4096, seed=99) + out = torch.empty_like(x) + op(x, out) + op.compare(out, op.reference(x, torch.empty_like(x))) + wrong = _control_quant(x, **variant) + n_diff = (_bits(out) != _bits(wrong)).sum().item() + print(f"{variant}: {n_diff} of {x.numel()} elements differ from the op") + assert n_diff > 0 + with pytest.raises(AssertionError, match="differ bitwise"): + op.compare(out, wrong) + + +# ─── rejected inputs (the op's own checks) ──────────────────────────────────── + + +def _raises(match: str, *args) -> None: + with pytest.raises(ValueError, match=match): + op(*args) + + +def test_rejected_inputs() -> None: + x = _window_kv_like(16, seed=4) + out = torch.empty_like(x) + pool = torch.empty(2, 128, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda") + codes = torch.empty(16, 512, dtype=torch.uint8, device="cuda") + scales = torch.empty(16, 16, dtype=torch.uint8, device="cuda") + _raises("x must be bfloat16", x.float(), out) + _raises("x must be bfloat16", x.half(), out) + _raises("x must be a CUDA tensor", x.cpu(), out) + _raises("out must be on cuda", x, out.cpu()) + _raises("positive multiple of 32", x[:, :48], out[:, :48]) + _raises("unit stride along its last dim", x.t().contiguous().t(), out) + _raises("rank >= 2", x[0], out) + _raises("out's last dim must be 512", x, out[:, :256]) + _raises("out's last dim must be 512 with unit stride", x, out.t().contiguous().t()) + _raises("out must be bfloat16, float8_e4m3fn or uint8", x, out.float(), None, scales) + _raises("needs out_scale", x, codes) + _raises("a bfloat16 out takes None", x, out, None, scales) + _raises("out_scale must be", x, codes, None, scales[:, :8]) + _raises("out_scale must be uint8 or float8_e8m0fnu", x, codes, None, scales.to(torch.int32)) + _raises("slot_mapping must be int32 or int64", x, pool, torch.zeros(16, device="cuda")) + _raises( + "slot_mapping must be \\[16\\]", x, pool, torch.zeros(15, dtype=torch.int32, device="cuda") + ) + _raises("needs >= 16 slots", x, out[:8]) + _raises("must be \\[slots, C\\] or \\[pages, page_size, C\\]", x, pool.view(2, 1, 128, 512)) + + +# ─── the largest SWA pool an engine can allocate ────────────────────────────── + +# The SWA pool (one pool shared by all 40 layers) is sized by the V4.1 +# DeepseekV4CacheManager from the per-rank KV quota, which an engine derives as +# kv_cache_fraction (0.5 at the gates) x the device's FREE memory after weights +# and activations. Free memory never exceeds total memory, so a manager built at +# 0.5 x the device's TOTAL memory (148,510,474,240 bytes on this GB300) allocates +# the largest SWA pool an engine of that configuration can have. At the gate +# settings (max_seq_len 16384, max_num_tokens 8192, max_batch_size {8, 32}, +# DSpark or not; test_largest_engine_pools_from_the_cache_manager) the largest +# is [491000, 128, 512] bf16. Over the engine envelope (max_num_tokens +# 2048..16384, max_batch_size up to 256, max_seq_len up to 1,048,576, DSpark or +# not) the unsplit manager's largest is [730880, 128, 512] bf16 at +# (max_num_tokens 16384, max_batch_size 1, no DSpark), whatever max_seq_len (the +# pool scan; test_envelope_largest_swa_pools_from_the_cache_manager re-measures it). +# A PP2 stage manager holds only its 20 layers' pools, so its SWA pool is larger: +# stage 1 alone reaches [958720, 128, 512] at (16384, 2, 0, 16384) (the PP2 +# stage scan; test_envelope_pp2_stage_swa_pools_from_the_cache_manager re-measures it): +# 122,716,160 slots, 6.3e10 elements, 125.7 GB. That stage-alone maximum is +# ENGINE_SWA_PAGES. test_pool_size_invariance shows the op's result does not +# depend on B, so the far-end cases below run at that B_max. Past slot 4,194,304 +# every element offset exceeds 2^31, so any 32-bit +# offset in the kernel's address arithmetic would land a row somewhere else (or +# fault). The pool is allocated at that exact shape, with one guard page past +# its end, and every element is set to a NaN bit pattern the op never writes +# for a finite row: after the call, the named slots must hold the reference and +# every other element of the pool and the guard page must still hold the +# sentinel. +GATES_SWA_PAGES = 491000 +UNSPLIT_SWA_PAGES = 730880 +ENGINE_SWA_PAGES = 958720 +ENGINE_PAGE = 128 +SENTINEL = 0x7FA5 # bf16 NaN bit pattern (as int16) + + +def _hazard_slots(num_slots: int, width: int, page_size: int) -> list[int]: + """Distinct slots on both sides of every 32-bit offset hazard of a [slots, width] bf16 pool, and both ends. + + Element offsets 2^30 (2^31 bytes), 2^31 (2^32 bytes), 2^32, 2^33, 2^34, 2^35 + and 2^36 (where inside the pool), the first page, the last page and the last slot. + """ + cands = {0, 1, page_size - 1, page_size} + for limit in (2**30, 2**31, 2**32, 2**33, 2**34, 2**35, 2**36): + s = limit // width + cands |= {s - 1, s} + cands |= {num_slots - page_size - 1, num_slots - page_size, num_slots - 2, num_slots - 1} + return sorted(c for c in cands if 0 <= c < num_slots) + + +def _engine_pool(pages: int, page_size: int, width: int = WINDOW_KV_WIDTH): + """(buffer with one guard page, [pages, page_size, width] bf16 pool view), every element SENTINEL.""" + buf = torch.empty(pages + 1, page_size, width, dtype=torch.bfloat16, device="cuda") + buf.view(torch.int16).fill_(SENTINEL) + return buf, buf[:pages] + + +def _check_engine_pool( + buf: torch.Tensor, pool: torch.Tensor, slots: torch.Tensor, y: torch.Tensor +) -> None: + """Named slots hold y bitwise; every other element of the pool and the guard page is still SENTINEL. + + Restores the named slots to SENTINEL, so the pool can be reused. + """ + num_slots = pool.shape[0] * pool.shape[1] + flat = pool.view(num_slots, -1) + s64 = slots.to(torch.int64) + live = (s64 >= 0) & (s64 < num_slots) + idx = s64[live] + _assert_bitwise(flat[idx], y[live]) + flat.view(torch.int16)[idx] = SENTINEL + whole = buf.view(torch.int16).view(-1) + step = 1 << 30 + n_bad = sum(int((whole[i : i + step] != SENTINEL).sum()) for i in range(0, whole.numel(), step)) + assert n_bad == 0, f"{n_bad} elements outside the named slots were written" + + +def _free_engine_pool() -> None: + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +@pytest.mark.parametrize("slot_dtype", [torch.int32, torch.int64]) +def test_engine_swa_pool_far_end(slot_dtype: torch.dtype) -> None: + """Decode-like: one row per 32-bit hazard boundary and at both ends of the engine-size pool, plus skips.""" + buf, pool = _engine_pool(ENGINE_SWA_PAGES, ENGINE_PAGE) + num_slots = ENGINE_SWA_PAGES * ENGINE_PAGE + hazard = _hazard_slots(num_slots, WINDOW_KV_WIDTH, ENGINE_PAGE) + assert hazard[-1] == num_slots - 1 and 2**35 // WINDOW_KV_WIDTH in hazard + # -1 and the first two slots past the end (which would land in the guard page) are skipped. + slots = torch.tensor( + hazard + [-1, num_slots, num_slots + ENGINE_PAGE - 1], dtype=torch.int64, device="cuda" + ) + x = _window_kv_like(slots.numel(), seed=31) + op(x, pool, slots.to(slot_dtype)) + _check_engine_pool(buf, pool, slots, _fake_quant(x)[0]) + del buf, pool + _free_engine_pool() + + +def test_engine_swa_pool_prefill_spread() -> None: + """Prefill-like, T = 32768 (the decoder pass's largest step): the last page written in order, every hazard + slot, the rest spread over the pool.""" + t = _MAX_ROWS + buf, pool = _engine_pool(ENGINE_SWA_PAGES, ENGINE_PAGE) + num_slots = ENGINE_SWA_PAGES * ENGINE_PAGE + last_page = torch.arange(num_slots - ENGINE_PAGE, num_slots, device="cuda") + hazard = torch.tensor(_hazard_slots(num_slots, WINDOW_KV_WIDTH, ENGINE_PAGE), device="cuda") + fixed = torch.unique(torch.cat([last_page, hazard])) + g = torch.Generator(device="cuda").manual_seed(8192) + perm = torch.randperm(num_slots, device="cuda", generator=g) + rest = perm[~torch.isin(perm, fixed)][: t - fixed.numel()] + slots = torch.cat([last_page, fixed[~torch.isin(fixed, last_page)], rest]).to(torch.int64) + assert slots.numel() == t and torch.unique(slots).numel() == t + slots[torch.randperm(t, device="cuda", generator=g)[: t // 16]] = -1 + x = _fast_rows(t, seed=8192) + op(x, pool, slots.to(torch.int32)) + _check_engine_pool(buf, pool, slots, _fake_quant(x)[0]) + del buf, pool + _free_engine_pool() + + +def test_engine_swa_pool_cuda_graph_far_end() -> None: + """Captured on the engine-size pool after a warm-up call, replayed with far-end and boundary slots.""" + t = 32 + buf, pool = _engine_pool(ENGINE_SWA_PAGES, ENGINE_PAGE) + num_slots = ENGINE_SWA_PAGES * ENGINE_PAGE + hazard = _hazard_slots(num_slots, WINDOW_KV_WIDTH, ENGINE_PAGE) + x = torch.empty(t, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda") + slots = torch.empty(t, dtype=torch.int32, device="cuda") + x.copy_(_window_kv_like(t, seed=40)) + slots.copy_(torch.arange(t, device="cuda")) + op(x, pool, slots) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(x, pool, slots) + buf.view(torch.int16).fill_(SENTINEL) + + g = torch.Generator(device="cuda").manual_seed(41) + hz = torch.tensor(hazard, device="cuda") + near_end = num_slots - 1 - torch.arange(4 * ENGINE_PAGE, device="cuda") + near_end = near_end[~torch.isin(near_end, hz)] + for seed in (42, 43): + pick = torch.randperm(near_end.numel(), device="cuda", generator=g)[: t - len(hazard) - 1] + new_slots = torch.cat([hz, torch.tensor([-1], device="cuda"), near_end[pick]]).to( + torch.int64 + ) + assert torch.unique(new_slots).numel() == t + new_x = _window_kv_like(t, seed=seed) + x.copy_(new_x) + slots.copy_(new_slots) + graph.replay() + torch.cuda.synchronize() + _check_engine_pool(buf, pool, new_slots, _fake_quant(new_x)[0]) + del graph, buf, pool + _free_engine_pool() + + +def test_engine_swa_pool_harness_sees_a_wrapped_offset() -> None: + """Control: the far-end check fails when rows land where a 32-bit (mod 2^32) element offset would put them. + + The op is handed the wrapped slots and the checker the true ones; both the + named-slot comparison and the outside-the-slots sentinel scan must see it. + """ + buf, pool = _engine_pool(ENGINE_SWA_PAGES, ENGINE_PAGE) + num_slots = ENGINE_SWA_PAGES * ENGINE_PAGE + true = torch.tensor( + [ + s + for s in _hazard_slots(num_slots, WINDOW_KV_WIDTH, ENGINE_PAGE) + if s * WINDOW_KV_WIDTH >= 2**32 + ], + dtype=torch.int64, + device="cuda", + ) + wrapped = (true * WINDOW_KV_WIDTH % 2**32) // WINDOW_KV_WIDTH + assert true.numel() >= 5 and not torch.isin(wrapped, true).any() + x = _window_kv_like(true.numel(), seed=51) + y = _fake_quant(x)[0] + op(x, pool, wrapped) + with pytest.raises(AssertionError, match="elements differ bitwise"): + _check_engine_pool(buf, pool, true, y) + flat = pool.view(num_slots, -1) + flat.view(torch.int16)[true] = SENTINEL + with pytest.raises(AssertionError, match="outside the named slots were written"): + _check_engine_pool(buf, pool, torch.tensor([], dtype=torch.int64, device="cuda"), y[:0]) + del buf, pool, flat + _free_engine_pool() + + +# ─── every row count 1..32768 in the core target's paged-scatter call ──────── +# +# The DeepSeek-V4.1 core target writes each layer's window KV with +# mxfp8_fake_quantize(kv [T, 512], swa_pool [pages, 128, 512], slots[:T]), +# where T is the rank's token count this step (any integer up to the engine's +# max_num_tokens, at most 16384; in the decoder bounded-replay pass of layers +# 21-39 up to 128 x 256 = 32768 rows) and slots is a prefix view of one row of +# the metadata's int64 [num_local_layers, max_num_tokens] slot table (-1 = no +# slot). Rows reach the kernel only as the size of grid axis 0: program r loads +# row r and slot_mapping[r] and writes one slot; no kernel argument carries T. + +import tensorrt_llm._torch.custom_ops.mxfp8_fake_quantize as _mx_op # noqa: E402 (the Triton kernel object, for the variant probe) + +_MAX_ROWS = 32768 +_SWEEP_PAGES = 320 # 40960 slots of 128: plainly larger than the 32768 rows swept + + +def _jit_variants(fn) -> int: + """Number of compiled variants Triton holds for a @triton.jit function, over all devices.""" + dc = getattr(fn, "device_caches", None) + if dc is not None: + return sum(len(v[0]) if isinstance(v, tuple) else len(v) for v in dc.values()) + return sum(len(v) for v in fn.cache.values()) + + +def _sweep_slots(t: int, num_slots: int, seed: int) -> torch.Tensor: + """Row 1 of an int64 [2, t] table (the metadata's layout): distinct random slots, 1/16 of them -1.""" + g = torch.Generator(device="cuda").manual_seed(seed) + table = torch.full((2, t), -1, dtype=torch.int64, device="cuda") + slots = table[1] + slots.copy_(torch.randperm(num_slots, device="cuda", generator=g)[:t]) + slots[torch.randperm(t, device="cuda", generator=g)[: t // 16]] = -1 + slots[1] = -1 # a skipped row inside the smallest prefixes too + return slots + + +def test_every_row_count_1_to_32768_in_the_paged_scatter_form() -> None: + """Every prefix (x[:r], slots[:r]), r = 1..32768, writes exactly the first r rows of the 32768-row call. + + The 32768-row call is held bitwise to the native reference and to the + sentinel scan. Then, into the same sentinel-filled pool, the op runs at + every r in increasing order. After each call the slots named by the live + entries of slots[:r] hold the reference rows bit for bit, and the number of + pool elements that are not the sentinel is exactly 512 per live row: a + skipped row, a row written to another slot, or a write past row r - 1 all + change one of the two counts. The mismatch counts accumulate on the device; + one sync at the end. The probe: rows are only grid axis 0, so the sweep must + compile no Triton variant beyond the one the 32768-row call compiled. + """ + t = _MAX_ROWS + num_slots = _SWEEP_PAGES * ENGINE_PAGE + buf, pool = _engine_pool(_SWEEP_PAGES, ENGINE_PAGE) + flat = pool.view(num_slots, -1).view(torch.int16) + whole = buf.view(torch.int16) + slots = _sweep_slots(t, num_slots, seed=8193) + assert slots.stride(0) == 1 + x = _window_kv_like(t, seed=8193) + y = _fake_quant(x)[0] + + kernel = _mx_op._mxfp8_fake_quantize_kernel + n0 = _jit_variants(kernel) + op(x, pool, slots) + n1 = _jit_variants(kernel) + _check_engine_pool( + buf, pool, slots, y + ) # the N_max call against the reference; restores the sentinel + + live = slots >= 0 + live_pos = torch.nonzero(live).squeeze(1) + live_slots = slots[live_pos] + y_live = y[live_pos].view(torch.int16) + cnt = torch.cumsum(live.to(torch.int64), 0).tolist() + bad_rows = torch.zeros((), dtype=torch.int64, device="cuda") + bad_count = torch.zeros((), dtype=torch.int64, device="cuda") + for r in range(1, t + 1): + op(x[:r], pool, slots[:r]) + k = cnt[r - 1] + bad_rows += (flat[live_slots[:k]] != y_live[:k]).sum() + bad_count += ((whole != SENTINEL).sum() - k * WINDOW_KV_WIDTH).abs() + n2 = _jit_variants(kernel) + print( + f"\n[rows 1..{t}] live rows {cnt[-1]}; Triton variants: {n0} before the {t}-row call, {n1} after, " + f"{n2} after the sweep; element mismatches {int(bad_rows)}, non-sentinel count error {int(bad_count)}" + ) + assert int(bad_rows) == 0, ( + f"{int(bad_rows)} elements of the named slots differ from the {t}-row result" + ) + assert int(bad_count) == 0, ( + f"non-sentinel element count off by {int(bad_count)} summed over the sweep" + ) + assert n1 - n0 <= 1 and n2 == n1, f"row counts compiled new variants: {n0} -> {n1} -> {n2}" + del buf, pool, flat, whole + _free_engine_pool() + + +def test_row_sweep_harness_sees_a_row_or_slot_shift() -> None: + """Control for the sweep: its two checks reject a read one row off and a write one slot off, at every r tried.""" + num_slots = _SWEEP_PAGES * ENGINE_PAGE + buf, pool = _engine_pool(_SWEEP_PAGES, ENGINE_PAGE) + flat = pool.view(num_slots, -1).view(torch.int16) + whole = buf.view(torch.int16) + g = torch.Generator(device="cuda").manual_seed(66) + slots = torch.randperm(num_slots, device="cuda", generator=g)[:66].to( + torch.int64 + ) # no -1: every row lands + x = _window_kv_like(66, seed=66) + y = _fake_quant(x)[0].view(torch.int16) + for r in (1, 2, 7, 63, 64): + for name, xa, sa in ( + ("row +1", x[1 : r + 1], slots[:r]), + ("slot +1", x[:r], slots[1 : r + 1]), + ): + whole.fill_(SENTINEL) + op(xa, pool, sa) + n_bad = int((flat[slots[:r]] != y[:r]).sum()) + n_count = int((whole != SENTINEL).sum()) - r * WINDOW_KV_WIDTH + print( + f"[control {name}] r={r}: {n_bad} / {r * WINDOW_KV_WIDTH} elements differ, count error {n_count}" + ) + assert n_bad > r * WINDOW_KV_WIDTH // 2 + del buf, pool, flat, whole + _free_engine_pool() + + +# ─── CUDA-graph capture at every graph batch the engines capture ───────────── +# +# The engine envelope captures every graph batch G of _ENVELOPE_GSET: every G +# 1..32 (padding off captures every batch up to max_batch_size) and TRT-LLM's +# padded list for cuda_graph_config max_batch_size 256 (1, 2, 4, 8..128 step 8, +# 192, 256). The target's SWA write runs T = G (decode) and T = 6G (DSpark +# target verify: 1 + 5 drafts per request) rows under capture; the drafter +# writes its stage rings at M = G (bonus) and 5G (back-fill) and its draft +# block at N = 5G (the ring and block tests below). _GRAPH_BS is the gates' +# earlier set {1..8, 16, 24, 32}, kept for the 99-page ring case. +_ENVELOPE_GSET = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +assert len(_ENVELOPE_GSET) == 46 +_GRAPH_BS = (1, 2, 3, 4, 5, 6, 7, 8, 16, 24, 32) +_SWA_GRAPH_T = tuple( + sorted(set(_ENVELOPE_GSET) | {6 * g for g in _ENVELOPE_GSET}) +) # 82 row counts, 1..1536 +_RING_GRAPH_T = tuple(sorted({3 * b for b in _GRAPH_BS} | {5 * b for b in _GRAPH_BS})) +# The DSpark worker's ring, _kv_windows [max_bs + 1, 3, 128, 512] bf16 at max_bs 32, +# viewed as [(32 + 1) * 3, 128, 512]. +_RING_PAGES = (32 + 1) * 3 + + +def _far_end_slots(num_slots: int, page_size: int, t: int, g: torch.Generator) -> torch.Tensor: + """t distinct slots: the pool's last slot, then a random pick from the 32-bit hazard slots and the last + 4 pages (and, when t needs more, slots spread uniformly over the pool); for t >= 2 one of them (not the + last slot) replaced by -1.""" + last = num_slots - 1 + hazard = torch.tensor( + [s for s in _hazard_slots(num_slots, WINDOW_KV_WIDTH, page_size) if s != last], + device="cuda", + ) + near_end = last - 1 - torch.arange(4 * page_size - 1, device="cuda") + near_end = near_end[~torch.isin(near_end, hazard)] + rest = torch.cat([hazard, near_end]) + rest = rest[torch.randperm(rest.numel(), device="cuda", generator=g)] + if rest.numel() < t - 1: + spread = torch.randint(0, last, (4 * t,), device="cuda", generator=g).unique() + spread = spread[~torch.isin(spread, rest)] + rest = torch.cat([rest, spread[torch.randperm(spread.numel(), device="cuda", generator=g)]]) + slots = torch.cat([torch.tensor([last], device="cuda"), rest[: t - 1]]).to(torch.int64) + assert slots.numel() == t and torch.unique(slots).numel() == t + if t >= 2: + slots[1 + int(torch.randint(0, t - 1, (1,), device="cuda", generator=g))] = -1 + return slots + + +def _capture_sweep(pages: int, page_size: int, ts, seed: int) -> None: + """Per T: eager warm-up at that shape, capture, then 2 replays on new rows and new far-end slots. + + Every replay must write exactly the reference into the named slots and + nothing else (sentinel scan of the whole pool plus a guard page), and equal + bit for bit the same call run eagerly on the same inputs. + """ + buf, pool = _engine_pool(pages, page_size) + num_slots = pages * page_size + flat = pool.view(num_slots, -1) + g = torch.Generator(device="cuda").manual_seed(seed) + table = torch.full( + (2, max(ts)), -1, dtype=torch.int64, device="cuda" + ) # the metadata's int64 slot table + for t in ts: + x = torch.empty(t, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda") + slots = table[1, :t] + x.copy_(_window_kv_like(t, seed=seed + t)) + slots.copy_(_far_end_slots(num_slots, page_size, t, g)) + op(x, pool, slots) # eager warm-up at this shape, as the engine's warm-up step runs it + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(x, pool, slots) + buf.view(torch.int16).fill_(SENTINEL) + for rep in (1, 2): + new_x = _window_kv_like(t, seed=seed + 1000 * rep + t) + new_slots = _far_end_slots(num_slots, page_size, t, g) + x.copy_(new_x) + slots.copy_(new_slots) + graph.replay() + torch.cuda.synchronize() + live = new_slots[new_slots >= 0] + replayed = flat[live].view(torch.int16).clone() + ref = _fake_quant(new_x)[0] + _check_engine_pool(buf, pool, new_slots, ref) + op(x, pool, slots) # the same call, eagerly, on the same inputs + torch.cuda.synchronize() + eager = flat[live].view(torch.int16).clone() + _check_engine_pool(buf, pool, new_slots, ref) + assert torch.equal(replayed, eager), ( + f"T={t} replay {rep}: {int((replayed != eager).sum())} elements" + ) + del graph + print( + f"\n[graphs on [{pages}, {page_size}, 512]] T = {list(ts)}: every replay bitwise == eager == reference" + ) + del buf, pool, flat + _free_engine_pool() + + +def test_cuda_graph_every_graph_batch_on_the_engine_swa_pool() -> None: + """T = G (decode) and 6G (DSpark verify) for every G of _ENVELOPE_GSET (82 row counts, 1..1536), on the B_max + bf16 SWA pool [958720, 128, 512] (a PP2 stage alone).""" + assert len(_SWA_GRAPH_T) == 82 and max(_SWA_GRAPH_T) == 1536 + _capture_sweep(ENGINE_SWA_PAGES, ENGINE_PAGE, _SWA_GRAPH_T, seed=700) + + +def test_cuda_graph_dspark_ring_and_block_rows() -> None: + """T in 3 x and 5 x {1..8, 16, 24, 32} (stage-ring rows, draft block) on a ring-shaped [99, 128, 512] pool.""" + _capture_sweep(_RING_PAGES, ENGINE_PAGE, _RING_GRAPH_T, seed=900) + + +# ─── the DSpark drafter's stage-ring writes, in the drafter's own form ─────── +# +# The drafter's _ring_write writes each stage s = 0, 1, 2 with +# mxfp8_fake_quantize(kv_s [M, 512], ring3, idx3[:, s]): ring3 is the worker's +# _kv_windows [max_batch_size + 1, 3, 128, 512] bf16 viewed [(max_batch_size + 1) +# * 3, 128, 512], and the slots are column s of an int64 [M, 3] table (a stride-3 +# view: split(idx3, 1, 1)[s] reshaped to [M]) holding (ring_row * 3 + s) * 128 + +# (position & 127), or -1 for a masked row. M = G on every draft step's bonus +# write (forward_batched) and 5G on the multi-accept back-fill +# (write_context_windows_batched), both captured; n <= 128 on the prefill +# seeding (write_context_windows, eager), where ring3 is one ring row's +# [3, 128, 512] view. Triton specializes an integer argument equal to 1 and a +# pointer's 16-byte alignment, so the stride-3 slot map runs different compiled +# variants from the target's contiguous slot rows (stage 1's column starts 8 +# bytes into the table: two variants over the three stages): these tests run them. +_ENVELOPE_RING_ROWS = 256 + 1 # the worker's ring rows at max_batch_size 256 +_ENVELOPE_RING_PAGES = 3 * _ENVELOPE_RING_ROWS # 771 +_ENVELOPE_RING_M = tuple( + sorted(set(_ENVELOPE_GSET) | {5 * g for g in _ENVELOPE_GSET}) +) # 83 row counts, 1..1280 + + +def _ring_idx3(m: int, ring_rows: int, g: torch.Generator, masked: bool = True) -> torch.Tensor: + """int64 [m, 3] like the drafter's idx3: G requests on distinct ring rows (the last row among them), + k = m // G consecutive positions each (k = 1, the bonus write, while m < ring_rows; else k = 5, the + back-fill), entry (row * 3 + s) * 128 + (position & 127); with `masked`, m // 8 whole rows -1.""" + k = 1 if m < ring_rows else 5 + assert m % k == 0 and m // k < ring_rows, (m, k) + n_req = m // k + rows = torch.cat( + [ + torch.tensor([ring_rows - 1], device="cuda"), + torch.randperm(ring_rows - 1, device="cuda", generator=g)[: n_req - 1], + ] + ) + p0 = torch.randint(0, 2**20, (n_req, 1), device="cuda", generator=g) + off = (p0 + torch.arange(k, device="cuda").view(1, k)) & (ENGINE_PAGE - 1) # [G, k] + stage = torch.arange(3, device="cuda").view(1, 1, 3) + idx3 = ( + ((rows.view(-1, 1, 1) * 3 + stage) * ENGINE_PAGE + off.unsqueeze(-1)) + .reshape(m, 3) + .contiguous() + ) + if masked and m >= 8: + idx3[torch.randperm(m, device="cuda", generator=g)[: m // 8]] = -1 + return idx3.to(torch.int64) + + +def _ring_cols(idx3: torch.Tensor) -> list[torch.Tensor]: + """The drafter's per-stage slot maps: split(idx3, 1, 1)[s] reshaped to [M], stride-3 views.""" + cols = [c.reshape(idx3.shape[0]) for c in torch.split(idx3, 1, 1)] + assert idx3.shape[0] < 2 or all( + c.stride(0) == 3 and c.data_ptr() == idx3.data_ptr() + 8 * s for s, c in enumerate(cols) + ) + return cols + + +def test_envelope_dspark_ring_write_captured_at_every_graph_batch() -> None: + """The drafter's ring write captured at M = G and 5G for every G of _ENVELOPE_GSET (83 row counts, 1..1280). + + As the drafter does it: three calls per graph (stage 0, 1, 2), each with its own contiguous [M, 512] + input and a stride-3 column of one int64 [M, 3] table, into ring3 = the [257, 3, 128, 512] ring (max + batch 256) viewed [771, 128, 512], with a guard page past it. Per M: an eager warm-up (the three calls), + the capture, then 2 replays after new rows and a new table (new ring rows and positions, the ring's + last row among them, 1/8 of the rows masked) are written in place. Every replay writes the reference + into the named slots of every stage and nothing else (the whole ring and its guard page hold the + sentinel), and equals the same three calls run eagerly. + """ + buf, ring3 = _engine_pool(_ENVELOPE_RING_PAGES, ENGINE_PAGE) + num_slots = _ENVELOPE_RING_PAGES * ENGINE_PAGE + flat = ring3.view(num_slots, -1) + g = torch.Generator(device="cuda").manual_seed(9100) + graphs = 0 + for m in _ENVELOPE_RING_M: + xs = [ + torch.empty(m, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda") for _ in range(3) + ] + idx3 = torch.empty(m, 3, dtype=torch.int64, device="cuda") + cols = _ring_cols(idx3) + for s in range(3): + xs[s].copy_(_fast_rows(m, seed=9100 + 3 * m + s)) + idx3.copy_(_ring_idx3(m, _ENVELOPE_RING_ROWS, g)) + for s in range(3): # eager warm-up at this shape + op(xs[s], ring3, cols[s]) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for s in range(3): + op(xs[s], ring3, cols[s]) + graphs += 1 + buf.view(torch.int16).fill_(SENTINEL) + for rep in (1, 2): + new_x = [_fast_rows(m, seed=9100 + 100_000 * rep + 3 * m + s) for s in range(3)] + new_idx = _ring_idx3(m, _ENVELOPE_RING_ROWS, g) + for s in range(3): + xs[s].copy_(new_x[s]) + idx3.copy_(new_idx) + graph.replay() + torch.cuda.synchronize() + all_slots = new_idx.t().reshape(-1) + live = all_slots[all_slots >= 0] + assert torch.unique(live).numel() == live.numel() + ref = torch.cat([_fake_quant(new_x[s])[0] for s in range(3)]) + replayed = flat[live].view(torch.int16).clone() + _check_engine_pool(buf, ring3, all_slots, ref) + for s in range(3): # the same three calls, eagerly, on the same inputs + op(xs[s], ring3, cols[s]) + torch.cuda.synchronize() + eager = flat[live].view(torch.int16).clone() + _check_engine_pool(buf, ring3, all_slots, ref) + assert torch.equal(replayed, eager), ( + f"M={m} replay {rep}: {int((replayed != eager).sum())} elements" + ) + del graph + print( + f"\n[dspark ring writes on [{_ENVELOPE_RING_PAGES}, 128, 512], stride-3 slot columns] {graphs} graphs of 3 " + f"calls, M = {list(_ENVELOPE_RING_M)}: every replay bitwise == eager == reference, nothing else written" + ) + assert graphs == 83 and max(_ENVELOPE_RING_M) == 1280 + del buf, ring3, flat + _free_engine_pool() + + +def test_envelope_dspark_ring_write_every_row_count() -> None: + """Every row count of the drafter's stride-3 ring write, eager: every M in 1..1280 (5 x 256, the largest the + ring write takes), and the prefill seeding's n in 1..128 into one ring row's [3, 128, 512] view. + + The 1280-row calls (three stages, a [1280, 3] back-fill table over 256 ring rows, the last among them, + 1/8 of the rows masked) are held to the reference and the sentinel scan of the whole ring plus a guard + page. Then per M: the three calls with (x_s[:M], cols_s[:M]) (stride-3 prefix views); the live slots of + the first M rows hold the reference bit for bit, the count of non-sentinel elements is 512 per live + (row, stage), and the named slots are restored, so each M is checked on its own. The probe: after the + first stride-3 call, no M compiles a new variant. Then the seeding form: for n in 1..128, the rows at + positions p .. p + n - 1 of ring row 256 (idx3 = s * 128 + (p & 127) into kv_windows[256], a [3, 128, + 512] view), checked the same way over the whole ring. + """ + m_max = 5 * 256 + buf, ring3 = _engine_pool(_ENVELOPE_RING_PAGES, ENGINE_PAGE) + num_slots = _ENVELOPE_RING_PAGES * ENGINE_PAGE + flat = ring3.view(num_slots, -1).view(torch.int16) + whole = buf.view(torch.int16) + g = torch.Generator(device="cuda").manual_seed(9200) + idx3 = _ring_idx3(m_max, _ENVELOPE_RING_ROWS, g) + cols = _ring_cols(idx3) + xs = [_window_kv_like(m_max, seed=9200 + s) for s in range(3)] + ys = [_fake_quant(x)[0].view(torch.int16) for x in xs] + kernel = _mx_op._mxfp8_fake_quantize_kernel + n0 = _jit_variants(kernel) + for s in range(3): + op(xs[s], ring3, cols[s]) + n1 = _jit_variants(kernel) + all_slots = idx3.t().reshape(-1) + _check_engine_pool(buf, ring3, all_slots, torch.cat([_fake_quant(x)[0] for x in xs])) + + bad_rows = torch.zeros((), dtype=torch.int64, device="cuda") + bad_count = torch.zeros((), dtype=torch.int64, device="cuda") + live = idx3[:, 0] >= 0 + cnt = torch.cumsum(live.to(torch.int64), 0).tolist() + for m in range(1, m_max + 1): + for s in range(3): + op(xs[s][:m], ring3, cols[s][:m]) + sel = live[:m] + for s in range(3): + bad_rows += (flat[cols[s][:m][sel]] != ys[s][:m][sel]).sum() + bad_count += ((whole != SENTINEL).sum() - 3 * cnt[m - 1] * WINDOW_KV_WIDTH).abs() + flat[idx3[:m][sel].reshape(-1)] = SENTINEL + n2 = _jit_variants(kernel) + + # The prefill seeding: n rows at positions p .. p + n - 1 into ring row 256's [3, 128, 512] view. + window = ring3.view(_ENVELOPE_RING_ROWS, 3, ENGINE_PAGE, WINDOW_KV_WIDTH)[ + _ENVELOPE_RING_ROWS - 1 + ] + wflat = window.reshape(3 * ENGINE_PAGE, -1).view(torch.int16) + assert window.data_ptr() == ring3[-3:].data_ptr() + bad_seed = torch.zeros((), dtype=torch.int64, device="cuda") + for n in range(1, ENGINE_PAGE + 1): + p = 1000 + 37 * n + off = (p + torch.arange(n, device="cuda")) & (ENGINE_PAGE - 1) + widx = ( + off.view(n, 1) + ENGINE_PAGE * torch.arange(3, device="cuda").view(1, 3) + ).contiguous() + wcols = _ring_cols(widx) + for s in range(3): + op(xs[s][:n], window, wcols[s]) + for s in range(3): + bad_seed += (wflat[wcols[s]] != ys[s][:n]).sum() + bad_seed += ((whole != SENTINEL).sum() - 3 * n * WINDOW_KV_WIDTH).abs() + wflat[widx.reshape(-1)] = SENTINEL + n3 = _jit_variants(kernel) + print( + f"\n[dspark ring write, stride-3 slots, M 1..{m_max}] {cnt[-1]} live rows x 3 stages; variants {n0} -> " + f"{n1} (first stride-3 calls) -> {n2} (M sweep) -> {n3} (seeding n 1..128); element mismatches " + f"{int(bad_rows)}, count error {int(bad_count)}; seeding mismatches + count error {int(bad_seed)}" + ) + assert int(bad_rows) == 0 and int(bad_count) == 0 and int(bad_seed) == 0 + assert n1 - n0 <= 1 and n3 == n2 == n1, ( + f"row counts compiled new variants: {n0} -> {n1} -> {n2} -> {n3}" + ) + del buf, ring3, flat, whole, window, wflat + _free_engine_pool() + + +# ─── the pool bound: any SWA pool B <= B_max ────────────────────────────────── +# +# The engine picks the SWA pool's page count B from free memory, so no single +# dumped B is "the" engine's. Two tests make the coverage an upper bound: +# (1) the cache manager at the largest quota a gate-configured engine can reach +# gives B_max = ENGINE_SWA_PAGES (and the op is exact at the far end of the +# manager's own buffer), and (2) the same rows into the same slots give the same +# bits at B = the smallest pool that holds them, sizes across the 2^31 / 2^32 +# element-count lines, a mid size and B_max (views of one allocation), and in a +# separate exact-size allocation. + +# (max_draft_len, max_batch_size) -> SWA pages at 0.5 x this GB300's total memory. +_BMAX_DEVICE_TOTAL = 297_020_948_480 +_BMAX_SWA_PAGES = {(0, 32): 261_600, (5, 32): 323_840, (0, 8): 468_720, (5, 8): 491_000} + + +def _pp_mapping(pp_stage): + """The single-process mapping of a rank of PP2 stage `pp_stage` (tp 2, as the PP2 stage scan builds + it; the V4.1 pools do not depend on tp), or of the unsplit manager when pp_stage is None.""" + from tensorrt_llm.mapping import Mapping + + if pp_stage is None: + return Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + return Mapping(world_size=4, rank=2 * pp_stage, tp_size=2, pp_size=2, gpus_per_node=4) + + +def _v41_manager( + quota: int, + max_draft_len: int, + max_batch_size: int, + max_num_tokens: int = 8192, + max_seq_len=None, + pp_stage=None, +): + """A DeepSeek-V4.1 DeepseekV4CacheManager with bf16-container KV storage and the given per-rank KV quota, + at the gate settings unless max_num_tokens / max_seq_len say otherwise (max_seq_len as the engine passes + it: + 14 with DSpark). The storage is pinned: when TRTLLM_DSV41_KV_STORAGE is unset the manager stores the + KV natively (the native section below builds those managers explicitly).""" + import os + from types import SimpleNamespace + + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, + KvCacheConfig, + ) + + text = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=128, + head_dim=512, + index_head_dim=128, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, + ) + config = SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **{k: v for k, v in text.items() if k != "model_type"}, + ) + saved = os.environ.get("TRTLLM_DSV41_KV_STORAGE") + os.environ["TRTLLM_DSV41_KV_STORAGE"] = "bf16" + try: + mgr = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + max_gpu_total_bytes=quota, + host_cache_size=0, + event_buffer_max_size=0, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=40, + num_kv_heads=1, + head_dim=512, + tokens_per_block=128, + # DSpark (max_draft_len 5, overlap on) grows the manager's max_seq_len by 5 + 4 + 5. + max_seq_len=max_seq_len or (16398 if max_draft_len else 16384), + max_batch_size=max_batch_size, + mapping=_pp_mapping(pp_stage), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=max_num_tokens, + sparse_attn_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ), + pretrained_config=config, + spec_config=DSparkDecodingConfig(max_draft_len=max_draft_len, block_size=max_draft_len) + if max_draft_len + else None, + max_cuda_graph_batch_size=max_batch_size, + disable_overlap_scheduler=False, + ) + finally: + if saved is None: + os.environ.pop("TRTLLM_DSV41_KV_STORAGE", None) + else: + os.environ["TRTLLM_DSV41_KV_STORAGE"] = saved + assert mgr.v41_layout.kv_storage == "bf16" + return mgr + + +def test_largest_engine_pools_from_the_cache_manager() -> None: + """B_max from the manager itself, and the far-end hazard call on the manager's own B_max buffer. + + The quota is 0.5 x the device's total memory, the upper bound of + kv_cache_fraction 0.5 x free memory. The largest configuration is built + last and kept for the far-end call: its SWA view (layer 0; one pool for all + 40 layers) is filled with the sentinel, the op writes rows to both sides of + every 2^31..2^34 element-offset line, the first and the last page and the + last slot, and the view must hold the reference there and the sentinel + everywhere else. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + total = torch.cuda.mem_get_info()[1] + assert total == _BMAX_DEVICE_TOTAL, ( + f"B_max constants are for a {_BMAX_DEVICE_TOTAL}-byte GB300, got {total}" + ) + quota = total // 2 + got = {} + order = sorted(_BMAX_SWA_PAGES, key=_BMAX_SWA_PAGES.get) # the largest last + for key in order: + mgr = _v41_manager(quota, *key) + pool = mgr.get_buffers(0, AT.SWA) + got[key] = pool.shape[0] + assert ( + tuple(pool.shape[1:]) == (ENGINE_PAGE, WINDOW_KV_WIDTH) and pool.dtype == torch.bfloat16 + ) + assert pool.stride() == (ENGINE_PAGE * WINDOW_KV_WIDTH, WINDOW_KV_WIDTH, 1), pool.stride() + if key != order[-1]: + mgr.shutdown() + del mgr, pool + _free_engine_pool() + print( + f"\n[quota {quota} B = 0.5 x device total] SWA pages per (max_draft_len, max_batch_size): {got}" + ) + assert got == _BMAX_SWA_PAGES, got + assert max(got.values()) == GATES_SWA_PAGES + try: + num_slots = pool.shape[0] * pool.shape[1] + pool.view(torch.int16).fill_(SENTINEL) + hazard = _hazard_slots(num_slots, WINDOW_KV_WIDTH, ENGINE_PAGE) + assert hazard[-1] == num_slots - 1 and 2**34 // WINDOW_KV_WIDTH in hazard + slots = torch.tensor(hazard + [-1, num_slots], dtype=torch.int64, device="cuda") + x = _window_kv_like(slots.numel(), seed=4910) + op(x, pool, slots) + _check_engine_pool(pool, pool, slots, _fake_quant(x)[0]) + print( + f"[manager's own SWA buffer {tuple(pool.shape)}] {len(hazard)} far-end / hazard rows == reference, " + f"rest sentinel" + ) + finally: + mgr.shutdown() + del mgr, pool + _free_engine_pool() + + +def _invariance_slots(t: int, page_size: int, seed: int): + """(smallest page count holding t slots, int64 [2, t] table row 1 of distinct slots in it): the last + slot of that smallest pool is named, 1/16 of the other rows -1.""" + small = -(-t // page_size) + last = small * page_size - 1 + g = torch.Generator(device="cuda").manual_seed(seed) + perm = torch.randperm(small * page_size, device="cuda", generator=g) + pick = torch.cat([torch.tensor([last], device="cuda"), perm[perm != last]])[:t] + pick = pick[torch.randperm(t, device="cuda", generator=g)] + table = torch.full((2, t), -1, dtype=torch.int64, device="cuda") + slots = table[1] + slots.copy_(pick) + skip = torch.randperm(t, device="cuda", generator=g)[: t // 16] + slots[skip[slots[skip] != last]] = -1 + return small, slots + + +@pytest.mark.parametrize("t", [32, 8192, 32768]) +def test_pool_size_invariance(t: int) -> None: + """The same t rows into the same slots: bitwise the same bits at every pool size B, and = reference. + + B = the smallest pool holding the slots, the page counts whose element count first passes 2^31, 2^32, + 2^33, 2^34 and 2^35 (32769, 65537, 131073, 262145, 524289), the gates' bound 491000, the unsplit + manager's envelope bound 730880, a mid size B_max // 2 + 1 and B_max (958720, a PP2 stage alone), + each a view pool[:B] of one sentinel-filled allocation of B_max pages plus a guard page, then a separate + allocation of exactly the smallest size. After every call the whole allocation (past B too) holds the + sentinel outside the named slots. + """ + small, slots = _invariance_slots(t, ENGINE_PAGE, seed=2300 + t) + x = _window_kv_like(t, seed=2300 + t) if t <= 8192 else _fast_rows(t, seed=2300 + t) + y = _fake_quant(x)[0] + live = slots[slots >= 0] + per_page = ENGINE_PAGE * WINDOW_KV_WIDTH + sizes = sorted( + { + small, + *(lim // per_page + 1 for lim in (2**31, 2**32, 2**33, 2**34, 2**35)), + GATES_SWA_PAGES, + UNSPLIT_SWA_PAGES, + ENGINE_SWA_PAGES // 2 + 1, + ENGINE_SWA_PAGES, + } + ) + assert {32769, 65537, 131073, 262145, 524289, 491000, 730880} < set(sizes) and sizes[ + -1 + ] == 958720 + buf, _ = _engine_pool(ENGINE_SWA_PAGES, ENGINE_PAGE) + first = None + for b in sizes: + pool = buf[:b] + op(x, pool, slots) + torch.cuda.synchronize() + got = pool.view(-1, WINDOW_KV_WIDTH)[live].view(torch.int16).clone() + _check_engine_pool(buf, pool, slots, y) + first = got if first is None else first + assert torch.equal(got, first), ( + f"B={b}: {int((got != first).sum())} elements differ from B={sizes[0]}" + ) + del buf, pool + _free_engine_pool() + buf, pool = _engine_pool(small, ENGINE_PAGE) + op(x, pool, slots) + torch.cuda.synchronize() + got = pool.view(-1, WINDOW_KV_WIDTH)[live].view(torch.int16).clone() + _check_engine_pool(buf, pool, slots, y) + assert torch.equal(got, first), "a separate exact-size pool differs" + print( + f"\n[t={t}] pools B = {sizes} pages (views of one allocation) and a separate [{small}, 128, 512]: " + f"{live.numel()} rows bitwise identical and == reference, nothing else written" + ) + del buf, pool + _free_engine_pool() + + +# ─── native storage on the live DeepSeek-V4.1 cache manager ────────────────── +# +# With TRTLLM_DSV41_KV_STORAGE=native the V4.1 DeepseekV4CacheManager keeps the +# SWA window KV natively: one uint8 pool [B, 128, 528] shared by all 40 layers, +# each 528-byte row the 512 E4M3 codes then the 16 UE8M0 bytes, and +# get_v41_native_kv_views(L, SWA) returns (codes [B, 128, 512], scales +# [B, 128, 16]), two uint8 views of that pool with strides (67584, 528, 1). The +# native-storage target writes it with this op's native form: +# mxfp8_fake_quantize(kv [T, 512], codes, md.v41_swa_slot_mapping[local(L)][:T], scales). +# B comes from free memory again. The table is the manager's native SWA pages at +# 0.5 x this GB300's total memory (the ceiling of kv_cache_fraction 0.5 x free +# memory) for every (max_num_tokens, max_batch_size, DSpark max_draft_len) the +# runs use or could use: the native gates run (2048, 8, 5) (boot) and +# (8192, 32, 5) (accuracy), 549,080 and 732,280 pages; over max_batch_size +# {8, 32} the largest is (8192, 8, 5), 854,440 pages. Over every max_batch_size +# 1..32 (the native-storage witness runs at 1) the page count is +# not monotone in max_batch_size, and B_max = 941,800 pages at (8192, 1, 0) +# (120,550,400 slots, 63.6 GB). Over the engine envelope (max_num_tokens up +# to 16384, max_batch_size up to 256, max_seq_len up to 1,048,576) the unsplit +# manager's largest is 1,098,720 pages at (16384, 1, 0) (140,636,160 slots, +# 74.3 GB), and a PP2 stage-1 manager alone reaches 1,648,320 at (16384, 1, 0) +# (210,984,960 slots, 111.4 GB; ENVELOPE_NATIVE_SWA_BMAX, the B_max the far-end, +# invariance and capture cases use). Every byte-offset hazard up to 2^35 lies +# inside all of them, and 2^36 inside the two largest. +# +# Round trips: rows the op writes natively are read back with the paired reader +# (the raw op torch.ops.trtllm.mxfp8_dequantize) and held bitwise to the op's own +# bf16-container output of the same rows, which is itself held to the recipe +# built here from native torch ops. The reader is part of what is checked, not +# the reference. + +import os # noqa: E402 +import random # noqa: E402 +from contextlib import contextmanager # noqa: E402 + +import tensorrt_llm._torch.custom_ops.mxfp8_dequantize # noqa: E402,F401 -- registers the paired reader + +_STORAGE_ENV = "TRTLLM_DSV41_KV_STORAGE" +NATIVE_ROW = 528 # 512 E4M3 code bytes, then 16 UE8M0 bytes +NATIVE_CODES = 512 +# 0xFF is E4M3 -NaN and the UE8M0 NaN: for a finite row the op writes neither +# (codes are clamped to +-448; the scale byte is e + 127 with e in -22..120). +NATIVE_SENTINEL = 0xFF +_NATIVE_SWA_PAGES = { + (2048, 8, 0): 399_640, + (2048, 8, 5): 549_080, + (2048, 32, 0): 366_400, + (2048, 32, 5): 569_440, + (8192, 8, 0): 776_000, + (8192, 8, 5): 854_440, + (8192, 32, 0): 585_840, + (8192, 32, 5): 732_280, +} +_NATIVE_GATES = ((2048, 8, 5), (8192, 32, 5)) +MBS_AXIS = tuple(range(1, 33)) # every max_batch_size the engine configuration permits (1..32) +# Native SWA pages at 0.5 x device total per (max_num_tokens, DSpark max_draft_len), for max_batch_size 1..32 +# (measured on this GB300; re-measured by test_native_swa_pools_from_the_cache_manager). +_SWA_PAGES_BY_MBS = { + (2048, 0): ( + 549_600, + 549_600, + 549_600, + 549_600, + 549_600, + 488_040, + 439_360, + 399_640, + 399_640, + 549_080, + 506_880, + 506_880, + 506_880, + 506_880, + 506_880, + 439_360, + 439_360, + 439_360, + 439_360, + 412_080, + 387_720, + 387_720, + 462_720, + 462_720, + 439_360, + 439_360, + 439_360, + 418_520, + 418_520, + 382_280, + 366_400, + 366_400, + ), + (2048, 5): ( + 549_600, + 549_600, + 549_600, + 549_600, + 732_280, + 599_240, + 599_240, + 549_080, + 549_080, + 549_080, + 506_880, + 628_040, + 628_040, + 628_040, + 549_080, + 549_080, + 549_080, + 646_400, + 646_400, + 610_160, + 578_400, + 578_400, + 549_080, + 628_040, + 599_240, + 599_240, + 573_440, + 573_440, + 573_440, + 591_800, + 569_440, + 569_440, + ), + (8192, 0): ( + 941_800, + 879_240, + 879_240, + 879_240, + 879_240, + 824_160, + 824_160, + 776_000, + 776_000, + 732_280, + 732_280, + 809_760, + 809_760, + 809_760, + 732_280, + 732_280, + 732_280, + 732_280, + 732_280, + 699_040, + 668_760, + 640_960, + 615_120, + 676_200, + 650_880, + 650_880, + 650_880, + 650_880, + 606_200, + 606_200, + 585_840, + 585_840, + ), + (8192, 5): ( + 879_240, + 879_240, + 879_240, + 879_240, + 879_240, + 905_080, + 905_080, + 854_440, + 854_440, + 809_760, + 809_760, + 879_240, + 879_240, + 879_240, + 798_840, + 798_840, + 798_840, + 798_840, + 859_880, + 824_160, + 790_880, + 732_280, + 732_280, + 732_280, + 757_600, + 757_600, + 757_600, + 757_600, + 708_960, + 708_960, + 732_280, + 732_280, + ), +} +_NATIVE_SWA_PAGES_AXIS = { + (t, s, d): pages[s - 1] for (t, d), pages in _SWA_PAGES_BY_MBS.items() for s in MBS_AXIS +} +assert all(_NATIVE_SWA_PAGES_AXIS[k] == v for k, v in _NATIVE_SWA_PAGES.items()) +NATIVE_SWA_BMAX = 941_800 # over every max_batch_size 1..32 +NATIVE_SWA_BMAX_KEY = (8192, 1, 0) # the configuration that reaches it +NATIVE_SWA_BMAX_8_32 = 854_440 # the bound over max_batch_size {8, 32} only +NATIVE_SWA_BMAX_8_32_KEY = (8192, 8, 5) +NATIVE_SWA_GATE_MAX = 732_280 +NATIVE_WITNESS_KEY = (2048, 1, 5) # the native-storage witness configuration +# The engine envelope (max_num_tokens 2048..16384, max_batch_size up to 256, max_seq_len up to 1,048,576, +# DSpark or not): the pool scan's largest native SWA pool, at (16384, 1, 0) whatever max_seq_len; +# re-measured by test_envelope_largest_swa_pools_from_the_cache_manager. +UNSPLIT_NATIVE_SWA_BMAX = 1_098_720 +ENVELOPE_SWA_BMAX_KEY = (16384, 1, 0) +# PP2 (the unified target's stage split at layer 20): a stage manager holds only its 20 layers' pools. The +# PP2 stage scan's largest native SWA pool is stage 1 alone at (16384, 1, 0), whatever max_seq_len (the +# stage alone is the upper bound: in an engine the stages' token capacities are cut to the smaller one); +# re-measured by test_envelope_pp2_stage_swa_pools_from_the_cache_manager. +ENVELOPE_NATIVE_SWA_BMAX = 1_648_320 +ENVELOPE_PP2_SWA_STAGE = 1 +_NATIVE_FAR_PAGES = (UNSPLIT_NATIVE_SWA_BMAX, ENVELOPE_NATIVE_SWA_BMAX) +_NATIVE_PAGE_BYTES = ENGINE_PAGE * NATIVE_ROW # 67584 + + +@contextmanager +def _native_storage(): + saved = os.environ.get(_STORAGE_ENV) + os.environ[_STORAGE_ENV] = "native" + try: + yield + finally: + if saved is None: + os.environ.pop(_STORAGE_ENV, None) + else: + os.environ[_STORAGE_ENV] = saved + + +def _v41_native_manager( + max_num_tokens: int, max_batch_size: int, max_draft_len: int, max_seq_len=None, pp_stage=None +): + """A V4.1 manager with native KV storage and 0.5 x the device's total memory, at the gates' max_seq_len + unless max_seq_len says otherwise (as the engine passes it: + 14 with DSpark).""" + from types import SimpleNamespace + + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, + KvCacheConfig, + ) + + total = torch.cuda.mem_get_info()[1] + assert total == _BMAX_DEVICE_TOTAL, ( + f"B_max constants are for a {_BMAX_DEVICE_TOTAL}-byte GB300, got {total}" + ) + text = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=128, + head_dim=512, + index_head_dim=128, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, + ) + config = SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **{k: v for k, v in text.items() if k != "model_type"}, + ) + with _native_storage(): + mgr = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + max_gpu_total_bytes=total // 2, + host_cache_size=0, + event_buffer_max_size=0, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=40, + num_kv_heads=1, + head_dim=512, + tokens_per_block=128, + max_seq_len=max_seq_len or (16398 if max_draft_len else 16384), + max_batch_size=max_batch_size, + mapping=_pp_mapping(pp_stage), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=max_num_tokens, + sparse_attn_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ), + pretrained_config=config, + spec_config=DSparkDecodingConfig(max_draft_len=max_draft_len, block_size=max_draft_len) + if max_draft_len + else None, + max_cuda_graph_batch_size=max_batch_size, + disable_overlap_scheduler=False, + ) + assert mgr.v41_layout.kv_storage == "native" + return mgr + + +def _native_swa(mgr, layer: int = 0): + """(the manager's SWA pool [B, 128, 528] uint8, codes view, scales view) through `layer` (a layer of the + manager's stage), geometry asserted.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + pool = mgr.get_buffers(layer, AT.SWA) + codes, scales = mgr.get_v41_native_kv_views(layer, AT.SWA) + b = pool.shape[0] + assert pool.dtype == torch.uint8 and tuple(pool.shape) == (b, ENGINE_PAGE, NATIVE_ROW), ( + pool.shape + ) + assert pool.stride() == (_NATIVE_PAGE_BYTES, NATIVE_ROW, 1), pool.stride() + assert tuple(codes.shape) == (b, ENGINE_PAGE, NATIVE_CODES) and codes.stride() == pool.stride() + assert tuple(scales.shape) == (b, ENGINE_PAGE, 16) and scales.stride() == pool.stride() + assert ( + codes.data_ptr() == pool.data_ptr() and scales.data_ptr() == pool.data_ptr() + NATIVE_CODES + ) + return pool, codes, scales + + +def _native_rows(x: torch.Tensor): + """(uint8 [T, 528] native rows = codes then UE8M0 bytes, bf16 [T, 512] container values) of the recipe.""" + y, codes, sbytes = _fake_quant(x) + return torch.cat([codes, sbytes], dim=1), y + + +def _decode_native(rows: torch.Tensor) -> torch.Tensor: + """bf16 [n, 512] of native rows, in native torch: e4m3(code) * 2^(byte - 127).""" + v = rows[:, :NATIVE_CODES].contiguous().view(torch.float8_e4m3fn).float().view(-1, 16, 32) + s = torch.exp2(rows[:, NATIVE_CODES:].float() - 127).unsqueeze(-1) + return (v * s).view(-1, NATIVE_CODES).to(torch.bfloat16) + + +def _round_trip( + codes: torch.Tensor, scales: torch.Tensor, slots: torch.Tensor, container: torch.Tensor +) -> int: + """Read the rows at `slots` back with the paired reader (compact: out row r = pool slot slots[r]) into a + NaN-filled bf16 out, and hold every row whose slot is in the pool bitwise to `container` (the op's own + bf16-container output of the same rows). Returns the number of rows checked.""" + num_slots = codes.shape[0] * codes.shape[1] + out = torch.empty(slots.numel(), NATIVE_CODES, dtype=torch.bfloat16, device="cuda") + out.view(torch.int16).fill_(SENTINEL) + torch.ops.trtllm.mxfp8_dequantize(codes, scales, out, slots) + s64 = slots.to(torch.int64) + live = (s64 >= 0) & (s64 < num_slots) + _assert_bitwise(out[live], container[live]) + return int(live.sum()) + + +def _native_hazard_slots(num_slots: int, page_size: int = ENGINE_PAGE) -> list[int]: + """Distinct slots of a 528-byte-row pool on both sides of every byte offset 2^31..2^36 inside it (the row + holding the offset and its neighbours), plus the first page, the last page and the last slot.""" + cands = {0, 1, page_size - 1, page_size} + for limit in (2**31, 2**32, 2**33, 2**34, 2**35, 2**36): + s = limit // NATIVE_ROW + cands |= {s - 1, s, s + 1} + cands |= {num_slots - page_size - 1, num_slots - page_size, num_slots - 2, num_slots - 1} + return sorted(c for c in cands if 0 <= c < num_slots) + + +def _count_not_sentinel(buf: torch.Tensor) -> torch.Tensor: + """Device count of bytes of `buf` (uint8, any shape, contiguous) that are not NATIVE_SENTINEL.""" + whole = buf.view(-1) + step = ( + 1 << 26 + ) # a bool reduction materializes an int64 copy (8 bytes per element): keep chunks small + n = torch.zeros((), dtype=torch.int64, device="cuda") + for i in range(0, whole.numel(), step): + n += (whole[i : i + step] != NATIVE_SENTINEL).sum() + return n + + +def _check_native_pool( + buf: torch.Tensor, num_slots: int, slots: torch.Tensor, rows: torch.Tensor +) -> None: + """buf: uint8 [pages(+guard), 128, 528] whose first num_slots slots are the pool. The named slots hold + `rows` bitwise; every other byte of buf is the sentinel. Restores the named slots to the sentinel.""" + flat = buf.view(-1, NATIVE_ROW) + s64 = slots.to(torch.int64) + live = (s64 >= 0) & (s64 < num_slots) + idx = s64[live] + got = flat[idx] + n_bad = int((got != rows[live]).sum()) + assert n_bad == 0, f"{n_bad} of {got.numel()} bytes in the named slots differ bitwise" + flat[idx] = NATIVE_SENTINEL + n_out = int(_count_not_sentinel(buf)) + assert n_out == 0, f"{n_out} bytes outside the named slots were written" + + +def _native_engine_pool(pages: int): + """(uint8 [pages + 1, 128, 528] buffer with a guard page, codes view, scales view), all bytes sentinel.""" + buf = torch.full( + (pages + 1, ENGINE_PAGE, NATIVE_ROW), NATIVE_SENTINEL, dtype=torch.uint8, device="cuda" + ) + return buf, buf[:pages, :, :NATIVE_CODES], buf[:pages, :, NATIVE_CODES:] + + +def _fast_rows(t: int, seed: int) -> torch.Tensor: + """bf16 [t, 512] window-KV-like rows, vectorized: one magnitude 2^-30..2^20 per 32-block, some blocks + zero and some under the 1e-4 floor (scale byte 105).""" + g = torch.Generator(device="cuda").manual_seed(seed) + mag = torch.randint(-30, 21, (t, 16, 1), device="cuda", generator=g).float() + x = torch.randn(t, 16, BLOCK, device="cuda", generator=g) * torch.exp2(mag) + pick = torch.randint(0, 16, (t, 1, 1), device="cuda", generator=g) + x = torch.where(torch.arange(16, device="cuda").view(1, 16, 1) == pick, 0.0, x) + return x.view(t, WINDOW_KV_WIDTH).to(torch.bfloat16) + + +def _manager_far_end_round_trip(pool, codes, scales, seed: int) -> None: + """On a manager's own native SWA pool: sentinel-fill it, write rows through the manager's views to both + sides of every 2^31..2^35 byte offset, the first and the last page and the last slot (plus a -1 and a + one-past-the-end slot, skipped); the named rows hold the recipe's codes and bytes, decode in torch and + read back through the paired reader to the op's own bf16-container output, the rest is the sentinel.""" + num_slots = pool.shape[0] * ENGINE_PAGE + pool.fill_(NATIVE_SENTINEL) + hazard = _native_hazard_slots(num_slots) + assert hazard[-1] == num_slots - 1 and 2**35 // NATIVE_ROW in hazard + slots = torch.tensor(hazard + [-1, num_slots], dtype=torch.int64, device="cuda") + x = _window_kv_like(slots.numel(), seed=seed) + rows, y = _native_rows(x) + op(x, codes, slots, scales) + live = slots[(slots >= 0) & (slots < num_slots)] + container = torch.empty_like(x) + op(x, container) + _assert_bitwise(_decode_native(pool.view(-1, NATIVE_ROW)[live]), container[: live.numel()]) + _assert_bitwise(container, y) + assert _round_trip(codes, scales, slots, container) == len(hazard) + _check_native_pool(pool, num_slots, slots, rows) + print( + f"[manager's own native SWA pool {tuple(pool.shape)}] {len(hazard)} far-end / hazard rows == " + f"recipe codes + bytes; torch decode == mxfp8_dequantize read-back == the bf16 container; rest sentinel" + ) + + +def test_native_swa_pools_from_the_cache_manager() -> None: + """Every native configuration's SWA pool from the manager itself, and the hazard call on the largest. + + For each (max_num_tokens, max_batch_size, max_draft_len) of {2048, 8192} x + every max_batch_size 1..32 x {0, 5} (128 managers) at 0.5 x total memory: + the pool is [B, 128, 528] uint8 for all 40 layers (one base), the views + are its code / scale halves, and 890 bytes per token; the pages equal the + table. B_max over the axis is 941,800 at (8192, 1, 0); over max_batch_size + {8, 32} it is 854,440 at (8192, 8, 5); 732,280 at the gates. On both of + those managers' own pools the far-end round trip runs: filled with the + sentinel, the op's native form writes rows to both sides of every + 2^31..2^35 byte offset, the first and the last page and the last slot + through the manager's views, and the pool must hold the recipe's codes and + bytes there, decoding (in torch, and through mxfp8_dequantize) to the op's + own bf16-container output, and the sentinel everywhere else. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + got = {} + order = sorted(_NATIVE_SWA_PAGES_AXIS, key=_NATIVE_SWA_PAGES_AXIS.get) # the largest last + assert order[-1] == NATIVE_SWA_BMAX_KEY + for key in order: + mgr = _v41_native_manager(*key) + pool, codes, scales = _native_swa(mgr) + assert mgr.max_batch_size == key[1] and mgr.max_seq_len == (16398 if key[2] else 16384) + assert mgr.get_cache_bytes_per_token() == 890, (key, mgr.get_cache_bytes_per_token()) + for layer in range(40): + v = mgr.get_buffers(layer, AT.SWA) + assert v.data_ptr() == pool.data_ptr() and v.shape == pool.shape, (key, layer) + got[key] = pool.shape[0] + if key == NATIVE_SWA_BMAX_8_32_KEY: + assert got[key] == NATIVE_SWA_BMAX_8_32 + _manager_far_end_round_trip(pool, codes, scales, seed=8544) + if key != order[-1]: + mgr.shutdown() + del mgr, pool, codes, scales, v + _free_engine_pool() + over = {k: v for k, v in sorted(got.items()) if v > NATIVE_SWA_BMAX_8_32} + print( + f"\n[native, quota = 0.5 x device total] SWA pages per (max_num_tokens, max_batch_size, " + f"max_draft_len), 128 configurations: {dict(sorted(got.items()))}\n" + f"[above the max_batch_size {{8, 32}} bound {NATIVE_SWA_BMAX_8_32}] {len(over)}: {over}" + ) + assert got == _NATIVE_SWA_PAGES_AXIS, { + k: (v, _NATIVE_SWA_PAGES_AXIS[k]) for k, v in got.items() if v != _NATIVE_SWA_PAGES_AXIS[k] + } + assert max(got.values()) == NATIVE_SWA_BMAX and got[NATIVE_SWA_BMAX_KEY] == NATIVE_SWA_BMAX + assert max(got[k] for k in _NATIVE_SWA_PAGES) == NATIVE_SWA_BMAX_8_32 + assert max(got[k] for k in _NATIVE_GATES) == NATIVE_SWA_GATE_MAX + try: + _manager_far_end_round_trip(pool, codes, scales, seed=9418) + finally: + mgr.shutdown() + del mgr, pool, codes, scales + _free_engine_pool() + + +# (storage, max_num_tokens, max_batch_size, DSpark max_draft_len, max_seq_len as the engine passes it) -> SWA +# pages at 0.5 x this GB300's total memory: the envelope ceilings (the pool scan's maxima over 1920 +# managers per storage: max_num_tokens 2048..16384, max_batch_size 1..256 sampled, DSpark 0 / 5, max_seq_len +# 16384..1,048,576) at two max_seq_len, and the max_batch_size-256 configurations beside them. +_ENVELOPE_SWA_PAGES = { + ("bf16", 16384, 1, 0, 16384): 730_880, + ("bf16", 16384, 1, 0, 1_048_576): 730_880, + ("bf16", 16384, 256, 0, 16384): 161_520, + ("bf16", 16384, 256, 5, 16398): 244_200, + ("native", 16384, 1, 0, 16384): 1_098_720, + ("native", 16384, 1, 0, 1_048_576): 1_098_720, + ("native", 16384, 256, 0, 16384): 420_520, + ("native", 16384, 256, 5, 16398): 617_600, +} + + +def test_envelope_largest_swa_pools_from_the_cache_manager() -> None: + """The envelope SWA ceilings from the manager itself, and the far-end call on each ceiling manager's + own pool. + + Each configuration of _ENVELOPE_SWA_PAGES is built at 0.5 x the device's total memory in its storage mode, the + pool geometry asserted ([B, 128, 512] bf16 with strides (65536, 512, 1); [B, 128, 528] uint8 with its + code / scale views) and B pinned: 730,880 (bf16) and 1,098,720 (native) at (16384, 1, 0) whatever + max_seq_len. On the (16384, 1, 0, 16384) managers' own pools (filled with the sentinel first) the op + writes rows to both sides of every 32-bit hazard, the first and the last page and the last slot: + bf16 element offsets 2^30..2^35 (hazard rows == reference, rest sentinel); native byte offsets + 2^31..2^36 (codes + bytes == recipe, torch decode and mxfp8_dequantize read-back == the bf16 container, + rest sentinel). + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + total = torch.cuda.mem_get_info()[1] + assert total == _BMAX_DEVICE_TOTAL, ( + f"B_max constants are for a {_BMAX_DEVICE_TOTAL}-byte GB300, got {total}" + ) + got = {} + for key in _ENVELOPE_SWA_PAGES: + storage, max_t, seqs, draft, lm = key + if storage == "bf16": + mgr = _v41_manager(total // 2, draft, seqs, max_num_tokens=max_t, max_seq_len=lm) + else: + mgr = _v41_native_manager(max_t, seqs, draft, max_seq_len=lm) + try: + assert mgr.max_seq_len == lm and mgr.max_batch_size == seqs + far = (max_t, seqs, draft, lm) == ENVELOPE_SWA_BMAX_KEY + (16384,) + if storage == "bf16": + pool = mgr.get_buffers(0, AT.SWA) + assert ( + tuple(pool.shape[1:]) == (ENGINE_PAGE, WINDOW_KV_WIDTH) + and pool.dtype == torch.bfloat16 + ) + assert pool.stride() == (ENGINE_PAGE * WINDOW_KV_WIDTH, WINDOW_KV_WIDTH, 1), ( + pool.stride() + ) + got[key] = pool.shape[0] + if far: + num_slots = pool.shape[0] * pool.shape[1] + pool.view(torch.int16).fill_(SENTINEL) + hazard = _hazard_slots(num_slots, WINDOW_KV_WIDTH, ENGINE_PAGE) + assert hazard[-1] == num_slots - 1 and 2**35 // WINDOW_KV_WIDTH in hazard + slots = torch.tensor(hazard + [-1, num_slots], dtype=torch.int64, device="cuda") + x = _window_kv_like(slots.numel(), seed=7308) + op(x, pool, slots) + _check_engine_pool(pool, pool, slots, _fake_quant(x)[0]) + print( + f"\n[bf16 manager {key}, own SWA buffer {tuple(pool.shape)}] {len(hazard)} far-end / " + f"hazard rows == reference, rest sentinel" + ) + del pool + else: + pool, codes, scales = _native_swa(mgr) + got[key] = pool.shape[0] + if far: + assert 2**36 // NATIVE_ROW in _native_hazard_slots(pool.shape[0] * ENGINE_PAGE) + _manager_far_end_round_trip(pool, codes, scales, seed=10987) + del pool, codes, scales + finally: + mgr.shutdown() + del mgr + _free_engine_pool() + print(f"\n[envelope SWA pages at 0.5 x device total] {got}") + assert got == _ENVELOPE_SWA_PAGES, { + k: (v, _ENVELOPE_SWA_PAGES[k]) for k, v in got.items() if v != _ENVELOPE_SWA_PAGES[k] + } + assert max(v for k, v in got.items() if k[0] == "bf16") == UNSPLIT_SWA_PAGES + assert max(v for k, v in got.items() if k[0] == "native") == UNSPLIT_NATIVE_SWA_BMAX + + +# (storage, PP2 stage, max_num_tokens, max_batch_size, DSpark max_draft_len, max_seq_len) -> SWA pages of the +# stage manager built alone at 0.5 x this GB300's total memory: each storage's stage-alone maximum over the +# engine envelope (stage 1) and the other stage at its own maximum's tuple. +_ENVELOPE_PP2_SWA_PAGES = { + ("bf16", 0, 16384, 1, 0, 16384): 596_220, + ("bf16", 1, 16384, 2, 0, 16384): 958_720, + ("native", 0, 16384, 1, 0, 16384): 878_760, + ("native", 1, 16384, 1, 0, 16384): 1_648_320, +} + + +def test_envelope_pp2_stage_swa_pools_from_the_cache_manager() -> None: + """The PP2 stage managers' SWA pools, and the far-end call on each storage's largest stage pool. + + Each configuration of _ENVELOPE_PP2_SWA_PAGES is built alone (one process) with the mapping of a rank of that + stage (pp_size 2, tp 2) at 0.5 x the device's total memory: the stage holds layers 0-19 or 20-39 + (pp_layers asserted), every layer of the stage has the same SWA view on one pool base, and B is pinned + (bf16 [B, 128, 512] with strides (65536, 512, 1); native [B, 128, 528] with its code / scale views). On + the two stage-1 managers' own pools (filled with the sentinel first) the op writes rows to both sides of + every 32-bit hazard, the first and the last page and the last slot through a stage-1 layer's views: + bf16 element offsets 2^30..2^35 (rows == reference, rest sentinel); native byte offsets 2^31..2^36 (codes + + bytes == recipe, torch decode and mxfp8_dequantize read-back == the bf16 container, rest sentinel). + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + total = torch.cuda.mem_get_info()[1] + assert total == _BMAX_DEVICE_TOTAL, ( + f"B_max constants are for a {_BMAX_DEVICE_TOTAL}-byte GB300, got {total}" + ) + got = {} + for key in _ENVELOPE_PP2_SWA_PAGES: + storage, stage, max_t, seqs, draft, lm = key + if storage == "bf16": + mgr = _v41_manager( + total // 2, draft, seqs, max_num_tokens=max_t, max_seq_len=lm, pp_stage=stage + ) + else: + mgr = _v41_native_manager(max_t, seqs, draft, max_seq_len=lm, pp_stage=stage) + try: + layers = list(range(20)) if stage == 0 else list(range(20, 40)) + assert ( + list(mgr.pp_layers) == layers + and mgr.max_seq_len == lm + and mgr.max_batch_size == seqs + ) + views = [mgr.get_buffers(layer, AT.SWA) for layer in layers] + assert all( + v.data_ptr() == views[0].data_ptr() and v.shape == views[0].shape for v in views + ), key + got[key] = views[0].shape[0] + if storage == "bf16": + pool = views[0] + assert ( + tuple(pool.shape[1:]) == (ENGINE_PAGE, WINDOW_KV_WIDTH) + and pool.dtype == torch.bfloat16 + ) + assert pool.stride() == (ENGINE_PAGE * WINDOW_KV_WIDTH, WINDOW_KV_WIDTH, 1), ( + pool.stride() + ) + if stage == 1: + num_slots = pool.shape[0] * pool.shape[1] + pool.view(torch.int16).fill_(SENTINEL) + hazard = _hazard_slots(num_slots, WINDOW_KV_WIDTH, ENGINE_PAGE) + assert hazard[-1] == num_slots - 1 and 2**35 // WINDOW_KV_WIDTH in hazard + slots = torch.tensor(hazard + [-1, num_slots], dtype=torch.int64, device="cuda") + x = _window_kv_like(slots.numel(), seed=9587) + op(x, pool, slots) + _check_engine_pool(pool, pool, slots, _fake_quant(x)[0]) + print( + f"\n[bf16 PP2 stage-1 manager {key}, own SWA buffer {tuple(pool.shape)}] {len(hazard)} " + f"far-end / hazard rows == reference, rest sentinel" + ) + del pool + else: + pool, codes, scales = _native_swa(mgr, layers[0]) + if stage == 1: + assert 2**36 // NATIVE_ROW in _native_hazard_slots(pool.shape[0] * ENGINE_PAGE) + _manager_far_end_round_trip(pool, codes, scales, seed=16483) + del pool, codes, scales + del views + finally: + mgr.shutdown() + del mgr + _free_engine_pool() + print(f"\n[PP2 stage SWA pages at 0.5 x device total, stage alone] {got}") + assert got == _ENVELOPE_PP2_SWA_PAGES, { + k: (v, _ENVELOPE_PP2_SWA_PAGES[k]) + for k, v in got.items() + if v != _ENVELOPE_PP2_SWA_PAGES[k] + } + assert max(v for k, v in got.items() if k[0] == "bf16") == ENGINE_SWA_PAGES + assert max(v for k, v in got.items() if k[0] == "native") == ENVELOPE_NATIVE_SWA_BMAX + + +@pytest.mark.parametrize("pages", _NATIVE_FAR_PAGES) +@pytest.mark.parametrize("slot_dtype", [torch.int32, torch.int64]) +def test_native_engine_swa_pool_far_end(slot_dtype: torch.dtype, pages: int) -> None: + """Decode-like on a [pages, 128, 528] pool with a guard page, at the envelope B_max (1,648,320, a PP2 + stage alone) and the unsplit manager's (1,098,720): hazard rows, -1 and past-the-end skipped; the rows read + back through mxfp8_dequantize equal the op's bf16-container output.""" + buf, codes, scales = _native_engine_pool(pages) + num_slots = pages * ENGINE_PAGE + hazard = _native_hazard_slots(num_slots) + assert 2**35 // NATIVE_ROW in hazard and hazard[-1] == num_slots - 1 + slots = torch.tensor( + hazard + [-1, num_slots, num_slots + ENGINE_PAGE - 1], dtype=torch.int64, device="cuda" + ) + x = _window_kv_like(slots.numel(), seed=8545) + op(x, codes, slots.to(slot_dtype), scales) + container = torch.empty_like(x) + op(x, container) + _assert_bitwise(container, _native_rows(x)[1]) + assert _round_trip(codes, scales, slots.to(slot_dtype), container) == len(hazard) + _check_native_pool(buf, num_slots, slots, _native_rows(x)[0]) + del buf, codes, scales + _free_engine_pool() + + +@pytest.mark.parametrize("pages", _NATIVE_FAR_PAGES) +def test_native_engine_swa_pool_prefill_spread(pages: int) -> None: + """Prefill-like, T = 32768 (the decoder pass's largest step) on a [pages, 128, 528] pool with a guard page + (1,648,320 and 1,098,720): the last page in order, every hazard slot, the rest spread uniformly, 1/16 of the + rows -1; the rows read back through mxfp8_dequantize equal the op's bf16-container output.""" + t = _MAX_ROWS + buf, codes, scales = _native_engine_pool(pages) + num_slots = pages * ENGINE_PAGE + last_page = torch.arange(num_slots - ENGINE_PAGE, num_slots, device="cuda") + hazard = torch.tensor(_native_hazard_slots(num_slots), device="cuda") + fixed = torch.unique(torch.cat([last_page, hazard])) + g = torch.Generator(device="cuda").manual_seed(8546) + rest = torch.randint(0, num_slots, (3 * t,), device="cuda", generator=g).unique() + rest = rest[~torch.isin(rest, fixed)] + rest = rest[torch.randperm(rest.numel(), device="cuda", generator=g)][: t - fixed.numel()] + slots = torch.cat([last_page, fixed[~torch.isin(fixed, last_page)], rest]).to(torch.int64) + assert slots.numel() == t and torch.unique(slots).numel() == t + slots[torch.randperm(t, device="cuda", generator=g)[: t // 16]] = -1 + x = _fast_rows(t, seed=8546) + op(x, codes, slots.to(torch.int32), scales) + container = torch.empty_like(x) + op(x, container) + _assert_bitwise(container, _native_rows(x)[1]) + assert _round_trip(codes, scales, slots.to(torch.int32), container) == t - t // 16 + _check_native_pool(buf, num_slots, slots, _native_rows(x)[0]) + del buf, codes, scales + _free_engine_pool() + + +@pytest.mark.parametrize("pages", _NATIVE_FAR_PAGES) +def test_native_engine_swa_pool_harness_sees_a_wrapped_offset(pages: int) -> None: + """Control, on [pages, 128, 528] (1,648,320 and 1,098,720): rows written where a mod-2^32 byte offset would + put them fail both halves of the pool check (the named rows and the sentinel scan) and the round trip + (the reader's rows at the true slots are not the container).""" + buf, codes, scales = _native_engine_pool(pages) + num_slots = pages * ENGINE_PAGE + true = torch.tensor( + [s for s in _native_hazard_slots(num_slots) if s * NATIVE_ROW >= 2**32], + dtype=torch.int64, + device="cuda", + ) + wrapped = (true * NATIVE_ROW % 2**32) // NATIVE_ROW + assert true.numel() >= 9 and not torch.isin(wrapped, true).any() + x = _window_kv_like(true.numel(), seed=8547) + rows = _native_rows(x)[0] + op(x, codes, wrapped, scales) + container = torch.empty_like(x) + op(x, container) + with pytest.raises(AssertionError, match="elements differ bitwise"): + _round_trip(codes, scales, true, container) + with pytest.raises(AssertionError, match="bytes in the named slots differ"): + _check_native_pool(buf, num_slots, true, rows) + buf.view(-1, NATIVE_ROW)[true] = NATIVE_SENTINEL + with pytest.raises(AssertionError, match="outside the named slots were written"): + _check_native_pool( + buf, num_slots, torch.tensor([], dtype=torch.int64, device="cuda"), rows[:0] + ) + del buf, codes, scales + _free_engine_pool() + + +@pytest.mark.parametrize("t", [32, 8192, 32768]) +def test_native_pool_size_invariance(t: int) -> None: + """The same rows into the same slots through views codes[:B] / scales[:B] of the live PP2 stage-1 + manager's native SWA pool (16384, 1, 0, 16384; 1,648,320 pages, the envelope B_max): the same bytes at + every B, equal to the recipe, read back through mxfp8_dequantize (codes[:B] / scales[:B]) bitwise as the + op's bf16-container output, nothing else written; then a separate allocation of exactly the smallest size. + + B = the smallest pool holding the slots, the page counts whose byte count first passes 2^31, 2^32, + 2^33, 2^34, 2^35 and 2^36, a mid size, the gates' largest (732,280), the max_batch_size {8, 32} bound + (854,440), B_max over max_batch_size 1..32 at the gates' max_num_tokens (941,800), the unsplit + manager's envelope B_max (1,098,720) and the stage's (1,648,320). + """ + mgr = _v41_native_manager(*ENVELOPE_SWA_BMAX_KEY, pp_stage=ENVELOPE_PP2_SWA_STAGE) + try: + assert list(mgr.pp_layers) == list(range(20, 40)) + pool, codes, scales = _native_swa(mgr, 20) + assert pool.shape[0] == ENVELOPE_NATIVE_SWA_BMAX + pool.fill_(NATIVE_SENTINEL) + small, slots = _invariance_slots(t, ENGINE_PAGE, seed=8548 + t) + x = _window_kv_like(t, seed=8548 + t) if t <= 8192 else _fast_rows(t, seed=8548 + t) + rows, y = _native_rows(x) + container = torch.empty_like(x) + op(x, container) + _assert_bitwise(container, y) + live = slots[slots >= 0] + sizes = sorted( + { + small, + *( + lim // _NATIVE_PAGE_BYTES + 1 + for lim in (2**31, 2**32, 2**33, 2**34, 2**35, 2**36) + ), + ENVELOPE_NATIVE_SWA_BMAX // 2 + 1, + NATIVE_SWA_GATE_MAX, + NATIVE_SWA_BMAX_8_32, + NATIVE_SWA_BMAX, + UNSPLIT_NATIVE_SWA_BMAX, + ENVELOPE_NATIVE_SWA_BMAX, + } + ) + first = None + for b in sizes: + op(x, codes[:b], slots, scales[:b]) + torch.cuda.synchronize() + got = pool.view(-1, NATIVE_ROW)[live].clone() + assert _round_trip(codes[:b], scales[:b], slots, container) == live.numel() + _check_native_pool(pool, b * ENGINE_PAGE, slots, rows) + first = got if first is None else first + assert torch.equal(got, first), ( + f"B={b}: {int((got != first).sum())} bytes differ from B={sizes[0]}" + ) + finally: + mgr.shutdown() + del mgr, pool, codes, scales + _free_engine_pool() + buf, codes, scales = _native_engine_pool(small) + op(x, codes, slots, scales) + torch.cuda.synchronize() + got = buf.view(-1, NATIVE_ROW)[live].clone() + assert _round_trip(codes, scales, slots, container) == live.numel() + _check_native_pool(buf, small * ENGINE_PAGE, slots, rows) + assert torch.equal(got, first), "a separate exact-size native pool differs" + print( + f"\n[native t={t}] views B = {sizes} of the manager's pool and a separate [{small}, 128, 528]: " + f"{live.numel()} rows bitwise identical and == recipe, mxfp8_dequantize read-back == bf16 container " + f"at every B, nothing else written" + ) + del buf, codes, scales + _free_engine_pool() + + +def test_native_every_row_count_1_to_32768() -> None: + """Every prefix (x[:r], slots[:r]), r = 1..32768, of the target's native call writes exactly the first r + rows of the 32768-row call, and nothing else. + + A native [320, 128, 528] pool plus a guard page, all bytes sentinel. The + 32768-row call is held to the recipe and the sentinel scan. Then per r: the + call, the named live slots against the recipe's rows, the count of + non-sentinel bytes against 528 per live row, and the named slots restored + to the sentinel, so each call's writes are checked on their own. One sync + at the end. The probe: the sweep compiles no variant beyond the 32768-row + call's. + """ + t = _MAX_ROWS + num_slots = _SWEEP_PAGES * ENGINE_PAGE + buf, codes, scales = _native_engine_pool(_SWEEP_PAGES) + flat = buf.view(-1, NATIVE_ROW) + slots = _sweep_slots(t, num_slots, seed=8549) + x = _window_kv_like(t, seed=8549) + rows = _native_rows(x)[0] + kernel = _mx_op._mxfp8_fake_quantize_kernel + n0 = _jit_variants(kernel) + op(x, codes, slots, scales) + n1 = _jit_variants(kernel) + _check_native_pool(buf, num_slots, slots, rows) + + live = slots >= 0 + live_pos = torch.nonzero(live).squeeze(1) + live_slots = slots[live_pos] + rows_live = rows[live_pos] + cnt = torch.cumsum(live.to(torch.int64), 0).tolist() + bad_rows = torch.zeros((), dtype=torch.int64, device="cuda") + bad_count = torch.zeros((), dtype=torch.int64, device="cuda") + for r in range(1, t + 1): + op(x[:r], codes, slots[:r], scales) + k = cnt[r - 1] + bad_rows += (flat[live_slots[:k]] != rows_live[:k]).sum() + bad_count += ((buf != NATIVE_SENTINEL).sum() - k * NATIVE_ROW).abs() + flat[live_slots[:k]] = NATIVE_SENTINEL + n2 = _jit_variants(kernel) + print( + f"\n[native rows 1..{t}] live rows {cnt[-1]}; Triton variants: {n0} before the {t}-row call, {n1} after, " + f"{n2} after the sweep; byte mismatches {int(bad_rows)}, non-sentinel count error {int(bad_count)}" + ) + assert int(bad_rows) == 0 and int(bad_count) == 0 + assert n1 - n0 <= 1 and n2 == n1, f"row counts compiled new variants: {n0} -> {n1} -> {n2}" + # Control: the same two checks see a read one row off and a write one slot off. + for r in (1, 2, 7, 64): + for name, xa, sa in ( + ("row +1", x[1 : r + 1], live_slots[:r]), + ("slot +1", x[:r], live_slots[1 : r + 1]), + ): + op(xa, codes, sa, scales) + n_bad = int((flat[live_slots[:r]] != _native_rows(x[:r])[0]).sum()) + print(f"[native control {name}] r={r}: {n_bad} of {r * NATIVE_ROW} bytes differ") + assert n_bad > r * NATIVE_ROW // 2 + buf.fill_(NATIVE_SENTINEL) + del buf, codes, scales, flat + _free_engine_pool() + + +# Captured native row counts: every graph batch G of the engine envelope (_ENVELOPE_GSET: every G 1..32 -- +# at max_batch_size s <= 32 padding off captures every G in 1..s -- and the padded list to 256), 1 token per +# request, or 6 with DSpark verify: T = G and 6G, 82 row counts up to 1536. The gates' {1..32} and +# {6, 12, ..., 192} are part of it. +_NATIVE_GRAPH_T = _SWA_GRAPH_T +assert set(MBS_AXIS) | {6 * g for g in MBS_AXIS} <= set(_NATIVE_GRAPH_T) + + +def test_native_cuda_graph_every_graph_batch_on_the_bmax_pool() -> None: + """T = G and 6G for every G of _ENVELOPE_GSET (82 row counts, 1..1536; decode and DSpark verify), native form, + on the envelope B_max native SWA pool (1,648,320 pages, a PP2 stage alone): per T an eager warm-up, + the capture, then 2 + replays on new rows and new far-end slots; every replay writes exactly the recipe into the named slots + and nothing else, reads back through mxfp8_dequantize as the op's bf16 container, and equals the same + call run eagerly.""" + buf, codes, scales = _native_engine_pool(ENVELOPE_NATIVE_SWA_BMAX) + num_slots = ENVELOPE_NATIVE_SWA_BMAX * ENGINE_PAGE + flat = buf.view(-1, NATIVE_ROW) + g = torch.Generator(device="cuda").manual_seed(8550) + hazard = torch.tensor(_native_hazard_slots(num_slots), device="cuda") + last = num_slots - 1 + table = torch.full((2, max(_NATIVE_GRAPH_T)), -1, dtype=torch.int64, device="cuda") + + def far_slots(t): + near = last - 1 - torch.arange(4 * ENGINE_PAGE - 1, device="cuda") + pool_ = torch.unique(torch.cat([hazard[hazard != last], near])) + if pool_.numel() < t - 1: # large graph batches: the rest spread uniformly over the pool + spread = torch.randint(0, last, (4 * t,), device="cuda", generator=g).unique() + pool_ = torch.cat([pool_, spread[~torch.isin(spread, pool_)]]) + pool_ = pool_[torch.randperm(pool_.numel(), device="cuda", generator=g)] + s = torch.cat([torch.tensor([last], device="cuda"), pool_[: t - 1]]).to(torch.int64) + assert torch.unique(s).numel() == t + if t >= 2: + s[1 + int(torch.randint(0, t - 1, (1,), device="cuda", generator=g))] = -1 + return s + + n_round_trip = 0 + for t in _NATIVE_GRAPH_T: + x = torch.empty(t, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda") + slots = table[1, :t] + x.copy_(_window_kv_like(t, seed=8550 + t)) + slots.copy_(far_slots(t)) + op(x, codes, slots, scales) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(x, codes, slots, scales) + buf.fill_(NATIVE_SENTINEL) + for rep in (1, 2): + new_x = _window_kv_like(t, seed=8550 + 1000 * rep + t) + new_slots = far_slots(t) + x.copy_(new_x) + slots.copy_(new_slots) + graph.replay() + torch.cuda.synchronize() + live = new_slots[new_slots >= 0] + replayed = flat[live].clone() + rows, y = _native_rows(new_x) + container = torch.empty_like(new_x) + op(new_x, container) + _assert_bitwise(container, y) + n_round_trip += _round_trip(codes, scales, new_slots, container) + _check_native_pool(buf, num_slots, new_slots, rows) + op(x, codes, slots, scales) + torch.cuda.synchronize() + eager = flat[live].clone() + _check_native_pool(buf, num_slots, new_slots, rows) + assert torch.equal(replayed, eager), ( + f"T={t} replay {rep}: {int((replayed != eager).sum())} bytes" + ) + del graph + print( + f"\n[native graphs on [{ENVELOPE_NATIVE_SWA_BMAX}, 128, 528]] {len(_NATIVE_GRAPH_T)} graphs, T = " + f"{list(_NATIVE_GRAPH_T)}: every replay bitwise == eager == recipe; {n_round_trip} replayed rows read " + f"back through mxfp8_dequantize == the bf16 container" + ) + del buf, codes, scales, flat + _free_engine_pool() + + +# ─── native form over random multi-request steps on the live manager ─── + + +def _random_step( + rng: random.Random, max_tokens: int = 8192, max_requests: int = 32, max_len: int = 16384 +): + """([(cached, new)] contexts first (random order) then generation requests, num_contexts, gen_len). + + k in 1..32 requests over T <= 8192 tokens; contexts get random lengths and a + random cached prefix (40% none); generation requests all carry gen_len + tokens (1 decode, 6 DSpark verify) at random positions. + """ + k = rng.randint(1, max_requests) + gen_len = rng.choice([1, 6]) + g = rng.randint(0, k) if rng.random() < 0.7 else 0 + c = k - g + budget = max_tokens - g * gen_len + t_ctx = rng.randint(c, budget) if c else 0 + cuts = sorted(rng.sample(range(1, t_ctx), c - 1)) if c > 1 else [] + news = [b - a for a, b in zip([0] + cuts, cuts + [t_ctx])] if c else [] + reqs = [(0 if rng.random() < 0.4 else rng.randint(0, max_len - n), n) for n in news] + rng.shuffle(reqs) + reqs += [(rng.randint(0, max_len - gen_len), gen_len) for _ in range(g)] + return reqs, c, gen_len + + +# Hand-picked steps at the envelope's edges, then the seeded random ones. +_EDGE_STEPS = [ + ([(0, 8192)], 1, 1), + ([(0, 1)], 1, 1), + ([(16383, 1)], 1, 1), + ([(0, 256)] * 32, 32, 1), + ([(16384 - 256, 256)] * 32, 32, 1), + ([(16383, 1)] * 32, 0, 1), + ([(16378, 6)] * 32, 0, 6), + ([(513, 7000)] + [(16000, 6)] * 31, 1, 6), + ([(0, 127), (0, 128), (0, 129), (127, 1), (128, 1), (8000, 192)], 6, 1), +] +_NUM_RANDOM_STEPS = 200 + + +def _run_step(mgr, reqs, num_contexts: int, gen_len: int, uid: int, max_num_tokens: int = 8192): + """Allocate the step's requests on the manager and build its DeepSeek-V4.1 attention metadata. + + Each request's first token is unique over the test (uid + request index), so + block reuse never shares a page between requests of different steps. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + DeepseekV4TrtllmAttentionMetadata, + ) + from tensorrt_llm._torch.metadata import KVCacheParams + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from tensorrt_llm.bindings import SamplingConfig + from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig + from tensorrt_llm.mapping import Mapping + + requests = [] + for rid, (c, n) in enumerate(reqs): + r = LlmRequest( + request_id=rid, + max_new_tokens=64, + input_tokens=[uid + rid] + list(range(1, c + n)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + assert mgr.prepare_context(r) and mgr.resize_context(r, r.context_chunk_size), (rid, c, n) + requests.append(r) + md = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor([n for _, n in reqs], dtype=torch.int), + request_ids=[r.py_request_id for r in requests], + max_num_requests=mgr.max_batch_size, + num_contexts=num_contexts, + prompt_lens=[c + n for c, n in reqs], + max_num_tokens=max_num_tokens, + kv_cache_manager=mgr, + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=[c for c, _ in reqs] + ), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + sparse_attention_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ), + ) + if num_contexts < len(reqs) and gen_len > 1: + md.update_spec_dec_param( + batch_size=len(reqs), + is_spec_decoding_enabled=False, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, + max_draft_len=gen_len - 1, + max_total_draft_tokens=gen_len - 1, + ) + md.prepare() + return requests, md + + +def test_native_writes_over_random_multi_request_steps() -> None: + """Random multi-request steps on the live accuracy-gate manager (8192, 32, DSpark 5; native SWA 732,280 pages). + + Per step (9 edge steps + 200 seeded random partitions of T <= 8192 into + k <= 32 requests, contexts with random cached prefixes, decode or DSpark + verify generation requests): the requests are allocated on the manager, + the metadata is prepared, and for layers 0, 1, 20 and 39 the op's native + form writes random rows at md.v41_swa_slot_mapping[local(L)][:T] through + get_v41_native_kv_views(L, SWA): the slots hold the recipe's codes and + bytes, decoding to the op's own bf16-container output of the same rows; + after the step every other byte of the pool is still the sentinel. The + op compiles no variant past its first native call. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + mgr = _v41_native_manager(8192, 32, 5) + rng = random.Random(8551) + steps = list(_EDGE_STEPS) + [_random_step(rng) for _ in range(_NUM_RANDOM_STEPS)] + kernel = _mx_op._mxfp8_fake_quantize_kernel + try: + pool, _, _ = _native_swa(mgr) + num_slots = pool.shape[0] * ENGINE_PAGE + assert pool.shape[0] == NATIVE_SWA_GATE_MAX + pool.fill_(NATIVE_SENTINEL) + flat = pool.view(-1, NATIVE_ROW) + container_buf = torch.empty(8192, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda") + n_variants = None + stats = dict(steps=0, rows=0, max_t=0, max_k=0, max_slot=0) + for i, (reqs, nctx, gen_len) in enumerate(steps): + requests, md = _run_step(mgr, reqs, nctx, gen_len, uid=200_000 + 32 * i) + try: + t = sum(n for _, n in reqs) + for layer in (0, 1, 20, 39): + codes, scales = mgr.get_v41_native_kv_views(layer, AT.SWA) + slots = md.v41_swa_slot_mapping[mgr.layer_offsets[layer]][:t] + x = _fast_rows(t, seed=100_000 + 4 * i + layer) + op(x, codes, slots, scales) + rows, y = _native_rows(x) + # A fixed 8192-slot out: in the identity form the slot count is an integer + # argument, and Triton specializes it (== 1, % 16 == 0) per call. + op(x, container_buf) + container = container_buf[:t] + if n_variants is None: + n_variants = _jit_variants(kernel) + s64 = slots.to(torch.int64) + assert int(s64.min()) >= 0 and int(s64.max()) < num_slots + got = flat[s64] + assert torch.equal(got, rows), ( + f"step {i} layer {layer}: {int((got != rows).sum())} bytes" + ) + _assert_bitwise(_decode_native(got), container) + _assert_bitwise(container, y) + flat[s64] = NATIVE_SENTINEL + stats["max_slot"] = max(stats["max_slot"], int(s64.max())) + n_out = int(_count_not_sentinel(pool)) + assert n_out == 0, f"step {i}: {n_out} bytes outside the step's slots were written" + stats["steps"] += 1 + stats["rows"] += 4 * t + stats["max_t"] = max(stats["max_t"], t) + stats["max_k"] = max(stats["max_k"], len(reqs)) + finally: + for r in requests: + mgr.free_resources(r) + assert _jit_variants(kernel) == n_variants, "a step shape compiled a new variant" + print( + f"\n[native multi-request steps on the live manager, SWA {tuple(pool.shape)}] {stats}; " + f"one compiled variant" + ) + assert stats["steps"] == len(_EDGE_STEPS) + _NUM_RANDOM_STEPS and stats["max_t"] == 8192 + assert stats["max_k"] == 32 + finally: + mgr.shutdown() + del mgr + _free_engine_pool() + + +# The envelope's step axes: T <= max_num_tokens 16384, k <= max_batch_size 256 requests per step. A step's KV +# tokens (cached + new, over its requests) are capped at _ENVELOPE_STEP_TOKENS: the native manager at (16384, 256, +# DSpark 5, max_seq_len 32,782) held 22 requests of 32,768 tokens and refused the 23rd, so 256 requests cannot +# all sit at the far end of max_seq_len at once. +_ENVELOPE_STEP_TOKENS = 560_000 +_ENVELOPE_STEP_L = 32768 + + +def _envelope_random_step(rng: random.Random, max_tokens: int, max_requests: int, seq_len: int): + """_random_step at the envelope's step axes with DSpark: k in 1..max_requests, contexts end at <= seq_len, + decode requests (next_n 1) at <= seq_len, DSpark-verify requests (next_n 6) at <= seq_len + 14.""" + k = rng.randint(1, max_requests) + gen_len = rng.choice([1, 6]) + g = rng.randint(0, k) if rng.random() < 0.7 else 0 + c = k - g + budget = max_tokens - g * gen_len + t_ctx = rng.randint(c, budget) if c else 0 + cuts = sorted(rng.sample(range(1, t_ctx), c - 1)) if c > 1 else [] + news = [b - a for a, b in zip([0] + cuts, cuts + [t_ctx])] if c else [] + reqs = [(0 if rng.random() < 0.4 else rng.randint(0, seq_len - n), n) for n in news] + rng.shuffle(reqs) + end = seq_len + 14 if gen_len == 6 else seq_len + reqs += [(rng.randint(0, end - gen_len), gen_len) for _ in range(g)] + return reqs, c, gen_len + + +def _envelope_fit(step, rng: random.Random): + """The step with its KV tokens capped at _ENVELOPE_STEP_TOKENS: in a random order, a request keeps its cached + prefix while the rest can still fit, else gets a random prefix below 1024. Order and new tokens kept.""" + reqs, nctx, gen_len = step + if sum(a + n for a, n in reqs) <= _ENVELOPE_STEP_TOKENS: + return step + out, used = list(reqs), 0 + reserve = sum(n + 1024 for _, n in reqs) + order = list(range(len(reqs))) + rng.shuffle(order) + for i in order: + a, n = reqs[i] + reserve -= n + 1024 + if used + a + n + reserve > _ENVELOPE_STEP_TOKENS: + a = rng.randint(0, 1023) + out[i] = (a, n) + used += a + n + assert used <= _ENVELOPE_STEP_TOKENS + return out, nctx, gen_len + + +def _envelope_edge_steps(max_t: int, seqs: int, L: int): + """A full max_t step from position 0 and one ending at L; seqs decode requests (8 at L - 1, the rest below + position 1152); seqs DSpark verifies (8 ending at L + 13, the rest below 1130); a context filling the step + beside seqs - 1 verifies (7 at the far end).""" + lm, rest = L + 14, seqs - 8 + steps = [ + ([(0, max_t)], 1, 1), + ([(L - max_t, max_t)], 1, 1), + ([(L - 1, 1)] * 8 + [(127 + (61 * i) % 1024, 1) for i in range(rest)], 0, 1), + ([(lm - 6, 6)] * 8 + [(100 + (37 * i) % 1024, 6) for i in range(rest)], 0, 6), + ( + [(0, max_t - 6 * (seqs - 1))] + + [(lm - 6, 6)] * 7 + + [(100 + (37 * i) % 1024, 6) for i in range(rest)], + 1, + 6, + ), + ] + for reqs, _, _ in steps: + assert sum(a + b for a, b in reqs) <= _ENVELOPE_STEP_TOKENS + return steps + + +def test_envelope_native_writes_over_random_steps_256_requests() -> None: + """Random multi-request steps at the envelope's step axes, on the live native manager of max_num_tokens 16384, + max_batch_size 256, DSpark 5, max_seq_len 32768 + 14: 5 edge steps and 200 seeded random partitions of + T <= 16384 into k <= 256 requests (contexts first with random cached prefixes, decode at positions + < 32,768 or DSpark verify at < 32,782), each step's KV tokens capped (_envelope_fit). Per step, exactly as + test_native_writes_over_random_multi_request_steps: layers 0, 1, 20 and 39 write random rows at + md.v41_swa_slot_mapping[local(L)][:T] through get_v41_native_kv_views(L, SWA); the slots hold the + recipe's codes and bytes, decoding to the op's own bf16-container output, and after the step every other + byte of the pool is still the sentinel; no compiled variant past the first native call.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + max_t, seqs, L = 16384, 256, _ENVELOPE_STEP_L + rng = random.Random(16_384 + seqs + L) + edges = _envelope_edge_steps(max_t, seqs, L) + steps = edges + [ + _envelope_fit(_envelope_random_step(rng, max_t, seqs, L), rng) + for _ in range(_NUM_RANDOM_STEPS) + ] + mgr = _v41_native_manager(max_t, seqs, 5, max_seq_len=L + 14) + kernel = _mx_op._mxfp8_fake_quantize_kernel + try: + assert mgr.max_seq_len == L + 14 and mgr.max_batch_size == seqs + pool, _, _ = _native_swa(mgr) + num_slots = pool.shape[0] * ENGINE_PAGE + pool.fill_(NATIVE_SENTINEL) + flat = pool.view(-1, NATIVE_ROW) + container_buf = torch.empty(max_t, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda") + n_variants = None + stats = dict(steps=0, rows=0, max_t=0, max_k=0, max_pos=0, max_slot=0) + for i, (reqs, nctx, gen_len) in enumerate(steps): + requests, md = _run_step( + mgr, reqs, nctx, gen_len, uid=5_000_000 + 256 * i, max_num_tokens=max_t + ) + try: + t = sum(n for _, n in reqs) + for layer in (0, 1, 20, 39): + codes, scales = mgr.get_v41_native_kv_views(layer, AT.SWA) + slots = md.v41_swa_slot_mapping[mgr.layer_offsets[layer]][:t] + x = _fast_rows(t, seed=500_000 + 4 * i + layer) + op(x, codes, slots, scales) + rows, y = _native_rows(x) + op(x, container_buf) # a fixed 16384-slot out (one identity variant) + container = container_buf[:t] + if n_variants is None: + n_variants = _jit_variants(kernel) + s64 = slots.to(torch.int64) + assert int(s64.min()) >= 0 and int(s64.max()) < num_slots + got = flat[s64] + assert torch.equal(got, rows), ( + f"step {i} layer {layer}: {int((got != rows).sum())} bytes" + ) + _assert_bitwise(_decode_native(got), container) + _assert_bitwise(container, y) + flat[s64] = NATIVE_SENTINEL + stats["max_slot"] = max(stats["max_slot"], int(s64.max())) + n_out = int(_count_not_sentinel(pool)) + assert n_out == 0, f"step {i}: {n_out} bytes outside the step's slots were written" + stats["steps"] += 1 + stats["rows"] += 4 * t + stats["max_t"] = max(stats["max_t"], t) + stats["max_k"] = max(stats["max_k"], len(reqs)) + stats["max_pos"] = max(stats["max_pos"], max(c + n - 1 for c, n in reqs)) + finally: + for r in requests: + mgr.free_resources(r) + assert _jit_variants(kernel) == n_variants, "a step shape compiled a new variant" + print( + f"\n[envelope native steps, ({max_t}, {seqs}, 5) at max_seq_len {L + 14}, SWA {tuple(pool.shape)}] " + f"{len(edges)} edge + {_NUM_RANDOM_STEPS} random steps: {stats}; one compiled variant" + ) + assert stats["steps"] == len(steps) and stats["max_t"] == max_t and stats["max_k"] == seqs + assert stats["max_pos"] == L + 13 + finally: + mgr.shutdown() + del mgr + _free_engine_pool() + + +# The max_batch_size axis: the target's native write on the live manager whose pool is +# B_max over every max_batch_size 1..32, (8192, 1, 0), and on the native-storage witness's (2048, 1, 5). +_MBS1_CONFIGS = (NATIVE_SWA_BMAX_KEY, NATIVE_WITNESS_KEY) +_MBS1_RANDOM_STEPS = 60 + + +def _single_request_steps(max_t: int, draft: int, rng: random.Random, n_random: int): + """Steps of one request (max_batch_size 1) that fit (max_num_tokens, DSpark): a full step from position 0, + one ending at position 16383, 1 new token at 0 and at 16383 (context), a decode at 16383, and with DSpark + a 6-token verify ending at 16397 (the manager's max_seq_len is 16398); then seeded random ones: a context + of 1..max_t new tokens with a random cached prefix (40% none), or a decode (or, with DSpark, a verify) at + a random position.""" + steps = [ + ([(0, max_t)], 1, 1), + ([(16384 - max_t, max_t)], 1, 1), + ([(0, 1)], 1, 1), + ([(16383, 1)], 1, 1), + ([(16383, 1)], 0, 1), + ] + if draft: + steps.append(([(16392, 6)], 0, 6)) + for _ in range(n_random): + if rng.random() < 0.5: + n = rng.randint(1, max_t) + steps.append(([(0 if rng.random() < 0.4 else rng.randint(0, 16384 - n), n)], 1, 1)) + else: + gen_len = 6 if draft and rng.random() < 0.5 else 1 + steps.append( + ([(rng.randint(0, (16398 if draft else 16384) - gen_len), gen_len)], 0, gen_len) + ) + return steps + + +@pytest.mark.parametrize("cfg", _MBS1_CONFIGS, ids=[f"{t}x{s}x{d}" for t, s, d in _MBS1_CONFIGS]) +def test_native_writes_at_max_batch_size_1(cfg) -> None: + """Random steps at max_batch_size 1, on the live manager of the configuration that gives B_max over every + max_batch_size 1..32 ((8192, 1, 0): 941,800 pages) and of the native-storage witness ((2048, 1, 5): 549,600). + + Per step (edge steps reaching T = max_num_tokens and position 16383, 16397 with DSpark, + 60 seeded + random single-request steps): the request is allocated on the manager, the metadata prepared with the + configuration's max_num_tokens, and for layers 0, 1, 20 and 39 the op's native form writes random rows + at md.v41_swa_slot_mapping[local(L)][:T] through get_v41_native_kv_views(L, SWA): the slots hold the + recipe's codes and bytes, decoding in torch and reading back through mxfp8_dequantize to the op's own + bf16-container output; after the step every other byte of the pool is still the sentinel. The op + compiles no variant past its first native call. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + max_t, seqs, draft = cfg + mgr = _v41_native_manager(*cfg) + rng = random.Random(9418 + max_t + draft) + steps = _single_request_steps(max_t, draft, rng, _MBS1_RANDOM_STEPS) + kernel = _mx_op._mxfp8_fake_quantize_kernel + try: + pool, _, _ = _native_swa(mgr) + num_slots = pool.shape[0] * ENGINE_PAGE + assert pool.shape[0] == _NATIVE_SWA_PAGES_AXIS[cfg] and mgr.max_batch_size == seqs == 1 + pool.fill_(NATIVE_SENTINEL) + flat = pool.view(-1, NATIVE_ROW) + container_buf = torch.empty(max_t, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda") + n_variants = None + stats = dict(steps=0, rows=0, round_trip_rows=0, max_t=0, max_pos=0, max_slot=0) + for i, (reqs, nctx, gen_len) in enumerate(steps): + requests, md = _run_step( + mgr, reqs, nctx, gen_len, uid=700_000 + 8 * i, max_num_tokens=max_t + ) + try: + t = sum(n for _, n in reqs) + for layer in (0, 1, 20, 39): + codes, scales = mgr.get_v41_native_kv_views(layer, AT.SWA) + slots = md.v41_swa_slot_mapping[mgr.layer_offsets[layer]][:t] + x = _fast_rows(t, seed=300_000 + 4 * i + layer) + op(x, codes, slots, scales) + rows, y = _native_rows(x) + op(x, container_buf) # a fixed max_t-slot out (one identity variant) + container = container_buf[:t] + if n_variants is None: + n_variants = _jit_variants(kernel) + s64 = slots.to(torch.int64) + assert int(s64.min()) >= 0 and int(s64.max()) < num_slots + got = flat[s64] + assert torch.equal(got, rows), ( + f"step {i} layer {layer}: {int((got != rows).sum())} bytes" + ) + _assert_bitwise(_decode_native(got), container) + _assert_bitwise(container, y) + stats["round_trip_rows"] += _round_trip(codes, scales, slots, container) + flat[s64] = NATIVE_SENTINEL + stats["max_slot"] = max(stats["max_slot"], int(s64.max())) + n_out = int(_count_not_sentinel(pool)) + assert n_out == 0, f"step {i}: {n_out} bytes outside the step's slots were written" + stats["steps"] += 1 + stats["rows"] += 4 * t + stats["max_t"] = max(stats["max_t"], t) + stats["max_pos"] = max(stats["max_pos"], max(c + n - 1 for c, n in reqs)) + finally: + for r in requests: + mgr.free_resources(r) + assert _jit_variants(kernel) == n_variants, "a step shape compiled a new variant" + print( + f"\n[native steps at max_batch_size 1, {cfg}, SWA {tuple(pool.shape)}] {stats}; one compiled variant" + ) + assert stats["steps"] == len(steps) and stats["max_t"] == max_t + assert stats["max_pos"] == (16397 if draft else 16383) + assert stats["round_trip_rows"] == stats["rows"] + finally: + mgr.shutdown() + del mgr + _free_engine_pool() + + +def test_native_variant_map() -> None: + """Which calls compile a new Triton variant: the native form into a manager-geometry pool compiles one + variant whatever T (rows are grid axis 0; the pool's slot count, a multiple of 128, is fixed); the + identity form passes the row count as the slot count, which Triton specializes (== 1, % 16 == 0, + other), so it compiles one variant per class and none for a second member of a class.""" + kernel = _mx_op._mxfp8_fake_quantize_kernel + buf, codes, scales = _native_engine_pool(64) + added = {} + for t in (8192, 1, 2, 16, 17, 4096, 8191): + slots = torch.randperm(64 * ENGINE_PAGE, device="cuda")[:t] + n0 = _jit_variants(kernel) + op(_fast_rows(t, seed=t), codes, slots, scales) + added[("native", t)] = _jit_variants(kernel) - n0 + classes = {"1": (1, 1), "16k": (16, 32, 48, 8192), "other": (2, 17, 33, 8191)} + for name, ts in classes.items(): + for i, t in enumerate(ts): + n0 = _jit_variants(kernel) + op( + _fast_rows(t, seed=t), + torch.empty(t, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda"), + ) + added[(f"identity[{name}]#{i}", t)] = _jit_variants(kernel) - n0 + # The drafter's ring write: stride-3 slot columns into the [771, 128, 512] bf16 ring. + ring = torch.zeros( + _ENVELOPE_RING_PAGES, ENGINE_PAGE, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda" + ) + g = torch.Generator(device="cuda").manual_seed(2201) + for m in (1280, 1, 2, 16, 17, 256, 1275): + cols = _ring_cols(_ring_idx3(m, _ENVELOPE_RING_ROWS, g, masked=False)) + assert all(c.stride(0) == (1 if m == 1 else 3) for c in cols), ( + m, + [c.stride() for c in cols], + ) + n0 = _jit_variants(kernel) + for s in range(3): + op(_fast_rows(m, seed=m + s), ring, cols[s]) + added[("ring stride 3", m)] = _jit_variants(kernel) - n0 + print(f"\n[variant map] new variants per call: {added}") + assert all(v == 0 for (form, t), v in added.items() if form == "native" and t != 8192) + # The 1280-row call compiles the stride-3 pair (stage 1's column is 8 bytes into the table, so not + # 16-byte aligned). M = 1: the one-element column comes out of the reshape with stride 1, the target's + # contiguous pair (already compiled in a whole-file run). So only M >= 2 is held to no new variant. + assert all( + v == 0 for (form, m), v in added.items() if form == "ring stride 3" and m not in (1280, 1) + ) + for name, ts in classes.items(): + assert all(added[(f"identity[{name}]#{i}", t)] == 0 for i, t in enumerate(ts) if i), name + del buf, codes, scales, ring + + +# ─── The DSpark drafter's block K/V write: the identity form ───────────────── +# +# The drafter's attention writes its draft block's K/V with +# mxfp8_fake_quantize(kv, blk_kv): `kv` a contiguous [N, 512] bf16 tensor, +# N = 5G for G <= max_batch_size (256 in the engine envelope), and `blk_kv` = +# split(buf, [N, rest], 0)[0], the first N rows of a [ring_rows * 5 + 128, 512] +# bf16 buffer built once per ring size (ring_rows = max_batch_size + 1). No +# slot_mapping: the identity form, whose slot count is the view's row count N. +# Rows reach the kernel as grid axis 0; the slot count N is a kernel argument +# that Triton specializes by class (== 1, % 16 == 0, other), so the whole row +# axis runs on three compiled variants. +_ID_N_MAX = 8192 +_ID_PAD_ROWS = 128 # the drafter's buffer carries 128 rows past the largest block +_DSPARK_RING_ROWS = 256 + 1 +_DSPARK_BLK_ROWS = _DSPARK_RING_ROWS * 5 + _ID_PAD_ROWS # 1413 +_DSPARK_BLOCK_T = tuple( + 5 * g for g in _ENVELOPE_GSET +) # N = 5G for every G of _ENVELOPE_GSET: 46 blocks, 5..1280 + + +def _sentinel_buffer(rows: int) -> torch.Tensor: + buf = torch.empty(rows, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda") + buf.view(torch.int16).fill_(SENTINEL) + return buf + + +def test_identity_form_every_row_count_1_to_8192() -> None: + """Every row count of the identity form in the drafter's call shape: out = the first T rows of a longer buffer. + + The 8192-row call is bitwise equal to the reference and leaves the 128 rows + past it at the sentinel. Then for every T = 1..8192 the first T rows are + reset to the sentinel and the op runs on (x[:T], buf[:T]): those rows hold + the first T rows of the 8192-row result bit for bit, and every element past + row T - 1 is still the sentinel. After one call per row-count class, the + sweep compiles no Triton variant. Mismatch counts accumulate on the device; + one sync at the end. + """ + t = _ID_N_MAX + x = _window_kv_like(t, seed=5160) + y = _fake_quant(x)[0].view(torch.int16) + buf = _sentinel_buffer(t + _ID_PAD_ROWS) + whole = buf.view(torch.int16) + kernel = _mx_op._mxfp8_fake_quantize_kernel + n0 = _jit_variants(kernel) + op(x, buf[:t]) + assert torch.equal(whole[:t], y), "the 8192-row identity call differs from the reference" + assert bool((whole[t:] == SENTINEL).all()), "the 8192-row identity call wrote past its view" + for r in (1, 2): # the two other row-count classes + op(x[:r], buf[:r]) + n1 = _jit_variants(kernel) + whole.fill_(SENTINEL) + + bad_rows = torch.zeros((), dtype=torch.int64, device="cuda") + bad_tail = torch.zeros((), dtype=torch.int64, device="cuda") + for r in range(1, t + 1): + whole[:r].fill_(SENTINEL) + op(x[:r], buf[:r]) + bad_rows += (whole[:r] != y[:r]).sum() + bad_tail += (whole[r:] != SENTINEL).sum() + n2 = _jit_variants(kernel) + print( + f"\n[identity rows 1..{t}] Triton variants: {n0} before, {n1} after one call per class, {n2} after the " + f"sweep; element mismatches {int(bad_rows)}, elements written past row T - 1: {int(bad_tail)}" + ) + assert int(bad_rows) == 0, f"{int(bad_rows)} elements differ from the 8192-row result's rows" + assert int(bad_tail) == 0, f"{int(bad_tail)} elements past the view were written" + assert n1 - n0 <= 3 and n2 == n1, f"row counts compiled new variants: {n0} -> {n1} -> {n2}" + + +def test_identity_row_sweep_harness_sees_a_shift_or_an_extra_row() -> None: + """Controls for the identity sweep: a read one row off, a write one row off and one row too many are all seen.""" + x = _window_kv_like(66, seed=5161) + y = _fake_quant(x)[0].view(torch.int16) + buf = _sentinel_buffer(66 + _ID_PAD_ROWS) + whole = buf.view(torch.int16) + for r in (1, 2, 7, 63, 64): + cases = ( + ("row +1", x[1 : r + 1], buf[:r]), + ("slot +1", x[:r], buf[1 : r + 1]), + ("one row too many", x[: r + 1], buf[: r + 1]), + ) + for name, xa, oa in cases: + whole.fill_(SENTINEL) + op(xa, oa) + n_bad = int((whole[:r] != y[:r]).sum()) + n_tail = int((whole[r:] != SENTINEL).sum()) + print( + f"[identity control {name}] r={r}: {n_bad} / {r * WINDOW_KV_WIDTH} elements differ, " + f"{n_tail} written past row {r - 1}" + ) + if name != "one row too many": # the rows check sees a shifted read or write + assert n_bad > r * WINDOW_KV_WIDTH // 2 + if name != "row +1": # the tail check sees a write at row r + assert n_tail == WINDOW_KV_WIDTH + + +def test_identity_form_dspark_block_captured_at_every_g() -> None: + """The drafter's block K/V write captured at N = 5G for every G of _ENVELOPE_GSET (46 blocks, 5..1280), + replayed on new rows. + + As the drafter does it at max_batch_size 256: one shared [257 * 5 + 128, 512] buffer, each graph's + `out` the first N rows of it, an eager call at each shape before its + capture, all graphs in one pool, largest first. Three rounds of interleaved + replays, each on new rows written into the captured input with the whole + buffer reset to the sentinel: the replay writes the reference into rows + 0..N-1 bit for bit, equals the same call run eagerly, differs from the + previous round's, and leaves every row from N on at the sentinel. A control + round writes the new rows where the graph does not read them; the replay + then matches only the stale rows. + """ + blk = _sentinel_buffer(_DSPARK_BLK_ROWS) + whole = blk.view(torch.int16) + pool = torch.cuda.graph_pool_handle() + graphs = {} + for n in sorted(_DSPARK_BLOCK_T, reverse=True): + x = torch.empty(n, WINDOW_KV_WIDTH, dtype=torch.bfloat16, device="cuda") + x.copy_(_window_kv_like(n, seed=n)) + out = torch.split(blk, [n, _DSPARK_BLK_ROWS - n], 0)[0] + op(x, out) # the engine's warm-up: one eager call at this shape + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + op(x, out) + graphs[n] = (graph, x, out) + torch.cuda.synchronize() + + order = _DSPARK_BLOCK_T[::2] + _DSPARK_BLOCK_T[1::2][::-1] + previous = {} + for rnd in range(3): + for n in order: + graph, x, out = graphs[n] + new = _window_kv_like(n, seed=10_000 * (rnd + 1) + n) + x.copy_(new) + whole.fill_(SENTINEL) + graph.replay() + torch.cuda.synchronize() + got = whole[:n].clone() + assert torch.equal(got, _fake_quant(new)[0].view(torch.int16)), (n, rnd) + assert bool((whole[n:] == SENTINEL).all()), (n, rnd) + eager = torch.empty_like(x) + op(x, eager) + assert torch.equal(got, eager.view(torch.int16)), (n, rnd) + if n in previous: + assert not torch.equal(got, previous[n]), (n, rnd) + previous[n] = got + + for n in order: + graph, x, out = graphs[n] + elsewhere = _window_kv_like(n, seed=99_999 + n) + whole.fill_(SENTINEL) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(whole[:n], _fake_quant(x)[0].view(torch.int16)), n + assert not torch.equal(whole[:n], _fake_quant(elsewhere)[0].view(torch.int16)), n + assert len(_DSPARK_BLOCK_T) == 46 and max(_DSPARK_BLOCK_T) == 1280 + print( + f"\n[dspark block] {len(_DSPARK_BLOCK_T)} graphs, N {_DSPARK_BLOCK_T}, 3 rounds: every replay bitwise " + f"the reference and eager, rows past N untouched; stale-input control rejected at every N" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_mxfp8_to_e4m3.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_mxfp8_to_e4m3.py new file mode 100644 index 000000000000..55b0a711205f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_mxfp8_to_e4m3.py @@ -0,0 +1,305 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the mxfp8_to_e4m3 catalog entry. + +The entry owns what is certified: `CELLS` is the list (the compact gather at each inv_scale from a +paged and a flat pool, the engine-sized pool into the compact window buffer, the identity form, +dst_slots scatter), `reference` and `compare` are the bitwise gate. This file turns a cell's spec +into tensors and keeps the checks a cell cannot carry: normal-range exactness, saturation, signed +zeros and NaN, CUDA-graph replay, the wrong-variant controls and the op's own rejections. + +Every comparison is bitwise (uint8 codes): the entry's reference decodes with the float8_e4m3fn -> +fp64 cast of the code, torch.ldexp(code, byte - 127) in fp64, the multiply by the fp32 inverse scale, +clamp to +-448, then torch's round-to-nearest-even cast to float8_e4m3fn. The kernel builds the scale +from the IEEE exponent bits and multiplies in fp32, so the two derivations are independent. + +The native pool is DeepSeek-V4.1's SWA pool: 528-byte rows (512 E4M3 codes then 16 UE8M0 bytes) in +128-slot pages, passed as two strided views, with the pool sized from the engine's cache manager (an +engine runtime-value dump: [188160, 128, 528]) in the engine-pool cells and far above any batch +elsewhere. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.quantization.mxfp8_to_e4m3 import ( # noqa: E501 + mxfp8_to_e4m3 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "mxfp8_to_e4m3 requires a CUDA device" + +BLOCK = 32 +E4M3_MAX = 448.0 +AMAX_FLOOR = 1e-4 +W = 512 +NB = W // BLOCK +PAGE = 128 +ROW_BYTES = W + NB # 528 +DEV = "cuda" +POOL_PAGES = 320 # 40960 slots, far above any batch here +E4M3_NAN_CODES = (0x7F, 0xFF) + + +def _gen(seed): + return torch.Generator(device=DEV).manual_seed(seed) + + +def _fake_quant_native(x): + """Checkpoint-recipe MXFP8 of bf16 [T, W]: (fp32 values, uint8 codes, uint8 UE8M0 bytes).""" + t, w = x.shape + xf = x.float().view(t, w // BLOCK, BLOCK) + a = xf.abs().amax(-1).clamp_min(AMAX_FLOOR) + m, e = torch.frexp(a.double() / E4M3_MAX) + e = torch.where(m == 0.5, e - 1, e).to(torch.int32) + s = torch.exp2(e.double()).float().unsqueeze(-1) + q = (xf / s).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) + y = (q.float() * s).view(t, w) + return y, q.view(t, w).view(torch.uint8), (e + 127).to(torch.uint8) + + +def _window_kv_like(t, seed=0, lo=-12, hi=4): + """bf16 [t, 512] rows whose 32-blocks have magnitudes 2^lo .. 2^hi (DeepSeek-V4.1's normed window KV + lies in about 2^-10 .. 2^7), plus all-zero and sub-floor blocks.""" + g = _gen(seed) + mag = torch.randint(lo, hi + 1, (t, NB, 1), device=DEV, generator=g).float() + x = torch.randn(t, NB, BLOCK, device=DEV, generator=g) * torch.exp2(mag) + for r in range(0, t, 5): + x[r, r % NB] = 0.0 + x[r, (r + 1) % NB] = torch.randn(BLOCK, device=DEV, generator=g) * 1e-6 + return x.view(t, W).to(torch.bfloat16) + + +def _ref(codes, scales, inv_scale): + """The entry's reference transcode of every slot of a [N, W] pool -> uint8 E4M3 codes [N, W].""" + out = torch.empty(codes.shape[0], codes.shape[-1], dtype=torch.float8_e4m3fn, device=DEV) + return op.reference(codes, scales, out, inv_scale).view(torch.uint8) + + +def _control(codes, scales, inv_scale, *, code_dtype=torch.float8_e4m3fn, bias=127): + """A deliberately-wrong transcode of uint8 codes [N, W] + UE8M0 bytes [N, W/32] -> uint8 codes.""" + n, w = codes.shape + v = codes.contiguous().view(code_dtype).float().double().view(n, w // BLOCK, BLOCK) + e = (scales.long() - bias).double().unsqueeze(-1) + y = torch.ldexp(v, e).float() * torch.tensor( + inv_scale, dtype=torch.float32, device=codes.device + ) + y = y.clamp(-E4M3_MAX, E4M3_MAX).view(n, w) + return y.to(torch.float8_e4m3fn).view(torch.uint8) + + +def _native_pool(pages, seed=0): + """[pages, 128, 528] uint8 pool + its (codes, scales) views, filled with checkpoint-recipe rows.""" + pool = torch.zeros(pages, PAGE, ROW_BYTES, dtype=torch.uint8, device=DEV) + n = pages * PAGE + x = _window_kv_like(n, seed) + _, codes, sc = _fake_quant_native(x) + flat = pool.view(n, ROW_BYTES) + flat[:, :W] = codes + flat[:, W:] = sc + return pool, pool[..., :W], pool[..., W:] + + +def _assert_bitwise(actual, expected): + bad = (actual.reshape(-1) != expected.reshape(-1)).sum().item() + assert bad == 0, f"{bad} of {actual.numel()} codes differ" + + +# ─── certified cells ────────────────────────────────────────────────────────── + + +def _build(spec, seed: int): + """A cell's spec -> (codes, scales, out, inv_scale, src_slots, dst_slots, extra). + + `extra` carries what the structural checks need: `pool` (whose bytes must not change), `tail` + (rows past `out` in its parent buffer that must keep their fill), `last_page_start` (the + engine pool's last page must be read) and `exactness` (the normal-range check). + """ + g = _gen(seed) + kind, inv_scale = spec["kind"], spec["inv_scale"] + src = dst = None + extra = {} + if kind == "compact": + # The target's call: every query's 128 window columns (with -1 tails) gathered to rows t*128+c. + pool, codes, scales = _native_pool(POOL_PAGES, seed=seed) + if not spec["paged"]: + codes = codes.reshape(-1, W) + scales = scales.reshape(-1, NB) + T = spec["queries"] + src = torch.randint( + 0, POOL_PAGES * PAGE, (T * 128,), device=DEV, generator=g, dtype=torch.int32 + ) + src[torch.randint(0, T * 128, (T * 20,), device=DEV, generator=g)] = -1 + buf = torch.full((T * 128 + 64, W), 0xAB, dtype=torch.uint8, device=DEV).view( + torch.float8_e4m3fn + ) + out = buf[: T * 128] + extra["tail"] = buf[T * 128 :] + elif kind == "engine_pool": + # The engine's SWA pool and the native-storage compact window buffer [max_num_tokens * 128, 512] + # (8192 tokens); rows named from the pool's first and last pages. + pages, T = spec["pool_pages"], spec["queries"] + pool = torch.zeros(pages, PAGE, ROW_BYTES, dtype=torch.uint8, device=DEV) + codes, scales = pool[..., :W], pool[..., W:] + n_fill = min(T * 128, 65536) + fill_rows = torch.cat( + [ + torch.arange(0, n_fill // 2, device=DEV), + torch.arange(pages * PAGE - n_fill // 2, pages * PAGE, device=DEV), + ] + ) + x = _window_kv_like(fill_rows.numel(), seed=seed) + _, c, s = _fake_quant_native(x) + flat = pool.view(-1, ROW_BYTES) + flat[fill_rows, :W] = c + flat[fill_rows, W:] = s + src = fill_rows[ + torch.randint(0, fill_rows.numel(), (T * 128,), device=DEV, generator=g) + ].to(torch.int32) + src[::13] = -1 + buf = torch.zeros(8192 * 128, W, dtype=torch.float8_e4m3fn, device=DEV) + out = buf[: T * 128] + extra["last_page_start"] = pages * PAGE - n_fill // 2 + elif kind == "identity": + pool, codes, scales = _native_pool(spec["pool_pages"], seed=seed) + out = torch.empty(spec["pool_pages"] * PAGE, W, dtype=torch.float8_e4m3fn, device=DEV) + extra["exactness"] = True + else: + assert kind == "scatter", kind + pool, codes, scales = _native_pool(8, seed=seed) + rows = 500 + src = torch.randint(0, 8 * PAGE, (rows,), device=DEV, generator=g, dtype=torch.int64) + dst = torch.randperm(1024, device=DEV, generator=g)[:rows] + dst[::7] = -1 + dst[1::11] = 5000 + out = torch.full((1024, W), 0x11, dtype=torch.uint8, device=DEV).view(torch.float8_e4m3fn) + extra["pool"] = pool + return codes, scales, out, inv_scale, src, dst, extra + + +def _check_normal_range_exact(codes, scales, out, inv_scale) -> None: + """With a power-of-two inv_scale, every value whose scaled magnitude is in [2^-6, 448] is reproduced + exactly: out * (1 / inv_scale) == the MXFP8 value. Only E4M3-subnormal results round.""" + n = out.shape[0] + flat_codes, flat_scales = codes.reshape(n, W), scales.reshape(n, NB) + # 2^(b - 127) built from the fp32 exponent bits (exact; torch.ldexp on the GPU goes through pow and + # can be one ulp off in fp64), times the code in fp32: exact, at most 4 significant bits. + p2 = (flat_scales.to(torch.int32) << 23).view(torch.float32) + v = ( + flat_codes.contiguous().view(torch.float8_e4m3fn).float().view(n, NB, BLOCK) + * p2.unsqueeze(-1) + ).view(n, W) + sv = (v * inv_scale).abs() + normal = (sv >= 2.0**-6) & (sv <= E4M3_MAX) + back = out.float() / inv_scale + frac = normal.float().mean().item() + n_bad = int((back[normal] != v[normal]).sum()) + print(f"inv_scale {inv_scale}: {frac:.3f} of values in the normal range, {n_bad} not exact") + assert frac > 0.2 + assert n_bad == 0, "a normal-range value was not reproduced exactly" + sub = (sv > 0) & (sv < 2.0**-6) + if sub.any(): + err = (back[sub] - v[sub]).abs().max().item() + assert err <= 2.0**-10 / inv_scale, err # half an E4M3 subnormal step + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + codes, scales, out, inv_scale, src, dst, extra = _build( + cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + ) + before = out.view(torch.uint8).clone().view(torch.float8_e4m3fn) + tail = None if "tail" not in extra else extra["tail"].view(torch.uint8).clone() + pool = extra["pool"] + pool_before = pool.clone() if pool.numel() <= 1 << 30 else None + result = op(codes, scales, out, inv_scale, src, dst) + assert result is out, "the op writes out in place and returns it" + op.compare(out, op.reference(codes, scales, before, inv_scale, src, dst)) + if pool_before is not None: + assert torch.equal(pool, pool_before), "the source pool must not be written" + if tail is not None: + assert torch.equal(extra["tail"].view(torch.uint8), tail), ( + "wrote past the rows it was given" + ) + if "last_page_start" in extra: + assert int(src.max()) >= extra["last_page_start"] # the pool's last page is read + if extra.get("exactness"): + _check_normal_range_exact(codes, scales, out, inv_scale) + + +# ─── special encodings, graphs, controls, rejections ────────────────────────── + + +def test_saturates_and_keeps_signed_zero_and_nan(): + codes = torch.tensor( + [[0x7E, 0xFE, 0x00, 0x80, 0x7F] + [0x38] * (W - 5)], dtype=torch.uint8, device=DEV + ) + scales = torch.full((1, NB), 127 + 4, dtype=torch.uint8, device=DEV) # x16: 448*16 saturates + out = torch.empty(1, W, dtype=torch.float8_e4m3fn, device=DEV) + op(codes, scales, out, 1.0) + o = out.view(torch.uint8)[0] + assert o[0].item() == 0x7E and o[1].item() == 0xFE, "448 * 16 must saturate to +-448" + assert o[2].item() == 0x00 and o[3].item() == 0x80, "signed zeros" + assert o[4].item() in E4M3_NAN_CODES, "an e4m3 NaN code stays NaN" + scales[0, 0] = 255 + op(codes, scales, out, 1.0) + assert all(out.view(torch.uint8)[0, i].item() in E4M3_NAN_CODES for i in range(32)), ( + "byte 255 -> NaN block" + ) + + +def test_cuda_graph_capture_replay(): + pool, codes, scales = _native_pool(32, seed=8) + T = 16 + src = torch.zeros(T * 128, dtype=torch.int32, device=DEV) + out = torch.zeros(T * 128, W, dtype=torch.float8_e4m3fn, device=DEV) + op(codes, scales, out, 32.0, src) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(codes, scales, out, 32.0, src) + flat = pool.view(-1, ROW_BYTES) + for seed in range(3): + g = _gen(100 + seed) + src.copy_( + torch.randint(-1, 32 * PAGE, (T * 128,), device=DEV, generator=g, dtype=torch.int32) + ) + graph.replay() + exp = _ref(flat[src.clamp_min(0).long(), :W], flat[src.clamp_min(0).long(), W:], 32.0) + exp[src < 0] = 0 + _assert_bitwise(out.view(torch.uint8), exp) + op.compare(out, op.reference(codes, scales, torch.zeros_like(out), 32.0, src)) + + +def test_wrong_variants_are_detected(): + """Controls: the bitwise gate sees a wrong scale bias, a scale not applied, and a wrong code format.""" + pool, codes, scales = _native_pool(4, seed=9) + out = torch.empty(4 * PAGE, W, dtype=torch.float8_e4m3fn, device=DEV) + op(codes, scales, out, 16.0) + op.compare(out, op.reference(codes, scales, torch.empty_like(out), 16.0)) + flat = pool.view(-1, ROW_BYTES) + got = out.view(torch.uint8) + for name, exp in ( + ("bias 126", _control(flat[:, :W], flat[:, W:], 16.0, bias=126)), + ("inv_scale ignored", _control(flat[:, :W], flat[:, W:], 1.0)), + ( + "codes read as e5m2", + _control(flat[:, :W], flat[:, W:], 16.0, code_dtype=torch.float8_e5m2), + ), + ): + frac = (got != exp).float().mean().item() + assert frac > 0.3, f"control '{name}' differs on only {frac:.3f} of codes" + + +def test_rejects_bad_inputs(): + pool, codes, scales = _native_pool(1, seed=10) + bf = torch.empty(PAGE, W, dtype=torch.bfloat16, device=DEV) + with pytest.raises(ValueError, match="float8_e4m3fn"): + op(codes, scales, bf, 1.0) + out = torch.empty(PAGE, W, dtype=torch.float8_e4m3fn, device=DEV) + with pytest.raises(ValueError, match="inv_scale"): + op(codes, scales, out, 0.0) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_nvfp4_fake_quantize.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_nvfp4_fake_quantize.py new file mode 100644 index 000000000000..340d44148038 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_nvfp4_fake_quantize.py @@ -0,0 +1,2690 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the nvfp4_fake_quantize entry. + +The reference (`op.reference`) is the model's `fp4_act_quant(x, 16, inplace, scale_dtype=float8_e4m3fn)` +recipe written out in native torch ops: amax floor 6*2^-9, scale = e4m3(amax / 6) (IEEE division, +round-to-nearest-even, saturating at 448), y = x / scale (IEEE division), E2M1 round-to-nearest-even with +the sign of y kept, dequantized value E2M1 * scale. Every gate is bit-exact (bf16 bit patterns, packed +bytes, scale bytes). + +`op.CELLS` holds the dense, layout and small-pool calls. This file keeps what a cell cannot carry: the +tie, floor and saturation properties with their deliberately-wrong controls, in place on the input, +CUDA-graph replay, the op's own rejections, the engine-size pools (far end at B_max, view-size +invariance, every row count 1..16384, every graph batch), the cache managers that size those pools, +live multi-request steps and the compiled-variant map. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.quantization.nvfp4_fake_quantize import ( # noqa: E501 + nvfp4_fake_quantize as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "nvfp4_fake_quantize requires a CUDA device" + +DEV = "cuda" +W = 512 # DeepSeek-V4.1 compressed-KV latent width (head_dim) +BLOCK = 16 +FLOOR = 6 * 2.0**-9 # reference amax floor for E4M3-scaled FP4 (kernel.py fp4_quant_kernel) +E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) # magnitude of code 0..7; bit 3 is the sign +MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +NUM_SLOTS = 4096 # pool rows: plainly larger than any batch below +_IDX = {"int32": torch.int32, "int64": torch.int64} + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _ref_rows(x: torch.Tensor): + """(dequantized bf16 [rows, W], packed uint8 [rows, W/2], scale bytes uint8 [rows, W/16]) of every row + of `x`: `op.reference` of the slot-r call into fresh dense buffers.""" + n, w = x.reshape(-1, x.shape[-1]).shape + deq = op.reference(x, torch.empty(n, w, dtype=torch.bfloat16, device=x.device)) + packed, s_bytes = op.reference( + x, + torch.empty(n, w // 2, dtype=torch.uint8, device=x.device), + None, + torch.empty(n, w // BLOCK, dtype=torch.uint8, device=x.device), + ) + return deq, packed, s_bytes + + +def _control( + x: torch.Tensor, *, floor=FLOOR, pow2_scale=False, ties_toward_zero=False, reciprocal=False +) -> torch.Tensor: + """Dequantized bf16 values of a deliberately-wrong variant of the recipe. + + The keyword variants: a power-of-two (UE8M0-style) scale, no amax floor (a zero scale reads back as + zero), exact midpoints rounded toward zero, and x * (1/scale) instead of x / scale. + """ + xf = x.float().unflatten(-1, (-1, BLOCK)) + amax = xf.abs().amax(-1, keepdim=True) + if floor is not None: + amax = amax.clamp_min(floor) + if pow2_scale: + u = (amax / 6.0).contiguous().view(torch.int32) + e = ((u >> 23) & 0xFF) + ((u & 0x7FFFFF) != 0).to(torch.int32) + s = (e << 23).view(torch.float32) + else: + s = (amax / 6.0).clamp(max=448.0).to(torch.float8_e4m3fn).float() + y = xf * (1.0 / s) if reciprocal else xf / s + y = torch.where(s == 0, xf * 0.0, y).clamp(-6.0, 6.0) + a = y.abs() + if ties_toward_zero: + idx = (a.unsqueeze(-1) > torch.tensor(MIDPOINTS, device=x.device)).sum(-1) + r = torch.tensor(E2M1, device=x.device)[idx] + else: # torch.round is round-half-to-even + r = torch.where( + a < 2, + torch.round(a * 2) / 2, + torch.where(a < 4, torch.round(a), torch.round(a / 2) * 2), + ) + return (torch.copysign(r, y) * s).flatten(-2).to(torch.bfloat16) + + +def _decode(packed: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Native pair -> bf16 via an E2M1 lookup table (independent of the op).""" + lut = torch.tensor(E2M1 + tuple(-v for v in E2M1), dtype=torch.float32, device=packed.device) + v = torch.stack([lut[(packed & 0xF).long()], lut[(packed >> 4).long()]], -1).flatten(-2) + s = scale.view(torch.float8_e4m3fn).float().repeat_interleave(BLOCK, -1) + return (v * s).to(torch.bfloat16) + + +def _run_core(x): + out = torch.full(x.shape, float("nan"), dtype=torch.bfloat16, device=DEV) + return op(x, out) + + +def _run_native(x): + n = x.shape[0] + pk = torch.full((n, x.shape[-1] // 2), 0xA5, dtype=torch.uint8, device=DEV) + sc = torch.full((n, x.shape[-1] // BLOCK), 0xA5, dtype=torch.uint8, device=DEV) + return op(x, pk, None, sc) + + +def _assert_bitexact(x): + """Both forms of the slot-r call equal the reference bitwise; the native pair decodes to the core output.""" + deq, packed, s_bytes = _ref_rows(x) + core = _run_core(x) + op.compare(core, deq) + pk, sc = _run_native(x) + op.compare((pk, sc), (packed, s_bytes)) + assert torch.equal(_bits(_decode(pk, sc)), _bits(core)), ( + "native pair does not decode to the core output" + ) + return core + + +def _random_blocks(t, width=W, lo=-16, hi=6, seed=0): + """bf16 [t, width], each 16-block scaled by 2^k, k uniform in [lo, hi] (the low end is below the floor).""" + g = torch.Generator(device=DEV).manual_seed(seed) + mag = torch.randint(lo, hi + 1, (t, width // BLOCK, 1), device=DEV, generator=g).float() + x = torch.randn(t, width // BLOCK, BLOCK, device=DEV, generator=g) * torch.exp2(mag) + return x.view(t, width).to(torch.bfloat16) + + +def _tie_blocks(): + """For every positive finite E4M3 scale s: a block [6s, +-t*s for t in MIDPOINTS, 0.1s], plus the same + blocks one bf16 ulp above and below every midpoint entry. 6s and t*s are exact in bf16, so the scale + comes out as exactly s and y = t is an exact midpoint.""" + s = torch.arange(1, 127, dtype=torch.uint8).view(torch.float8_e4m3fn).float() # 2^-9 .. 448 + pattern = torch.tensor([6.0, *MIDPOINTS, *(-m for m in MIDPOINTS), 0.1]) + ties = (s[:, None] * pattern[None, :]).to(torch.bfloat16) # [126, 16] + b = ties.view(torch.int16) + up, dn = b.clone(), b.clone() + up[:, 1:15] += 1 + dn[:, 1:15] -= 1 + blocks = torch.cat([ties, up.view(torch.bfloat16), dn.view(torch.bfloat16)], 0) # [378, 16] + pad = (-blocks.shape[0]) % (W // BLOCK) + blocks = torch.cat([blocks, torch.zeros(pad, BLOCK, dtype=torch.bfloat16)], 0) + return blocks.view(-1, W).to(DEV) + + +def _floor_blocks(): + """Blocks with amax below, at and above the floor, all-zero, all -0.0, tiny negatives.""" + g = torch.Generator().manual_seed(7) + rows = [] + for a in ( + 0.0, + 2.0**-40, + 2.0**-20, + 6 * 2.0**-12, + 6 * 2.0**-11, + 6 * 2.0**-10, + 9 * 2.0**-10, + 6 * 2.0**-9 * 0.75, + 6 * 2.0**-9 * 0.99, + 6 * 2.0**-9, + 6 * 2.0**-9 * 1.01, + 6 * 2.0**-8, + ): + blk = (torch.rand(32, BLOCK, generator=g) * 2 - 1) * a + blk[:, 0] = a + rows.append(blk) + rows.append(torch.zeros(32, BLOCK)) + rows.append(-torch.zeros(32, BLOCK)) + rows.append(-torch.rand(32, BLOCK, generator=g) * 2.0**-14) + mixed = torch.rand(32, BLOCK, generator=g) * 0.01 - 0.005 + mixed[:, 0] = 1.0 # scale 1/6 -> 0.25*s = 0.0417: every other element rounds to +-0 + rows.append(mixed) + return torch.cat(rows, 0).to(torch.bfloat16).view(-1, W).to(DEV) + + +def _every_bf16_amax(): + """One block per positive finite bf16 value a (0x0001..0x7F7F), block = a * fixed pattern (amax = a).""" + a = torch.arange(1, 0x7F80, dtype=torch.int32).to(torch.int16).view(torch.bfloat16).float() + pattern = torch.tensor( + [ + 1.0, + -0.3, + 0.77, + -0.51, + 0.13, + 0.9, + -0.66, + 0.2, + -0.99, + 0.45, + 0.06, + -0.12, + 0.35, + -0.8, + 0.58, + -0.25, + ] + ) + blocks = (a[:, None] * pattern[None, :]).to(torch.bfloat16) + pad = (-blocks.shape[0]) % (W // BLOCK) + blocks = torch.cat([blocks, blocks[:pad]], 0) + return blocks.view(-1, W).to(DEV), a.to(DEV) + + +def _slot_map(t, idx_dtype, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + slots = torch.randperm(NUM_SLOTS, device=DEV, generator=g)[:t].to(torch.int64) + slots[3] = -1 + slots[7] = NUM_SLOTS + 5 + slots[11] = -(2**40) if idx_dtype == torch.int64 else -(2**30) + return slots.to(idx_dtype) + + +def _build_x(spec): + kind = spec["kind"] + if kind == "random": + width = spec["width"] + return _random_blocks(spec["t"], width=width, seed=spec["t"] if width == W else width) + if kind == "ties": + return _tie_blocks() + if kind == "floor": + return _floor_blocks() + if kind == "every_amax": + return _every_bf16_amax()[0] + if kind == "leading_dims": + return _random_blocks(96, seed=3).view(96, 1, W) + if kind == "row_stride": + wide = torch.randn(96, W + 128, device=DEV).to(torch.bfloat16) + wide[:, :W] = _random_blocks(96, seed=3) + return wide[:, :W] + if kind == "bf16_pool": + return _random_blocks(37, seed=5) + if kind == "native_pool": + return _random_blocks(37, seed=9) + assert kind == "zero_rows", kind + return torch.empty(0, W, dtype=torch.bfloat16, device=DEV) + + +def _build(spec): + """(x, out, slot_mapping, out_scale) for a cell: sentinel-filled buffers, so `compare` also proves + every element the op must not write was left alone.""" + kind, layout = spec["kind"], spec.get("layout") + x = _build_x(spec) + rows, width = x.reshape(-1, x.shape[-1]).shape + if kind == "bf16_pool": + shape = { + "flat": (NUM_SLOTS, W), + "paged64": (NUM_SLOTS // 64, 64, W), + "paged128": (NUM_SLOTS // 128, 128, W), + }[layout] + out = torch.full(shape, -7.0, dtype=torch.bfloat16, device=DEV) + return x, out, _slot_map(rows, _IDX[spec["idx_dtype"]], seed=len(layout)), None + if kind == "native_pool": + slots = _slot_map(rows, _IDX[spec["idx_dtype"]], seed=13) + if layout == "split": + data = torch.full((NUM_SLOTS, W // 2), 0x5A, dtype=torch.uint8, device=DEV) + scale = torch.full((NUM_SLOTS, W // BLOCK), 0x5A, dtype=torch.uint8, device=DEV) + return x, data, slots, scale.view(torch.float8_e4m3fn) + pages = (NUM_SLOTS,) if layout == "row288" else (NUM_SLOTS // 128, 128) + buf = torch.full(pages + (W // 2 + W // BLOCK,), 0x5A, dtype=torch.uint8, device=DEV) + assert buf.view(NUM_SLOTS, -1).stride(0) == 288 + return x, buf[..., : W // 2], slots, buf[..., W // 2 :] + if kind == "zero_rows": + out = torch.full((8, W), 3.0, dtype=torch.bfloat16, device=DEV) + return x, out, torch.empty(0, dtype=torch.int32, device=DEV), None + if spec["form"] == "core": + return ( + x, + torch.full((rows, width), float("nan"), dtype=torch.bfloat16, device=DEV), + None, + None, + ) + pk = torch.full((rows, width // 2), 0xA5, dtype=torch.uint8, device=DEV) + sc = torch.full((rows, width // BLOCK), 0xA5, dtype=torch.uint8, device=DEV) + return x, pk, None, sc + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + x, out, slots, out_scale = _build(cell.spec) + expected = op.reference(x, out, slots, out_scale) + result = op(x, out, slots, out_scale) + if out_scale is None: + assert result is out, "the core form writes in place and returns out" + else: + assert result[0] is out and result[1] is out_scale, ( + "the native form returns (out, out_scale)" + ) + op.compare(result, expected) + rows = x.reshape(-1, x.shape[-1]).shape[0] + if slots is None and out_scale is not None: + assert torch.equal(_bits(_decode(out, out_scale)), _bits(_ref_rows(x)[0])), ( + "native pair does not decode to the core values" + ) + if slots is not None and rows: + s64 = slots.to(torch.int64) + assert int(((s64 >= 0) & (s64 < NUM_SLOTS)).sum()) == rows - 3, ( + "three slots are out of range" + ) + + +def test_exact_ties_round_to_even(): + x = _tie_blocks() + _, _, s_bytes = _ref_rows(x) + # the construction really produces exact midpoints (126 scales x 14 per block) + s = s_bytes.view(torch.float8_e4m3fn).float().repeat_interleave(BLOCK, -1) + y = (x.float() / s).abs() + n_ties = int((y.unsqueeze(-1) == torch.tensor(MIDPOINTS, device=DEV)).any(-1).sum()) + assert n_ties == 126 * 14, n_ties + core = _assert_bitexact(x) + # up-ties (0.75, 1.75, 3.5) go up, down-ties (0.25, 1.25, 2.5, 5) go down + tie_rows = core.float().view(-1, BLOCK)[:126] + s_rows = s.view(-1, BLOCK)[:126, :1] + got = (tie_rows[:, 1:8] / s_rows).abs() + want = torch.tensor([0.0, 1.0, 1.0, 2.0, 2.0, 4.0, 4.0], device=DEV).expand_as(got) + assert torch.equal(got, want) + # controls: ties toward zero, and a reciprocal multiply, both miss the reference here + tz = _control(x, ties_toward_zero=True) + rc = _control(x, reciprocal=True) + n_tz = int((_bits(tz) != _bits(core)).sum()) + n_rc = int((_bits(rc) != _bits(core)).sum()) + print( + f"ties: {n_ties} exact midpoints; control ties-toward-zero differs on {n_tz}, reciprocal on {n_rc}" + ) + assert n_tz == 126 * 6 # +-0.75, +-1.75, +-3.5 per scale + assert n_rc > 0 + + +def test_floor_region(): + x = _floor_blocks() + core = _assert_bitexact(x) + _, sc = _run_native(x) + amax = x.float().view(-1, BLOCK).abs().amax(-1).view(sc.shape) + below = amax <= FLOOR + # every block at or below the floor gets the floor scale 6*2^-9 / 6 = 2^-9 (E4M3 byte 0x01) + assert bool((sc[below] == 0x01).all()) + # -0.0 inputs and negatives that round to zero keep the sign (code 0x8, value -0.0) + neg_zero = _bits(core) == -32768 + assert int(neg_zero.sum()) > 0 + assert bool((torch.signbit(x) == torch.signbit(core))[core == 0].all()) + nf = _control(x, floor=None) + n_nf = int((nf.float() != core.float()).sum()) + print( + f"floor: {int(below.sum())} blocks at/below the floor; control no-floor differs on {n_nf} values" + ) + assert n_nf > 0 + + +def test_every_bf16_amax(): + x, a = _every_bf16_amax() + core = _assert_bitexact(x) + _, sc = _run_native(x) + sc = sc.view(-1)[: a.numel()] + # the scale saturates at 448 (byte 0x7E) once amax / 6 rounds past it, and never becomes NaN (0x7F) + assert not bool((sc == 0x7F).any()) + assert bool((sc[a >= 2688.0] == 0x7E).all()) + p2 = _control(x, pow2_scale=True) + n_p2 = int((p2.float() != core.float()).sum()) + print( + f"every bf16 amax: {a.numel()} blocks, {int((sc == 0x7E).sum())} saturated; " + f"control power-of-two scale differs on {n_p2} values" + ) + assert n_p2 > 0 + + +def test_controls_fail_on_random_data(): + x = torch.cat([_random_blocks(1024, seed=11), _tie_blocks(), _floor_blocks()], 0) + core = _assert_bitexact(x) + total = core.numel() + assert torch.equal(_bits(_control(x)), _bits(core)), ( + "the control recipe without a variant is the op" + ) + for name, kw in [ + ("power-of-two scale", dict(pow2_scale=True)), + ("no floor", dict(floor=None)), + ("ties toward zero", dict(ties_toward_zero=True)), + ("reciprocal", dict(reciprocal=True)), + ]: + wrong = _control(x, **kw) + n = int((_bits(wrong) != _bits(core)).sum()) + print(f"control {name}: {n} of {total} elements differ") + assert n > 0, name + + +def test_in_place_on_input(): + """`out` may be `x` itself in the slot-r form: each program reads its whole row before writing it.""" + x = _random_blocks(256, seed=17) + y = x.clone() + expected = op.reference(y, y) + assert op(y, y) is y + op.compare(y, expected) + + +def test_cuda_graph_replay(): + t = 64 + x = torch.empty(t, W, dtype=torch.bfloat16, device=DEV) + pool = torch.zeros(NUM_SLOTS, W, dtype=torch.bfloat16, device=DEV) + slots = torch.empty(t, dtype=torch.int32, device=DEV) + x.copy_(_random_blocks(t, seed=1)) + slots.copy_(_slot_map(t, torch.int32, seed=1)) + op(x, pool, slots) # compile outside capture + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + op(x, pool, slots) + for seed in (2, 3): + x.copy_(_random_blocks(t, seed=seed)) + slots.copy_(_slot_map(t, torch.int32, seed=seed)) + g.replay() + ref, _, _ = _ref_rows(x) + s64 = slots.to(torch.int64) + valid = (s64 >= 0) & (s64 < NUM_SLOTS) + assert torch.equal(_bits(pool[s64[valid]]), _bits(ref[valid])) + + +@pytest.mark.parametrize( + "case, match", + [ + ("fp32_input", "x must be bfloat16"), + ("width_not_16", "positive multiple of 16"), + ("strided_last_dim", "unit stride along its last dim"), + ("native_without_scale", "needs out_scale"), + ("scale_wrong_shape", "out_scale must be"), + ("bf16_out_with_scale", "out_scale is only written"), + ("out_wrong_width", "out's last dim must be"), + ("slot_wrong_length", "slot_mapping must be"), + ("slot_float", "slot_mapping must be int32 or int64"), + ("too_few_slots", "without slot_mapping out needs"), + ("fp16_out", "out must be bfloat16 or uint8"), + ], +) +def test_rejections(case, match): + x = _random_blocks(8, seed=0) + out = torch.empty(8, W, dtype=torch.bfloat16, device=DEV) + pk = torch.empty(8, W // 2, dtype=torch.uint8, device=DEV) + sc = torch.empty(8, W // BLOCK, dtype=torch.uint8, device=DEV) + args = { + "fp32_input": (x.float(), out), + "width_not_16": (x[:, :24], torch.empty(8, 24, dtype=torch.bfloat16, device=DEV)), + "strided_last_dim": (torch.empty(W, 8, dtype=torch.bfloat16, device=DEV).t(), out), + "native_without_scale": (x, pk), + "scale_wrong_shape": (x, pk, None, sc[:, :16]), + "bf16_out_with_scale": (x, out, None, sc), + "out_wrong_width": (x, pk.view(8, -1)[:, :128]), + "slot_wrong_length": (x, out, torch.zeros(7, dtype=torch.int32, device=DEV)), + "slot_float": (x, out, torch.zeros(8, dtype=torch.float32, device=DEV)), + "too_few_slots": (x, out[:4]), + "fp16_out": (x, out.half()), + }[case] + with pytest.raises(ValueError, match=match): + op(*args) + + +# ─── the largest COMPRESS views an engine can allocate ──────────────────────── + +# The per-owner COMPRESS views (ratio 2 [B, 64, 512] for owners 2 / 8 / 14, +# ratio 1 [B, 128, 512] for owner 20, bf16) are sized by the V4.1 +# DeepseekV4CacheManager from the per-rank KV quota, which an engine derives as +# kv_cache_fraction (0.5 at the gates) x the device's FREE memory after weights +# and activations. Free memory never exceeds total memory, so a manager built at +# 0.5 x the device's TOTAL memory (148,510,474,240 bytes on this GB300) +# allocates the largest views an engine of that configuration can have. At the +# gate settings (max_seq_len 16384, max_num_tokens 8192, max_batch_size {8, 32}, +# DSpark or not; test_largest_engine_views_from_the_cache_manager) the largest +# are ratio 2 [862719, 64, 512] and ratio 1 [287573, 128, 512] (GATES_COMPRESS). +# Over the engine envelope (max_num_tokens 2048..16384, max_batch_size up +# to 256, max_seq_len up to 1,048,576, DSpark or not) the unsplit manager's +# largest are ratio 2 [1270272, 64, 512] and ratio 1 [423424, 128, 512], both at +# (max_num_tokens 2048, max_batch_size 48, no DSpark, max_seq_len 1,048,576) +# (UNSPLIT_COMPRESS, the pool scan; +# test_envelope_largest_compress_views_from_the_cache_manager re-measures them). A PP2 +# stage manager holds only its 20 layers' pools, so its views are larger: stage 0 +# alone (the ratio-2 sources) reaches [2118144, 64, 512] (135,561,216 slots, +# 6.9e10 elements, 138.8 GB) and stage 1 alone (the ratio-1 source) [1059584, +# 128, 512] (135,626,752 slots, 6.9e10 elements, 138.9 GB), both at (2048, 96, 0, +# 1,048,576) (the PP2 stage scan; +# test_envelope_pp2_stage_compress_views_from_the_cache_manager re-measures them). +# Those stage-alone maxima are ENGINE_COMPRESS. test_pool_size_invariance shows +# the op's result does not depend on B, so the far-end cases below run at that +# B_max. Past slot 4,194,304 every element offset +# exceeds 2^31, past 8,388,608 it exceeds 2^32 and past 67,108,864 it exceeds +# 2^35, so any 32-bit offset in the kernel's address arithmetic would land a +# row elsewhere (or fault). Each pool is allocated at the exact shape with one +# guard page past its end, every element set to a NaN bit pattern the op never +# writes for a finite row; after the call the named slots must hold the +# reference and every other element must still hold it. +GATES_COMPRESS = {"ratio2": (862719, 64), "ratio1": (287573, 128)} +UNSPLIT_COMPRESS = {"ratio2": (1270272, 64), "ratio1": (423424, 128)} +ENGINE_COMPRESS = {"ratio2": (2118144, 64), "ratio1": (1059584, 128)} +SENTINEL = 0x7FA5 # bf16 NaN bit pattern (as int16) + + +def _hazard_slots(num_slots, width, page_size): + """Distinct slots on both sides of every 32-bit offset hazard of a [slots, width] bf16 pool, and both ends. + + Element offsets 2^30 (2^31 bytes), 2^31 (2^32 bytes), 2^32, 2^33, 2^34, 2^35 and 2^36 (where inside the + pool), the first page, the last page and the last slot. + """ + cands = {0, 1, page_size - 1, page_size} + for limit in (2**30, 2**31, 2**32, 2**33, 2**34, 2**35, 2**36): + s = limit // width + cands |= {s - 1, s} + cands |= {num_slots - page_size - 1, num_slots - page_size, num_slots - 2, num_slots - 1} + return sorted(c for c in cands if 0 <= c < num_slots) + + +def _engine_pool(pages, page_size): + """(buffer with one guard page, [pages, page_size, W] bf16 pool view), every element SENTINEL.""" + buf = torch.empty(pages + 1, page_size, W, dtype=torch.bfloat16, device=DEV) + buf.view(torch.int16).fill_(SENTINEL) + return buf, buf[:pages] + + +def _check_engine_pool(buf, pool, slots, ref): + """Named slots hold ref bitwise; every other element of the pool and the guard page is still SENTINEL. + + Restores the named slots to SENTINEL, so the pool can be reused. + """ + num_slots = pool.shape[0] * pool.shape[1] + flat = pool.view(num_slots, W) + s64 = slots.to(torch.int64) + valid = (s64 >= 0) & (s64 < num_slots) + idx = s64[valid] + got = flat[idx] + n_bad = int((_bits(got) != _bits(ref[valid])).sum()) + assert n_bad == 0, f"{n_bad} of {got.numel()} elements in the named slots differ bitwise" + flat.view(torch.int16)[idx] = SENTINEL + whole = buf.view(torch.int16).view(-1) + step = 1 << 30 + n_out = sum(int((whole[i : i + step] != SENTINEL).sum()) for i in range(0, whole.numel(), step)) + assert n_out == 0, f"{n_out} elements outside the named slots were written" + + +def _free_engine_pool(): + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +@pytest.mark.parametrize("idx_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("ratio", ["ratio2", "ratio1"]) +def test_engine_pool_far_end(ratio, idx_dtype): + """Decode-like: one row per 32-bit hazard boundary and at both ends of the engine-size pool, plus skips.""" + pages, page_size = ENGINE_COMPRESS[ratio] + buf, pool = _engine_pool(pages, page_size) + num_slots = pages * page_size + hazard = _hazard_slots(num_slots, W, page_size) + assert hazard[-1] == num_slots - 1 and 2**34 // W in hazard + # -1 and the first two slots past the end (which would land in the guard page) are skipped. + slots = torch.tensor( + hazard + [-1, num_slots, num_slots + page_size - 1], dtype=torch.int64, device=DEV + ) + x = _random_blocks(slots.numel(), seed=31) + op(x, pool, slots.to(idx_dtype)) + _check_engine_pool(buf, pool, slots, _ref_rows(x)[0]) + del buf, pool + _free_engine_pool() + + +@pytest.mark.parametrize("ratio", ["ratio2", "ratio1"]) +def test_engine_pool_prefill_spread(ratio): + """Prefill-like, n = 16384 (the envelope's max_num_tokens): the last page written in order, every hazard + slot, the rest spread over the pool.""" + t = _MAX_ROWS + pages, page_size = ENGINE_COMPRESS[ratio] + buf, pool = _engine_pool(pages, page_size) + num_slots = pages * page_size + last_page = torch.arange(num_slots - page_size, num_slots, device=DEV) + hazard = torch.tensor(_hazard_slots(num_slots, W, page_size), device=DEV) + fixed = torch.unique(torch.cat([last_page, hazard])) + g = torch.Generator(device=DEV).manual_seed(8192) + perm = torch.randperm(num_slots, device=DEV, generator=g) + rest = perm[~torch.isin(perm, fixed)][: t - fixed.numel()] + slots = torch.cat([last_page, fixed[~torch.isin(fixed, last_page)], rest]).to(torch.int64) + assert slots.numel() == t and torch.unique(slots).numel() == t + slots[torch.randperm(t, device=DEV, generator=g)[: t // 16]] = -1 + x = _sweep_rows(t, seed=8192) + op(x, pool, slots.to(torch.int32)) + _check_engine_pool(buf, pool, slots, _ref_rows(x)[0]) + del buf, pool + _free_engine_pool() + + +def test_engine_pool_cuda_graph_far_end(): + """Captured on the ratio-2 engine-size pool after a warm-up call, replayed with far-end and boundary slots.""" + t = 32 + pages, page_size = ENGINE_COMPRESS["ratio2"] + buf, pool = _engine_pool(pages, page_size) + num_slots = pages * page_size + hazard = _hazard_slots(num_slots, W, page_size) + x = torch.empty(t, W, dtype=torch.bfloat16, device=DEV) + slots = torch.empty(t, dtype=torch.int32, device=DEV) + x.copy_(_random_blocks(t, seed=40)) + slots.copy_(torch.arange(t, device=DEV)) + op(x, pool, slots) # compile outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(x, pool, slots) + buf.view(torch.int16).fill_(SENTINEL) + + g = torch.Generator(device=DEV).manual_seed(41) + hz = torch.tensor(hazard, device=DEV) + near_end = num_slots - 1 - torch.arange(4 * page_size, device=DEV) + near_end = near_end[~torch.isin(near_end, hz)] + for seed in (42, 43): + pick = torch.randperm(near_end.numel(), device=DEV, generator=g)[: t - len(hazard) - 1] + new_slots = torch.cat([hz, torch.tensor([-1], device=DEV), near_end[pick]]).to(torch.int64) + assert torch.unique(new_slots).numel() == t + new_x = _random_blocks(t, seed=seed) + x.copy_(new_x) + slots.copy_(new_slots) + graph.replay() + torch.cuda.synchronize() + _check_engine_pool(buf, pool, new_slots, _ref_rows(new_x)[0]) + del graph, buf, pool + _free_engine_pool() + + +def test_engine_pool_harness_sees_a_wrapped_offset(): + """Control: the far-end check fails when rows land where a 32-bit (mod 2^32) element offset would put them. + + The op is handed the wrapped slots and the checker the true ones; both the named-slot comparison and the + outside-the-slots sentinel scan must see it. + """ + pages, page_size = ENGINE_COMPRESS["ratio2"] + buf, pool = _engine_pool(pages, page_size) + num_slots = pages * page_size + true = torch.tensor( + [s for s in _hazard_slots(num_slots, W, page_size) if s * W >= 2**32], + dtype=torch.int64, + device=DEV, + ) + wrapped = (true * W % 2**32) // W + assert true.numel() >= 5 and not torch.isin(wrapped, true).any() + x = _random_blocks(true.numel(), seed=51) + ref = _ref_rows(x)[0] + op(x, pool, wrapped) + with pytest.raises(AssertionError, match="elements in the named slots differ bitwise"): + _check_engine_pool(buf, pool, true, ref) + flat = pool.view(num_slots, W) + flat.view(torch.int16)[true] = SENTINEL + with pytest.raises(AssertionError, match="outside the named slots were written"): + _check_engine_pool(buf, pool, torch.tensor([], dtype=torch.int64, device=DEV), ref[:0]) + del buf, pool, flat + _free_engine_pool() + + +# ─── every row count 1..16384 in the core target's paged-scatter call ──────── +# +# The DeepSeek-V4.1 core target writes each KV source's compressed entries with +# nvfp4_fake_quantize(latent [n, 512], compress_view, slots[:n]): ratio-2 owners +# into a [pages, 64, 512] view, the ratio-1 owner into a [pages, 128, 512] view, +# n = the step's compressed-entry count (T for ratio 1; up to T / 2 plus the +# reserved decode rows for ratio 2; at most the engine's max_num_tokens, 16384 +# in the engine envelope; the KV sources are layers 2, 8, 14 and 20, so the +# decoder bounded-replay pass of layers 21-39 makes no such call), +# and slots a prefix view of the metadata's int64 [max_num_tokens] +# v41_compress_slot_mapping[r] (-1 = no slot). Rows reach the kernel only as the +# size of grid axis 0: program r loads row r and slot_mapping[r] and writes one +# slot; no kernel argument carries n. + +import tensorrt_llm._torch.custom_ops.nvfp4_fake_quantize as _nv_op # noqa: E402 (the Triton kernel object, for the variant probe) + +_MAX_ROWS = 16384 +_SWEEP_SLOTS = 20480 # plainly larger than the 16384 rows swept +_SWEEP_GEOMETRY = {"ratio2": 64, "ratio1": 128} # page size of each COMPRESS view + + +def _jit_variants(fn) -> int: + """Number of compiled variants Triton holds for a @triton.jit function, over all devices.""" + dc = getattr(fn, "device_caches", None) + if dc is not None: + return sum(len(v[0]) if isinstance(v, tuple) else len(v) for v in dc.values()) + return sum(len(v) for v in fn.cache.values()) + + +def _sweep_rows(t: int, seed: int) -> torch.Tensor: + """Random 16-blocks over 2^-16 .. 2^6, with the tie set and the floor set spliced in near the top.""" + x = _random_blocks(t, seed=seed) + ties, floors = _tie_blocks(), _floor_blocks() + x[2 : 2 + ties.shape[0]] = ties + x[40 : 40 + floors.shape[0]] = floors + return x + + +def _sweep_slots(t: int, num_slots: int, seed: int) -> torch.Tensor: + """A prefix-able int64 [t] view like v41_compress_slot_mapping[r]: distinct random slots, 1/16 of them -1.""" + g = torch.Generator(device=DEV).manual_seed(seed) + slots = torch.randperm(num_slots, device=DEV, generator=g)[:t].to(torch.int64) + slots[torch.randperm(t, device=DEV, generator=g)[: t // 16]] = -1 + slots[1] = -1 # a skipped row inside the smallest prefixes too + return slots + + +@pytest.mark.parametrize("ratio", list(_SWEEP_GEOMETRY)) +def test_every_row_count_1_to_16384_in_the_paged_scatter_form(ratio): + """Every prefix (x[:r], slots[:r]), r = 1..16384, writes exactly the first r rows of the 16384-row call. + + The 16384-row call is held bitwise to the native reference and to the + sentinel scan. Then, into the same sentinel-filled pool, the op runs at + every r in increasing order. After each call the slots named by the live + entries of slots[:r] hold the reference rows bit for bit, and the number of + pool elements that are not the sentinel is exactly 512 per live row: a + skipped row, a row written to another slot, or a write past row r - 1 all + change one of the two counts. The mismatch counts accumulate on the device; + one sync at the end. The probe: rows are only grid axis 0, so the sweep must + compile no Triton variant beyond the one the 16384-row call compiled. + """ + t = _MAX_ROWS + page_size = _SWEEP_GEOMETRY[ratio] + buf, pool = _engine_pool(_SWEEP_SLOTS // page_size, page_size) + flat = pool.view(_SWEEP_SLOTS, W).view(torch.int16) + whole = buf.view(torch.int16) + slots = _sweep_slots(t, _SWEEP_SLOTS, seed=8193 + page_size) + x = _sweep_rows(t, seed=8193 + page_size) + ref = _ref_rows(x)[0] + + kernel = _nv_op._nvfp4_fake_quantize_kernel + n0 = _jit_variants(kernel) + op(x, pool, slots) + n1 = _jit_variants(kernel) + _check_engine_pool( + buf, pool, slots, ref + ) # the N_max call against the reference; restores the sentinel + + live = slots >= 0 + live_pos = torch.nonzero(live).squeeze(1) + live_slots = slots[live_pos] + ref_live = _bits(ref[live_pos]) + cnt = torch.cumsum(live.to(torch.int64), 0).tolist() + bad_rows = torch.zeros((), dtype=torch.int64, device=DEV) + bad_count = torch.zeros((), dtype=torch.int64, device=DEV) + for r in range(1, t + 1): + op(x[:r], pool, slots[:r]) + k = cnt[r - 1] + bad_rows += (flat[live_slots[:k]] != ref_live[:k]).sum() + bad_count += ((whole != SENTINEL).sum() - k * W).abs() + n2 = _jit_variants(kernel) + print( + f"\n[{ratio} rows 1..{t}, page {page_size}] live rows {cnt[-1]}; Triton variants: {n0} before the " + f"{t}-row call, {n1} after, {n2} after the sweep; element mismatches {int(bad_rows)}, " + f"non-sentinel count error {int(bad_count)}" + ) + assert int(bad_rows) == 0, ( + f"{int(bad_rows)} elements of the named slots differ from the {t}-row result" + ) + assert int(bad_count) == 0, ( + f"non-sentinel element count off by {int(bad_count)} summed over the sweep" + ) + assert n1 - n0 <= 1 and n2 == n1, f"row counts compiled new variants: {n0} -> {n1} -> {n2}" + del buf, pool, flat, whole + _free_engine_pool() + + +def test_row_sweep_harness_sees_a_row_or_slot_shift(): + """Control for the sweep: its checks reject a read one row off and a write one slot off, at every r tried.""" + page_size = 64 + buf, pool = _engine_pool(_SWEEP_SLOTS // page_size, page_size) + flat = pool.view(_SWEEP_SLOTS, W).view(torch.int16) + whole = buf.view(torch.int16) + g = torch.Generator(device=DEV).manual_seed(66) + slots = torch.randperm(_SWEEP_SLOTS, device=DEV, generator=g)[:66].to( + torch.int64 + ) # no -1: every row lands + x = _random_blocks(66, seed=66) + ref = _bits(_ref_rows(x)[0]) + for r in (1, 2, 7, 63, 64): + for name, xa, sa in ( + ("row +1", x[1 : r + 1], slots[:r]), + ("slot +1", x[:r], slots[1 : r + 1]), + ): + whole.fill_(SENTINEL) + op(xa, pool, sa) + n_bad = int((flat[slots[:r]] != ref[:r]).sum()) + n_count = int((whole != SENTINEL).sum()) - r * W + print( + f"[control {name}] r={r}: {n_bad} / {r * W} elements differ, count error {n_count}" + ) + assert n_bad > r * W // 2 + del buf, pool, flat, whole + _free_engine_pool() + + +# ─── CUDA-graph capture at every graph batch the engines capture ───────────── +# +# The engine envelope captures every graph batch G of _ENVELOPE_GSET: every G +# 1..32 (padding off captures every batch up to max_batch_size) and TRT-LLM's +# padded list for cuda_graph_config max_batch_size 256 (1, 2, 4, 8..128 step 8, +# 192, 256). With DSpark the target-verify step is captured at 6 tokens per +# request. Rows per captured call: the ratio-1 owner writes one entry per token +# (n = G, and 6G for verify); a ratio-2 owner writes the runtime-reserved +# ceil(next_n / 2) rows per request (n = G, and 3G for verify), the reserved +# rows without a completed pair carrying slot -1. +_ENVELOPE_GSET = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +assert len(_ENVELOPE_GSET) == 46 +_GRAPH_N = { + "ratio1": tuple( + sorted(set(_ENVELOPE_GSET) | {6 * g for g in _ENVELOPE_GSET}) + ), # 82 row counts, 1..1536 + "ratio2": tuple( + sorted(set(_ENVELOPE_GSET) | {3 * g for g in _ENVELOPE_GSET}) + ), # 77 row counts, 1..768 +} + + +def _far_end_slots(num_slots, page_size, n, g): + """n distinct slots: the pool's last slot, then a random pick from the 32-bit hazard slots and the last + 4 pages (and, when n needs more, slots spread uniformly over the pool); for n >= 2 one of them (not the + last slot) replaced by -1.""" + last = num_slots - 1 + hazard = torch.tensor( + [s for s in _hazard_slots(num_slots, W, page_size) if s != last], device=DEV + ) + near_end = last - 1 - torch.arange(4 * page_size - 1, device=DEV) + near_end = near_end[~torch.isin(near_end, hazard)] + rest = torch.cat([hazard, near_end]) + rest = rest[torch.randperm(rest.numel(), device=DEV, generator=g)] + if rest.numel() < n - 1: + spread = torch.randint(0, last, (4 * n,), device=DEV, generator=g).unique() + spread = spread[~torch.isin(spread, rest)] + rest = torch.cat([rest, spread[torch.randperm(spread.numel(), device=DEV, generator=g)]]) + slots = torch.cat([torch.tensor([last], device=DEV), rest[: n - 1]]).to(torch.int64) + assert slots.numel() == n and torch.unique(slots).numel() == n + if n >= 2: + slots[1 + int(torch.randint(0, n - 1, (1,), device=DEV, generator=g))] = -1 + return slots + + +@pytest.mark.parametrize("ratio", ["ratio1", "ratio2"]) +def test_cuda_graph_every_graph_batch_on_the_engine_compress_view(ratio): + """Per n of that ratio's envelope capture set (ratio 1: G and 6G, 82 row counts to 1536; ratio 2: G and + 3G, 77 row counts to 768; G over _ENVELOPE_GSET): eager warm-up at that shape, capture, then 2 replays on new + rows and new far-end slots. + + On the B_max bf16 COMPRESS view of that ratio. Every replay must write + exactly the reference into the named slots and nothing else (sentinel scan + of the whole pool plus a guard page), and equal bit for bit the same call + run eagerly on the same inputs. + """ + pages, page_size = ENGINE_COMPRESS[ratio] + ns = _GRAPH_N[ratio] + buf, pool = _engine_pool(pages, page_size) + num_slots = pages * page_size + flat = pool.view(num_slots, W) + g = torch.Generator(device=DEV).manual_seed(700 + page_size) + table = torch.full( + (_MAX_ROWS,), -1, dtype=torch.int64, device=DEV + ) # like v41_compress_slot_mapping[r] + for n in ns: + x = torch.empty(n, W, dtype=torch.bfloat16, device=DEV) + slots = table[:n] + x.copy_(_random_blocks(n, seed=700 + n)) + slots.copy_(_far_end_slots(num_slots, page_size, n, g)) + op(x, pool, slots) # eager warm-up at this shape, as the engine's warm-up step runs it + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(x, pool, slots) + buf.view(torch.int16).fill_(SENTINEL) + for rep in (1, 2): + new_x = _random_blocks(n, seed=700 + 1000 * rep + n) + new_slots = _far_end_slots(num_slots, page_size, n, g) + x.copy_(new_x) + slots.copy_(new_slots) + graph.replay() + torch.cuda.synchronize() + live = new_slots[new_slots >= 0] + replayed = _bits(flat[live]).clone() + ref = _ref_rows(new_x)[0] + _check_engine_pool(buf, pool, new_slots, ref) + op(x, pool, slots) # the same call, eagerly, on the same inputs + torch.cuda.synchronize() + eager = _bits(flat[live]).clone() + _check_engine_pool(buf, pool, new_slots, ref) + assert torch.equal(replayed, eager), ( + f"n={n} replay {rep}: {int((replayed != eager).sum())} elements" + ) + del graph + print( + f"\n[{ratio} graphs on [{pages}, {page_size}, 512]] n = {list(ns)}: every replay bitwise == eager == reference" + ) + del buf, pool, flat + _free_engine_pool() + + +# ─── any COMPRESS view B <= B_max ───────────────────────────────────────────── +# +# The engine picks the COMPRESS views' page counts from free memory, so no single +# dumped B is "the" engine's. Two tests make the coverage an upper bound: +# (1) the cache manager at the largest quota a gate-configured engine can reach +# gives B_max = ENGINE_COMPRESS (and the op is exact at the far end of the +# manager's own views), and (2) the same rows into the same slots give the same +# bits at B = the smallest view that holds them, sizes across the 2^31 / 2^32 +# element-count lines, a mid size and B_max (views of one allocation), and in a +# separate exact-size allocation. + +# (max_draft_len, max_batch_size) -> (ratio-2 pages of owners 2 / 8 / 14, ratio-1 pages of owner 20) +# at 0.5 x this GB300's total memory. +_BMAX_DEVICE_TOTAL = 297_020_948_480 +_BMAX_COMPRESS_PAGES = { + (5, 8): ((509184, 509183, 509182), 169728), + (0, 8): ((527358, 527357, 527356), 175786), + (5, 32): ((801024, 801023, 801022), 267008), + (0, 32): ((862719, 862718, 862717), 287573), +} + + +def _pp_mapping(pp_stage): + """The single-process mapping of a rank of PP2 stage `pp_stage` (tp 2, as the PP2 stage scan builds + it; the V4.1 pools do not depend on tp), or of the unsplit manager when pp_stage is None.""" + from tensorrt_llm.mapping import Mapping + + if pp_stage is None: + return Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + return Mapping(world_size=4, rank=2 * pp_stage, tp_size=2, pp_size=2, gpus_per_node=4) + + +def _v41_manager( + quota, max_draft_len, max_batch_size, max_num_tokens=8192, max_seq_len=None, pp_stage=None +): + """A DeepSeek-V4.1 DeepseekV4CacheManager with bf16-container KV storage and the given per-rank KV quota, + at the gate settings unless max_num_tokens / max_seq_len say otherwise (max_seq_len as the engine passes + it: + 14 with DSpark). The storage is pinned: when TRTLLM_DSV41_KV_STORAGE is unset the manager stores the + KV natively (the native section below builds those managers explicitly).""" + import os + from types import SimpleNamespace + + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, + KvCacheConfig, + ) + + text = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=128, + head_dim=512, + index_head_dim=128, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, + ) + config = SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **{k: v for k, v in text.items() if k != "model_type"}, + ) + saved = os.environ.get("TRTLLM_DSV41_KV_STORAGE") + os.environ["TRTLLM_DSV41_KV_STORAGE"] = "bf16" + try: + mgr = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + max_gpu_total_bytes=quota, + host_cache_size=0, + event_buffer_max_size=0, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=40, + num_kv_heads=1, + head_dim=512, + tokens_per_block=128, + # DSpark (max_draft_len 5, overlap on) grows the manager's max_seq_len by 5 + 4 + 5. + max_seq_len=max_seq_len or (16398 if max_draft_len else 16384), + max_batch_size=max_batch_size, + mapping=_pp_mapping(pp_stage), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=max_num_tokens, + sparse_attn_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ), + pretrained_config=config, + spec_config=DSparkDecodingConfig(max_draft_len=max_draft_len, block_size=max_draft_len) + if max_draft_len + else None, + max_cuda_graph_batch_size=max_batch_size, + disable_overlap_scheduler=False, + ) + finally: + if saved is None: + os.environ.pop("TRTLLM_DSV41_KV_STORAGE", None) + else: + os.environ["TRTLLM_DSV41_KV_STORAGE"] = saved + assert mgr.v41_layout.kv_storage == "bf16" + return mgr + + +def test_largest_engine_views_from_the_cache_manager(): + """B_max from the manager itself, and the far-end hazard call on the manager's own B_max views. + + The quota is 0.5 x the device's total memory, the upper bound of kv_cache_fraction 0.5 x free memory. + The largest configuration (max_batch_size 32, no DSpark) is built last and kept: on owner 2's and + owner 14's ratio-2 views and owner 20's ratio-1 view (each filled with the sentinel first), the op + writes rows to both sides of every 2^31..2^34 element-offset line, the first and the last page and + the last slot; the view must hold the reference there and the sentinel everywhere else. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + total = torch.cuda.mem_get_info()[1] + assert total == _BMAX_DEVICE_TOTAL, ( + f"B_max constants are for a {_BMAX_DEVICE_TOTAL}-byte GB300, got {total}" + ) + quota = total // 2 + got = {} + keys = list(_BMAX_COMPRESS_PAGES) # the largest last + for key in keys: + mgr = _v41_manager(quota, *key) + views = {o: mgr.get_buffers(o, AT.COMPRESS) for o in (2, 8, 14, 20)} + for o, v in views.items(): + ps = 64 if o != 20 else 128 + assert tuple(v.shape[1:]) == (ps, W) and v.dtype == torch.bfloat16, ( + o, + v.shape, + v.dtype, + ) + assert v.stride() == (ps * W, W, 1), (o, v.stride()) + got[key] = (tuple(views[o].shape[0] for o in (2, 8, 14)), views[20].shape[0]) + if key != keys[-1]: + mgr.shutdown() + del mgr, views + _free_engine_pool() + print( + f"\n[quota {quota} B = 0.5 x device total] COMPRESS pages per (max_draft_len, max_batch_size): {got}" + ) + assert got == _BMAX_COMPRESS_PAGES, got + assert max(v[0][0] for v in got.values()) == GATES_COMPRESS["ratio2"][0] + assert max(v[1] for v in got.values()) == GATES_COMPRESS["ratio1"][0] + try: + for o in (2, 14, 20): + pool = views[o] + page_size = pool.shape[1] + num_slots = pool.shape[0] * page_size + pool.view(torch.int16).fill_(SENTINEL) + hazard = _hazard_slots(num_slots, W, page_size) + assert hazard[-1] == num_slots - 1 and 2**34 // W in hazard + slots = torch.tensor(hazard + [-1, num_slots], dtype=torch.int64, device=DEV) + x = _random_blocks(slots.numel(), seed=8620 + o) + op(x, pool, slots) + _check_engine_pool(pool, pool, slots, _ref_rows(x)[0]) + print( + f"[manager's own owner-{o} view {tuple(pool.shape)}] {len(hazard)} far-end / hazard rows == " + f"reference, rest of the view sentinel" + ) + finally: + mgr.shutdown() + del mgr, views + _free_engine_pool() + + +def _invariance_slots(n, page_size, seed): + """(smallest page count holding n slots, int64 [n] prefix of a [8192] table of distinct slots in it): + the last slot of that smallest view is named, 1/16 of the other rows -1.""" + small = -(-n // page_size) + last = small * page_size - 1 + g = torch.Generator(device=DEV).manual_seed(seed) + perm = torch.randperm(small * page_size, device=DEV, generator=g) + pick = torch.cat([torch.tensor([last], device=DEV), perm[perm != last]])[:n] + pick = pick[torch.randperm(n, device=DEV, generator=g)] + table = torch.full( + (_MAX_ROWS,), -1, dtype=torch.int64, device=DEV + ) # like v41_compress_slot_mapping[r] + slots = table[:n] + slots.copy_(pick) + skip = torch.randperm(n, device=DEV, generator=g)[: n // 16] + slots[skip[slots[skip] != last]] = -1 + return small, slots + + +@pytest.mark.parametrize("n", [32, 8192, 16384]) +@pytest.mark.parametrize("ratio", ["ratio2", "ratio1"]) +def test_pool_size_invariance(ratio, n): + """The same n rows into the same slots: bitwise the same bits at every view size B, and = reference. + + B = the smallest view holding the slots, the page counts whose element count first passes 2^31, 2^32, + 2^33, 2^34, 2^35 and 2^36 (where below B_max), the gates' bound (GATES_COMPRESS), the unsplit + manager's envelope bound (UNSPLIT_COMPRESS), a mid size B_max // 2 + 1 and the envelope B_max (a + PP2 stage alone), each a view pool[:B] of one sentinel-filled allocation of B_max pages plus a + guard page, then a separate allocation of exactly the smallest size. After every call the whole + allocation (past B too) holds the sentinel outside the named slots. + """ + pages, page_size = ENGINE_COMPRESS[ratio] + small, slots = _invariance_slots(n, page_size, seed=2300 + n + page_size) + x = _random_blocks(n, seed=2300 + n + page_size) + ref = _ref_rows(x)[0] + live = slots[slots >= 0] + per_page = page_size * W + sizes = sorted( + {small, pages // 2 + 1, GATES_COMPRESS[ratio][0], UNSPLIT_COMPRESS[ratio][0], pages} + | { + lim // per_page + 1 + for lim in (2**31, 2**32, 2**33, 2**34, 2**35, 2**36) + if lim // per_page + 1 < pages + } + ) + buf, _ = _engine_pool(pages, page_size) + first = None + for b in sizes: + pool = buf[:b] + op(x, pool, slots) + torch.cuda.synchronize() + got = _bits(pool.view(-1, W)[live]).clone() + _check_engine_pool(buf, pool, slots, ref) + first = got if first is None else first + assert torch.equal(got, first), ( + f"B={b}: {int((got != first).sum())} elements differ from B={sizes[0]}" + ) + del buf, pool + _free_engine_pool() + buf, pool = _engine_pool(small, page_size) + op(x, pool, slots) + torch.cuda.synchronize() + got = _bits(pool.view(-1, W)[live]).clone() + _check_engine_pool(buf, pool, slots, ref) + assert torch.equal(got, first), "a separate exact-size view differs" + print( + f"\n[{ratio} n={n}] views B = {sizes} pages of [B, {page_size}, 512] (one allocation) and a separate " + f"[{small}, {page_size}, 512]: {live.numel()} rows bitwise identical and == reference, nothing else written" + ) + del buf, pool + _free_engine_pool() + + +# ─── native storage on the live DeepSeek-V4.1 cache manager ────────────────── +# +# With TRTLLM_DSV41_KV_STORAGE=native the V4.1 DeepseekV4CacheManager keeps the +# compressed main KV natively: per ratio one uint8 buffer of 288-byte rows, the +# 256 packed E2M1 bytes then the 32 E4M3 bytes; the ratio-2 owners 2 / 8 / 14 +# view it as [B, 64, 288] with bases one page apart (so B, B - 1, B - 2 pages), +# owner 20 (ratio 1) as [B, 128, 288]. get_v41_native_kv_views(o, COMPRESS) +# returns (codes [..., 256], scales [..., 32]), two uint8 views with the pool's +# strides. The native-storage target writes them with this op's native form: +# nvfp4_fake_quantize(latent [n, 512], codes, md.v41_compress_slot_mapping[r][:n], scales). +# The table is the manager's native COMPRESS pages at 0.5 x this GB300's total +# memory (the ceiling of kv_cache_fraction 0.5 x free memory) for every +# (max_num_tokens, max_batch_size, DSpark max_draft_len) the runs use or could +# use: the native gates run (2048, 8, 5) and (8192, 32, 5) and reach 1,954,512 / +# 651,504 pages; that cross-product's largest, (2048, 32, 0), is 2,606,019 +# (ratio 2, 166.8 M rows, 48.0 GB) and 868,673 (ratio 1, 111.2 M rows, 32.0 GB) +# (GATES_NATIVE_COMPRESS). Over the engine envelope the largest native views +# are 3,886,647 / 1,295,549 pages at (2048, 128, 0, max_seq_len 1,048,576) +# (NATIVE_COMPRESS, the B_max the far-end, invariance and capture cases use). + +import os # noqa: E402 +import random # noqa: E402 +from contextlib import contextmanager # noqa: E402 + +_STORAGE_ENV = "TRTLLM_DSV41_KV_STORAGE" +NATIVE_ROW = 288 # 256 packed E2M1 bytes, then 32 E4M3 bytes +NATIVE_CODES = 256 +# 0xFF is the E4M3 NaN: the op never writes it as a scale byte (scales are +# positive, saturating at 0x7E). A packed code byte can be 0xFF (two -6 +# nibbles), so the checks compare the named rows first and then restore them. +NATIVE_SENTINEL = 0xFF +_NATIVE_COMPRESS_PAGES = { + (2048, 8, 0): ((2_131_740, 2_131_739, 2_131_738), 710_580), + (2048, 8, 5): ((1_954_512, 1_954_511, 1_954_510), 651_504), + (2048, 32, 0): ((2_606_019, 2_606_018, 2_606_017), 868_673), + (2048, 32, 5): ((2_315_604, 2_315_603, 2_315_602), 771_868), + (8192, 8, 0): ((1_378_986, 1_378_985, 1_378_984), 459_662), + (8192, 8, 5): ((1_302_528, 1_302_527, 1_302_526), 434_176), + (8192, 32, 0): ((2_083_497, 2_083_496, 2_083_495), 694_499), + (8192, 32, 5): ((1_895_082, 1_895_081, 1_895_080), 631_694), +} +_NATIVE_GATES = ((2048, 8, 5), (8192, 32, 5)) +# ratio -> (B_max pages over max_batch_size 1..32 at max_num_tokens {2048, 8192}, page size, gates' largest) +GATES_NATIVE_COMPRESS = {"ratio2": (2_606_019, 64, 1_954_512), "ratio1": (868_673, 128, 651_504)} +# ratio -> (the unsplit manager's B_max pages over the engine envelope, page size, gates' largest): the +# pool scan's largest native COMPRESS views, both at (max_num_tokens 2048, max_batch_size 128, no DSpark, +# max_seq_len 1,048,576); re-measured by test_envelope_largest_compress_views_from_the_cache_manager. +UNSPLIT_NATIVE_COMPRESS = { + "ratio2": (3_886_647, 64, 1_954_512), + "ratio1": (1_295_549, 128, 651_504), +} +ENVELOPE_NATIVE_COMPRESS_KEY = ( + 2048, + 128, + 0, + 1_048_576, +) # (max_num_tokens, max_batch_size, max_draft_len, max_seq_len) +# PP2 (the unified target's stage split at layer 20): a stage manager holds only its 20 layers' pools. The +# PP2 stage scan's largest native COMPRESS views are stage 0 alone (ratio-2 owners 2 / 8 / 14) and stage 1 +# alone (ratio-1 owner 20), both at (2048, 256, 0, 1,048,576) +# (the PP2 stage scan; the stage alone is +# the upper bound: in an engine the stages' token capacities are cut to the smaller one); re-measured by +# test_envelope_pp2_stage_compress_views_from_the_cache_manager. ratio -> (B_max pages, page size, gates' largest). +NATIVE_COMPRESS = {"ratio2": (6_480_780, 64, 1_954_512), "ratio1": (3_245_959, 128, 651_504)} +# ratio -> (the PP2 stage holding that ratio's views, the stage manager's (max_num_tokens, max_batch_size, +# max_draft_len, max_seq_len) at the native stage-alone maximum) +ENVELOPE_PP2_NATIVE_COMPRESS = { + "ratio2": (0, (2048, 256, 0, 1_048_576)), + "ratio1": (1, (2048, 256, 0, 1_048_576)), +} +_OWNERS = {"ratio2": (2, 8, 14), "ratio1": (20,)} + + +@contextmanager +def _native_storage(): + saved = os.environ.get(_STORAGE_ENV) + os.environ[_STORAGE_ENV] = "native" + try: + yield + finally: + if saved is None: + os.environ.pop(_STORAGE_ENV, None) + else: + os.environ[_STORAGE_ENV] = saved + + +def _v41_native_manager( + max_num_tokens, max_batch_size, max_draft_len, max_seq_len=None, pp_stage=None +): + """A V4.1 manager with native KV storage and 0.5 x the device's total memory, at the gates' max_seq_len + unless max_seq_len says otherwise (as the engine passes it: + 14 with DSpark).""" + from types import SimpleNamespace + + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, + KvCacheConfig, + ) + + total = torch.cuda.mem_get_info()[1] + assert total == _BMAX_DEVICE_TOTAL, ( + f"B_max constants are for a {_BMAX_DEVICE_TOTAL}-byte GB300, got {total}" + ) + text = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=128, + head_dim=512, + index_head_dim=128, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, + ) + config = SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **{k: v for k, v in text.items() if k != "model_type"}, + ) + with _native_storage(): + mgr = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + max_gpu_total_bytes=total // 2, + host_cache_size=0, + event_buffer_max_size=0, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=40, + num_kv_heads=1, + head_dim=512, + tokens_per_block=128, + max_seq_len=max_seq_len or (16398 if max_draft_len else 16384), + max_batch_size=max_batch_size, + mapping=_pp_mapping(pp_stage), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=max_num_tokens, + sparse_attn_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ), + pretrained_config=config, + spec_config=DSparkDecodingConfig(max_draft_len=max_draft_len, block_size=max_draft_len) + if max_draft_len + else None, + max_cuda_graph_batch_size=max_batch_size, + disable_overlap_scheduler=False, + ) + assert mgr.v41_layout.kv_storage == "native" + return mgr + + +def _native_compress(mgr, owner): + """(owner's COMPRESS pool [B, page, 288] uint8, codes view, scales view), geometry asserted.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + page_size = 128 if owner == 20 else 64 + pool = mgr.get_buffers(owner, AT.COMPRESS) + codes, scales = mgr.get_v41_native_kv_views(owner, AT.COMPRESS) + b = pool.shape[0] + assert pool.dtype == torch.uint8 and tuple(pool.shape) == (b, page_size, NATIVE_ROW), ( + owner, + pool.shape, + ) + assert pool.stride() == (page_size * NATIVE_ROW, NATIVE_ROW, 1), (owner, pool.stride()) + assert tuple(codes.shape) == (b, page_size, NATIVE_CODES) and codes.stride() == pool.stride() + assert tuple(scales.shape) == (b, page_size, 32) and scales.stride() == pool.stride() + assert ( + codes.data_ptr() == pool.data_ptr() and scales.data_ptr() == pool.data_ptr() + NATIVE_CODES + ) + return pool, codes, scales + + +def _native_rows(x): + """(uint8 [n, 288] native rows = packed codes then E4M3 bytes, bf16 [n, 512] core values) of the reference.""" + deq, packed, s_bytes = _ref_rows(x) + return torch.cat([packed, s_bytes], dim=1), deq + + +def _native_hazard_slots(num_slots, page_size): + """Distinct slots of a 288-byte-row pool on both sides of every byte offset 2^31..2^36 inside it (the row + holding the offset and its neighbours), plus the first page, the last page and the last slot.""" + cands = {0, 1, page_size - 1, page_size} + for limit in (2**31, 2**32, 2**33, 2**34, 2**35, 2**36): + s = limit // NATIVE_ROW + cands |= {s - 1, s, s + 1} + cands |= {num_slots - page_size - 1, num_slots - page_size, num_slots - 2, num_slots - 1} + return sorted(c for c in cands if 0 <= c < num_slots) + + +def _count_not_sentinel(buf): + whole = buf.view(-1) + step = ( + 1 << 26 + ) # a bool reduction materializes an int64 copy (8 bytes per element): keep chunks small + n = torch.zeros((), dtype=torch.int64, device=DEV) + for i in range(0, whole.numel(), step): + n += (whole[i : i + step] != NATIVE_SENTINEL).sum() + return n + + +def _check_native_pool(buf, flat, num_slots, slots, rows): + """`flat`: the pool as [slots(+guard), 288] rows inside `buf`. The named slots hold `rows` bitwise; after + they are restored to the sentinel, every byte of `buf` is the sentinel.""" + s64 = slots.to(torch.int64) + live = (s64 >= 0) & (s64 < num_slots) + idx = s64[live] + got = flat[idx] + n_bad = int((got != rows[live]).sum()) + assert n_bad == 0, f"{n_bad} of {got.numel()} bytes in the named slots differ bitwise" + flat[idx] = NATIVE_SENTINEL + n_out = int(_count_not_sentinel(buf)) + assert n_out == 0, f"{n_out} bytes outside the named slots were written" + + +def _native_engine_pool(pages, page_size): + """(uint8 [pages + 1, page_size, 288] buffer with a guard page, codes view, scales view), all sentinel.""" + buf = torch.full( + (pages + 1, page_size, NATIVE_ROW), NATIVE_SENTINEL, dtype=torch.uint8, device=DEV + ) + return buf, buf[:pages, :, :NATIVE_CODES], buf[:pages, :, NATIVE_CODES:] + + +def test_native_compress_pools_from_the_cache_manager(): + """Every native configuration's COMPRESS views from the manager itself, and the hazard call on the largest. + + For each (max_num_tokens, max_batch_size, max_draft_len) at 0.5 x total + memory: the owners' views are [B, 64, 288] (owners 2 / 8 / 14, bases one + page apart) and [B, 128, 288] (owner 20), uint8, with the code / scale + halves as the native views. The largest, (2048, 32, 0), is kept: on owner + 2's and owner 14's ratio-2 views and owner 20's ratio-1 view (each buffer + filled with the sentinel first) the op's native form writes rows to both + sides of every 2^31..2^35 byte offset, the first and the last page and the + last slot; the buffer must hold the reference's packed codes and scale + bytes there, decoding to the op's own core output, and the sentinel + everywhere else. + """ + got = {} + keys = sorted( + _NATIVE_COMPRESS_PAGES, key=lambda k: _NATIVE_COMPRESS_PAGES[k][0][0] + ) # the largest last + for key in keys: + mgr = _v41_native_manager(*key) + views = {o: _native_compress(mgr, o) for o in (2, 8, 14, 20)} + base = views[2][0].data_ptr() + for i, o in enumerate((2, 8, 14)): + assert views[o][0].data_ptr() == base + i * 64 * NATIVE_ROW, o + got[key] = (tuple(views[o][0].shape[0] for o in (2, 8, 14)), views[20][0].shape[0]) + if key != keys[-1]: + mgr.shutdown() + del mgr, views + _free_engine_pool() + print( + f"\n[native, quota = 0.5 x device total] COMPRESS pages per (max_num_tokens, max_batch_size, " + f"max_draft_len): {got}" + ) + assert got == _NATIVE_COMPRESS_PAGES, got + for ratio, i in (("ratio2", 0), ("ratio1", 1)): + bmax, _, gate_max = GATES_NATIVE_COMPRESS[ratio] + first = (lambda v: v[0][0]) if ratio == "ratio2" else (lambda v: v[1]) + assert max(first(v) for v in got.values()) == bmax + assert max(first(got[k]) for k in _NATIVE_GATES) == gate_max + try: + for o in (2, 14, 20): + pool, codes, scales = views[o] + whole = ( + views[2][0] if o != 20 else pool + ) # owner 2's view spans the whole ratio-2 buffer + page_size = pool.shape[1] + num_slots = pool.shape[0] * page_size + whole.fill_(NATIVE_SENTINEL) + hazard = _native_hazard_slots(num_slots, page_size) + assert hazard[-1] == num_slots - 1 + slots = torch.tensor(hazard + [-1, num_slots], dtype=torch.int64, device=DEV) + x = _random_blocks(slots.numel(), seed=26060 + o) + rows, deq = _native_rows(x) + op(x, codes, slots, scales) + flat = pool.view(-1, NATIVE_ROW) + core = _run_core(x) + assert torch.equal(_bits(core), _bits(deq)) + got_rows = flat[slots[: len(hazard)]] + assert torch.equal( + _bits(_decode(got_rows[:, :NATIVE_CODES], got_rows[:, NATIVE_CODES:])), + _bits(core[: len(hazard)]), + ) + _check_native_pool(whole, flat, num_slots, slots, rows) + print( + f"[manager's own native owner-{o} view {tuple(pool.shape)}] {len(hazard)} far-end / hazard rows " + f"== reference codes + bytes, decode == core output, rest of the buffer sentinel" + ) + finally: + mgr.shutdown() + del mgr, views + _free_engine_pool() + + +# The max_batch_size axis: native COMPRESS pages at 0.5 x device total per (max_num_tokens, +# DSpark max_draft_len), (ratio 2 owner 2, ratio 1 owner 20) for max_batch_size 1..32 (measured on this +# GB300; owners 8 / 14 are 1 / 2 pages less than owner 2 in every configuration). +MBS_AXIS = tuple(range(1, 33)) +_COMPRESS_PAGES_BY_MBS = { + (2048, 0): ( + ( + 1_954_512, + 1_954_512, + 1_954_512, + 1_954_512, + 1_954_512, + 2_171_790, + 1_954_512, + 2_131_740, + 2_131_740, + 1_954_512, + 2_104_860, + 2_104_860, + 2_104_860, + 2_104_860, + 2_104_860, + 2_345_640, + 2_345_640, + 2_345_640, + 2_345_640, + 2_443_035, + 2_528_916, + 2_528_916, + 2_468_520, + 2_468_520, + 2_541_339, + 2_541_339, + 2_541_339, + 2_419_368, + 2_419_368, + 2_550_441, + 2_606_019, + 2_606_019, + ), + ( + 651_504, + 651_504, + 651_504, + 651_504, + 651_504, + 723_930, + 651_504, + 710_580, + 710_580, + 651_504, + 701_620, + 701_620, + 701_620, + 701_620, + 701_620, + 781_880, + 781_880, + 781_880, + 781_880, + 814_345, + 842_972, + 842_972, + 822_840, + 822_840, + 847_113, + 847_113, + 847_113, + 806_456, + 806_456, + 850_147, + 868_673, + 868_673, + ), + ), + (2048, 5): ( + ( + 1_954_512, + 1_954_512, + 1_954_512, + 1_954_512, + 1_736_703, + 1_776_753, + 1_776_753, + 1_954_512, + 1_954_512, + 1_954_512, + 2_104_860, + 1_954_512, + 1_954_512, + 1_954_512, + 2_198_184, + 2_198_184, + 2_198_184, + 2_069_844, + 2_069_844, + 2_171_790, + 2_262_810, + 2_262_810, + 2_345_640, + 2_233_683, + 2_310_144, + 2_310_144, + 2_209_107, + 2_209_107, + 2_209_107, + 2_255_208, + 2_315_604, + 2_315_604, + ), + ( + 651_504, + 651_504, + 651_504, + 651_504, + 578_901, + 592_251, + 592_251, + 651_504, + 651_504, + 651_504, + 701_620, + 651_504, + 651_504, + 651_504, + 732_728, + 732_728, + 732_728, + 689_948, + 689_948, + 723_930, + 754_270, + 754_270, + 781_880, + 744_561, + 770_048, + 770_048, + 736_369, + 736_369, + 736_369, + 751_736, + 771_868, + 771_868, + ), + ), + (8192, 0): ( + ( + 1_116_840, + 1_040_865, + 1_040_865, + 1_040_865, + 1_040_865, + 1_220_607, + 1_220_607, + 1_378_986, + 1_378_986, + 1_518_891, + 1_518_891, + 1_439_061, + 1_439_061, + 1_439_061, + 1_674_807, + 1_674_807, + 1_674_807, + 1_674_807, + 1_674_807, + 1_776_753, + 1_869_594, + 1_954_512, + 1_875_966, + 1_804_167, + 1_881_429, + 1_881_429, + 1_881_429, + 1_881_429, + 2_020_692, + 2_020_692, + 2_083_497, + 2_083_497, + ), + ( + 372_280, + 346_955, + 346_955, + 346_955, + 346_955, + 406_869, + 406_869, + 459_662, + 459_662, + 506_297, + 506_297, + 479_687, + 479_687, + 479_687, + 558_269, + 558_269, + 558_269, + 558_269, + 558_269, + 592_251, + 623_198, + 651_504, + 625_322, + 601_389, + 627_143, + 627_143, + 627_143, + 627_143, + 673_564, + 673_564, + 694_499, + 694_499, + ), + ), + (8192, 5): ( + ( + 1_040_865, + 1_040_865, + 1_040_865, + 1_040_865, + 1_040_865, + 1_148_805, + 1_148_805, + 1_302_528, + 1_302_528, + 1_439_061, + 1_439_061, + 1_367_151, + 1_367_151, + 1_367_151, + 1_598_349, + 1_598_349, + 1_598_349, + 1_598_349, + 1_529_172, + 1_627_476, + 1_720_320, + 1_736_703, + 1_736_703, + 1_736_703, + 1_751_265, + 1_751_265, + 1_751_265, + 1_751_265, + 1_889_619, + 1_889_619, + 1_895_082, + 1_895_082, + ), + ( + 346_955, + 346_955, + 346_955, + 346_955, + 346_955, + 382_935, + 382_935, + 434_176, + 434_176, + 479_687, + 479_687, + 455_717, + 455_717, + 455_717, + 532_783, + 532_783, + 532_783, + 532_783, + 509_724, + 542_492, + 573_440, + 578_901, + 578_901, + 578_901, + 583_755, + 583_755, + 583_755, + 583_755, + 629_873, + 629_873, + 631_694, + 631_694, + ), + ), +} +_NATIVE_COMPRESS_PAGES_AXIS = { + (t, s, d): ((c2[s - 1], c2[s - 1] - 1, c2[s - 1] - 2), c20[s - 1]) + for (t, d), (c2, c20) in _COMPRESS_PAGES_BY_MBS.items() + for s in MBS_AXIS +} +assert all(_NATIVE_COMPRESS_PAGES_AXIS[k] == v for k, v in _NATIVE_COMPRESS_PAGES.items()) + + +def test_native_compress_pools_over_every_max_batch_size(): + """B_max over the whole max_batch_size axis: the manager in native mode at 0.5 x total memory, for every + (max_num_tokens, max_batch_size, max_draft_len) of {2048, 8192} x 1..32 x {0, 5} (128 managers; + max_seq_len 16384 without DSpark, 16398 with it). The owners' views are [B, 64, 288] (2 / 8 / 14, bases + one page apart, so B, B - 1, B - 2 pages) and [B, 128, 288] (20); the pages equal the table. The + largest is B_max itself, 2,606,019 / 868,673 pages, reached at (2048, 31, 0) and (2048, 32, 0): no + max_batch_size in 1..32 gives a larger native COMPRESS view, so the coverage below at B_max holds for + every pool a manager at max_batch_size 1..32 allocates.""" + got = {} + for key in _NATIVE_COMPRESS_PAGES_AXIS: + mgr = _v41_native_manager(*key) + try: + assert mgr.max_batch_size == key[1] and mgr.max_seq_len == (16398 if key[2] else 16384) + views = {o: _native_compress(mgr, o) for o in (2, 8, 14, 20)} + base = views[2][0].data_ptr() + for i, o in enumerate((2, 8, 14)): + assert views[o][0].data_ptr() == base + i * 64 * NATIVE_ROW, (key, o) + got[key] = (tuple(views[o][0].shape[0] for o in (2, 8, 14)), views[20][0].shape[0]) + del views + finally: + mgr.shutdown() + del mgr + _free_engine_pool() + top2 = max(v[0][0] for v in got.values()) + top1 = max(v[1] for v in got.values()) + at = sorted(k for k, v in got.items() if v[0][0] == top2 or v[1] == top1) + print( + f"\n[native COMPRESS pages at 0.5 x device total, 128 configurations] {dict(sorted(got.items()))}\n" + f"[axis maximum] ratio 2 {top2}, ratio 1 {top1}, at {at}" + ) + assert got == _NATIVE_COMPRESS_PAGES_AXIS, { + k: (v, _NATIVE_COMPRESS_PAGES_AXIS[k]) + for k, v in got.items() + if v != _NATIVE_COMPRESS_PAGES_AXIS[k] + } + assert (top2, top1) == (GATES_NATIVE_COMPRESS["ratio2"][0], GATES_NATIVE_COMPRESS["ratio1"][0]) + assert at == [(2048, 31, 0), (2048, 32, 0)] + + +# (storage, max_num_tokens, max_batch_size, DSpark max_draft_len, max_seq_len as the engine passes it) -> +# (ratio-2 pages of owners 2 / 8 / 14, ratio-1 pages of owner 20) at 0.5 x this GB300's total memory: the +# envelope ceilings (the pool scan's maxima over 1920 managers per storage: max_num_tokens 2048..16384, +# max_batch_size 1..256 sampled, DSpark 0 / 5, max_seq_len 16384..1,048,576) and the max_batch_size-256 +# configurations at max_num_tokens 16384. +_ENVELOPE_COMPRESS_PAGES = { + ("bf16", 2048, 48, 0, 1_048_576): ((1_270_272, 1_270_271, 1_270_270), 423_424), + ("bf16", 16384, 256, 0, 16384): ((1_035_264, 1_035_263, 1_035_262), 345_088), + ("bf16", 16384, 256, 5, 16398): ((947_712, 947_711, 947_710), 315_904), + ("native", 2048, 128, 0, 1_048_576): ((3_886_647, 3_886_646, 3_886_645), 1_295_549), + ("native", 16384, 256, 0, 16384): ((2_703_360, 2_703_359, 2_703_358), 901_120), + ("native", 16384, 256, 5, 16398): ((2_408_448, 2_408_447, 2_408_446), 802_816), +} + + +def test_envelope_largest_compress_views_from_the_cache_manager(): + """The envelope COMPRESS ceilings from the manager itself, and the far-end call on each ceiling + manager's own views. + + Each configuration of _ENVELOPE_COMPRESS_PAGES is built at 0.5 x the device's total memory in its storage mode; + the owners' views are asserted (bf16: [B, 64, 512] / [B, 128, 512] with strides (page * 512, 512, 1); + native: [B, 64, 288] / [B, 128, 288] uint8 with code / scale halves; owners 2 / 8 / 14 one page apart) + and their pages pinned. On the two ceiling managers (2048, 48 / 128, 0, 1,048,576), through owner 2's and + owner 14's ratio-2 views and owner 20's ratio-1 view (each buffer filled with the sentinel first), the op + writes rows to both sides of every 32-bit hazard (bf16 element offsets 2^30..2^35; native byte offsets + 2^31..2^36), the first and the last page and the last slot: the rows equal the reference (native: packed + codes + scale bytes, decoding to the op's own core output) and nothing else in the buffer is written. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + total = torch.cuda.mem_get_info()[1] + assert total == _BMAX_DEVICE_TOTAL, ( + f"B_max constants are for a {_BMAX_DEVICE_TOTAL}-byte GB300, got {total}" + ) + got = {} + ceilings = ( + ("bf16",) + ENVELOPE_NATIVE_COMPRESS_KEY[:1] + (48,) + ENVELOPE_NATIVE_COMPRESS_KEY[2:], + ("native",) + ENVELOPE_NATIVE_COMPRESS_KEY, + ) + for key in _ENVELOPE_COMPRESS_PAGES: + storage, max_t, seqs, draft, lm = key + if storage == "bf16": + mgr = _v41_manager(total // 2, draft, seqs, max_num_tokens=max_t, max_seq_len=lm) + else: + mgr = _v41_native_manager(max_t, seqs, draft, max_seq_len=lm) + try: + assert mgr.max_seq_len == lm and mgr.max_batch_size == seqs + if storage == "bf16": + views = {o: mgr.get_buffers(o, AT.COMPRESS) for o in (2, 8, 14, 20)} + for o, v in views.items(): + ps = 64 if o != 20 else 128 + assert tuple(v.shape[1:]) == (ps, W) and v.dtype == torch.bfloat16, ( + o, + v.shape, + v.dtype, + ) + assert v.stride() == (ps * W, W, 1), (o, v.stride()) + base = views[2].data_ptr() + for i, o in enumerate((2, 8, 14)): + assert views[o].data_ptr() == base + i * 64 * W * 2, o + got[key] = (tuple(views[o].shape[0] for o in (2, 8, 14)), views[20].shape[0]) + if key in ceilings: + for o in (2, 14, 20): + pool = views[o] + page_size = pool.shape[1] + num_slots = pool.shape[0] * page_size + whole = views[2] if o != 20 else pool + whole.view(torch.int16).fill_(SENTINEL) + hazard = _hazard_slots(num_slots, W, page_size) + assert hazard[-1] == num_slots - 1 and 2**34 // W in hazard + assert o == 20 or 2**35 // W in hazard + slots = torch.tensor( + hazard + [-1, num_slots], dtype=torch.int64, device=DEV + ) + x = _random_blocks(slots.numel(), seed=12700 + o) + op(x, pool, slots) + _check_engine_pool(whole, pool, slots, _ref_rows(x)[0]) + print( + f"\n[bf16 manager {key}, own owner-{o} view {tuple(pool.shape)}] {len(hazard)} far-end " + f"/ hazard rows == reference, rest of the buffer sentinel" + ) + del pool, whole + del views + else: + views = {o: _native_compress(mgr, o) for o in (2, 8, 14, 20)} + base = views[2][0].data_ptr() + for i, o in enumerate((2, 8, 14)): + assert views[o][0].data_ptr() == base + i * 64 * NATIVE_ROW, o + got[key] = (tuple(views[o][0].shape[0] for o in (2, 8, 14)), views[20][0].shape[0]) + if key in ceilings: + for o in (2, 14, 20): + pool, codes, scales = views[o] + whole = views[2][0] if o != 20 else pool + page_size = pool.shape[1] + num_slots = pool.shape[0] * page_size + whole.fill_(NATIVE_SENTINEL) + hazard = _native_hazard_slots(num_slots, page_size) + assert hazard[-1] == num_slots - 1 and 2**35 // NATIVE_ROW in hazard + assert o == 20 or 2**36 // NATIVE_ROW in hazard + slots = torch.tensor( + hazard + [-1, num_slots], dtype=torch.int64, device=DEV + ) + x = _random_blocks(slots.numel(), seed=38860 + o) + rows, deq = _native_rows(x) + op(x, codes, slots, scales) + flat = pool.view(-1, NATIVE_ROW) + core = _run_core(x) + assert torch.equal(_bits(core), _bits(deq)) + got_rows = flat[slots[: len(hazard)]] + assert torch.equal( + _bits(_decode(got_rows[:, :NATIVE_CODES], got_rows[:, NATIVE_CODES:])), + _bits(core[: len(hazard)]), + ) + _check_native_pool(whole, flat, num_slots, slots, rows) + print( + f"\n[native manager {key}, own owner-{o} view {tuple(pool.shape)}] {len(hazard)} " + f"far-end / hazard rows == reference codes + bytes, decode == core output, rest of " + f"the buffer sentinel" + ) + del pool, codes, scales, whole, flat + del views + finally: + mgr.shutdown() + del mgr + _free_engine_pool() + print(f"\n[envelope COMPRESS pages at 0.5 x device total] {got}") + assert got == _ENVELOPE_COMPRESS_PAGES, { + k: (v, _ENVELOPE_COMPRESS_PAGES[k]) + for k, v in got.items() + if v != _ENVELOPE_COMPRESS_PAGES[k] + } + for storage, table in (("bf16", UNSPLIT_COMPRESS), ("native", UNSPLIT_NATIVE_COMPRESS)): + assert max(v[0][0] for k, v in got.items() if k[0] == storage) == table["ratio2"][0] + assert max(v[1] for k, v in got.items() if k[0] == storage) == table["ratio1"][0] + + +# (storage, PP2 stage, max_num_tokens, max_batch_size, DSpark max_draft_len, max_seq_len) -> {owner: COMPRESS +# pages} of the stage manager built alone at 0.5 x this GB300's total memory: each storage's stage-alone maxima +# over the engine envelope. Stage 0 holds the ratio-2 owners 2 / 8 / 14 (one +# shared ratio-2 buffer, views one page apart), stage 1 the ratio-1 owner 20. +_ENVELOPE_PP2_COMPRESS_PAGES = { + ("bf16", 0, 2048, 96, 0, 1_048_576): {2: 2_118_144, 8: 2_118_143, 14: 2_118_142}, + ("bf16", 1, 2048, 96, 0, 1_048_576): {20: 1_059_584}, + ("native", 0, 2048, 256, 0, 1_048_576): {2: 6_480_780, 8: 6_480_779, 14: 6_480_778}, + ("native", 1, 2048, 256, 0, 1_048_576): {20: 3_245_959}, +} + + +def test_envelope_pp2_stage_compress_views_from_the_cache_manager(): + """The PP2 stage managers' COMPRESS views at their stage-alone maxima, and the far-end call on each + manager's own views. + + Each configuration of _ENVELOPE_PP2_COMPRESS_PAGES is built alone (one process) with the mapping of a rank of + that stage (pp_size 2, tp 2) at 0.5 x the device's total memory: pp_layers asserted, the stage's owners' + views asserted (bf16 [B, 64 | 128, 512] with strides (page * 512, 512, 1); native [B, 64 | 128, 288] + with code / scale halves; owners 2 / 8 / 14 one page apart) and their pages pinned. Through owner 2's + and owner 14's views (stage 0) or owner 20's (stage 1), each buffer filled with the sentinel first, the + op writes rows to both sides of every 32-bit hazard (bf16 element offsets 2^30..2^36; native byte offsets + 2^31..2^36), the first and the last page and the last slot: the rows equal the reference (native: packed + codes + scale bytes, decoding to the op's own core output) and nothing else in the buffer is written. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + total = torch.cuda.mem_get_info()[1] + assert total == _BMAX_DEVICE_TOTAL, ( + f"B_max constants are for a {_BMAX_DEVICE_TOTAL}-byte GB300, got {total}" + ) + got = {} + for key, want in _ENVELOPE_PP2_COMPRESS_PAGES.items(): + storage, stage, max_t, seqs, draft, lm = key + if storage == "bf16": + mgr = _v41_manager( + total // 2, draft, seqs, max_num_tokens=max_t, max_seq_len=lm, pp_stage=stage + ) + else: + mgr = _v41_native_manager(max_t, seqs, draft, max_seq_len=lm, pp_stage=stage) + try: + assert list(mgr.pp_layers) == (list(range(20)) if stage == 0 else list(range(20, 40))) + assert mgr.max_seq_len == lm and mgr.max_batch_size == seqs + owners = tuple(want) + if storage == "bf16": + views = {o: mgr.get_buffers(o, AT.COMPRESS) for o in owners} + for o, v in views.items(): + ps = 64 if o != 20 else 128 + assert tuple(v.shape[1:]) == (ps, W) and v.dtype == torch.bfloat16, ( + o, + v.shape, + v.dtype, + ) + assert v.stride() == (ps * W, W, 1), (o, v.stride()) + if stage == 0: + for i, o in enumerate((2, 8, 14)): + assert views[o].data_ptr() == views[2].data_ptr() + i * 64 * W * 2, o + got[key] = {o: v.shape[0] for o, v in views.items()} + for o in (2, 14) if stage == 0 else (20,): + pool = views[o] + page_size = pool.shape[1] + num_slots = pool.shape[0] * page_size + whole = views[2] if o != 20 else pool + whole.view(torch.int16).fill_(SENTINEL) + hazard = _hazard_slots(num_slots, W, page_size) + assert hazard[-1] == num_slots - 1 and 2**36 // W in hazard + slots = torch.tensor(hazard + [-1, num_slots], dtype=torch.int64, device=DEV) + x = _random_blocks(slots.numel(), seed=21180 + o) + op(x, pool, slots) + _check_engine_pool(whole, pool, slots, _ref_rows(x)[0]) + print( + f"\n[bf16 PP2 stage-{stage} manager {key}, own owner-{o} view {tuple(pool.shape)}] " + f"{len(hazard)} far-end / hazard rows == reference, rest of the buffer sentinel" + ) + del pool, whole + del views + else: + views = {o: _native_compress(mgr, o) for o in owners} + if stage == 0: + for i, o in enumerate((2, 8, 14)): + assert ( + views[o][0].data_ptr() == views[2][0].data_ptr() + i * 64 * NATIVE_ROW + ), o + got[key] = {o: v[0].shape[0] for o, v in views.items()} + for o in (2, 14) if stage == 0 else (20,): + pool, codes, scales = views[o] + whole = views[2][0] if o != 20 else pool + page_size = pool.shape[1] + num_slots = pool.shape[0] * page_size + whole.fill_(NATIVE_SENTINEL) + hazard = _native_hazard_slots(num_slots, page_size) + assert hazard[-1] == num_slots - 1 and 2**36 // NATIVE_ROW in hazard + slots = torch.tensor(hazard + [-1, num_slots], dtype=torch.int64, device=DEV) + x = _random_blocks(slots.numel(), seed=64807 + o) + rows, deq = _native_rows(x) + op(x, codes, slots, scales) + flat = pool.view(-1, NATIVE_ROW) + core = _run_core(x) + assert torch.equal(_bits(core), _bits(deq)) + got_rows = flat[slots[: len(hazard)]] + assert torch.equal( + _bits(_decode(got_rows[:, :NATIVE_CODES], got_rows[:, NATIVE_CODES:])), + _bits(core[: len(hazard)]), + ) + _check_native_pool(whole, flat, num_slots, slots, rows) + print( + f"\n[native PP2 stage-{stage} manager {key}, own owner-{o} view {tuple(pool.shape)}] " + f"{len(hazard)} far-end / hazard rows == reference codes + bytes, decode == core output, " + f"rest of the buffer sentinel" + ) + del pool, codes, scales, whole, flat + del views + finally: + mgr.shutdown() + del mgr + _free_engine_pool() + print(f"\n[PP2 stage COMPRESS pages at 0.5 x device total, stage alone] {got}") + assert got == _ENVELOPE_PP2_COMPRESS_PAGES, { + k: (v, _ENVELOPE_PP2_COMPRESS_PAGES[k]) + for k, v in got.items() + if v != _ENVELOPE_PP2_COMPRESS_PAGES[k] + } + for storage, table in (("bf16", ENGINE_COMPRESS), ("native", NATIVE_COMPRESS)): + assert max(v.get(2, 0) for k, v in got.items() if k[0] == storage) == table["ratio2"][0] + assert max(v.get(20, 0) for k, v in got.items() if k[0] == storage) == table["ratio1"][0] + + +@pytest.mark.parametrize("idx_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("ratio", ["ratio2", "ratio1"]) +def test_native_engine_pool_far_end(ratio, idx_dtype): + """Decode-like on a B_max native pool with a guard page: hazard rows; -1 and past-the-end skipped.""" + pages, page_size, _ = NATIVE_COMPRESS[ratio] + buf, codes, scales = _native_engine_pool(pages, page_size) + num_slots = pages * page_size + hazard = _native_hazard_slots(num_slots, page_size) + slots = torch.tensor( + hazard + [-1, num_slots, num_slots + page_size - 1], dtype=torch.int64, device=DEV + ) + x = _random_blocks(slots.numel(), seed=26061 + page_size) + op(x, codes, slots.to(idx_dtype), scales) + _check_native_pool(buf, buf.view(-1, NATIVE_ROW), num_slots, slots, _native_rows(x)[0]) + del buf, codes, scales + _free_engine_pool() + + +@pytest.mark.parametrize("ratio", ["ratio2", "ratio1"]) +def test_native_engine_pool_prefill_spread(ratio): + """Prefill-like, n = 8192 on a B_max native pool: the last page in order, every hazard slot, the rest + spread uniformly, 1/16 of the rows -1.""" + n = _MAX_ROWS + pages, page_size, _ = NATIVE_COMPRESS[ratio] + buf, codes, scales = _native_engine_pool(pages, page_size) + num_slots = pages * page_size + last_page = torch.arange(num_slots - page_size, num_slots, device=DEV) + hazard = torch.tensor(_native_hazard_slots(num_slots, page_size), device=DEV) + fixed = torch.unique(torch.cat([last_page, hazard])) + g = torch.Generator(device=DEV).manual_seed(26062 + page_size) + rest = torch.randint(0, num_slots, (3 * n,), device=DEV, generator=g).unique() + rest = rest[~torch.isin(rest, fixed)] + rest = rest[torch.randperm(rest.numel(), device=DEV, generator=g)][: n - fixed.numel()] + slots = torch.cat([last_page, fixed[~torch.isin(fixed, last_page)], rest]).to(torch.int64) + assert slots.numel() == n and torch.unique(slots).numel() == n + slots[torch.randperm(n, device=DEV, generator=g)[: n // 16]] = -1 + x = _sweep_rows(n, seed=26062 + page_size) + op(x, codes, slots.to(torch.int32), scales) + _check_native_pool(buf, buf.view(-1, NATIVE_ROW), num_slots, slots, _native_rows(x)[0]) + del buf, codes, scales + _free_engine_pool() + + +def test_native_engine_pool_harness_sees_a_wrapped_offset(): + """Control: rows written where a mod-2^32 byte offset would put them fail both halves of the check.""" + pages, page_size, _ = NATIVE_COMPRESS["ratio2"] + buf, codes, scales = _native_engine_pool(pages, page_size) + num_slots = pages * page_size + flat = buf.view(-1, NATIVE_ROW) + true = torch.tensor( + [s for s in _native_hazard_slots(num_slots, page_size) if s * NATIVE_ROW >= 2**32], + dtype=torch.int64, + device=DEV, + ) + wrapped = (true * NATIVE_ROW % 2**32) // NATIVE_ROW + assert true.numel() >= 9 and not torch.isin(wrapped, true).any() + x = _random_blocks(true.numel(), seed=26063) + rows = _native_rows(x)[0] + op(x, codes, wrapped, scales) + with pytest.raises(AssertionError, match="bytes in the named slots differ"): + _check_native_pool(buf, flat, num_slots, true, rows) + flat[true] = NATIVE_SENTINEL + with pytest.raises(AssertionError, match="outside the named slots were written"): + _check_native_pool( + buf, flat, num_slots, torch.tensor([], dtype=torch.int64, device=DEV), rows[:0] + ) + del buf, codes, scales, flat + _free_engine_pool() + + +@pytest.mark.parametrize("n", [32, 8192, 16384]) +@pytest.mark.parametrize("ratio", ["ratio2", "ratio1"]) +def test_native_pool_size_invariance(ratio, n): + """The same rows into the same slots through views codes[:B] / scales[:B] of the live PP2 stage + manager's envelope B_max view (owner 2 of stage 0, or owner 20 of stage 1, each built alone at + (2048, 256, 0, 1,048,576)): the same bytes at every B, equal to the reference, nothing else in the buffer + written; then a separate allocation of exactly the smallest size. + + B = the smallest view holding the slots, the page counts whose byte count first passes 2^31, 2^32, + 2^33, 2^34, 2^35 and 2^36 (where below B_max), a mid size, the gates' largest, the max_batch_size 1..32 + bound (GATES_NATIVE_COMPRESS), the unsplit manager's envelope bound (UNSPLIT_NATIVE_COMPRESS) and the + stage's B_max. + """ + pages, page_size, gate_max = NATIVE_COMPRESS[ratio] + owner = 20 if ratio == "ratio1" else 2 + page_bytes = page_size * NATIVE_ROW + stage, (max_t, seqs, draft, lm) = ENVELOPE_PP2_NATIVE_COMPRESS[ratio] + mgr = _v41_native_manager(max_t, seqs, draft, max_seq_len=lm, pp_stage=stage) + try: + assert list(mgr.pp_layers) == (list(range(20)) if stage == 0 else list(range(20, 40))) + pool, codes, scales = _native_compress(mgr, owner) + assert pool.shape[0] == pages + pool.fill_(NATIVE_SENTINEL) + flat = pool.view(-1, NATIVE_ROW) + small, slots = _invariance_slots(n, page_size, seed=26064 + n + page_size) + x = ( + _sweep_rows(n, seed=26064 + n + page_size) + if n >= 200 + else _random_blocks(n, seed=26064 + page_size) + ) + rows = _native_rows(x)[0] + live = slots[slots >= 0] + sizes = sorted( + { + small, + pages // 2 + 1, + gate_max, + GATES_NATIVE_COMPRESS[ratio][0], + UNSPLIT_NATIVE_COMPRESS[ratio][0], + pages, + } + | { + lim // page_bytes + 1 + for lim in (2**31, 2**32, 2**33, 2**34, 2**35, 2**36) + if lim // page_bytes + 1 < pages + } + ) + first = None + for b in sizes: + op(x, codes[:b], slots, scales[:b]) + torch.cuda.synchronize() + got = flat[live].clone() + _check_native_pool(pool, flat, b * page_size, slots, rows) + first = got if first is None else first + assert torch.equal(got, first), ( + f"B={b}: {int((got != first).sum())} bytes differ from B={sizes[0]}" + ) + finally: + mgr.shutdown() + del mgr, pool, codes, scales, flat + _free_engine_pool() + buf, codes, scales = _native_engine_pool(small, page_size) + op(x, codes, slots, scales) + torch.cuda.synchronize() + flat = buf.view(-1, NATIVE_ROW) + got = flat[live].clone() + _check_native_pool(buf, flat, small * page_size, slots, rows) + assert torch.equal(got, first), "a separate exact-size native view differs" + print( + f"\n[native {ratio} n={n}] views B = {sizes} of the manager's owner-{owner} view and a separate " + f"[{small}, {page_size}, 288]: {live.numel()} rows bitwise identical and == reference, nothing else written" + ) + del buf, codes, scales, flat + _free_engine_pool() + + +@pytest.mark.parametrize("ratio", list(_SWEEP_GEOMETRY)) +def test_native_every_row_count_1_to_16384(ratio): + """Every prefix (x[:r], slots[:r]), r = 1..16384, of the target's native call writes exactly the first r + rows of the 16384-row call, and nothing else. + + A native pool of 20,480 slots in that ratio's page geometry plus a guard + page, all bytes sentinel. The 16384-row call is held to the reference and + the sentinel scan. Then per r: the call, the named live slots against the + reference rows, the count of non-sentinel bytes against the reference's + own count for the live rows, and the named slots restored to the sentinel, + so each call's writes are checked on their own. One sync at the end. The + probe: the sweep compiles no variant beyond the 8192-row call's. + """ + t = _MAX_ROWS + page_size = _SWEEP_GEOMETRY[ratio] + buf, codes, scales = _native_engine_pool(_SWEEP_SLOTS // page_size, page_size) + flat = buf.view(-1, NATIVE_ROW) + slots = _sweep_slots(t, _SWEEP_SLOTS, seed=26065 + page_size) + x = _sweep_rows(t, seed=26065 + page_size) + rows = _native_rows(x)[0] + kernel = _nv_op._nvfp4_fake_quantize_kernel + n0 = _jit_variants(kernel) + op(x, codes, slots, scales) + n1 = _jit_variants(kernel) + _check_native_pool(buf, flat, _SWEEP_SLOTS, slots, rows) + + live = slots >= 0 + live_pos = torch.nonzero(live).squeeze(1) + live_slots = slots[live_pos] + rows_live = rows[live_pos] + cnt = torch.cumsum(live.to(torch.int64), 0).tolist() + # A packed code byte can equal the sentinel, so the expected count is the reference rows' own. + expect = torch.cumsum((rows_live != NATIVE_SENTINEL).sum(1), 0) + bad_rows = torch.zeros((), dtype=torch.int64, device=DEV) + bad_count = torch.zeros((), dtype=torch.int64, device=DEV) + for r in range(1, t + 1): + op(x[:r], codes, slots[:r], scales) + k = cnt[r - 1] + bad_rows += (flat[live_slots[:k]] != rows_live[:k]).sum() + bad_count += ((buf != NATIVE_SENTINEL).sum() - (expect[k - 1] if k else 0)).abs() + flat[live_slots[:k]] = NATIVE_SENTINEL + n2 = _jit_variants(kernel) + print( + f"\n[native {ratio} rows 1..{t}, page {page_size}] live rows {cnt[-1]}; Triton variants: {n0} before the " + f"{t}-row call, {n1} after, {n2} after the sweep; byte mismatches {int(bad_rows)}, " + f"non-sentinel count error {int(bad_count)}" + ) + assert int(bad_rows) == 0 and int(bad_count) == 0 + assert n1 - n0 <= 1 and n2 == n1, f"row counts compiled new variants: {n0} -> {n1} -> {n2}" + # Control: the same checks see a read one row off and a write one slot off. + for r in (1, 2, 7, 64): + for name, xa, sa in ( + ("row +1", x[1 : r + 1], live_slots[:r]), + ("slot +1", x[:r], live_slots[1 : r + 1]), + ): + op(xa, codes, sa, scales) + n_bad = int((flat[live_slots[:r]] != _native_rows(x[:r])[0]).sum()) + print(f"[native control {name}] r={r}: {n_bad} of {r * NATIVE_ROW} bytes differ") + assert n_bad > r * NATIVE_ROW // 4 + buf.fill_(NATIVE_SENTINEL) + del buf, codes, scales, flat + _free_engine_pool() + + +@pytest.mark.parametrize("ratio", ["ratio1", "ratio2"]) +def test_native_cuda_graph_every_graph_batch_on_the_bmax_pool(ratio): + """Per n of that ratio's envelope capture set (ratio 1: G and 6G, 82 row counts to 1536; ratio 2: G and + 3G, 77 row counts to 768), native form on the envelope B_max native pool (3,886,647 / 1,295,549 + pages): an eager warm-up, the capture, then 2 replays on new rows and new far-end slots; every replay + writes exactly the reference into the named slots and nothing else, and equals the same call run + eagerly.""" + pages, page_size, _ = NATIVE_COMPRESS[ratio] + ns = _GRAPH_N[ratio] + buf, codes, scales = _native_engine_pool(pages, page_size) + num_slots = pages * page_size + flat = buf.view(-1, NATIVE_ROW) + g = torch.Generator(device=DEV).manual_seed(26066 + page_size) + hazard = torch.tensor(_native_hazard_slots(num_slots, page_size), device=DEV) + last = num_slots - 1 + table = torch.full( + (_MAX_ROWS,), -1, dtype=torch.int64, device=DEV + ) # like v41_compress_slot_mapping[r] + + def far_slots(n): + near = last - 1 - torch.arange(4 * page_size - 1, device=DEV) + cand = torch.unique(torch.cat([hazard[hazard != last], near])) + if cand.numel() < n - 1: # large graph batches: the rest spread uniformly over the pool + spread = torch.randint(0, last, (4 * n,), device=DEV, generator=g).unique() + cand = torch.cat([cand, spread[~torch.isin(spread, cand)]]) + cand = cand[torch.randperm(cand.numel(), device=DEV, generator=g)] + s = torch.cat([torch.tensor([last], device=DEV), cand[: n - 1]]).to(torch.int64) + assert torch.unique(s).numel() == n + if n >= 2: + s[1 + int(torch.randint(0, n - 1, (1,), device=DEV, generator=g))] = -1 + return s + + for n in ns: + x = torch.empty(n, W, dtype=torch.bfloat16, device=DEV) + slots = table[:n] + x.copy_(_random_blocks(n, seed=26066 + n)) + slots.copy_(far_slots(n)) + op(x, codes, slots, scales) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(x, codes, slots, scales) + buf.fill_(NATIVE_SENTINEL) + for rep in (1, 2): + new_x = _random_blocks(n, seed=26066 + 1000 * rep + n) + new_slots = far_slots(n) + x.copy_(new_x) + slots.copy_(new_slots) + graph.replay() + torch.cuda.synchronize() + live = new_slots[new_slots >= 0] + replayed = flat[live].clone() + rows = _native_rows(new_x)[0] + _check_native_pool(buf, flat, num_slots, new_slots, rows) + op(x, codes, slots, scales) + torch.cuda.synchronize() + eager = flat[live].clone() + _check_native_pool(buf, flat, num_slots, new_slots, rows) + assert torch.equal(replayed, eager), ( + f"n={n} replay {rep}: {int((replayed != eager).sum())} bytes" + ) + del graph + print( + f"\n[native {ratio} graphs on [{pages}, {page_size}, 288]] n = {list(ns)}: every replay bitwise == eager " + f"== reference" + ) + del buf, codes, scales, flat + _free_engine_pool() + + +# ─── native form over random multi-request steps on the live manager ─── + + +def _random_step(rng, max_tokens=8192, max_requests=32, max_len=16384): + """([(cached, new)] contexts first (random order) then generation requests, num_contexts, gen_len).""" + k = rng.randint(1, max_requests) + gen_len = rng.choice([1, 6]) + g = rng.randint(0, k) if rng.random() < 0.7 else 0 + c = k - g + budget = max_tokens - g * gen_len + t_ctx = rng.randint(c, budget) if c else 0 + cuts = sorted(rng.sample(range(1, t_ctx), c - 1)) if c > 1 else [] + news = [b - a for a, b in zip([0] + cuts, cuts + [t_ctx])] if c else [] + reqs = [(0 if rng.random() < 0.4 else rng.randint(0, max_len - n), n) for n in news] + rng.shuffle(reqs) + reqs += [(rng.randint(0, max_len - gen_len), gen_len) for _ in range(g)] + return reqs, c, gen_len + + +_EDGE_STEPS = [ + ([(0, 8192)], 1, 1), + ([(0, 1)], 1, 1), + ([(16383, 1)], 1, 1), + ([(0, 256)] * 32, 32, 1), + ([(16384 - 256, 256)] * 32, 32, 1), + ([(16383, 1)] * 32, 0, 1), + ([(16378, 6)] * 32, 0, 6), + ([(513, 7000)] + [(16000, 6)] * 31, 1, 6), + ([(0, 127), (0, 128), (0, 129), (127, 1), (128, 1), (8000, 192)], 6, 1), +] +_NUM_RANDOM_STEPS = 200 + + +def _run_step(mgr, reqs, num_contexts, gen_len, uid, max_num_tokens=8192): + """Allocate the step's requests on the manager and build its DeepSeek-V4.1 attention metadata. + + Each request's first token is unique over the test (uid + request index), so + block reuse never shares a page between requests of different steps. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + DeepseekV4TrtllmAttentionMetadata, + ) + from tensorrt_llm._torch.metadata import KVCacheParams + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from tensorrt_llm.bindings import SamplingConfig + from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig + from tensorrt_llm.mapping import Mapping + + requests = [] + for rid, (c, n) in enumerate(reqs): + r = LlmRequest( + request_id=rid, + max_new_tokens=64, + input_tokens=[uid + rid] + list(range(1, c + n)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + assert mgr.prepare_context(r) and mgr.resize_context(r, r.context_chunk_size), (rid, c, n) + requests.append(r) + md = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor([n for _, n in reqs], dtype=torch.int), + request_ids=[r.py_request_id for r in requests], + max_num_requests=mgr.max_batch_size, + num_contexts=num_contexts, + prompt_lens=[c + n for c, n in reqs], + max_num_tokens=max_num_tokens, + kv_cache_manager=mgr, + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=[c for c, _ in reqs] + ), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + sparse_attention_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ), + ) + if num_contexts < len(reqs) and gen_len > 1: + md.update_spec_dec_param( + batch_size=len(reqs), + is_spec_decoding_enabled=False, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, + max_draft_len=gen_len - 1, + max_total_draft_tokens=gen_len - 1, + ) + md.prepare() + return requests, md + + +def test_native_writes_over_random_multi_request_steps(): + """On the live accuracy-gate manager (8192, 32, DSpark 5; native COMPRESS 1,895,082 / 631,694). + + Per step (9 edge steps + 200 seeded random partitions of T <= 8192 into + k <= 32 requests, contexts with random cached prefixes, decode or DSpark + verify generation requests): the requests are allocated on the manager and + the metadata is prepared; per ratio, every owner (2, 8, 14 at ratio 2, 20 + at ratio 1) gets random latents written by the op's native form at + md.v41_compress_slot_mapping[r][:n] (reserved generation rows -1) through + get_v41_native_kv_views(o, COMPRESS): the slots hold the reference's + packed codes and scale bytes, decoding to the op's own core output of the + same rows; after the step every other byte of both ratio buffers is still + the sentinel. The op compiles no variant past its first native call. + """ + mgr = _v41_native_manager(8192, 32, 5) + rng = random.Random(26067) + steps = list(_EDGE_STEPS) + [_random_step(rng) for _ in range(_NUM_RANDOM_STEPS)] + kernel = _nv_op._nvfp4_fake_quantize_kernel + try: + views = {o: _native_compress(mgr, o) for o in (2, 8, 14, 20)} + whole = {"ratio2": views[2][0], "ratio1": views[20][0]} + for w in whole.values(): + w.fill_(NATIVE_SENTINEL) + core_buf = torch.empty(_MAX_ROWS, W, dtype=torch.bfloat16, device=DEV) + n_variants = None + stats = dict(steps=0, rows=0, skipped=0, max_t=0, max_k=0, max_n={1: 0, 2: 0}) + for i, (reqs, nctx, gen_len) in enumerate(steps): + requests, md = _run_step(mgr, reqs, nctx, gen_len, uid=300_000 + 32 * i) + try: + for ratio, r in (("ratio2", 2), ("ratio1", 1)): + n = md.num_total_compressed_tokens[r] + if n == 0: + continue + slots = md.v41_compress_slot_mapping[r][:n] + s64 = slots.to(torch.int64) + live = s64 >= 0 + for o in _OWNERS[ratio]: + pool, codes, scales = views[o] + flat = pool.view(-1, NATIVE_ROW) + x = _random_blocks(n, seed=400_000 + 4 * i + o) + op(x, codes, slots, scales) + rows, deq = _native_rows(x) + # A fixed 8192-slot out: in the identity form the slot count is an integer + # argument, and Triton specializes it (== 1, % 16 == 0) per call. + op(x, core_buf) + core = core_buf[:n] + if n_variants is None: + n_variants = _jit_variants(kernel) + assert int(s64.max()) < pool.shape[0] * pool.shape[1] + got = flat[s64[live]] + assert torch.equal(got, rows[live]), ( + f"step {i} owner {o}: {int((got != rows[live]).sum())}" + ) + assert torch.equal( + _bits(_decode(got[:, :NATIVE_CODES], got[:, NATIVE_CODES:])), + _bits(core[live]), + ) + assert torch.equal(_bits(core), _bits(deq)) + flat[s64[live]] = NATIVE_SENTINEL + stats["rows"] += int(live.sum()) + stats["skipped"] += int((~live).sum()) + stats["max_n"][r] = max(stats["max_n"][r], n) + for ratio, w in whole.items(): + n_out = int(_count_not_sentinel(w)) + assert n_out == 0, ( + f"step {i} {ratio}: {n_out} bytes outside the step's slots were written" + ) + stats["steps"] += 1 + stats["max_t"] = max(stats["max_t"], sum(n for _, n in reqs)) + stats["max_k"] = max(stats["max_k"], len(reqs)) + finally: + for r_ in requests: + mgr.free_resources(r_) + assert _jit_variants(kernel) == n_variants, "a step shape compiled a new variant" + print( + f"\n[native multi-request steps on the live manager, COMPRESS {tuple(views[2][0].shape)} / " + f"{tuple(views[20][0].shape)}] {stats}; one compiled variant" + ) + assert stats["steps"] == len(_EDGE_STEPS) + _NUM_RANDOM_STEPS and stats["max_t"] == 8192 + assert stats["max_k"] == 32 and stats["skipped"] > 0 + finally: + mgr.shutdown() + del mgr, views, whole + _free_engine_pool() + + +# The envelope's step axes: T <= max_num_tokens 16384, k <= max_batch_size 256 requests per step. A step's KV +# tokens (cached + new, over its requests) are capped at _ENVELOPE_STEP_TOKENS: a native manager at (16384, 256, +# DSpark 5, max_seq_len 32,782) held 22 requests of 32,768 tokens and refused the 23rd, so 256 requests cannot +# all sit at the far end of max_seq_len at once. +_ENVELOPE_STEP_TOKENS = 560_000 +_ENVELOPE_STEP_L = 32768 + + +def _envelope_random_step(rng, max_tokens, max_requests, seq_len): + """_random_step at the envelope's step axes with DSpark: k in 1..max_requests, contexts end at <= seq_len, + decode requests (next_n 1) at <= seq_len, DSpark-verify requests (next_n 6) at <= seq_len + 14.""" + k = rng.randint(1, max_requests) + gen_len = rng.choice([1, 6]) + g = rng.randint(0, k) if rng.random() < 0.7 else 0 + c = k - g + budget = max_tokens - g * gen_len + t_ctx = rng.randint(c, budget) if c else 0 + cuts = sorted(rng.sample(range(1, t_ctx), c - 1)) if c > 1 else [] + news = [b - a for a, b in zip([0] + cuts, cuts + [t_ctx])] if c else [] + reqs = [(0 if rng.random() < 0.4 else rng.randint(0, seq_len - n), n) for n in news] + rng.shuffle(reqs) + end = seq_len + 14 if gen_len == 6 else seq_len + reqs += [(rng.randint(0, end - gen_len), gen_len) for _ in range(g)] + return reqs, c, gen_len + + +def _envelope_fit(step, rng): + """The step with its KV tokens capped at _ENVELOPE_STEP_TOKENS: in a random order, a request keeps its cached + prefix while the rest can still fit, else gets a random prefix below 1024. Order and new tokens kept.""" + reqs, nctx, gen_len = step + if sum(a + n for a, n in reqs) <= _ENVELOPE_STEP_TOKENS: + return step + out, used = list(reqs), 0 + reserve = sum(n + 1024 for _, n in reqs) + order = list(range(len(reqs))) + rng.shuffle(order) + for i in order: + a, n = reqs[i] + reserve -= n + 1024 + if used + a + n + reserve > _ENVELOPE_STEP_TOKENS: + a = rng.randint(0, 1023) + out[i] = (a, n) + used += a + n + assert used <= _ENVELOPE_STEP_TOKENS + return out, nctx, gen_len + + +def _envelope_edge_steps(max_t, seqs, L): + """A full max_t step from position 0 and one ending at L; seqs decode requests (8 at L - 1, the rest below + position 1152); seqs DSpark verifies (8 ending at L + 13, the rest below 1130); a context filling the step + beside seqs - 1 verifies (7 at the far end).""" + lm, rest = L + 14, seqs - 8 + steps = [ + ([(0, max_t)], 1, 1), + ([(L - max_t, max_t)], 1, 1), + ([(L - 1, 1)] * 8 + [(127 + (61 * i) % 1024, 1) for i in range(rest)], 0, 1), + ([(lm - 6, 6)] * 8 + [(100 + (37 * i) % 1024, 6) for i in range(rest)], 0, 6), + ( + [(0, max_t - 6 * (seqs - 1))] + + [(lm - 6, 6)] * 7 + + [(100 + (37 * i) % 1024, 6) for i in range(rest)], + 1, + 6, + ), + ] + for reqs, _, _ in steps: + assert sum(a + b for a, b in reqs) <= _ENVELOPE_STEP_TOKENS + return steps + + +def test_envelope_native_writes_over_random_steps_256_requests(): + """At the envelope's step axes, on the live native manager of max_num_tokens 16384, + max_batch_size 256, DSpark 5, max_seq_len 32768 + 14: 5 edge steps and 200 seeded random partitions of + T <= 16384 into k <= 256 requests (contexts first with random cached prefixes, decode at positions + < 32,768 or DSpark verify at < 32,782), each step's KV tokens capped (_envelope_fit). Per step, exactly as + test_native_writes_over_random_multi_request_steps: per ratio, every owner (2, 8, 14 at ratio 2, 20 at + ratio 1) gets random latents written by the op's native form at md.v41_compress_slot_mapping[r][:n] + (reserved generation rows -1) through get_v41_native_kv_views(o, COMPRESS); the slots hold the + reference's packed codes and scale bytes, decoding to the op's own core output of the same rows, and + after the step every other byte of both ratio buffers is still the sentinel; no compiled variant past the + first native call.""" + max_t, seqs, L = 16384, 256, _ENVELOPE_STEP_L + rng = random.Random(26_068 + seqs + L) + edges = _envelope_edge_steps(max_t, seqs, L) + steps = edges + [ + _envelope_fit(_envelope_random_step(rng, max_t, seqs, L), rng) + for _ in range(_NUM_RANDOM_STEPS) + ] + mgr = _v41_native_manager(max_t, seqs, 5, max_seq_len=L + 14) + kernel = _nv_op._nvfp4_fake_quantize_kernel + try: + assert mgr.max_seq_len == L + 14 and mgr.max_batch_size == seqs + views = {o: _native_compress(mgr, o) for o in (2, 8, 14, 20)} + whole = {"ratio2": views[2][0], "ratio1": views[20][0]} + for w in whole.values(): + w.fill_(NATIVE_SENTINEL) + core_buf = torch.empty(max_t, W, dtype=torch.bfloat16, device=DEV) + n_variants = None + stats = dict(steps=0, rows=0, skipped=0, max_t=0, max_k=0, max_pos=0, max_n={1: 0, 2: 0}) + for i, (reqs, nctx, gen_len) in enumerate(steps): + requests, md = _run_step( + mgr, reqs, nctx, gen_len, uid=6_000_000 + 256 * i, max_num_tokens=max_t + ) + try: + for ratio, r in (("ratio2", 2), ("ratio1", 1)): + n = md.num_total_compressed_tokens[r] + if n == 0: + continue + slots = md.v41_compress_slot_mapping[r][:n] + s64 = slots.to(torch.int64) + live = s64 >= 0 + for o in _OWNERS[ratio]: + pool, codes, scales = views[o] + flat = pool.view(-1, NATIVE_ROW) + x = _random_blocks(n, seed=600_000 + 4 * i + o) + op(x, codes, slots, scales) + rows, deq = _native_rows(x) + op(x, core_buf) # a fixed 16384-slot out (one identity variant) + core = core_buf[:n] + if n_variants is None: + n_variants = _jit_variants(kernel) + assert int(s64.max()) < pool.shape[0] * pool.shape[1] + got = flat[s64[live]] + assert torch.equal(got, rows[live]), ( + f"step {i} owner {o}: {int((got != rows[live]).sum())}" + ) + assert torch.equal( + _bits(_decode(got[:, :NATIVE_CODES], got[:, NATIVE_CODES:])), + _bits(core[live]), + ) + assert torch.equal(_bits(core), _bits(deq)) + flat[s64[live]] = NATIVE_SENTINEL + stats["rows"] += int(live.sum()) + stats["skipped"] += int((~live).sum()) + stats["max_n"][r] = max(stats["max_n"][r], n) + for ratio, w in whole.items(): + n_out = int(_count_not_sentinel(w)) + assert n_out == 0, ( + f"step {i} {ratio}: {n_out} bytes outside the step's slots were written" + ) + stats["steps"] += 1 + stats["max_t"] = max(stats["max_t"], sum(n for _, n in reqs)) + stats["max_k"] = max(stats["max_k"], len(reqs)) + stats["max_pos"] = max(stats["max_pos"], max(c + n - 1 for c, n in reqs)) + finally: + for r_ in requests: + mgr.free_resources(r_) + assert _jit_variants(kernel) == n_variants, "a step shape compiled a new variant" + print( + f"\n[envelope native steps, ({max_t}, {seqs}, 5) at max_seq_len {L + 14}, COMPRESS " + f"{tuple(views[2][0].shape)} / {tuple(views[20][0].shape)}] {len(edges)} edge + {_NUM_RANDOM_STEPS} " + f"random steps: {stats}; one compiled variant" + ) + assert stats["steps"] == len(steps) and stats["max_t"] == max_t and stats["max_k"] == seqs + assert stats["max_pos"] == L + 13 and stats["skipped"] > 0 and stats["max_n"][1] == max_t + finally: + mgr.shutdown() + del mgr, views, whole + _free_engine_pool() + + +def test_native_variant_map(): + """Which calls compile a new Triton variant: the native form into a manager-geometry view compiles one + variant whatever n (rows are grid axis 0; the view's slot count, a multiple of 64, is fixed); the + identity (core) form passes the row count as the slot count, which Triton specializes (== 1, + % 16 == 0, other), so it compiles one variant per class and none for a second member of a class.""" + kernel = _nv_op._nvfp4_fake_quantize_kernel + added = {} + for ratio, page_size in (("ratio2", 64), ("ratio1", 128)): + buf, codes, scales = _native_engine_pool(10240 // page_size, page_size) + for n in (8192, 1, 2, 16, 17, 4096, 8191): + slots = torch.randperm(10240, device=DEV)[:n] + n0 = _jit_variants(kernel) + op(_random_blocks(n, seed=n), codes, slots, scales) + added[(f"native[{ratio}]", n)] = _jit_variants(kernel) - n0 + del buf, codes, scales + classes = {"1": (1, 1), "16k": (16, 32, 48, 8192), "other": (2, 17, 33, 8191)} + for name, ns in classes.items(): + for i, n in enumerate(ns): + n0 = _jit_variants(kernel) + op(_random_blocks(n, seed=n), torch.empty(n, W, dtype=torch.bfloat16, device=DEV)) + added[(f"identity[{name}]#{i}", n)] = _jit_variants(kernel) - n0 + print(f"\n[variant map] new variants per call: {added}") + assert all(v == 0 for (form, n), v in added.items() if form.startswith("native") and n != 8192) + for name, ns in classes.items(): + assert all(added[(f"identity[{name}]#{i}", n)] == 0 for i, n in enumerate(ns) if i), name diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_nvfp4_kv_to_bf16.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_nvfp4_kv_to_bf16.py new file mode 100644 index 000000000000..e16930cdb7d5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_nvfp4_kv_to_bf16.py @@ -0,0 +1,1929 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the nvfp4_kv_to_bf16 entry: the bitwise twin of FlashInfer's nvfp4_kv_dequantize_paged. + +The pool is DeepSeek-V4.1's native compressed-KV layout: one 288-byte row per entry, 256 packed E2M1 +bytes (even element in the low nibble) then 32 E4M3 scale bytes, read as two strided views. Every gate +is bit-exact on bf16 bit patterns, against BOTH FlashInfer's `nvfp4_kv_dequantize_paged` on the same +inputs (`_twin`, called as the DeepSeek-V4.1 target's FlashInfer form calls it) and the native-torch E2M1 +lookup-table reference (`op.reference`); entries the op must not write are checked against a NaN sentinel +that both ops must leave in place. + +Shapes (`op.CELLS`): the target's per-query (top-k) form -- page_size 1 over the flattened pool and an +[R, 512] row table with -1 padding, out-of-range ids and short rows, R from 1 to 8192 -- and its +per-request form -- page_size 64 / 128 page tables of whole contexts up to 16384 entries, table holes, +request lengths 0 / 1 / page +- 1 / L, the compact [B, 132 * 128 / r, 1, 512] views and the engine's +pool sizes. This file keeps what a cell cannot carry: the quantizer round trip (-0.0 kept), a global +scale != 1, CUDA-graph capture and replay, wrong variants that must fail, the op's own rejections, and +every case on the live DeepSeek-V4.1 cache manager's native views. +""" + +import pytest +import torch +from flashinfer import nvfp4_kv_dequantize_paged + +import tensorrt_llm._torch.custom_ops.nvfp4_fake_quantize # noqa: F401 -- registers torch.ops.trtllm.nvfp4_fake_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.quantization.nvfp4_kv_to_bf16 import ( # noqa: E501 + nvfp4_kv_to_bf16 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "nvfp4_kv_to_bf16 requires a CUDA device" + +DEV = "cuda" +W = 512 +BLOCK = 16 +CODE_BYTES = W // 2 +ROW = CODE_BYTES + W // BLOCK # 288 +NUM_SLOTS = 16384 +E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +SENTINEL = -1 # int16 view of a bf16 NaN (0xFFFF) +_TABLE_DTYPE = {"int32": torch.int32, "int64": torch.int64} + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _random_pool(page_size: int, seed: int, slots: int = NUM_SLOTS) -> torch.Tensor: + g = torch.Generator(device=DEV).manual_seed(seed) + codes = torch.randint(0, 256, (slots, CODE_BYTES), generator=g, device=DEV, dtype=torch.int32) + scales = torch.randint(0, 0x7F, (slots, W // BLOCK), generator=g, device=DEV, dtype=torch.int32) + flat = torch.cat([codes, scales], dim=1).to(torch.uint8) + return flat.view(slots // page_size, page_size, ROW) + + +def _page_views(pool: torch.Tensor): + return pool[..., :CODE_BYTES].unsqueeze(2), pool[..., CODE_BYTES:].unsqueeze(2) + + +def _row_views(pool: torch.Tensor): + flat = pool.view(-1, ROW) + return flat[:, :CODE_BYTES][:, None, None], flat[:, CODE_BYTES:][:, None, None] + + +def _gs(v: float = 1.0) -> torch.Tensor: + return torch.full((1,), v, dtype=torch.float32, device=DEV) + + +def _sentinel_out(*shape) -> torch.Tensor: + return torch.full(shape, SENTINEL, dtype=torch.int16, device=DEV).view(torch.bfloat16) + + +def _ref_rows(flat: torch.Tensor, rows: torch.Tensor, global_scale: float = 1.0) -> torch.Tensor: + """bf16 [n, W] of pool rows `rows` (int64 [n], all inside the pool) of the flat [S, 288] pool: + `op.reference` of a one-request row gather over pages of one row.""" + n = rows.numel() + out = torch.empty(1, n, 1, W, dtype=torch.bfloat16, device=flat.device) + ref = op.reference( + flat[:, :CODE_BYTES][:, None, None], + flat[:, CODE_BYTES:][:, None, None], + rows.reshape(1, n), + torch.full((1,), n, dtype=torch.int32, device=flat.device), + _gs(global_scale), + out, + ) + return ref.view(n, W) + + +def _twin(kv_cache, kv_cache_sf, block_tables, seq_lens, global_scale, out): + """FlashInfer's paged NVFP4 dequantize on the same operands, one call made the way the target's + FlashInfer form makes it: the latent has no V, so a 16-element view of the same rows (8 code bytes, + 1 scale byte) is the V operand, into a scratch that is discarded. Writes `out`, returns it.""" + v_scratch = torch.empty(out.shape[:-1] + (16,), dtype=out.dtype, device=out.device) + nvfp4_kv_dequantize_paged( + (kv_cache, kv_cache[..., :8]), + (kv_cache_sf, kv_cache_sf[..., :1]), + block_tables, + seq_lens, + global_scale, + global_scale, + out, + v_scratch, + "NHD", + ) + return out + + +def _both(kv, sf, table, lens, gs, shape): + """Run the op and the FlashInfer twin on identical sentinel-filled outputs.""" + a, b = _sentinel_out(*shape), _sentinel_out(*shape) + op(kv, sf, table, lens, gs, a) + _twin(kv, sf, table, lens, gs, b) + torch.cuda.synchronize() + return a, b + + +def _row_table(R: int, K: int, seed: int, slots: int, table_dtype=torch.int32): + g = torch.Generator(device=DEV).manual_seed(seed) + tab = torch.randint(0, slots, (R, K), generator=g, device=DEV, dtype=torch.int64) + if R > 1: + tab[1, K // 2 :] = -1 + tab[0, 3] = -1 + tab[0, 5] = slots + tab[0, 7] = slots + 1000 + sl = torch.full((R,), K, dtype=torch.int32, device=DEV) + if R > 2: + sl[2] = K - 37 + valid = (tab >= 0) & (tab < slots) & (torch.arange(K, device=DEV)[None] < sl[:, None].long()) + return tab.to(table_dtype), sl, valid, tab.clamp(0, slots - 1) + + +def _i32(*values) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int32, device=DEV) + + +def _build(spec): + """(kv_cache, kv_cache_sf, block_tables, seq_lens, global_scale, out) for a cell; `out` is + sentinel-filled, so `compare` also proves every entry the op must skip was left alone.""" + kind = spec["kind"] + if kind == "topk": + R, K = spec["R"], 512 + kv, sf = _row_views(_random_pool(64, seed=R)) + dtype = _TABLE_DTYPE[spec["table_dtype"]] + tab, sl, _, _ = _row_table(R, K, seed=100 + R, slots=NUM_SLOTS, table_dtype=dtype) + return kv, sf, tab, sl, _gs(), _sentinel_out(R, K, 1, W) + if kind == "request": + page_size = spec["page_size"] + kv, sf = _page_views(_random_pool(page_size, seed=page_size)) + pages = NUM_SLOTS // page_size + L = 8192 + lens = [0, 1, page_size - 1, page_size + 1, 700, 4097, L] + B = len(lens) + g = torch.Generator(device=DEV).manual_seed(7) + bt = torch.stack( + [torch.randperm(pages, generator=g, device=DEV)[: L // page_size] for _ in range(B)] + ) + bt[4, 2] = -1 + bt[5, 1] = pages + 3 + bt = bt.to(_TABLE_DTYPE[spec["table_dtype"]]) + return kv, sf, bt, _i32(*lens), _gs(), _sentinel_out(B, L, 1, W) + if kind == "prefill_16k": + kv, sf = _page_views(_random_pool(128, seed=11)) + bt = torch.randperm(NUM_SLOTS // 128, device=DEV).view(1, -1).to(torch.int32) + return kv, sf, bt, _i32(NUM_SLOTS), _gs(), _sentinel_out(1, NUM_SLOTS, 1, W) + if kind == "compact": + r, B = spec["ratio"], spec["B"] + page_size = 128 // r + kv, sf = _page_views(_random_pool(page_size, seed=20 + r)) + L = 132 * 128 // r + g = torch.Generator(device=DEV).manual_seed(B) + bt = torch.randint( + 0, NUM_SLOTS // page_size, (B, 132), generator=g, device=DEV, dtype=torch.int32 + ) + sl = torch.randint(0, L + 1, (B,), generator=g, device=DEV, dtype=torch.int32) + return kv, sf, bt, sl, _gs(), _sentinel_out(B, L, 1, W) + assert kind == "engine_pool", kind + rows, page_size = spec["rows"], spec["page_size"] + pages = rows // page_size + pool = _random_pool(page_size, seed=pages, slots=rows) + if spec["form"] == "per_request": + kv, sf = _page_views(pool) + bt = _i32(0, pages - 1, pages // 2, 1).view(1, 4) + L = 4 * page_size + return kv, sf, bt, _i32(L), _gs(), _sentinel_out(1, L, 1, W) + kv, sf = _row_views(pool) + return kv, sf, _i32(0, rows - 1, 12345).view(1, 3), _i32(3), _gs(), _sentinel_out(1, 3, 1, W) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + args = _build(cell.spec) + out = args[5] + expected = op.reference(*args) + assert op(*args) is out, "the op writes in place and returns what it wrote" + twin = _twin(*args[:5], _sentinel_out(*out.shape)) + torch.cuda.synchronize() + op.compare(out, expected) + assert torch.equal(_bits(out), _bits(twin)), "the op and FlashInfer's twin differ" + written = int((_bits(out) != SENTINEL).any(-1).sum()) + assert written > 0, "the cell wrote nothing, so it certifies nothing" + + +def test_quantizer_round_trip_and_negative_zero(): + """Rows written by nvfp4_fake_quantize's native mode read back equal to its bf16 output.""" + n = 4096 + g = torch.Generator(device=DEV).manual_seed(3) + x = ( + torch.randn(n, W, generator=g, device=DEV) * torch.logspace(-4, 4, n, device=DEV)[:, None] + ).to(torch.bfloat16) + x[:8] = 0.0 + x[8:16] = -0.0 + x[16, :BLOCK] = 3000.0 + slots = torch.randperm(NUM_SLOTS, device=DEV)[:n] + pool = torch.full((NUM_SLOTS // 64, 64, ROW), 0xAB, dtype=torch.uint8, device=DEV) + torch.ops.trtllm.nvfp4_fake_quantize(x, pool[..., :CODE_BYTES], slots, pool[..., CODE_BYTES:]) + core = torch.full((n, W), float("nan"), dtype=torch.bfloat16, device=DEV) + torch.ops.trtllm.nvfp4_fake_quantize(x, core, None, None) + kv, sf = _row_views(pool) + a, b = _both(kv, sf, slots.view(1, n).to(torch.int32), _i32(n), _gs(), (1, n, 1, W)) + assert torch.equal(_bits(a), _bits(b)) + assert torch.equal(_bits(a[0, :, 0]), _bits(core)) + assert int((_bits(core) == -32768).sum()) > 0 + + +@pytest.mark.parametrize("gs", [0.5, 2.0, 1.0 / 3.0]) +def test_global_scale_matches_the_flashinfer_twin(gs): + """A global scale other than 1 (the target passes 1.0): the same bits as FlashInfer's twin.""" + pool = _random_pool(64, seed=5) + kv, sf = _row_views(pool) + tab, sl, _, _ = _row_table(64, 512, seed=9, slots=NUM_SLOTS) + a, b = _both(kv, sf, tab, sl, _gs(gs), (64, 512, 1, W)) + assert torch.equal(_bits(a), _bits(b)) + + +@pytest.mark.parametrize("R", [1, 6, 64, 384, 1536]) +def test_cuda_graph_capture_replay(R): + """Captured once, replayed on new tables and lengths: equal to eager FlashInfer each time.""" + K = 512 + pool = _random_pool(64, seed=R + 1) + kv, sf = _row_views(pool) + tab = torch.zeros((R, K), dtype=torch.int32, device=DEV) + sl = torch.zeros((R,), dtype=torch.int32, device=DEV) + out = _sentinel_out(R, K, 1, W) + gs = _gs() + op(kv, sf, tab, sl, gs, out) # warm-up (JIT) outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(kv, sf, tab, sl, gs, out) + for rep in range(3): + t, seq_lens, _, _ = _row_table(R, K, seed=1000 * rep + R, slots=NUM_SLOTS) + tab.copy_(t) + sl.copy_(seq_lens) + out.view(torch.int16).fill_(SENTINEL) + graph.replay() + ref = _sentinel_out(R, K, 1, W) + _twin(kv, sf, tab, sl, gs, ref) + torch.cuda.synchronize() + assert torch.equal(_bits(out), _bits(ref)), rep + + +def test_wrong_variants_fail_the_gate(): + """Nibble order swapped, one scale per 32: each must differ from the op's output.""" + pool = _random_pool(64, seed=31) + flat = pool.view(-1, ROW) + kv, sf = _row_views(pool) + rows = torch.randint(0, NUM_SLOTS, (256,), device=DEV) + a, _ = _both(kv, sf, rows.view(1, -1).to(torch.int32), _i32(256), _gs(), (1, 256, 1, W)) + good = _bits(a[0, :, 0]) + lut = torch.tensor(E2M1 + tuple(-v for v in E2M1), dtype=torch.float32, device=DEV) + r = flat[rows] + codes = r[:, :CODE_BYTES] + s16 = r[:, CODE_BYTES:].contiguous().view(torch.float8_e4m3fn).float() + swapped = torch.stack([lut[(codes >> 4).long()], lut[(codes & 0xF).long()]], -1).flatten(-2) + per32 = s16[:, ::2].repeat_interleave(32, -1) + straight = torch.stack([lut[(codes & 0xF).long()], lut[(codes >> 4).long()]], -1).flatten(-2) + for bad in ( + (swapped * s16.repeat_interleave(BLOCK, -1)).to(torch.bfloat16), + (straight * per32).to(torch.bfloat16), + ): + assert (good != _bits(bad)).float().mean() > 0.3 + + +def test_rejects_wrong_operands(): + pool = _random_pool(64, seed=2) + kv, sf = _row_views(pool) + tab = torch.zeros((2, 8), dtype=torch.int32, device=DEV) + sl = torch.full((2,), 8, dtype=torch.int32, device=DEV) + with pytest.raises(ValueError): + op(kv, sf, tab, sl, _gs(), torch.empty((2, 8, 1, W), dtype=torch.float16, device=DEV)) + with pytest.raises(ValueError): + op(kv, sf, tab, sl.long(), _gs(), _sentinel_out(2, 8, 1, W)) + with pytest.raises(ValueError): + op(kv, sf, tab, sl, _gs().double(), _sentinel_out(2, 8, 1, W)) + + +# ─── native storage on the live DeepSeek-V4.1 cache manager, at the engine envelope ───── +# +# The target calls this entry at two sites (unified sm_103 modeling, native KV storage): +# per request: nvfp4_kv_to_bf16(codes4, scales4, md.compress_block_tables[r][:B], md.compressed_kv_lens_cuda[r][:B], +# one, comp_out[r][:B]) +# codes4 / scales4 the KV source's native COMPRESS views [pages, 128 / r, 1, 256] / [.., 32]; +# comp_out[r] the [seqs, cols * 128 / r, 1, 512] view of one buffer (seqs = max_batch_size, +# cols = the manager's max_blocks_per_seq); +# per query: nvfp4_kv_to_bf16(codes1, scales1, topk_slots[:T], topk_cols[:T], one, topk_out[:T]) +# codes1 / scales1 the same views re-viewed with page size 1 ([pages * 128 / r, 1, 1, 256]), +# topk_slots the index op's [T, 512] int32 pool slots (-1 past each query's valid entries), +# topk_cols 512 per query, topk_out the [T, 512, 1, 512] view of a [max_num_tokens * 512, 512] +# buffer. +# The tests below are the FlashInfer form's native-storage and envelope tests run on this op; the reference is +# the writer's own core output or the LUT decode of the pool bytes, and the full calls are also held bitwise to +# FlashInfer's nvfp4_kv_dequantize_paged (the twin, `_twin`). The Triton kernel specializes on integer +# arguments; _jit_variants counts its compiled variants so the tests can show that a row count, a batch or +# a capture never compiles a new one. + +import os # noqa: E402 +import random # noqa: E402 +from contextlib import contextmanager # noqa: E402 + +import tensorrt_llm._torch.custom_ops.nvfp4_kv_to_bf16 as _kvb_op # noqa: E402 (the Triton kernel, variant probe) + + +def _jit_variants() -> int: + """Number of compiled variants Triton holds for the entry's kernel, over all devices.""" + fn = _kvb_op._nvfp4_kv_to_bf16_kernel + dc = getattr(fn, "device_caches", None) + if dc is not None: + return sum(len(v[0]) if isinstance(v, tuple) else len(v) for v in dc.values()) + return sum(len(v) for v in fn.cache.values()) + + +_STORAGE_ENV = "TRTLLM_DSV41_KV_STORAGE" +BYTE_SENTINEL = 0xFF # never a scale byte the writer produces (E4M3 NaN) +TOPK = 512 # the target's compressed top-k width (index_topk) +_BMAX_DEVICE_TOTAL = 297_020_948_480 +# ratio -> (B_max pages over the gates' axis {2048, 8192} x max_batch_size 1..32 x DSpark {0, 5} at max_seq_len +# 16384, page size, gates' largest, owner whose view spans the ratio's buffer) +NATIVE_COMPRESS_AXIS = { + "ratio2": (2_606_019, 64, 1_954_512, 2), + "ratio1": (868_673, 128, 651_504, 20), +} +# ratio -> (ceiling pages, page size, gates' largest, owner). The ceiling is the measured maximum of the native +# COMPRESS views over the engine envelope's scanned tuples at 0.5 x the device's memory, PP2 stage managers +# included (stage alone): max_num_tokens {2048, 4096, 8192, 16384} x max_batch_size {1..8, 12, 16, 24, 32, 48, +# 64, 96, 128, 160, 192, 224, 256} x DSpark {0, 5} x max_seq_len {16K, 32K, 64K, 128K, 256K, 1M}, unsplit +# and per PP2 stage. A stage holds half the layers, so its pools are the largest: ratio 2 on stage 0 (layers +# 0-19), ratio 1 on stage 1 (layers 20-39), both at CEILING_KEY. Quotas above 0.5 x and unsampled +# max_batch_size values are not covered. +NATIVE_COMPRESS = {"ratio2": (6_480_780, 64, 1_954_512, 2), "ratio1": (3_245_959, 128, 651_504, 20)} +CEILING_KEY = ( + 2048, + 256, + 0, + 1_048_576, +) # (max_num_tokens, max_batch_size, DSpark max_draft_len, max_seq_len) +CEILING_STAGE = { + "ratio2": 0, + "ratio1": 1, +} # the PP2 stage whose manager holds the ratio's ceiling view +NATIVE_COMPRESS_UNSPLIT = { + "ratio2": 3_886_647, + "ratio1": 1_295_549, +} # pp_size 1, at (2048, 128, 0, 1M) +# Pool sizes the invariance test also runs: the gates' axis B_max, envelope managers' views (the engine dump's +# (16384, 256, 5, 32768) views, the (16384, 256, 5, 1M) manager's), the unsplit ceiling, stage 0's largest at +# max_seq_len 16384 and stage 1's largest in a coupled PP2 engine. +_EXTRA_B = { + "ratio2": (2_606_019, 1_807_701, 2_974_605, 3_869_352, 3_886_647, 4_747_719), + "ratio1": (868_673, 602_567, 991_535, 1_289_784, 1_295_549, 2_159_957), +} + + +def _write(x, pool, slots): + """The writer's native form into a [.., 288] pool (codes then scales), rows at `slots`.""" + torch.ops.trtllm.nvfp4_fake_quantize(x, pool[..., :CODE_BYTES], slots, pool[..., CODE_BYTES:]) + + +def _core(x): + """The writer's own core (bf16) output of x, through a fixed-size out (one compiled variant).""" + out = torch.empty(max(8192, x.shape[0]), W, dtype=torch.bfloat16, device=DEV) + torch.ops.trtllm.nvfp4_fake_quantize(x, out, None, None) + return out[: x.shape[0]] + + +def _latents(n: int, seed: int) -> torch.Tensor: + """bf16 [n, 512], one magnitude 2^-16..2^6 per 16-block (below the amax floor up to large), a few -0.0.""" + g = torch.Generator(device=DEV).manual_seed(seed) + mag = torch.randint(-16, 7, (n, W // BLOCK, 1), device=DEV, generator=g).float() + x = (torch.randn(n, W // BLOCK, BLOCK, device=DEV, generator=g) * torch.exp2(mag)).view(n, W) + x[:, 5] = -0.0 + return x.to(torch.bfloat16) + + +@contextmanager +def _native_storage(): + saved = os.environ.get(_STORAGE_ENV) + os.environ[_STORAGE_ENV] = "native" + try: + yield + finally: + if saved is None: + os.environ.pop(_STORAGE_ENV, None) + else: + os.environ[_STORAGE_ENV] = saved + + +def _v41_native_manager( + max_num_tokens, max_batch_size, max_draft_len, max_seq_len=None, quota_div=2, pp_stage=None +): + """A V4.1 manager with native KV storage at the gate settings and total / quota_div of the device's memory + (0.5 by default). max_seq_len defaults to the gates' 16384, as the engine passes it: + 14 with DSpark.""" + from types import SimpleNamespace + + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager + from tensorrt_llm.bindings import DataType + from tensorrt_llm.bindings.internal.batch_manager import CacheType + from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, + KvCacheConfig, + ) + from tensorrt_llm.mapping import Mapping + + total = torch.cuda.mem_get_info()[1] + assert total == _BMAX_DEVICE_TOTAL, ( + f"B_max constants are for a {_BMAX_DEVICE_TOTAL}-byte GB300, got {total}" + ) + text = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=128, + head_dim=512, + index_head_dim=128, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, + ) + config = SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**text), + **{k: v for k, v in text.items() if k != "model_type"}, + ) + with _native_storage(): + mgr = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + max_gpu_total_bytes=total // quota_div, + host_cache_size=0, + event_buffer_max_size=0, + ), + kv_cache_type=CacheType.SELFKONLY, + num_layers=40, + num_kv_heads=1, + head_dim=512, + tokens_per_block=128, + max_seq_len=max_seq_len or (16398 if max_draft_len else 16384), + max_batch_size=max_batch_size, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + if pp_stage is None + else Mapping(world_size=2, rank=pp_stage, tp_size=1, pp_size=2, gpus_per_node=4), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=max_num_tokens, + sparse_attn_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ), + pretrained_config=config, + spec_config=DSparkDecodingConfig(max_draft_len=max_draft_len, block_size=max_draft_len) + if max_draft_len + else None, + max_cuda_graph_batch_size=max_batch_size, + disable_overlap_scheduler=False, + ) + assert mgr.v41_layout.kv_storage == "native" + return mgr + + +def _ceiling_manager(ratio): + """The PP2 stage manager whose native COMPRESS view of `ratio` is the envelope's ceiling (NATIVE_COMPRESS), + built as the stage probe builds it (Mapping(world_size 2, rank = stage, tp_size 1, pp_size 2)); the page + count is asserted.""" + t, s, d, L = CEILING_KEY + pages, _, _, owner = NATIVE_COMPRESS[ratio] + mgr = _v41_native_manager(t, s, d, max_seq_len=L, pp_stage=CEILING_STAGE[ratio]) + try: + got = _native_view(mgr, owner).shape[0] + assert got == pages, (ratio, got, pages) + except BaseException: + mgr.shutdown() + raise + return mgr + + +def _native_view(mgr, owner): + """The owner's native COMPRESS view [B, page, 288] uint8, geometry and native views asserted.""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + page_size = 128 if owner == 20 else 64 + pool = mgr.get_buffers(owner, AT.COMPRESS) + codes, scales = mgr.get_v41_native_kv_views(owner, AT.COMPRESS) + b = pool.shape[0] + assert pool.dtype == torch.uint8 and tuple(pool.shape) == (b, page_size, ROW), ( + owner, + pool.shape, + ) + assert pool.stride() == (page_size * ROW, ROW, 1), (owner, pool.stride()) + assert codes.data_ptr() == pool.data_ptr() and tuple(codes.shape) == (b, page_size, CODE_BYTES) + assert scales.data_ptr() == pool.data_ptr() + CODE_BYTES and tuple(scales.shape) == ( + b, + page_size, + W // BLOCK, + ) + return pool + + +def _n_ne(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Device count of differing elements of two same-shape tensors, in chunks (a bool reduction + materializes an int64 copy, 8 bytes per element).""" + fa, fb = a.reshape(-1), b.reshape(-1) + n = torch.zeros((), dtype=torch.int64, device=DEV) + for i in range(0, fa.numel(), 1 << 26): + n += (fa[i : i + (1 << 26)] != fb[i : i + (1 << 26)]).sum() + return n + + +def _count_not(buf: torch.Tensor, value: int) -> torch.Tensor: + whole = buf.view(-1) + step = 1 << 26 + n = torch.zeros((), dtype=torch.int64, device=DEV) + for i in range(0, whole.numel(), step): + n += (whole[i : i + step] != value).sum() + return n + + +def _hazard_rows(num_rows: int, per_page: int) -> list[int]: + """Rows of a 288-byte-row pool on both sides of every byte offset 2^31..2^36, both ends, the last page.""" + cands = {0, 1, per_page - 1, per_page} + for limit in (2**31, 2**32, 2**33, 2**34, 2**35, 2**36): + s = limit // ROW + cands |= {s - 1, s, s + 1} + cands |= set(range(num_rows - per_page - 1, num_rows)) + return sorted(c for c in cands if 0 <= c < num_rows) + + +def _hazard_pages(pages: int, per_page: int) -> list[int]: + """Pages on both sides of every byte offset 2^31..2^36 of the [pages, per_page, 288] view, and the ends.""" + cands = {0, 1, pages - 2, pages - 1} + for limit in (2**31, 2**32, 2**33, 2**34, 2**35, 2**36): + p = limit // (per_page * ROW) + cands |= {p - 1, p, p + 1} + return sorted(c for c in cands if 0 <= c < pages) + + +def _check_gather(out, rows_table, seq_lens, expected_of_row): + """out [R, k, 1, W]: entry (r, j) with j < seq_lens[r] and a valid table row equals expected_of_row(row); + every other entry is still the sentinel.""" + R, k = rows_table.shape + j = torch.arange(k, device=DEV)[None] + valid = (rows_table >= 0) & (j < seq_lens[:, None].long()) + got = _bits(out[:, :, 0]) + rows = rows_table[valid].long() + exp = _bits(expected_of_row(rows)) + n_bad = int(_n_ne(got[valid], exp)) + assert n_bad == 0, f"{n_bad} of {exp.numel()} gathered elements differ bitwise" + assert bool((got[~valid] == SENTINEL).all()), ( + "an entry outside the table's valid rows was written" + ) + return int(valid.sum()) + + +@pytest.mark.parametrize("table_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("owner", [2, 14, 20]) +def test_native_round_trip_far_end_on_the_manager(owner, table_dtype): + """On the live PP2 stage managers' ceiling views (CEILING_KEY; stage 0: owner 2 [6480780, 64, 288], owner + 14 two pages shorter; stage 1: owner 20 [3245959, 128, 288]): the writer's native form puts rows at the far end + (both sides of every 2^31..2^36 byte offset, the ends, the last page) and fills whole pages there + (the last pages, pages across every byte-offset line); this op reads them back + + - row gather: a [6, 512] table of those rows with -1 padding, ids one past the pool and far past it, + and one row whose length ends early; + - whole context: 3 requests of 8192 entries each through [3, n] page tables of those pages, lengths + 8192 / 8191 / 1; + + every read bitwise equal to the writer's core output of the row (and to the LUT reference of the pool + bytes), unwritten entries untouched, and the pool left exactly as the writer left it. + """ + mgr = _ceiling_manager("ratio1" if owner == 20 else "ratio2") + try: + pool = _native_view(mgr, owner) + spans = ( + _native_view(mgr, 2) if owner != 20 else pool + ) # owner 2's view spans the ratio-2 buffer + spans.fill_(BYTE_SENTINEL) + pages, per_page = pool.shape[0], pool.shape[1] + num_rows = pages * per_page + flat = pool.view(-1, ROW) + g = torch.Generator(device=DEV).manual_seed(26070 + owner) + # Whole pages for the context form: the last pages and pages across every byte-offset line. + L = 8192 + n_pages = L // per_page + tail = list(range(pages - n_pages, pages)) + haz = _hazard_pages(pages, per_page) + mid = [ + p + for p in torch.randperm(pages, device=DEV, generator=g)[: 4 * n_pages].tolist() + if p not in tail and p not in haz + ][: n_pages - len(haz)] + page_lists = [tail, haz + mid, [0] + mid[: n_pages - 1]] + page_rows = torch.unique( + torch.cat( + [ + torch.tensor(p, device=DEV)[:, None] * per_page + + torch.arange(per_page, device=DEV)[None] + for p in page_lists + ] + ).view(-1) + ) + rows = torch.unique( + torch.cat([page_rows, torch.tensor(_hazard_rows(num_rows, per_page), device=DEV)]) + ) + x = _latents(rows.numel(), seed=26071 + owner) + _write(x, pool, rows) + y = _core(x) + + def expected_of_row(r): + idx = torch.searchsorted(rows, r) + assert torch.equal(rows[idx], r) + return y[idx] + + gs = _gs() + # Row gather: far-end rows, -1, past-the-end ids. + hz = torch.tensor(_hazard_rows(num_rows, per_page), device=DEV) + R = 6 + tab = torch.full((R, TOPK), -1, dtype=torch.int64, device=DEV) + pick = torch.cat( + [hz, rows[torch.randperm(rows.numel(), device=DEV, generator=g)][: R * TOPK]] + ) + tab.view(-1)[: R * TOPK] = pick[: R * TOPK] + tab[0, 3], tab[0, 5], tab[0, 7] = -1, num_rows, num_rows + 1000 + tab[1, TOPK // 2 :] = -1 + sl = torch.full((R,), TOPK, dtype=torch.int32, device=DEV) + sl[2] = TOPK - 37 + kv, sf = _row_views(pool) + out = _sentinel_out(R, TOPK, 1, W) + op(kv, sf, tab.to(table_dtype), sl, gs, out) + tab_ok = torch.where((tab >= 0) & (tab < num_rows), tab, -1) + n_rg = _check_gather(out, tab_ok, sl, expected_of_row) + twin = _sentinel_out(R, TOPK, 1, W) + _twin(kv, sf, tab.to(table_dtype), sl, gs, twin) + assert torch.equal(_bits(out), _bits(twin)), "row gather differs from the FlashInfer twin" + assert torch.equal( + _bits( + out[:, :, 0][(tab_ok >= 0) & (torch.arange(TOPK, device=DEV)[None] < sl[:, None])] + ), + _bits( + _ref_rows( + flat, + tab_ok[ + (tab_ok >= 0) & (torch.arange(TOPK, device=DEV)[None] < sl[:, None]) + ].long(), + ) + ), + ) + # Whole context through page tables of those pages. + bt = torch.tensor( + [p[:n_pages] + [-1] * (n_pages - len(p[:n_pages])) for p in page_lists], device=DEV + ) + lens = torch.tensor([L, L - 1, 1], dtype=torch.int32, device=DEV) + kvp, sfp = _page_views(pool) + out2 = _sentinel_out(3, L, 1, W) + op(kvp, sfp, bt.to(table_dtype), lens, gs, out2) + t = torch.arange(L, device=DEV) + ctx_rows = torch.where( + bt[:, t // per_page] >= 0, bt[:, t // per_page] * per_page + t % per_page, -1 + ) + n_ctx = _check_gather(out2, ctx_rows, lens, expected_of_row) + twin = _sentinel_out(3, L, 1, W) + _twin(kvp, sfp, bt.to(table_dtype), lens, gs, twin) + assert torch.equal(_bits(out2), _bits(twin)), ( + "whole context differs from the FlashInfer twin" + ) + # The pool is exactly as the writer left it. + _, packed, s_bytes = _reference_native(x) + assert torch.equal(flat[rows], torch.cat([packed, s_bytes], 1)) + flat[rows] = BYTE_SENTINEL + assert int(_count_not(spans, BYTE_SENTINEL)) == 0, "a read changed the native pool" + print( + f"\n[owner {owner} view {tuple(pool.shape)}, {table_dtype}] {rows.numel()} far-end rows written; " + f"row gather {n_rg} entries and whole context {n_ctx} entries == the writer's core output" + ) + finally: + mgr.shutdown() + del mgr + torch.cuda.empty_cache() + + +def _reference_native(x): + """(bf16 core values, packed codes [n, 256], E4M3 bytes [n, 32]) of the checkpoint's NVFP4 recipe + (amax floor 6 * 2^-9, scale e4m3(amax / 6), E2M1 round-to-nearest-even), in native torch.""" + xf = x.float().unflatten(-1, (-1, BLOCK)) + amax = xf.abs().amax(-1, keepdim=True).clamp_min(6 * 2.0**-9) + s8 = (amax / 6.0).clamp(max=448.0).to(torch.float8_e4m3fn) + s = s8.float() + y = (xf / s).clamp(-6.0, 6.0) + a = y.abs() + r = torch.where( + a < 2, torch.round(a * 2) / 2, torch.where(a < 4, torch.round(a), torch.round(a / 2) * 2) + ) + deq = (torch.copysign(r, y) * s).flatten(-2).to(torch.bfloat16) + code = torch.where(r < 2, r * 2, torch.where(r == 6, 7.0, r + 2)).to(torch.int32) + nib = (code | (torch.signbit(y).to(torch.int32) << 3)).flatten(-2) + packed = (nib[..., 0::2] | (nib[..., 1::2] << 4)).to(torch.uint8) + return deq, packed, s8.view(torch.uint8).squeeze(-1) + + +def test_native_harness_sees_a_wrapped_offset(): + """Control: tables naming the rows / pages where a mod-2^32 byte offset would land read other rows, and + the gate sees it in both forms.""" + pages, per_page, _, owner = NATIVE_COMPRESS["ratio2"] + buf = torch.full((pages, per_page, ROW), BYTE_SENTINEL, dtype=torch.uint8, device=DEV) + num_rows = pages * per_page + true = torch.tensor( + [s for s in _hazard_rows(num_rows, per_page) if s * ROW >= 2**32], device=DEV + ) + wrapped = (true * ROW % 2**32) // ROW + assert true.numel() >= 9 and not torch.isin(wrapped, true).any() + x = _latents(2 * true.numel(), seed=26072) + _write(x, buf, torch.cat([true, wrapped])) + y = _core(x)[: true.numel()] + kv, sf = _row_views(buf) + out = _sentinel_out(1, true.numel(), 1, W) + op( + kv, + sf, + wrapped.view(1, -1).to(torch.int32), + torch.tensor([true.numel()], dtype=torch.int32, device=DEV), + _gs(), + out, + ) + n_bad = int((_bits(out[0, :, 0]) != _bits(y)).sum()) + tp = torch.tensor( + [p for p in _hazard_pages(pages, per_page) if p * per_page * ROW >= 2**32], device=DEV + ) + wp = (tp * per_page * ROW % 2**32) // (per_page * ROW) + xp = _latents(2 * tp.numel() * per_page, seed=26073) + prow = torch.cat([tp, wp])[:, None] * per_page + torch.arange(per_page, device=DEV)[None] + _write(xp, buf, prow.view(-1)) + yp = _core(xp)[: tp.numel() * per_page] + kvp, sfp = _page_views(buf) + out2 = _sentinel_out(1, tp.numel() * per_page, 1, W) + op( + kvp, + sfp, + wp.view(1, -1).to(torch.int32), + torch.tensor([tp.numel() * per_page], dtype=torch.int32, device=DEV), + _gs(), + out2, + ) + n_bad2 = int((_bits(out2[0, :, 0]) != _bits(yp)).sum()) + print( + f"\n[wrapped control] row gather {n_bad} of {out[0, :, 0].numel()}, whole context {n_bad2} of " + f"{out2[0, :, 0].numel()} elements differ from the true rows" + ) + assert n_bad > out[0, :, 0].numel() // 2 and n_bad2 > out2[0, :, 0].numel() // 2 + del buf + torch.cuda.empty_cache() + + +@pytest.mark.parametrize("n", [32, 8192]) +@pytest.mark.parametrize("ratio", ["ratio2", "ratio1"]) +def test_native_pool_size_invariance(ratio, n): + """The same rows at the same page indices read through views [:B] of the live PP2 stage manager's ceiling + view (CEILING_KEY on CEILING_STAGE): the same bits at every B (row gather over pages of one row -- + the per-query form's page-size-1 view -- and whole context over the real pages, the per-request + form) and equal to the writer's core output; then from a separate pool of exactly the smallest size. + + B = the smallest view holding the rows, the page counts whose byte count first passes 2^31..2^36 + (where below the ceiling), a mid size, the gates' largest, _EXTRA_B (the gates' axis B_max and envelope + managers' views) and the ceiling. + """ + pages, per_page, gate_max, owner = NATIVE_COMPRESS[ratio] + page_bytes = per_page * ROW + mgr = _ceiling_manager(ratio) + try: + pool = _native_view(mgr, owner) + pool.fill_(BYTE_SENTINEL) + small = -(-n // per_page) + g = torch.Generator(device=DEV).manual_seed(26074 + n + per_page) + order = torch.randperm(small * per_page, device=DEV, generator=g) + rows = torch.cat( + [torch.tensor([small * per_page - 1], device=DEV), order[order != small * per_page - 1]] + )[:n] + x = _latents(n, seed=26074 + n) + _write(x, pool, rows) + y = _core(x) + tab = rows.view(1, n).to(torch.int32) + sl = torch.tensor([n], dtype=torch.int32, device=DEV) + bt = torch.randperm(small, device=DEV, generator=g).view(1, -1).to(torch.int32) + L = small * per_page + t = torch.arange(L, device=DEV) + ctx_rows = bt[0].long()[t // per_page] * per_page + t % per_page + sizes = sorted( + {small, pages // 2 + 1, gate_max, pages} + | set(_EXTRA_B[ratio]) + | { + lim // page_bytes + 1 + for lim in (2**31, 2**32, 2**33, 2**34, 2**35, 2**36) + if lim // page_bytes + 1 < pages + } + ) + first = None + for b in sizes: + kv, sf = _row_views(pool[:b]) + out = _sentinel_out(1, n, 1, W) + op(kv, sf, tab, sl, _gs(), out) + kvp, sfp = _page_views(pool[:b]) + out2 = _sentinel_out(1, L, 1, W) + op(kvp, sfp, bt, torch.tensor([L], dtype=torch.int32, device=DEV), _gs(), out2) + got = torch.cat([_bits(out[0, :, 0]), _bits(out2[0, :, 0])]) + assert torch.equal(_bits(out[0, :, 0]), _bits(y)) + assert torch.equal(_bits(out2[0, :, 0]), _bits(_ref_rows(pool.view(-1, ROW), ctx_rows))) + first = got if first is None else first + assert torch.equal(got, first), f"B={b} differs from B={sizes[0]}" + exact = pool.view(-1, ROW)[: small * per_page].clone().view(small, per_page, ROW) + finally: + mgr.shutdown() + del mgr + torch.cuda.empty_cache() + kv, sf = _row_views(exact) + out = _sentinel_out(1, n, 1, W) + op(kv, sf, tab, sl, _gs(), out) + kvp, sfp = _page_views(exact) + out2 = _sentinel_out(1, L, 1, W) + op(kvp, sfp, bt, torch.tensor([L], dtype=torch.int32, device=DEV), _gs(), out2) + assert torch.equal(torch.cat([_bits(out[0, :, 0]), _bits(out2[0, :, 0])]), first) + print( + f"\n[native {ratio} n={n}] B = {sizes} (views of owner {owner}'s view) and a separate [{small}, " + f"{per_page}, 288]: both forms bitwise identical and == the writer's core output" + ) + + +def test_native_every_row_count_1_to_16384_row_gather(): + """Over the query rows of the row-gather form (the per-query form's T, up to the envelope's + max_num_tokens 16,384): every R in 1..16384 at the target's width 512 reads exactly the first R rows of + the 16384-row call. + + The pool: 16,384 native rows written by the writer. The table: [16384, 512] random rows with -1 + padding (every 5th row shortened), one row per query; seq_lens 512 (every 7th row shorter). Before + each call out[:R] is reset to the sentinel; after it out[:R] must equal the first R rows of the full + call (which equals the writer's core output where the table names a row and the sentinel elsewhere). + """ + g = torch.Generator(device=DEV).manual_seed(26075) + pool = torch.full((NUM_SLOTS // 64, 64, ROW), BYTE_SENTINEL, dtype=torch.uint8, device=DEV) + x = _latents(NUM_SLOTS, seed=26075) + _write(x, pool, torch.arange(NUM_SLOTS, device=DEV)) + y = _core(x) + rmax = 16384 + tab = torch.randint(0, NUM_SLOTS, (rmax, TOPK), device=DEV, generator=g, dtype=torch.int32) + cut = torch.randint(0, TOPK, (rmax,), device=DEV, generator=g) + short = (torch.arange(rmax, device=DEV) % 5 == 0)[:, None] & ( + torch.arange(TOPK, device=DEV)[None] >= cut[:, None] + ) + tab[short] = -1 + sl = torch.full((rmax,), TOPK, dtype=torch.int32, device=DEV) + sl[::7] = torch.randint(1, TOPK, (sl[::7].numel(),), device=DEV, generator=g, dtype=torch.int32) + kv, sf = _row_views(pool) + out = _sentinel_out(rmax, TOPK, 1, W) + op(kv, sf, tab, sl, _gs(), out) + _check_gather(out, tab, sl, lambda r: y[r]) + full = out.clone() + out.view(torch.int16).fill_(SENTINEL) + _twin(kv, sf, tab, sl, _gs(), out) + assert torch.equal(_bits(out), _bits(full)), ( + "the 16384-row call differs from the FlashInfer twin" + ) + n1 = _jit_variants() + o16, f16 = out.view(torch.int16), full.view(torch.int16) + mism = torch.zeros((), dtype=torch.int64, device=DEV) + for r in range(1, rmax + 1): + o16[:r] = SENTINEL + op(kv, sf, tab[:r], sl[:r], _gs(), out[:r]) + mism += _n_ne(o16[:r], f16[:r]) + n2 = _jit_variants() + print( + f"\n[row gather R = 1..{rmax} x {TOPK}] prefix mismatches {int(mism)}; compiled variants {n1} -> {n2}" + ) + assert int(mism) == 0 and n2 == n1 + + +_ENVELOPE_GSET = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_ENVELOPE_GRAPH_R = tuple(sorted(set(_ENVELOPE_GSET) | {6 * b for b in _ENVELOPE_GSET})) + + +@pytest.mark.parametrize("ratio", ["ratio2", "ratio1"]) +def test_native_cuda_graph_every_graph_batch_on_the_bmax_view(ratio): + """Row gather (the per-query form) at every query-row count a captured step makes over the engine + envelope, R in G and 6 G for every G of _ENVELOPE_GSET (1..1536; this holds the gates' {1..8, 16, 24, 32} and 6 x + those), width 512, on the live PP2 stage manager's ceiling view: per R an eager warm-up, the capture, then 2 replays + after the writer has put new rows at new far-end rows and the table (-1 entries included), the lengths + and the pool changed in place; every replay equals the writer's core output of the named rows, leaves + the rest of out (the -1 entries' rows included) untouched, and equals the eager call. The view is the + CEILING_KEY manager's.""" + pages, per_page, _, owner = NATIVE_COMPRESS[ratio] + mgr = _ceiling_manager(ratio) + try: + pool = _native_view(mgr, owner) + num_rows = pages * per_page + g = torch.Generator(device=DEV).manual_seed(26077 + per_page) + cand = torch.unique( + torch.cat( + [ + torch.tensor(_hazard_rows(num_rows, per_page), device=DEV), + num_rows - 1 - torch.arange(64 * per_page * 8, device=DEV), + ] + ) + ) + kv, sf = _row_views(pool) + gs = _gs() + + def fresh(R): + rows = cand[torch.randperm(cand.numel(), device=DEV, generator=g)][: R * TOPK // 4] + x = _latents( + rows.numel(), seed=int(torch.randint(0, 2**30, (1,), generator=g, device=DEV)) + ) + _write(x, pool, rows) + y = _core(x) + srt, idx = torch.sort(rows) + tab = rows[torch.randint(0, rows.numel(), (R, TOPK), device=DEV, generator=g)] + tab[:, TOPK - 13 :] = -1 + sl = torch.full((R,), TOPK, dtype=torch.int32, device=DEV) + sl[0] = TOPK - 100 + return tab.to(torch.int32), sl, lambda r: y[idx[torch.searchsorted(srt, r)]] + + n1 = None + for R in _ENVELOPE_GRAPH_R: + tab, sl, _ = fresh(R) + out = _sentinel_out(R, TOPK, 1, W) + op(kv, sf, tab, sl, gs, out) + torch.cuda.synchronize() + n1 = _jit_variants() if n1 is None else n1 + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(kv, sf, tab, sl, gs, out) + for rep in (1, 2): + new_tab, new_sl, exp = fresh(R) + tab.copy_(new_tab) + sl.copy_(new_sl) + out.view(torch.int16).fill_(SENTINEL) + graph.replay() + torch.cuda.synchronize() + _check_gather(out, tab, sl, exp) + replayed = out.clone() + out.view(torch.int16).fill_(SENTINEL) + op(kv, sf, tab, sl, gs, out) + assert torch.equal(_bits(out), _bits(replayed)), ( + f"R={R} replay {rep} differs from eager" + ) + del graph + n2 = _jit_variants() + print( + f"\n[{ratio} graphs, owner {owner} view {tuple(pool.shape)}] R = {list(_ENVELOPE_GRAPH_R)} x {TOPK} " + f"({len(_ENVELOPE_GRAPH_R)} graphs): every replay == eager == the writer's core output; compiled variants " + f"after the first eager call {n1} -> {n2}" + ) + assert n2 == n1, "a row count or a capture compiled a new variant" + finally: + mgr.shutdown() + del mgr + torch.cuda.empty_cache() + + +_NUM_RANDOM_STEPS = 200 + + +def _run_step(mgr, reqs, num_contexts, gen_len, uid, max_num_tokens=8192): + """Allocate the step's requests on the manager and build its DeepSeek-V4.1 attention metadata as the + engine builds it (max_num_requests = the manager's max_batch_size, the engine's max_num_tokens; each + request's first token unique over the test, so block reuse shares no page across steps).""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( + DeepseekV4TrtllmAttentionMetadata, + ) + from tensorrt_llm._torch.metadata import KVCacheParams + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from tensorrt_llm.bindings import SamplingConfig + from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig + from tensorrt_llm.mapping import Mapping + + requests = [] + for rid, (c, n) in enumerate(reqs): + r = LlmRequest( + request_id=rid, + max_new_tokens=64, + input_tokens=[uid + rid] + list(range(1, c + n)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + assert mgr.prepare_context(r) and mgr.resize_context(r, r.context_chunk_size), (rid, c, n) + requests.append(r) + md = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor([n for _, n in reqs], dtype=torch.int), + request_ids=[r.py_request_id for r in requests], + max_num_requests=mgr.max_batch_size, + num_contexts=num_contexts, + prompt_lens=[c + n for c, n in reqs], + max_num_tokens=max_num_tokens, + kv_cache_manager=mgr, + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=[c for c, _ in reqs] + ), + mapping=Mapping(world_size=1, tp_size=1, rank=0), + sparse_attention_config=DeepSeekV4SparseAttentionConfig( + index_head_dim=128, window_size=128, compress_ratios=[1] * 43 + ), + ) + if num_contexts < len(reqs) and gen_len > 1: + md.update_spec_dec_param( + batch_size=len(reqs), + is_spec_decoding_enabled=False, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, + max_draft_len=gen_len - 1, + max_total_draft_tokens=gen_len - 1, + ) + md.prepare() + return requests, md + + +_KVQ_COLS = ( + 132 # md.compress_block_tables[r].shape[1] (max_blocks_per_seq) in every envelope configuration +) +_KVQ_L = {2: _KVQ_COLS * 64, 1: _KVQ_COLS * 128} # 8,448 / 16,896 entries per request in the buffer +_KVQ_OWNERS = ( + (2, 2), + (8, 2), + (14, 2), + (20, 1), +) # the KV sources in forward order, with their ratio +_KVQ_DSPARK_ENTRIES = { + 2: 16398 // 2, + 1: 16398, +} # entries at the engine's max_seq_len 16398 (DSpark) +# The max_batch_size axis. The compressed read buffer is max_batch_size * 132 * 128 rows +# (16,896 per request), the only native-storage object whose size follows max_batch_size; every value 1..32 the +# engine configuration permits is run, the heavier sweeps at the listed values, and the native-storage witness's +# own configuration (max_num_tokens 2048, max_batch_size 1, DSpark) explicitly. + + +def _kvq_random_step(rng, max_tokens=8192, max_requests=32, seq_len=16384): + """_random_step over the engine's envelope with DSpark: contexts end at <= seq_len tokens, decode + requests (next_n 1) at <= seq_len, DSpark-verify requests (next_n 6) at <= seq_len + 14 (16398 at the + gates' 16384).""" + k = rng.randint(1, max_requests) + gen_len = rng.choice([1, 6]) + g = rng.randint(0, k) if rng.random() < 0.7 else 0 + c = k - g + budget = max_tokens - g * gen_len + t_ctx = rng.randint(c, budget) if c else 0 + cuts = sorted(rng.sample(range(1, t_ctx), c - 1)) if c > 1 else [] + news = [b - a for a, b in zip([0] + cuts, cuts + [t_ctx])] if c else [] + reqs = [(0 if rng.random() < 0.4 else rng.randint(0, seq_len - n), n) for n in news] + rng.shuffle(reqs) + end = seq_len + 14 if gen_len == 6 else seq_len + reqs += [(rng.randint(0, end - gen_len), gen_len) for _ in range(g)] + return reqs, c, gen_len + + +def _kvq_codes_scales(mgr, owner): + """The target's operands: the owner's native (codes, scales) COMPRESS views reshaped to + [pages, 128 / r, 1, 256] and [pages, 128 / r, 1, 32] (views, asserted).""" + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType as AT, + ) + + codes, scales = mgr.get_v41_native_kv_views(owner, AT.COMPRESS) + c4 = codes.reshape(codes.shape[0], codes.shape[1], 1, CODE_BYTES) + s4 = scales.reshape(scales.shape[0], scales.shape[1], 1, W // BLOCK) + assert c4.data_ptr() == codes.data_ptr() and s4.data_ptr() == scales.data_ptr() + return c4, s4 + + +def _kvq_page_candidates(pages, per_page, n, g): + """n distinct pages of a [pages, per_page, 288] view: every page across a 2^31..2^35 byte line and both + ends in random order, then the view's last pages and random pages alternately.""" + hz = torch.tensor(_hazard_pages(pages, per_page), device=DEV) + hz = hz[torch.randperm(hz.numel(), device=DEV, generator=g)] + last = pages - 1 - torch.arange(n, device=DEV) + last = last[~torch.isin(last, hz)] + rnd = torch.randint(0, pages, (4 * n + 64,), device=DEV, generator=g).unique() + rnd = rnd[~torch.isin(rnd, torch.cat([hz, last]))] + rnd = rnd[torch.randperm(rnd.numel(), device=DEV, generator=g)][: last.numel()] + tail = torch.stack([last[: rnd.numel()], rnd], 1).reshape(-1) + cand = torch.cat([hz, tail, last[rnd.numel() :]])[:n] + assert cand.numel() == n and torch.unique(cand).numel() == n + return cand.to(torch.int64) + + +def _mbs_edge_steps(max_t, seqs): + """Edge steps of at most seqs requests: a full-T context, 1 token at 0 and at 16383, one DSpark verify + at kv_len 16398, then full batches of seqs requests (contexts ending at 16384 with T = seqs * + (max_t // seqs), decodes at 16384, verifies at 16398) and, when seqs >= 2, one context filling T with + seqs - 1 verifies at 16398.""" + n = max_t // seqs + steps = [ + ([(0, max_t)], 1, 1), + ([(0, 1)], 1, 1), + ([(16383, 1)], 1, 1), + ([(16392, 6)], 0, 6), + ([(16384 - n, n)] * seqs, seqs, 1), + ([(16383, 1)] * seqs, 0, 1), + ([(16392, 6)] * seqs, 0, 6), + ] + if seqs > 1: + c = max_t - 6 * (seqs - 1) + steps.append(([(16384 - c, c)] + [(16392, 6)] * (seqs - 1), 1, 6)) + return steps + + +def _kvq_per_query_ops(mgr, owner): + """The per-query form's operands: the owner's native (codes, scales) COMPRESS views re-viewed with page + size 1, [pages * (128 / r), 1, 1, 256] and [.., 1, 1, 32], so a pool slot page * (128 / r) + e is its + own page index (views, asserted).""" + c4, s4 = _kvq_codes_scales(mgr, owner) + c1 = c4.view(c4.shape[0] * c4.shape[1], 1, 1, CODE_BYTES) + s1 = s4.view(s4.shape[0] * s4.shape[1], 1, 1, W // BLOCK) + assert c1.data_ptr() == c4.data_ptr() and s1.data_ptr() == s4.data_ptr() and c1.stride(0) == ROW + return c1, s1 + + +def _kvq_renumbered_steps(mgr, steps, max_t, seqs, seed, cols=_KVQ_COLS, per_query=False): + """The body of test_kvq_compressed_read_renumbered_over_random_steps for one manager: per step and per + source in forward order, the writer's native writes, the target's read into ONE buffer sized from the + step's metadata (asserted [seqs, cols]; cols 132 at the gates' max_seq_len) and kept across sources and + steps, the buffer checked against the bf16 container and the renumbered top-k against the bf16 path, and + the three controls. With per_query, each source is also read in the per-query form (see + test_envelope_compressed_read_renumbered_over_random_steps_256_requests). Returns (stats, controls + [source-steps detected, applicable, elements], buffer shape).""" + Lr = {r: cols * 128 // r for r in (1, 2)} + comp_c = q_out = None + views = {o: _native_view(mgr, o) for o, _ in _KVQ_OWNERS} + ops = {o: _kvq_codes_scales(mgr, o) for o, _ in _KVQ_OWNERS} + ops1 = {o: _kvq_per_query_ops(mgr, o) for o, _ in _KVQ_OWNERS} if per_query else None + one = _gs() + stats = dict(steps=0, calls=0, max_t=0, max_k=0, entries=0, topk_rows=0, max_len={1: 0, 2: 0}) + ctrl = { + "renumbered off by one": [0, 0, 0], + "other ratio's row stride": [0, 0, 0], + "stale buffer row": [0, 0, 0], + } + if per_query: + q_out = torch.zeros((max_t, TOPK, 1, W), dtype=torch.bfloat16, device=DEV) + q_cols = torch.full((max_t,), TOPK, dtype=torch.int32, device=DEV) + stats.update(pq_calls=0, pq_rows=0, pq_minus1=0, pq_max_t=0) + ctrl.update({"per-query row off by one": [0, 0, 0], "per-query stale row": [0, 0, 0]}) + chunk = 1 << 20 + try: + for i, (reqs, nctx, gen_len) in enumerate(steps): + requests, md = _run_step( + mgr, reqs, nctx, gen_len, uid=1_000_000 + 64 * i, max_num_tokens=max_t + ) + try: + s_, c_ = md.compress_block_tables[1].shape + assert (s_, c_) == (seqs, cols) and tuple(md.compress_block_tables[2].shape) == ( + s_, + c_, + ) + if comp_c is None: + comp_c = torch.zeros((s_ * cols * 128, W), dtype=torch.bfloat16, device=DEV) + comp_out = {r: comp_c[: s_ * Lr[r]].view(s_, Lr[r], 1, W) for r in (2, 1)} + comp_tbl = torch.arange(s_ * cols, dtype=torch.int32, device=DEV).view(s_, cols) + t, k = sum(n for _, n in reqs), len(reqs) + B = md.num_seqs + assert B == k <= seqs and t <= max_t + req = md.req_idx_per_token[:t].long() + gsd = torch.Generator(device=DEV).manual_seed(seed + i) + for o, r in _KVQ_OWNERS: + bk, L = 128 // r, Lr[r] + ent = [(c + n) // r for c, n in reqs] + lens = md.compressed_kv_lens_cuda[r][:B] + assert lens.tolist() == ent, (i, r) + cbt = md.compress_block_tables[r] + lmax = max(max(ent), 1) + e = torch.arange(lmax, device=DEV) + page = cbt[:B, (e // bk).clamp(max=cols - 1)].long() + row_of = page * bk + e % bk + has = e[None] < torch.tensor(ent, device=DEV)[:, None] + assert bool((page[has] >= 0).all()), f"step {i} r={r}: an entry without a page" + rows = row_of[has] + y = torch.empty((0, W), dtype=torch.bfloat16, device=DEV) + if rows.numel(): + x = _latents(rows.numel(), seed=seed + 1_000_000 + 4 * i + o) + _write(x, views[o], rows) + y = _core(x) + srt, idx = torch.sort(rows) + assert rows.numel() == 0 or bool((srt[1:] != srt[:-1]).all()) + before = comp_c.clone() + op(*ops[o], cbt[:B], lens, one, comp_out[r][:B]) + expected = before.clone() + expected[(torch.arange(B, device=DEV)[:, None] * L + e[None])[has]] = y + n_bad = int(_n_ne(_bits(comp_c), _bits(expected))) + assert n_bad == 0, f"step {i} source {o}: {n_bad} buffer elements differ" + del expected + vis = md.compressed_visible_lens_cuda[r][:t].long() + kk = min(TOPK, lmax) + scores = torch.rand(t, lmax, device=DEV, generator=gsd) + scores[e[None] >= vis[:, None]] = -1.0 + top_s, top_e = torch.topk(scores, kk, dim=1) + ok = top_s >= 0 + ren = comp_tbl[req[:, None], top_e // bk].long() * bk + top_e % bk + assert torch.equal(ren[ok], (req[:, None] * L + top_e)[ok]) + ren_ok, eb = ren[ok], req[:, None].expand_as(top_e)[ok] + counts = { + name: 0 + for name in ( + "renumbered off by one", + "other ratio's row stride", + "stale buffer row", + ) + } + if ren_ok.numel(): + ref_idx = idx[torch.searchsorted(srt, row_of[eb, top_e[ok]])] + other = eb * Lr[3 - r] + top_e[ok] + for s0 in range(0, ren_ok.numel(), chunk): + sl = slice(s0, s0 + chunk) + ref = _bits(y[ref_idx[sl]]) + n_bad = int(_n_ne(_bits(comp_c[ren_ok[sl]]), ref)) + assert n_bad == 0, ( + f"step {i} source {o}: {n_bad} top-k elements differ from the bf16 path" + ) + counts["renumbered off by one"] += int( + _n_ne(_bits(comp_c[ren_ok[sl] + 1]), ref) + ) + counts["other ratio's row stride"] += int( + _n_ne(_bits(comp_c[other[sl].clamp(max=comp_c.shape[0] - 1)]), ref) + ) + counts["stale buffer row"] += int(_n_ne(_bits(before[ren_ok[sl]]), ref)) + applicable = { + "renumbered off by one": ren_ok.numel() > 0, + "other ratio's row stride": bool((eb > 0).any()), + "stale buffer row": ren_ok.numel() > 0, + } + for name, c in counts.items(): + ctrl[name][0] += int(applicable[name] and c > 0) + ctrl[name][1] += int(applicable[name]) + ctrl[name][2] += c + del before + if per_query: + # The per-query form: the index op's pool slots, -1 where a query has no entry, as the + # page table of the page-size-1 view, 512 per query, into ONE persistent zero-filled + # [max_t, 512, 1, 512] buffer; query q's column j lands at row q * 512 + j. + slots = torch.full((t, TOPK), -1, dtype=torch.int32, device=DEV) + slots[:, :kk] = torch.where(ok, row_of[req[:, None], top_e], -1).to( + torch.int32 + ) + qf = q_out.view(-1, W) + q_before = qf.clone() + op(*ops1[o], slots, q_cols[:t], one, q_out[:t]) + valid = torch.zeros((t, TOPK), dtype=torch.bool, device=DEV) + valid[:, :kk] = ok + qi = torch.arange(t * TOPK, device=DEV).view(t, TOPK)[valid] + pcnt = {"per-query row off by one": 0, "per-query stale row": 0} + for s0 in range(0, qi.numel(), chunk): + sl = slice(s0, s0 + chunk) + ref = _bits(y[ref_idx[sl]]) + n_bad = int(_n_ne(_bits(qf[qi[sl]]), ref)) + assert n_bad == 0, ( + f"step {i} source {o}: {n_bad} per-query elements differ from the bf16 path" + ) + pcnt["per-query row off by one"] += int( + _n_ne(_bits(qf[(qi[sl] + 1).clamp(max=qf.shape[0] - 1)]), ref) + ) + pcnt["per-query stale row"] += int(_n_ne(_bits(q_before[qi[sl]]), ref)) + minus = torch.arange(t * TOPK, device=DEV).view(t, TOPK)[~valid] + for s0 in range(0, minus.numel(), chunk): + sl = slice(s0, s0 + chunk) + n_bad = int(_n_ne(_bits(qf[minus[sl]]), _bits(q_before[minus[sl]]))) + assert n_bad == 0, ( + f"step {i} source {o}: {n_bad} elements of -1 entries' rows written" + ) + n_tail = int(_n_ne(_bits(qf[t * TOPK :]), _bits(q_before[t * TOPK :]))) + assert n_tail == 0, ( + f"step {i} source {o}: {n_tail} elements past query {t} written" + ) + for name, c in pcnt.items(): + ctrl[name][0] += int(qi.numel() > 0 and c > 0) + ctrl[name][1] += int(qi.numel() > 0) + ctrl[name][2] += c + stats["pq_calls"] += 1 + stats["pq_rows"] += qi.numel() + stats["pq_minus1"] += minus.numel() + stats["pq_max_t"] = max(stats["pq_max_t"], t) + del q_before, qi, minus, valid, slots + stats["calls"] += 1 + stats["entries"] += rows.numel() + stats["topk_rows"] += ren_ok.numel() + stats["max_len"][r] = max(stats["max_len"][r], max(ent)) + stats["steps"] += 1 + stats["max_t"] = max(stats["max_t"], t) + stats["max_k"] = max(stats["max_k"], k) + finally: + for r_ in requests: + mgr.free_resources(r_) + torch.cuda.empty_cache() + return stats, ctrl, tuple(comp_c.shape) + finally: + comp_c = q_out = None + + +# ─── the engine envelope: max_seq_len to 1,048,576, max_batch_size and requests per step to 256 ─── +# +# The target's call is the one above, unchanged: +# nvfp4_kv_to_bf16(codes4, scales4, md.compress_block_tables[r][:B], md.compressed_kv_lens_cuda[r][:B], one, +# comp_out[r][:B]) +# comp_out[r] is the [seqs, cols * 128 / r, 1, 512] view of one buffer of seqs * cols * 128 rows, with +# (seqs, cols) = md.compress_block_tables[r].shape. The envelope moves three of its numbers: seqs (the engine's +# max_batch_size) and B (the step's requests; under capture every G of the padded graph-batch list) up to 256, +# and cols (the manager's max_blocks_per_seq) with max_seq_len. The buffer is seqs * cols * 128 KiB: 256 GiB at +# (256, 8196), more than a GB300 holds beside a pool, so the 1M-width tests stop at 128 requests. + +_ENVELOPE_L = ( + 16384, + 20480, + 32768, + 65536, + 131072, + 262144, + 1048576, +) # max_seq_len; 20480 is the decoder-replay guard run's +_ENVELOPE_COLS = { + 16384: 132, + 20480: 164, + 32768: 260, + 65536: 516, + 131072: 1028, + 262144: 2052, + 1048576: 8196, +} +_ENVELOPE_SEQS = { + L: 128 if L == 1048576 else 256 for L in _ENVELOPE_L +} # the buffer's requests: <= 128.1 GiB +assert len(_ENVELOPE_GSET) == 46 +_ENVELOPE_CAND = 16384 # distinct pool pages a width test draws its tables from (every request a new permutation) +# Widths whose every request length l in 1..L_r is run at B = 1: all of them. The call at l = L_r, B = 1 stays +# under 1 ms at every width, so the sampled length set does not apply (the time is printed). +_ENVELOPE_EVERY_LENGTH = set(_ENVELOPE_L) +# Pages of the native COMPRESS views (owner 2 at ratio 2, owner 20 at ratio 1) of the manager at 0.5 x this +# GB300's total memory, per (DSpark max_draft_len, max_seq_len) at max_num_tokens 16384, max_batch_size 256. +_ENVELOPE_PAGES = { + (0, 16384): (2_703_360, 901_120), + (5, 16384): (2_408_448, 802_816), + (0, 20480): (2_880_852, 960_284), + (5, 20480): (2_606_019, 868_673), + (0, 32768): (3_192_147, 1_064_049), + (5, 32768): (2_974_605, 991_535), + (0, 65536): (3_515_277, 1_171_759), + (5, 65536): (3_371_463, 1_123_821), + (0, 131072): (3_700_854, 1_233_618), + (5, 131072): (3_619_899, 1_206_633), + (0, 262144): (3_802_908, 1_267_636), + (5, 262144): (3_759_216, 1_253_072), + (0, 1048576): (3_877_545, 1_292_515), + (5, 1048576): (3_869_352, 1_289_784), +} + + +def test_envelope_read_buffer_widths_from_the_metadata(): + """The table width of the target's call at every max_seq_len of the envelope, DSpark off and on. On the + native manager at (max_num_tokens 16384, max_batch_size 256, max_draft_len 0 or 5) with max_seq_len L + (L + 14 with DSpark, as the engine passes it): max_blocks_per_seq is cols(L), and metadata built as the + engine builds it -- the step's, and its CUDA-graph copies at G = 1 and G = 256 -- has + compress_block_tables[1] and [2] of [256, cols] int32 (contiguous) and compressed_kv_lens_cuda[r] of [256] + int32. So out is [B, cols * 128 / r, 1, 512] with the same cols DSpark on and off. The views' page counts + at 0.5 x the device's memory are pinned in _ENVELOPE_PAGES.""" + got, pages = {}, {} + for L in _ENVELOPE_L: + for draft in (0, 5): + lm = L + 14 if draft else L + mgr = _v41_native_manager(16384, 256, draft, max_seq_len=lm) + try: + assert mgr.max_seq_len == lm and mgr.max_batch_size == 256, ( + L, + draft, + mgr.max_seq_len, + ) + pages[(draft, L)] = (_native_view(mgr, 2).shape[0], _native_view(mgr, 20).shape[0]) + requests, md = _run_step(mgr, [(0, 100)], 1, 1, uid=970_000, max_num_tokens=16384) + try: + forms = {"step": md} + for G in (1, 256): + forms[f"graph G={G}"] = md.create_cuda_graph_metadata(G, False, draft) + for name, m in forms.items(): + for r in (1, 2): + tbl, lens = m.compress_block_tables[r], m.compressed_kv_lens_cuda[r] + assert tuple(tbl.shape) == (256, mgr.max_blocks_per_seq), ( + L, + draft, + name, + r, + tbl.shape, + ) + assert tbl.dtype == torch.int32 and tbl.is_contiguous(), ( + L, + draft, + name, + r, + ) + assert tuple(lens.shape) == (256,) and lens.dtype == torch.int32, ( + L, + draft, + name, + r, + ) + got[(draft, L)] = mgr.max_blocks_per_seq + finally: + for r_ in requests: + mgr.free_resources(r_) + finally: + mgr.shutdown() + del mgr + torch.cuda.empty_cache() + widths = {L: (got[(0, L)] * 128, got[(0, L)] * 64) for L in _ENVELOPE_L} + print( + f"\n[envelope widths] cols per (draft, max_seq_len) {got}; out widths (ratio 1, ratio 2) {widths}; " + f"COMPRESS pages at 0.5 x total {pages}" + ) + assert got == {(d, L): _ENVELOPE_COLS[L] for L in _ENVELOPE_L for d in (0, 5)} + assert pages == _ENVELOPE_PAGES + + +def _envelope_fill(pool, pages, g): + """Random native rows into `pages` of the [P, bk, 288] view: every code byte, scale bytes 0x00..0x7E.""" + bk = pool.shape[1] + for i in range(0, pages.numel(), 2048): + p = pages[i : i + 2048] + codes = torch.randint( + 0, 256, (p.numel(), bk, CODE_BYTES), generator=g, device=DEV, dtype=torch.int32 + ) + scales = torch.randint( + 0, 0x7F, (p.numel(), bk, W // BLOCK), generator=g, device=DEV, dtype=torch.int32 + ) + pool[p] = torch.cat([codes, scales], -1).to(torch.uint8) + + +def _envelope_ref(pool, pages): + """bf16 [n, bk, 512]: the LUT reference of every row of `pages` of the [P, bk, 288] view.""" + bk = pool.shape[1] + flat = pool.view(-1, ROW) + out = torch.empty((pages.numel(), bk, W), dtype=torch.bfloat16, device=DEV) + for i in range(0, pages.numel(), 1024): + p = pages[i : i + 1024] + rows = (p[:, None] * bk + torch.arange(bk, device=DEV)[None]).view(-1) + out[i : i + 1024] = _ref_rows(flat, rows).view(-1, bk, W) + return out + + +def _envelope_dirty(buf): + """Device count of the 8-byte words of the contiguous bf16 `buf` that are not four sentinels.""" + w = buf.view(-1).view(torch.int64) + n = torch.zeros((), dtype=torch.int64, device=DEV) + for i in range(0, w.numel(), 1 << 26): + n += (w[i : i + (1 << 26)] != -1).sum() + return int(n) + + +def _envelope_check(out, ci, lens, ref, bk): + """Entries (b, e < lens[b]) of out [seqs, L_r, 1, 512] must be bitwise ref[ci[b, e // bk], e % bk] (the + LUT reference of the row the table names); each is then reset to the sentinel, so _envelope_dirty(buffer) == 0 + afterwards means the call wrote nothing else. Returns the entries checked.""" + n_all = 0 + for b, n in enumerate(lens): + if n: + exp = _bits(ref[ci[b, : -(-n // bk)]].view(-1, W)[:n]) + got = _bits(out[b, :n, 0]) + if not torch.equal(got, exp): + raise AssertionError( + f"request {b} (length {n}): {int(_n_ne(got, exp))} elements differ" + ) + out[b, :n].view(torch.int16).fill_(SENTINEL) + n_all += n + return n_all + + +def _envelope_far(L, r): + """Far-end lengths of a request at max_seq_len L, ratio r: the out's width L_r and L_r - 1, the DSpark + manager's maximum (L + 14) // r and the graph metadata's (L + 19) // r, L // r and (L - 1) // r.""" + Lr = _ENVELOPE_COLS[L] * 128 // r + return sorted({Lr, Lr - 1, (L + 14) // r, (L + 19) // r, L // r, (L - 1) // r}) + + +def _envelope_lengths(L, r): + """The request lengths run at B = 1: every l in 1..L_r where the call stays under 1 ms, otherwise + every l in 1..256, every l = {0, 1, -1} mod the page size 128 / r, and the far-end lengths.""" + bk, Lr = 128 // r, _ENVELOPE_COLS[L] * 128 // r + if L in _ENVELOPE_EVERY_LENGTH: + return list(range(1, Lr + 1)) + ls = set(range(1, 257)) | set(_envelope_far(L, r)) + for k in range(1, Lr // bk + 1): + ls |= {k * bk - 1, k * bk, k * bk + 1} + return sorted(v for v in ls if 1 <= v <= Lr) + + +@pytest.mark.parametrize("L", _ENVELOPE_L) +def test_envelope_fixed_out_every_width(L): + """Every length and the envelope's widths for the target's call, per max_seq_len L. Operands: the native views + of a live manager (16384, 256, DSpark 5; owner 2 at ratio 2, owner 20 at ratio 1) built at 1/8 of the + device's memory so the widest buffer fits beside it; out the ratio's [seqs, L_r, 1, 512] view (L_r = cols * + 128 / r) of one sentinel-filled buffer of seqs * cols * 128 rows (seqs 256, 128 at the 1M width) with a + guard past it; the table [seqs, cols] int32 over a 16,384-page candidate set (the view's hazard and last + pages, random pages; every request a different permutation) whose rows hold random native bytes; the + reference the LUT decode of those bytes. Every call's named entries are bitwise the reference and nothing + else in the buffer or the guard is written: + + - the full call, every request at L_r (every row of the buffer), one launch of the entry's Triton + kernel; the FlashInfer twin on the same operands writes the same bits; + - B = seqs with ragged lengths: the far-end lengths (_envelope_far), the page edges, 0, 1, random (twin too); + - every B in 1..seqs, the last request at a far-end length and the others short (0..2 pages); + - every request length l in 1..L_r at B = 1: each call's first l entries bitwise the full call's, entry l + still the sentinel; + - no compiled variant is added after the full call. + """ + from torch.profiler import ProfilerActivity, profile + + cols, seqs = _ENVELOPE_COLS[L], _ENVELOPE_SEQS[L] + mgr = _v41_native_manager(16384, 256, 5, quota_div=8) + big = None + try: + rows_total = seqs * cols * 128 + big = torch.full((rows_total + 128, W), SENTINEL, dtype=torch.int16, device=DEV).view( + torch.bfloat16 + ) + one = _gs() + report = {} + for owner, r in ((2, 2), (20, 1)): + bk = 128 // r + Lr = cols * bk + out = big[: seqs * Lr].view(seqs, Lr, 1, W) + pool = _native_view(mgr, owner) + c4, s4 = _kvq_codes_scales(mgr, owner) + g = torch.Generator(device=DEV).manual_seed(16101 + L + r) + rng = random.Random(16101 + L + r) + cand = _kvq_page_candidates(pool.shape[0], bk, _ENVELOPE_CAND, g) + _envelope_fill(pool, cand, g) + ref = _envelope_ref(pool, cand) + ci = torch.stack( + [ + torch.randperm(_ENVELOPE_CAND, device=DEV, generator=g)[:cols] + for _ in range(seqs) + ] + ) + tbl = cand[ci].to(torch.int32) + far = _envelope_far(L, r) + # the full call + full = torch.full((seqs,), Lr, dtype=torch.int32, device=DEV) + with profile(activities=[ProfilerActivity.CUDA]) as prof: + op(c4, s4, tbl, full, one, out) + torch.cuda.synchronize() + kernels = [ + e.name + for e in prof.events() + if e.device_type.name == "CUDA" and "nvfp4_kv_to_bf16" in e.name + ] + assert len(kernels) == 1, kernels + n1 = _jit_variants() + n_full = _envelope_check(out, ci, [Lr] * seqs, ref, bk) + assert _envelope_dirty(big) == 0, "the full call wrote outside its entries" + _twin(c4, s4, tbl, full, one, out) + assert ( + _envelope_check(out, ci, [Lr] * seqs, ref, bk) == n_full + and _envelope_dirty(big) == 0 + ), "twin differs" + # ragged lengths at B = seqs + ragged = (far + [0, 1, bk - 1, bk, bk + 1] + [rng.randint(0, Lr) for _ in range(seqs)])[ + :seqs + ] + op(c4, s4, tbl, torch.tensor(ragged, dtype=torch.int32, device=DEV), one, out) + n_rag = _envelope_check(out, ci, ragged, ref, bk) + assert _envelope_dirty(big) == 0, "the ragged call wrote outside its entries" + _twin(c4, s4, tbl, torch.tensor(ragged, dtype=torch.int32, device=DEV), one, out) + assert ( + _envelope_check(out, ci, ragged, ref, bk) == n_rag and _envelope_dirty(big) == 0 + ), "twin differs (ragged)" + # every B, the last request at the far end + n_b = 0 + for B in range(1, seqs + 1): + lens_b = [rng.randint(0, 2 * bk) for _ in range(B - 1)] + [far[B % len(far)]] + op( + c4, + s4, + tbl[:B], + torch.tensor(lens_b, dtype=torch.int32, device=DEV), + one, + out[:B], + ) + n_b += _envelope_check(out, ci, lens_b, ref, bk) + d = _envelope_dirty(big) + assert d == 0, f"B={B}: {d} words written outside the named entries" + # request lengths at B = 1; the call at l = L_r timed first (median of 5 CUDA-event timings) + ls = _envelope_lengths(L, r) + e64 = _bits(ref[ci[0]].view(-1, W)[:Lr]).view(torch.int64) + o64 = out[0, :, 0].view(torch.int16).view(torch.int64) + neq = torch.empty(e64.shape, dtype=torch.bool, device=DEV) + mism = torch.zeros((), dtype=torch.int64, device=DEV) + past = torch.zeros((), dtype=torch.int64, device=DEV) + lens1 = torch.full((1,), Lr, dtype=torch.int32, device=DEV) + times = [] + for _ in range(5): + ev0, ev1 = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + ev0.record() + op(c4, s4, tbl[:1], lens1, one, out[:1]) + ev1.record() + torch.cuda.synchronize() + times.append(ev0.elapsed_time(ev1)) + assert torch.equal(o64, e64), "the B = 1 full-length call differs from the reference" + assert L in _ENVELOPE_EVERY_LENGTH or sorted(times)[2] > 1.0, ( + f"a sampled width's call takes {times} ms" + ) + o64.fill_(-1) + for length in ls: + lens1.fill_(length) + op(c4, s4, tbl[:1], lens1, one, out[:1]) + torch.ne(o64[:length], e64[:length], out=neq[:length]) + mism += neq[:length].any() + if length < Lr: + past += (o64[length] != -1).any() + o64[:length].fill_(-1) + assert int(mism) == 0 and int(past) == 0, (int(mism), int(past)) + assert _envelope_dirty(big) == 0 + n2 = _jit_variants() + assert n2 == n1, f"a batch or a length compiled a new variant ({n1} -> {n2})" + report[f"ratio {r}"] = dict( + L_r=Lr, + pages=pool.shape[0], + full=n_full, + ragged=n_rag, + every_B=n_b, + call_ms_B1_full=round(sorted(times)[2], 3), + lengths=f"{len(ls)} ({'every l' if L in _ENVELOPE_EVERY_LENGTH else 'sampled'})", + far=far, + ) + del ref, e64, neq + print( + f"\n[envelope width, max_seq_len {L}: cols {cols}, buffer [{rows_total}, 512] ({seqs} requests)] " + f"{report}" + ) + finally: + big = None + mgr.shutdown() + del mgr + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +@pytest.mark.parametrize("L", _ENVELOPE_L) +@pytest.mark.parametrize("ratio", ["ratio2", "ratio1"]) +def test_envelope_cuda_graph_every_gset_batch_every_width(ratio, L): + """The target's call under CUDA-graph capture at every graph batch G of the envelope (_ENVELOPE_GSET: 1..32 and + 40..128 step 8, 192, 256; up to 128 at the 1M width) x next_n {1, 6}, at the table width the graph metadata + uses at max_seq_len L. Operands as test_envelope_fixed_out_every_width: the ratio's owner view of the live + (16384, 256, 5) manager at 1/8 of the device's memory, out the [seqs, L_r, 1, 512] view of one + sentinel-filled seqs * cols * 128-row buffer (a guard past it), a persistent [seqs, cols] table and [seqs] + lengths. First one eager call at the buffer's largest shape (seqs requests of L_r entries): the only + warm-up. Then per (G, next_n): capture on the [:G] views, and 2 replays after the table rows, the lengths + and the bytes of pages request 0 names all change in place. Lengths are kv_len // r: request 0 and + request G - 1 at the far end (kv_len L for next_n 1, L + 14 for next_n 6), request 1 at next_n, the rest + random up to kv_len 4096. Every replay writes exactly the named entries, bitwise the LUT reference, and + nothing else in the buffer (requests G.. and the guard included); the same call run eagerly after it + gives the same bits. + """ + cols, seqs = _ENVELOPE_COLS[L], _ENVELOPE_SEQS[L] + owner, r = (2, 2) if ratio == "ratio2" else (20, 1) + bk = 128 // r + Lr = cols * bk + mgr = _v41_native_manager(16384, 256, 5, quota_div=8) + big = None + try: + pool = _native_view(mgr, owner) + c4, s4 = _kvq_codes_scales(mgr, owner) + rows_total = seqs * cols * 128 + big = torch.full((rows_total + 128, W), SENTINEL, dtype=torch.int16, device=DEV).view( + torch.bfloat16 + ) + out = big[: seqs * Lr].view(seqs, Lr, 1, W) + one = _gs() + g = torch.Generator(device=DEV).manual_seed(16201 + L + r) + rng = random.Random(16201 + L + r) + cand = _kvq_page_candidates(pool.shape[0], bk, _ENVELOPE_CAND, g) + _envelope_fill(pool, cand, g) + ref = _envelope_ref(pool, cand) + ci = torch.zeros((seqs, cols), dtype=torch.int64, device=DEV) + tbl = torch.full((seqs, cols), -1, dtype=torch.int32, device=DEV) + lens = torch.zeros((seqs,), dtype=torch.int32, device=DEV) + + def fresh(entries): + """New table rows for the first len(entries) requests, new bytes on up to 16 pages request 0 + names (its first and last 8), the lengths in place.""" + B = len(entries) + ci[:B] = torch.stack( + [torch.randperm(_ENVELOPE_CAND, device=DEV, generator=g)[:cols] for _ in range(B)] + ) + tbl[:B] = cand[ci[:B]].to(torch.int32) + n0 = -(-entries[0] // bk) + if n0: + touch = torch.unique(torch.cat([ci[0, : min(n0, 8)], ci[0, max(n0 - 8, 0) : n0]])) + _envelope_fill(pool, cand[touch], g) + ref[touch] = _envelope_ref(pool, cand[touch]) + lens[:B] = torch.tensor(entries, dtype=torch.int32, device=DEV) + return entries + + entries = fresh([Lr] * seqs) + op(c4, s4, tbl, lens, one, out) + n_warm = _envelope_check(out, ci, entries, ref, bk) + assert _envelope_dirty(big) == 0 + n1 = _jit_variants() + graphs, named = 0, 0 + for G in (v for v in _ENVELOPE_GSET if v <= seqs): + for next_n in (1, 6): + end = L + 14 if next_n == 6 else L + + def kv_lens(): + kv = [rng.randint(next_n, min(end, 4096)) for _ in range(G)] + kv[0] = end + if G > 1: + kv[1], kv[-1] = next_n, end + return [v // r for v in kv] + + fresh(kv_lens()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(c4, s4, tbl[:G], lens[:G], one, out[:G]) + graphs += 1 + assert _envelope_dirty(big) == 0, "capture wrote into the buffer" + for rep in (1, 2): + entries = fresh(kv_lens()) + graph.replay() + torch.cuda.synchronize() + named += _envelope_check(out, ci, entries, ref, bk) + d = _envelope_dirty(big) + assert d == 0, ( + f"G={G} next_n={next_n} replay {rep}: {d} words written outside the named entries" + ) + op(c4, s4, tbl[:G], lens[:G], one, out[:G]) + _envelope_check(out, ci, entries, ref, bk) + assert _envelope_dirty(big) == 0, ( + f"G={G} next_n={next_n} replay {rep}: eager differs" + ) + del graph + print( + f"\n[envelope graphs {ratio}, max_seq_len {L}: owner {owner} view {tuple(pool.shape)}, out " + f"[{seqs}, {Lr}, 1, 512]] eager warm-up {n_warm} entries; {graphs} graphs (G in the envelope's set up " + f"to {seqs} x next_n 1/6), {named} replayed entries: every replay == eager == the LUT reference, " + f"nothing else written" + ) + assert graphs == 2 * len([v for v in _ENVELOPE_GSET if v <= seqs]) + assert _jit_variants() == n1, "a capture compiled a new variant" + finally: + big = None + mgr.shutdown() + del mgr + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +# KV tokens one step of the 256-request tests may hold in total. The native manager at (16384, 256, DSpark +# 5, max_seq_len 32,782) held 22 requests of 32,768 tokens and refused the 23rd (about 720K tokens), so a +# step of 256 requests cannot put every request at the far end; steps keep their total under this bound. +_ENVELOPE_STEP_TOKENS = 560_000 + + +def _envelope_edge_steps(max_t, seqs, L): + """Edge steps of a (max_num_tokens, max_batch_size) configuration at max_seq_len L with DSpark, inside + _ENVELOPE_STEP_TOKENS: a full-T context from position 0 and one ending at L, one token at L - 1, one DSpark + verify at kv_len L + 14; full batches of seqs requests with 8 at the far end and the rest at positions + below 1024 -- contexts (8 ending at L), decodes (8 at L - 1), verifies (8 at kv_len L + 14) -- and a + context filling T beside seqs - 1 verifies (7 at kv_len L + 14).""" + n, lm = max_t // seqs, L + 14 + rest = seqs - 8 + steps = [ + ([(0, max_t)], 1, 1), + ([(L - max_t, max_t)], 1, 1), + ([(L - 1, 1)], 1, 1), + ([(lm - 6, 6)], 0, 6), + ([(L - n, n)] * 8 + [(0, n)] * rest, seqs, 1), + ([(L - 1, 1)] * 8 + [(127 + (61 * i) % 1024, 1) for i in range(rest)], 0, 1), + ([(lm - 6, 6)] * 8 + [(100 + (37 * i) % 1024, 6) for i in range(rest)], 0, 6), + ] + c = max_t - 6 * (seqs - 1) + steps.append( + ([(L - c, c)] + [(lm - 6, 6)] * 7 + [(100 + (37 * i) % 1024, 6) for i in range(rest)], 1, 6) + ) + for reqs, _, _ in steps: + assert sum(a + b for a, b in reqs) <= _ENVELOPE_STEP_TOKENS, reqs[:3] + return steps + + +def _envelope_fit(step, rng): + """A random step with its total KV tokens capped at _ENVELOPE_STEP_TOKENS: its requests, in a random order, + keep their cached prefix while the rest can still fit; the others get a random prefix below 1024. The + request order (contexts first) and every request's new tokens are unchanged.""" + reqs, nctx, gen_len = step + if sum(a + n for a, n in reqs) <= _ENVELOPE_STEP_TOKENS: + return step + out, used = list(reqs), 0 + reserve = sum(n + 1024 for _, n in reqs) + order = list(range(len(reqs))) + rng.shuffle(order) + for i in order: + a, n = reqs[i] + reserve -= n + 1024 + if used + a + n + reserve > _ENVELOPE_STEP_TOKENS: + a = rng.randint(0, 1023) + out[i] = (a, n) + used += a + n + assert used <= _ENVELOPE_STEP_TOKENS + return out, nctx, gen_len + + +def test_envelope_compressed_read_renumbered_over_random_steps_256_requests(): + """At the envelope's step axes, in the target's exact call, on the live native manager of the + engine configuration max_num_tokens 16384, max_batch_size 256, DSpark 5, max_seq_len 32768 + 14 (cols + 260: out widths 33,280 / 16,640). Steps: _envelope_edge_steps (a 16,384-token context, 256 requests of a full + batch, verifies at kv_len 32,782) and 200 seeded random partitions of T <= 16,384 into k <= 256 requests + (contexts first with random cached prefixes, then decode next_n 1 at kv_len <= 32,768 or DSpark verify + next_n 6 at <= 32,782), each step's KV tokens capped at _ENVELOPE_STEP_TOKENS (_envelope_fit). + Per step and per KV source in forward order (2, 8, 14 at ratio 2, 20 at ratio 1), + exactly as test_kvq_compressed_read_renumbered_over_random_steps: the writer fills every compressed entry + at its row of md.compress_block_tables[r]; the call reads md.compress_block_tables[r][:B] and + md.compressed_kv_lens_cuda[r][:B] into the ratio's view of ONE persistent zero-filled [256 * 260 * 128, + 512] buffer; every entry (b, e < length) equals the writer's bf16 container and nothing else changes; the + renumbered top-k is bitwise the bf16 path; the three controls differ wherever they apply. + + The same steps and writes also drive the per-query form (the target's _native_topk, used when + max_num_tokens * 512 rows is the smaller buffer): op(codes1, scales1, slots, cols, + one, out) with (codes1, scales1) the owner's views re-viewed with page size 1 ([pages * 128 / r, 1, 1, + 256] / [.., 32]), slots the [T, 512] int32 pool slots of each query's top-k (up to 512 random visible + entries, -1 after them, as the index op packs them), cols 512 per query, and out the [T, 512, 1, 512] + view of ONE persistent zero-filled [16384, 512, 1, 512] buffer. After each call every valid (q, j) row + q * 512 + j is bitwise the bf16-container row of that slot, every -1 entry's row and every row past query + T is exactly as before the call, and two controls (the row off by one, the buffer before the call) + differ wherever a row is named.""" + L, max_t, seqs = 32768, 16384, 256 + rng = random.Random(16301) + edges = _envelope_edge_steps(max_t, seqs, L) + steps = edges + [ + _envelope_fit(_kvq_random_step(rng, max_tokens=max_t, max_requests=seqs, seq_len=L), rng) + for _ in range(_NUM_RANDOM_STEPS) + ] + mgr = _v41_native_manager(max_t, seqs, 5, max_seq_len=L + 14) + try: + stats, ctrl, shape = _kvq_renumbered_steps( + mgr, steps, max_t, seqs, seed=16_302_000, cols=_ENVELOPE_COLS[L], per_query=True + ) + print( + f"\n[envelope renumbered read, (16384, 256, 5) at max_seq_len {L + 14}, COMPRESS " + f"{tuple(_native_view(mgr, 2).shape)} / {tuple(_native_view(mgr, 20).shape)}, buffer {shape}] " + f"{len(edges)} edge + {_NUM_RANDOM_STEPS} random steps: {stats}; controls [detected, applicable, " + f"elements]: {ctrl}" + ) + assert shape == (seqs * _ENVELOPE_COLS[L] * 128, W) + assert stats["steps"] == len(steps) and stats["max_t"] == max_t and stats["max_k"] == seqs + assert stats["max_len"] == {2: (L + 14) // 2, 1: L + 14} + assert ( + stats["pq_calls"] == stats["calls"] + and stats["pq_max_t"] == max_t + and stats["pq_minus1"] > 0 + ) + for name, (hit, applicable, _) in ctrl.items(): + assert applicable > 0 and hit == applicable, f"control '{name}': {hit} of {applicable}" + finally: + mgr.shutdown() + del mgr + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +def test_kvq_compressed_read_renumbered_over_random_steps_gate_config(): + """In the target's two calls on the live accuracy-gate manager (8192, 32, DSpark 5; cols 132): + _mbs_edge_steps(8192, 32) (a full-T context, a DSpark verify at kv_len 16,398, full batches of 32 + requests) and 200 seeded random partitions of T <= 8192 into k <= 32 requests (contexts with cached + prefixes, decode next_n 1, DSpark verify next_n 6). Per step and KV source in forward order, as + test_envelope_compressed_read_renumbered_over_random_steps_256_requests: the per-request read into ONE + [32 * 132 * 128, 512] buffer and the per-query read into ONE [8192, 512, 1, 512] buffer, both bitwise the + writer's bf16 container, nothing else written, and every control fires where it applies.""" + max_t, seqs = 8192, 32 + rng = random.Random(17301) + edges = _mbs_edge_steps(max_t, seqs) + steps = edges + [ + _kvq_random_step(rng, max_tokens=max_t, max_requests=seqs) for _ in range(_NUM_RANDOM_STEPS) + ] + mgr = _v41_native_manager(max_t, seqs, 5) + try: + stats, ctrl, shape = _kvq_renumbered_steps( + mgr, steps, max_t, seqs, seed=17_302_000, per_query=True + ) + print( + f"\n[renumbered read, gate manager (8192, 32, 5), COMPRESS {tuple(_native_view(mgr, 2).shape)} / " + f"{tuple(_native_view(mgr, 20).shape)}, buffer {shape}] {len(edges)} edge + {_NUM_RANDOM_STEPS} random " + f"steps: {stats}; controls [detected, applicable, elements]: {ctrl}" + ) + assert shape == (seqs * _KVQ_COLS * 128, W) + assert stats["steps"] == len(steps) and stats["max_t"] == max_t and stats["max_k"] == seqs + assert stats["max_len"] == _KVQ_DSPARK_ENTRIES + assert ( + stats["pq_calls"] == stats["calls"] + and stats["pq_max_t"] == max_t + and stats["pq_minus1"] > 0 + ) + for name, (hit, applicable, _) in ctrl.items(): + assert applicable > 0 and hit == applicable, f"control '{name}': {hit} of {applicable}" + finally: + mgr.shutdown() + del mgr + torch.cuda.synchronize() + torch.cuda.empty_cache() diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_nvfp4_kv_to_e4m3.py b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_nvfp4_kv_to_e4m3.py new file mode 100644 index 000000000000..9702cc4fd51a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/quantization/test_modeling_v2_nvfp4_kv_to_e4m3.py @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the nvfp4_kv_to_e4m3 catalog entry. + +The pool is DeepSeek-V4.1's native compressed-KV layout: one 288-byte row per entry, 256 packed E2M1 +bytes (even element in the low nibble) then 32 E4M3 scale bytes (one per 16 elements), no global scale, +read as two strided views. Every gate is bitwise on the uint8 codes, and `out` is pre-filled with a +sentinel byte so that `compare` also proves every entry the op must skip was left alone. + +The out buffer is the native-storage target's compact compressed buffer [max_batch_size * 132 * 128 / r, 512] +viewed as [B, L_r, 1, 512] with L_r = 132 * 128 / r (r = 2 and 1); the pools are the engine's page +geometries (64-row pages for ratio 2, 128-row pages for ratio 1). +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.quantization.nvfp4_kv_to_e4m3 import ( # noqa: E501 + nvfp4_kv_to_e4m3 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "nvfp4_kv_to_e4m3 requires a CUDA device" + +DEV = "cuda" +W = 512 +BLOCK = 16 +CODE_BYTES = W // 2 +ROW = CODE_BYTES + W // BLOCK # 288 +E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +E4M3_MAX = 448.0 +SENT = 0x5A +COLS = 132 # compress_block_tables width at max_seq_len 16398 (native-storage runtime dump) + + +def _gen(seed): + return torch.Generator(device=DEV).manual_seed(seed) + + +def _pool(pages, page_size, seed, scale_hi=0x7E): + """[pages, page_size, 288] uint8: random code bytes, E4M3 scale bytes 0x00..scale_hi (the quantizer + writes non-negative finite scales; its values are e4m3(amax/6) with amax >= 6*2^-9).""" + g = _gen(seed) + n = pages * page_size + codes = torch.randint(0, 256, (n, CODE_BYTES), generator=g, device=DEV, dtype=torch.int32) + scales = torch.randint( + 0, scale_hi + 1, (n, W // BLOCK), generator=g, device=DEV, dtype=torch.int32 + ) + return torch.cat([codes, scales], 1).to(torch.uint8).view(pages, page_size, ROW) + + +def _views(pool): + return pool[..., :CODE_BYTES].unsqueeze(2), pool[..., CODE_BYTES:].unsqueeze(2) + + +def _sentinel_out(*shape): + return torch.full(shape, SENT, dtype=torch.uint8, device=DEV).view(torch.float8_e4m3fn) + + +def _values(flat, rows): + """fp32 [n, W] NVFP4 values of pool rows `rows` (for the error bound and the controls).""" + lut = torch.tensor(E2M1 + tuple(-v for v in E2M1), dtype=torch.float32, device=DEV) + r = flat[rows] + c = r[:, :CODE_BYTES] + v = torch.stack([lut[(c & 0xF).long()], lut[(c >> 4).long()]], -1).flatten(-2) + s = ( + r[:, CODE_BYTES:] + .contiguous() + .view(torch.float8_e4m3fn) + .float() + .repeat_interleave(BLOCK, -1) + ) + return v * s + + +def _build_compact(spec): + """The target's call: every compressed entry of each request of the step, into [B, L_r, 1, 512].""" + ratio = spec["ratio"] + page_size = 128 // ratio + L = COLS * page_size + pages = 2048 + pool = _pool(pages, page_size, seed=ratio) + kv, sf = _views(pool) + B = 8 + g = _gen(10 + ratio) + bt = torch.stack([torch.randperm(pages, generator=g, device=DEV)[:COLS] for _ in range(B)]).to( + torch.int32 + ) + lens = [0, 1, page_size - 1, page_size + 1, 700 // ratio, 4097 // ratio, 16398 // ratio, L] + bt[5, 3] = -1 # a hole: its entries are skipped + sl = torch.tensor(lens, dtype=torch.int32, device=DEV) + return kv, sf, bt, sl, _sentinel_out(B, L, 1, W), spec["inv_scale"] + + +def _build_engine(spec): + """Engine page counts (ratio-2 [453120, 64, .], ratio-1 [151040, 128, .]) with rows named in the + first and last pages, into the compact buffer view of max_batch_size mbs.""" + ratio, pages, mbs = spec["ratio"], spec["pages"], spec["mbs"] + page_size = 128 // ratio + L = COLS * page_size + pool = torch.zeros(pages, page_size, ROW, dtype=torch.uint8, device=DEV) + g = _gen(20 + ratio + mbs) + B = mbs + bt = torch.randint(0, pages, (B, COLS), generator=g, device=DEV, dtype=torch.int32) + bt[0, 0] = pages - 1 + bt[B - 1, 1] = 0 + used = torch.unique(bt.long()) + pool[used] = _pool(used.numel(), page_size, seed=30 + ratio) + sl = torch.randint(1, 16398 // ratio + 1, (B,), generator=g, device=DEV, dtype=torch.int32) + sl[0] = 16398 // ratio + buf = torch.full((mbs * COLS * 128 // ratio, W), SENT, dtype=torch.uint8, device=DEV) + out = buf.view(torch.float8_e4m3fn).view(mbs, L, 1, W) + kv, sf = _views(pool) + return kv, sf, bt, sl, out, spec["inv_scale"] + + +def _build_per_query(spec): + """The target's per-query read: the pool viewed as pages of one row ([pages * bk, 1, 1, C], a slot + number is its own page id), a [T, 512] table of the slots each query's top-k names (valid entries + first, -1 after), seq_lens 512 for every query, into [T, 512, 1, 512]. Engine pool geometry + (ratio-2 453120 x 64, ratio-1 151040 x 128 rows), slots named in the first and last pages.""" + ratio, T = spec["ratio"], spec["T"] + page_size = 128 // ratio + pages = 453120 if ratio == 2 else 151040 + pool = torch.zeros(pages, page_size, ROW, dtype=torch.uint8, device=DEV) + g = _gen(80 + ratio + T) + n_slots = pages * page_size + lo = torch.randint(0, 4096, (T, 512), generator=g, device=DEV) + hi = torch.randint(n_slots - 4096, n_slots, (T, 512), generator=g, device=DEV) + tab = torch.where(torch.rand(T, 512, generator=g, device=DEV) < 0.5, lo, hi) + k = torch.randint(0, 513, (T,), generator=g, device=DEV) + k[0] = 512 + if T > 1: + k[1] = 0 + tab = ( + torch.where(torch.arange(512, device=DEV)[None] < k[:, None], tab, -1) + .to(torch.int32) + .contiguous() + ) + used = torch.unique(tab[tab >= 0].long()) + flat = pool.view(-1, ROW) + flat[used] = _pool(1, used.numel(), seed=90 + ratio).view(-1, ROW) + kv = pool[..., :CODE_BYTES].reshape(n_slots, 1, 1, CODE_BYTES) + sf = pool[..., CODE_BYTES:].reshape(n_slots, 1, 1, W // BLOCK) + assert kv.data_ptr() == pool.data_ptr(), "the per-query view must alias the pool" + assert int(tab.max()) >= n_slots - 4096 + cols = torch.full((T,), 512, dtype=torch.int32, device=DEV) + return kv, sf, tab, cols, _sentinel_out(T, 512, 1, W), spec["inv_scale"] + + +_BUILDERS = {"compact": _build_compact, "engine": _build_engine, "per_query": _build_per_query} + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + args = _BUILDERS[cell.spec["kind"]](cell.spec) + out = args[4] + expected = op.reference(*args) + assert op(*args) is out, "the op writes in place and returns what it wrote" + op.compare(out, expected) + written = (out.view(torch.uint8) != SENT).any(-1).sum().item() + assert written > 0, "the cell wrote nothing, so it certifies nothing" + + +def test_signed_zero_saturation_and_rounding(): + """-0.0 (code 0x8) stays -0; 6 * 448 saturates to 448; 6-significant-bit products round to E4M3 RNE.""" + pool = torch.zeros(1, 1, ROW, dtype=torch.uint8, device=DEV) + pool[0, 0, 0] = 0x87 # element 0: code 7 (6.0), element 1: code 8 (-0.0) + pool[0, 0, 1] = 0x35 # element 2: code 5 (3.0), element 3: code 3 (1.5) + pool[0, 0, CODE_BYTES] = 0x7E # scale 448 for block 0 + kv, sf = _views(pool) + args = ( + kv, + sf, + torch.zeros(1, 1, dtype=torch.int32, device=DEV), + torch.ones(1, dtype=torch.int32, device=DEV), + _sentinel_out(1, 1, 1, W), + 1.0, + ) + expected = op.reference(*args) + out = op(*args) + o = out.view(torch.uint8)[0, 0, 0] + assert o[0].item() == 0x7E, "6 * 448 saturates to 448" + assert o[1].item() == 0x80, "-0.0 keeps its sign" + op.compare(out, expected) + + +def test_transcode_is_lossy_by_at_most_half_an_e4m3_ulp(): + """E2M1 x E4M3 has up to 6 significant bits: the transcode rounds, by at most 2^-4 relative in the + normal range; measured here on random rows (the fraction that is inexact is reported).""" + pool = _pool(16, 64, seed=40, scale_hi=0x60) + flat = pool.view(-1, ROW) + kv, sf = _views(pool) + bt = torch.arange(16, dtype=torch.int32, device=DEV).view(1, 16) + out = torch.empty(1, 1024, 1, W, dtype=torch.float8_e4m3fn, device=DEV) + op(kv, sf, bt, torch.tensor([1024], dtype=torch.int32, device=DEV), out, 64.0) + v = _values(flat, torch.arange(1024, device=DEV)) * 64.0 + got = out[0, :, 0].float() + normal = (v.abs() >= 2.0**-6) & (v.abs() <= E4M3_MAX) + rel = ((got - v).abs() / v.abs().clamp_min(1e-30))[normal] + inexact = (got != v)[v != 0].float().mean().item() + print(f"inexact fraction {inexact:.3f}, max rel error in normal range {rel.max().item():.4f}") + assert rel.max().item() <= 2.0**-4 + 1e-7 + assert inexact > 0.05 # the value set is genuinely not E4M3 + + +def test_cuda_graph_capture_replay(): + page_size = 64 + L = COLS * page_size + pool = _pool(512, page_size, seed=50) + kv, sf = _views(pool) + B = 4 + bt = torch.zeros(B, COLS, dtype=torch.int32, device=DEV) + sl = torch.ones(B, dtype=torch.int32, device=DEV) + out = _sentinel_out(B, L, 1, W) + op(kv, sf, bt, sl, out, 32.0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(kv, sf, bt, sl, out, 32.0) + for seed in range(3): + g = _gen(60 + seed) + out.view(torch.uint8).fill_(SENT) + bt.copy_(torch.randint(0, 512, (B, COLS), generator=g, device=DEV, dtype=torch.int32)) + sl.copy_(torch.randint(0, 4000, (B,), generator=g, device=DEV, dtype=torch.int32)) + expected = op.reference(kv, sf, bt, sl, _sentinel_out(B, L, 1, W), 32.0) + graph.replay() + op.compare(out, expected) + + +def test_wrong_variants_are_detected(): + """Controls: nibble order swapped, scale per 32 instead of 16, inv_scale ignored.""" + pool = _pool(4, 64, seed=70) + flat = pool.view(-1, ROW) + kv, sf = _views(pool) + out = torch.empty(1, 256, 1, W, dtype=torch.float8_e4m3fn, device=DEV) + op( + kv, + sf, + torch.arange(4, dtype=torch.int32, device=DEV).view(1, 4), + torch.tensor([256], dtype=torch.int32, device=DEV), + out, + 16.0, + ) + got = out.view(torch.uint8)[0, :, 0] + rows = torch.arange(256, device=DEV) + lut = torch.tensor(E2M1 + tuple(-v for v in E2M1), dtype=torch.float32, device=DEV) + r = flat[rows] + c = r[:, :CODE_BYTES] + s16 = r[:, CODE_BYTES:].contiguous().view(torch.float8_e4m3fn).float() + swapped = torch.stack([lut[(c >> 4).long()], lut[(c & 0xF).long()]], -1).flatten(-2) + v = torch.stack([lut[(c & 0xF).long()], lut[(c >> 4).long()]], -1).flatten(-2) + for name, vals in ( + ("nibbles swapped", swapped * s16.repeat_interleave(BLOCK, -1) * 16.0), + ("scale per 32", v * s16[:, ::2].repeat_interleave(2 * BLOCK, -1) * 16.0), + ("inv_scale ignored", v * s16.repeat_interleave(BLOCK, -1)), + ): + exp = vals.clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn).view(torch.uint8) + frac = (got != exp).float().mean().item() + assert frac > 0.3, f"control '{name}' differs on only {frac:.3f}" diff --git a/tests/unittest/_torch/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/attention/test_modeling_v2_fused_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/attention/test_modeling_v2_fused_qk_norm_rope.py new file mode 100644 index 000000000000..d25c372c4dac --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/attention/test_modeling_v2_fused_qk_norm_rope.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fused_qk_norm_rope catalog entry.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.gpt_oss.gpt_oss_120b__sm_103__tp1.ops.attention.fused_qk_norm_rope import ( # noqa: E501 + fused_qk_norm_rope as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "fused_qk_norm_rope requires a CUDA device" + + +def _build(spec, seed): + """Materialize one cell as (kwargs, qkv), with qkv separate because it is + the buffer the op writes and the caller has to keep a copy of.""" + g = torch.Generator(device="cuda").manual_seed(seed) + tokens = spec["tokens"] + heads = spec["num_heads_q"] + spec["num_heads_k"] + spec["num_heads_v"] + width = heads * spec["head_dim"] + qkv = torch.randn(tokens, width, generator=g, device="cuda").to(torch.bfloat16) + weight = torch.randn(spec["head_dim"], generator=g, device="cuda").to(torch.bfloat16) + kwargs = {k: v for k, v in spec.items() if k != "tokens"} + kwargs.update( + eps=1e-6, + q_weight=weight, + k_weight=weight, + position_ids=torch.arange(tokens, dtype=torch.int32, device="cuda"), + ) + return qkv, kwargs + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + qkv, kwargs = _build(cell.spec, seed=abs(hash(cell.why)) % 2**31) + ref = op.reference(qkv, **kwargs) + with validating(op): + out = op(qkv, **kwargs) + assert out is qkv, "the op writes in place" + op.compare(qkv, ref) + + +def test_the_v_heads_are_left_alone() -> None: + """The op's boundary, and nothing about the output shape shows it. + + A kernel that rotated the v heads too would still return a buffer of the + right size, and every cell compares against a reference that makes the same + claim -- so this pins it against the input instead. + """ + spec = dict(op.CELLS[1].spec) + qkv, kwargs = _build(spec, seed=11) + before = qkv.clone() + op(qkv, **kwargs) + v_start = (kwargs["num_heads_q"] + kwargs["num_heads_k"]) * kwargs["head_dim"] + assert torch.equal(qkv[:, v_start:], before[:, v_start:]) + assert not torch.equal(qkv[:, :v_start], before[:, :v_start]), "q and k were not touched" + + +def _ok(**over): + spec = dict(op.CELLS[0].spec) + spec.update(over) + return _build(spec, seed=3) + + +def test_guard_refuses_a_width_that_disagrees_with_the_head_counts() -> None: + """Read as whatever the counts say; the heads come out shifted.""" + qkv, kwargs = _ok() + with pytest.raises(AssertionError, match="last dim must be"): + with validating(op): + op(qkv[:, :-64], **kwargs) + + +def test_guard_refuses_a_rotary_dim_past_the_head() -> None: + qkv, kwargs = _ok() + kwargs["rotary_dim"] = 128 + with pytest.raises(AssertionError, match="exceeds head_dim"): + with validating(op): + op(qkv, **kwargs) + + +def test_guard_refuses_a_position_id_count_that_is_not_the_token_count() -> None: + qkv, kwargs = _ok() + kwargs["position_ids"] = torch.arange(3, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="one position per token"): + with validating(op): + op(qkv, **kwargs) diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py index ae56f8a4e4ac..0a3f46287979 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py +++ b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py @@ -2,13 +2,20 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the fp4_block_scale_moe_runner catalog entry.""" +import zlib + +import pytest import torch +from torch.profiler import ProfilerActivity, profile from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( fp4_block_scale_moe_runner as moe, ) from tensorrt_llm._torch.autotuner import AutoTuner, autotune +__extra_import_path__ = [".."] +from _validating import validating # noqa: E402 — needs the path declared above + assert torch.cuda.is_available(), "fp4_block_scale_moe_runner requires a CUDA device" DEV = "cuda" @@ -274,36 +281,37 @@ def _call(data, sf, args, num_experts, top_k, scale_as_is=False, **kw): "gemm2_weights_scale", ): kw.setdefault(k, args[k]) - return moe( - kw.pop("routing_logits"), - kw.pop("routing_bias"), - data, - sf, - kw.pop("gemm1_weights"), - kw.pop("gemm1_weights_scale"), - kw.pop("gemm1_bias"), - kw.pop("gemm1_alpha"), - kw.pop("gemm1_beta"), - kw.pop("gemm1_clamp_limit"), - kw.pop("gemm2_weights"), - kw.pop("gemm2_weights_scale"), - kw.pop("gemm2_bias"), - kw.pop("output1_scale_scalar"), - kw.pop("output1_scale_gate_scalar"), - kw.pop("output2_scale_scalar"), - num_experts, - top_k, - kw.pop("n_group"), - kw.pop("topk_group"), - kw.pop("intermediate_size"), - kw.pop("local_expert_offset"), - kw.pop("local_num_experts"), - kw.pop("routed_scaling_factor"), - kw.pop("routing_method_type"), - kw.pop("do_finalize"), - kw.pop("act_type"), - **kw, - ) + with validating(moe): + return moe( + kw.pop("routing_logits"), + kw.pop("routing_bias"), + data, + sf, + kw.pop("gemm1_weights"), + kw.pop("gemm1_weights_scale"), + kw.pop("gemm1_bias"), + kw.pop("gemm1_alpha"), + kw.pop("gemm1_beta"), + kw.pop("gemm1_clamp_limit"), + kw.pop("gemm2_weights"), + kw.pop("gemm2_weights_scale"), + kw.pop("gemm2_bias"), + kw.pop("output1_scale_scalar"), + kw.pop("output1_scale_gate_scalar"), + kw.pop("output2_scale_scalar"), + num_experts, + top_k, + kw.pop("n_group"), + kw.pop("topk_group"), + kw.pop("intermediate_size"), + kw.pop("local_expert_offset"), + kw.pop("local_num_experts"), + kw.pop("routed_scaling_factor"), + kw.pop("routing_method_type"), + kw.pop("do_finalize"), + kw.pop("act_type"), + **kw, + ) # ── reference ───────────────────────────────────────────────────────────── @@ -486,14 +494,25 @@ def _dsr1(): def test_layout_helpers_match(): """The trtllm preprocessing ops reproduce the pure-torch layout exactly. - The last two scale shapes are the DeepSeek-R1 routed geometry's, at + Two scale shapes are the DeepSeek-R1 routed geometry's, at H = 7168 / I = 2048: FC1 `[E, 2*I, H/16] = [E, 4096, 448]` and FC2 - `[E, H, I/16] = [E, 7168, 128]`. Both satisfy the swizzle's `M % 128 == 0` - and `C % 4 == 0`, so `block_scale_interleave` returns exactly `E*M*C` bytes - and reshapes back — no padding enters either operand at this geometry. + `[E, H, I/16] = [E, 7168, 128]`; four are DeepSeek-V4.1's, at H = 5120: + FC1 `[E, 4608, 320]` / FC2 `[E, 5120, 144]` for I = 2304 and FC1 + `[E, 1152, 320]` / FC2 `[E, 5120, 36]` for a MoE-TP4 rank's I = 576. All + satisfy the swizzle's `M % 128 == 0` and `C % 4 == 0`, so + `block_scale_interleave` returns exactly `E*M*C` bytes and reshapes back — + no padding enters any operand at these geometries. """ gen = torch.Generator(device=DEV).manual_seed(201) - for rows, cols in ((256, 12), (4096, 448), (7168, 128)): + for rows, cols in ( + (256, 12), + (4096, 448), + (7168, 128), + (4608, 320), + (5120, 144), + (1152, 320), + (5120, 36), + ): x = torch.randint(0, 256, (4, rows, cols), dtype=torch.uint8, device=DEV, generator=gen) perm = _fc1_perm(rows) for e in range(x.shape[0]): @@ -1426,6 +1445,520 @@ def test_ep_window_64_of_256(): ) +# ── DeepSeek-V4.1-Flash routed geometry ─────────────────────────────────── +# +# 384 routed experts, H = 5120, I = 2304, top-6, and the clamped SwiGLU +# gate = min(x W_gate^T, 10), up = clamp(x W_up^T, -10, 10). Unlike the geometries +# above, every operand carries a PER-EXPERT NVFP4 global scale, as a modelopt +# checkpoint does: the FC1 and FC2 weight global scales `gw1[e]`, `gw2[e]` and +# the FC2-input global scale `g2[e]` (one per expert, as a checkpoint's +# `w2.input_scale` is), so the three scale scalars and the clamp limit are +# genuinely different numbers for every expert. + +V41_E, V41_H, V41_I, V41_K = 384, 5120, 2304, 6 +V41_LIMIT = 10.0 +V41_EP = 96 # local experts per rank under four-way expert parallelism +V41_TP_I = 576 # intermediate columns per rank under four-way MoE tensor parallelism +_V41_TOKENS = (1, 2, 3, 5, 8, 16, 32, 64, 128, 256, 1024, 4096, 8192) +_V41 = None + + +def _v41(): + """The V4.1 routed layer, built once: kernel-ready tensors, reference + operands, per-expert global scales, and a static per-expert g2.""" + global _V41 + if _V41 is None: + args, ref = _build(V41_E, V41_H, V41_I, seed=4141) + gen = torch.Generator(device=DEV).manual_seed(4142) + # FC1 weight global scales log-uniform in [1/8, 1/2]: with g1 = 137 the + # true-domain gate/up RMS then runs from ~4 to ~16 across experts, so + # the limit of 10 bites on anywhere from ~1% to ~50% of an expert's + # elements. FC2's in [1/2, 2]. + ref["gw1"] = torch.exp2(torch.empty(V41_E, device=DEV).uniform_(-3.0, -1.0, generator=gen)) + ref["gw2"] = torch.exp2(torch.empty(V41_E, device=DEV).uniform_(-1.0, 1.0, generator=gen)) + ref["g1"] = 137.0 + # One static FC2-input global scale per expert, fixed before any test + # batch is drawn (as a checkpoint's is): 448*6 / the amax of that + # expert's clamped activation over a 4096-token calibration batch, + # times a per-expert factor in [1/2, 1] so no two experts share it. + _, _, xc = _rand_activation(4096, V41_H, ref["g1"], gen) + idc, _ = _v41_routing(4096, gen) + amax = torch.zeros(V41_E, device=DEV) + for e in range(V41_E): + tok = (idc == e).any(dim=1).nonzero(as_tuple=True)[0] + up, gate = _v41_fc1(xc[tok], ref, e, V41_LIMIT) + amax[e] = (up * gate * torch.sigmoid(gate)).abs().max() + factor = torch.empty(V41_E, device=DEV).uniform_(0.5, 1.0, generator=gen) + ref["g2"] = E4M3_MAX * E2M1_MAX / amax * factor + _V41 = (args, ref) + return _V41 + + +def _v41_routing(num_tokens, gen): + """Distinct top-6 ids of 384 and V4.1-style weights (positive, summing to 1.5).""" + ids = torch.rand(num_tokens, V41_E, device=DEV, generator=gen).argsort(dim=1)[:, :V41_K] + w = torch.rand(num_tokens, V41_K, device=DEV, generator=gen) + 0.05 + w = w / w.sum(dim=1, keepdim=True) * 1.5 + return ids.to(torch.int32).contiguous(), w.to(torch.bfloat16).contiguous() + + +def _v41_fc1(xe, ref, e, limit, swap_gate_up=False): + """True-domain (up, gate) of expert `e` over rows `xe`, clamped when `limit`.""" + up_c, up_s = ref["up"] + gt_c, gt_s = ref["gate"] + wu = _dequant(up_c[e], up_s[e]) / ref["gw1"][e] + wg = _dequant(gt_c[e], gt_s[e]) / ref["gw1"][e] + if swap_gate_up: + wu, wg = wg, wu + up, gate = xe @ wu.t(), xe @ wg.t() + if limit is not None: + gate = torch.clamp(gate, max=limit) + up = torch.clamp(up, min=-limit, max=limit) + return up, gate + + +def _ref_v41(x, ids, wts, ref, limit=V41_LIMIT, swap_gate_up=False): + """Native-torch V4.1 routed MoE, fp32, over all 384 experts. + + Returns `(full, windows, slices)`: the whole layer, the four 96-expert + expert-parallel windows' shares, and the four 576-column MoE-TP slices' + shares. A slice's share is exact to take from the full computation: its + columns start on a multiple of 16, so the FC1-output NVFP4 blocks — and + with the expert's own g2 their scales — are the ones the full layer has. + Also returns the fraction of gate / up elements the limit clamped. + """ + num_tokens = x.shape[0] + dn_c, dn_s = ref["down"] + g2 = ref["g2"] + full = torch.zeros(num_tokens, V41_H, dtype=torch.float32, device=DEV) + windows = [torch.zeros_like(full) for _ in range(V41_E // V41_EP)] + slices = [torch.zeros_like(full) for _ in range(V41_I // V41_TP_I)] + clamped = torch.zeros(2, dtype=torch.float64, device=DEV) + count = 0 + for e in range(V41_E): + tok, slot = (ids == e).nonzero(as_tuple=True) + if tok.numel() == 0: + continue + up, gate = _v41_fc1(x[tok], ref, e, None, swap_gate_up) + if limit is not None: + clamped[0] += (gate > limit).sum() + clamped[1] += (up.abs() > limit).sum() + count += gate.numel() + gate = torch.clamp(gate, max=limit) + up = torch.clamp(up, min=-limit, max=limit) + act = up * gate * torch.sigmoid(gate) + g2e = float(g2[e]) + act = _q_nvfp4(act, g2e) / g2e + wd = _dequant(dn_c[e], dn_s[e]) / ref["gw2"][e] + w = wts[tok, slot].float().unsqueeze(1) + y = (act @ wd.t()) * w + full.index_add_(0, tok, y) + windows[e // V41_EP].index_add_(0, tok, y) + for s in range(len(slices)): + a, b = s * V41_TP_I, (s + 1) * V41_TP_I + slices[s].index_add_(0, tok, (act[:, a:b] @ wd[:, a:b].t()) * w) + frac = (clamped / max(count, 1)).tolist() + return full, windows, slices, frac + + +def _v41_scalars(ref, g2=None, gw1=None, limit=V41_LIMIT): + """The kernel's four per-expert fp32 tensors for the V4.1 layer. + + With `alpha1[e] = 1 / (g1 * gw1[e])` (the FC1 accumulator's dequant): + `output1_scale_gate_scalar = alpha1`, `output1_scale_scalar = alpha1 * g2`, + `output2_scale_scalar = 1 / (g2 * gw2)`, and the clamp in the accumulator + domain, `gemm1_clamp_limit = limit / alpha1`. + """ + g2 = ref["g2"] if g2 is None else g2 + gw1 = ref["gw1"] if gw1 is None else gw1 + a1 = 1.0 / (ref["g1"] * gw1) + return dict( + output1_scale_scalar=(a1 * g2).contiguous(), + output1_scale_gate_scalar=a1.contiguous(), + output2_scale_scalar=(1.0 / (g2 * ref["gw2"])).contiguous(), + gemm1_clamp_limit=(limit / a1).contiguous(), + ) + + +def _v41_slice(d, off: int, n: int): + return {k: v[off : off + n].contiguous() for k, v in d.items()} + + +_V41_TP_ARGS = {} + + +def _v41_tp_args(width: int = V41_TP_I): + """Kernel-ready operands of the MoE-TP ranks at `width` intermediate columns + per rank (576: the four MoE-TP4 ranks; 1152: the two MoE-TP2 ranks): rank r + holds intermediate columns [width r, width (r + 1)) of every expert — its + rows of up and gate, its columns of down — prepared exactly like a full + layer with I = width.""" + if width not in _V41_TP_ARGS: + _, ref = _v41() + up_p, gt_p, up_s, gt_s, dn_p, dn_s = ref["raw"] + ranks = [] + for r in range(V41_I // width): + a, b = r * width, (r + 1) * width + w1, s1 = _prep_fc1(up_p[:, a:b], gt_p[:, a:b], up_s[:, a:b], gt_s[:, a:b]) + w2, s2 = _prep_fc2( + dn_p[:, :, a // 2 : b // 2].contiguous(), dn_s[:, :, a // SV : b // SV].contiguous() + ) + ranks.append( + dict( + gemm1_weights=w1, + gemm1_weights_scale=s1, + gemm2_weights=w2, + gemm2_weights_scale=s2, + intermediate_size=width, + ) + ) + _V41_TP_ARGS[width] = ranks + return _V41_TP_ARGS[width] + + +def test_v41_clamp_is_in_the_accumulator_domain(): + """`gemm1_clamp_limit` clamps the FC1 accumulator, before any scalar. + + Read straight out of the kernel (identity FC2 with unit block scales and + `output2_scale_scalar = 1`, so the output IS the requantized intermediate, + bit for bit) with four experts whose `output1_scale_gate_scalar` (alpha1) + and FC2-input global scales g2 are all different, one g2 deliberately too + large so some blocks saturate the e4m3 scale at 448. Passing + `gemm1_clamp_limit[e] = L[e] / alpha1[e]` gives, bit-exactly on every + element, `gate = min(gate_true, L)`, `up = clamp(up_true, -L, L)` — gate + clamped from above only, up on both sides, both halves by the same alpha1 + (up's own scalar `output1_scale_scalar` plays no part). A true-domain limit + handed over as is clamps at `L * alpha1` instead. + """ + num_experts, hidden, top_k, num_tokens = 4, 256, 1, 256 + gen = torch.Generator(device=DEV).manual_seed(4321) + up_p, up_s, up_c = _rand_nvfp4(num_experts, hidden, hidden, gen) + gt_p, gt_s, gt_c = _rand_nvfp4(num_experts, hidden, hidden, gen) + w1, s1 = _prep_fc1(up_p, gt_p, up_s, gt_s) + eye = torch.zeros(num_experts, hidden, hidden, dtype=torch.uint8, device=DEV) + d = torch.arange(hidden, device=DEV) + eye[:, d, d] = 2 # e2m1 code 2 == +1.0 + dn_p = (eye[..., 0::2] | (eye[..., 1::2] << 4)).contiguous() + dn_s = torch.ones(num_experts, hidden, hidden // SV, device=DEV).to(torch.float8_e4m3fn) + w2, s2 = _prep_fc2(dn_p, dn_s) + args = dict( + gemm1_weights=w1, + gemm1_weights_scale=s1, + gemm2_weights=w2, + gemm2_weights_scale=s2, + intermediate_size=hidden, + ) + g1 = 137.0 + gw1 = torch.tensor([1.0, 0.37, 2.9, 0.0061], device=DEV) # per-expert FC1 weight global scales + alpha1 = 1.0 / (g1 * gw1) + data, sf, x = _rand_activation(num_tokens, hidden, g1, gen) + ids = torch.arange(num_tokens, device=DEV, dtype=torch.int32).remainder(num_experts) + ids = ids.unsqueeze(1).contiguous() + eid = ids[:, 0].long() + wts = torch.ones(num_tokens, 1, dtype=torch.bfloat16, device=DEV) + up_t = torch.zeros(num_tokens, hidden, device=DEV) + gate_t = torch.zeros(num_tokens, hidden, device=DEV) + for e in range(num_experts): + tok = (eid == e).nonzero(as_tuple=True)[0] + up_t[tok] = x[tok] @ _dequant(up_c[e], up_s[e]).t() / gw1[e] + gate_t[tok] = x[tok] @ _dequant(gt_c[e], gt_s[e]).t() / gw1[e] + # a true-domain limit per expert that bites on ~40% of its gate elements + lim = torch.stack([gate_t[eid == e].abs().flatten().quantile(0.6) for e in range(num_experts)]) + act_free = up_t * gate_t * torch.sigmoid(gate_t) + g2 = torch.stack( + [E4M3_MAX * E2M1_MAX / act_free[eid == e].abs().max() for e in range(num_experts)] + ) * torch.tensor([1.0, 0.5, 0.25, 4.0], device=DEV) + + def kernel(clamp): + return _call( + data, + sf, + args, + num_experts, + top_k, + topk_ids=ids, + topk_weights=wts, + output1_scale_scalar=(alpha1 * g2).contiguous(), + output1_scale_gate_scalar=alpha1.contiguous(), + output2_scale_scalar=torch.ones(num_experts, device=DEV), + gemm1_clamp_limit=clamp, + )[0].float() + + def expect(gate_lim, up_lim, gate_both=False, up_above_only=False): + g, u = gate_t.clone(), up_t.clone() + if gate_lim is not None: + gl = gate_lim[eid].unsqueeze(1) + g = torch.clamp(g, min=-gl, max=gl) if gate_both else torch.minimum(g, gl) + if up_lim is not None: + ul = up_lim[eid].unsqueeze(1) + u = torch.minimum(u, ul) if up_above_only else torch.clamp(u, min=-ul, max=ul) + a = u * g * torch.sigmoid(g) + out = torch.empty_like(a) + for e in range(num_experts): + r = eid == e + out[r] = _q_nvfp4(a[r], float(g2[e])) + return out + + assert torch.equal(kernel(None), expect(None, None)), "unclamped intermediate is not bit-exact" + saturated = sum( + int( + ( + (g2[e] * act_free[eid == e]).reshape(-1, SV).abs().amax(-1) / E2M1_MAX > E4M3_MAX + ).sum() + ) + for e in range(num_experts) + ) + assert saturated > 0, "no e4m3 block scale saturated — the g2 x4 expert did not exercise it" + + y = kernel((lim / alpha1).contiguous()) + assert torch.equal(y, expect(lim, lim)), ( + f"accumulator-domain clamp: {(y != expect(lim, lim)).sum().item()} of {y.numel()} " + "elements differ from gate<=L, |up|<=L" + ) + # the readings that are NOT what the kernel does, each far from bit-exact + fracs = {} + for name, cand in ( + ("clamp in the true domain", expect(lim / alpha1, lim / alpha1)), + ("gate clamped on both sides", expect(lim, lim, gate_both=True)), + ("up clamped from above only", expect(lim, lim, up_above_only=True)), + ("up clamped after output1_scale_scalar", expect(lim, lim / (alpha1 * g2))), + ("gate only", expect(lim, None)), + ("up only", expect(None, lim)), + ): + fracs[name] = (y == cand).float().mean().item() + assert fracs[name] < 0.9, f"'{name}' also matches {fracs[name]:.2%} — no discrimination" + # control: a true-domain limit handed over as is clamps at L * alpha1 + raw = kernel(lim.contiguous()) + assert torch.equal(raw, expect(lim * alpha1, lim * alpha1)) + wrong = (raw == expect(lim, lim)).float().mean().item() + assert wrong < 0.9, f"true-domain limit passed raw still matches {wrong:.2%}" + clamped = (gate_t > lim[eid].unsqueeze(1)).float().mean().item() + print( + f" test_v41_clamp_is_in_the_accumulator_domain OK (bit-exact on all {y.numel()} elements, " + f"gate clamped on {clamped:.1%}, {saturated} saturated block scales; other readings match " + f"{max(fracs.values()):.1%} at most; the true-domain limit passed raw matches {wrong:.1%})" + ) + + +def test_v41_routed_geometry(): + """The V4.1 routed expert layer: 384 experts, H=5120, I=2304, top-6, clamp. + + Three layouts, over the whole token column (1 through 8192 rows per call): + the full 384-expert stack at offset 0; the four EP4 windows (96 local + experts at offsets 0 / 96 / 192 / 288, `num_experts` 384 on every call); + and the four MoE-TP4 ranks (`I = 576` — a multiple of 64 but not of 128 — + with all 384 experts local). Every expert has its own FC1/FC2 weight global + scale, FC2-input global scale and clamp limit. + """ + args, ref = _v41() + tp_args = _v41_tp_args() + gen = torch.Generator(device=DEV).manual_seed(4143) + g1 = ref["g1"] + scal = _v41_scalars(ref) + worst = dict( + full=(0.0, 0.0), win=(0.0, 0.0), ep_sum=(0.0, 0.0), tp=(0.0, 0.0), tp_sum=(0.0, 0.0) + ) + worst_ep_vs_full = worst_tp_vs_full = 0.0 + zero_rows = 0 + clamp_frac = [1.0, 1.0] + top = {} + + def keep(name, y, exp): + dv = _dev(y, exp) + worst[name] = tuple(max(a, b) for a, b in zip(worst[name], dv)) + top[name] = tuple(max(a, b) for a, b in zip(top.get(name, (0.0, 0.0)), dv)) + + for num_tokens in _V41_TOKENS: + top = {} + data, sf, x = _rand_activation(num_tokens, V41_H, g1, gen) + ids, wts = _v41_routing(num_tokens, gen) + ref_full, ref_win, ref_tp, frac = _ref_v41(x, ids, wts, ref) + if num_tokens >= 256: + clamp_frac = [min(a, b) for a, b in zip(clamp_frac, frac)] + exp_full = ref_full.to(torch.bfloat16) + full = _call(data, sf, args, V41_E, V41_K, topk_ids=ids, topk_weights=wts, **scal)[0] + assert full.shape == (num_tokens, V41_H) and full.dtype == torch.bfloat16 + _assert_moe_close(full, exp_full) + keep("full", full, exp_full) + + parts = [] + for r in range(V41_E // V41_EP): + off = r * V41_EP + y = _call( + data, + sf, + _window_args(args, off, V41_EP), + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=V41_EP, + topk_ids=ids, + topk_weights=wts, + **_v41_slice(scal, off, V41_EP), + )[0] + exp = ref_win[r].to(torch.bfloat16) + _assert_moe_close(y, exp) + keep("win", y, exp) + outside = ~((ids >= off) & (ids < off + V41_EP)).any(dim=1) + assert torch.equal(y[outside], torch.zeros_like(y[outside])), ( + f"EP window {r} wrote rows routed entirely elsewhere at T={num_tokens}" + ) + zero_rows = max(zero_rows, int(outside.sum())) + parts.append(y) + ep_sum = sum(p.float() for p in parts).to(torch.bfloat16) + _assert_moe_close(ep_sum, exp_full) + keep("ep_sum", ep_sum, exp_full) + worst_ep_vs_full = max(worst_ep_vs_full, _dev(ep_sum, full.float())[0]) + + parts = [] + for r, ta in enumerate(tp_args): + y = _call(data, sf, ta, V41_E, V41_K, topk_ids=ids, topk_weights=wts, **scal)[0] + exp = ref_tp[r].to(torch.bfloat16) + _assert_moe_close(y, exp) + keep("tp", y, exp) + parts.append(y) + tp_sum = sum(p.float() for p in parts).to(torch.bfloat16) + _assert_moe_close(tp_sum, exp_full) + keep("tp_sum", tp_sum, exp_full) + worst_tp_vs_full = max(worst_tp_vs_full, _dev(tp_sum, full.float())[0]) + + # the clamp was genuinely exercised on both halves at every batch of >= 256 rows + assert clamp_frac[0] > 0.01 and clamp_frac[1] > 0.01, clamp_frac + fmt = lambda d: f"{d[0]:.2f} elt / {d[1]:.2f} rms" # noqa: E731 + print( + f" test_v41_routed_geometry OK (worst: full {fmt(worst['full'])}, EP4 window " + f"{fmt(worst['win'])}, EP4 sum {fmt(worst['ep_sum'])}, TP4 rank {fmt(worst['tp'])}, " + f"TP4 sum {fmt(worst['tp_sum'])} ulp; EP4 sum vs full call {worst_ep_vs_full:.2f} elt, " + f"TP4 sum vs full call {worst_tp_vs_full:.2f} elt; <= {zero_rows} bitwise-zero rows per " + f"window; clamp active on >= {clamp_frac[0]:.1%} of gate / {clamp_frac[1]:.1%} of up " + f"elements; at T={_V41_TOKENS[-1]}: full {fmt(top['full'])}, window {fmt(top['win'])}, " + f"TP4 rank {fmt(top['tp'])})" + ) + + +def test_v41_wrong_operands_are_discriminated(): + """The V4.1 layer's plausible mistakes, each run and measured against the gate. + + At T = 1024 on the V4.1 geometry: the clamp limit given in the true domain + (10, not 10 / alpha1); FC1 stacked `[gate ; up]` (the dense-MLP order) + instead of `[up ; gate]`; an EP window issued at `local_expert_offset = 0` + (tokens routed to its experts are silently dropped, others computed with the + wrong expert's weights); per-expert scalars shifted by one expert; one TP4 + rank left out of the sum. Each must land far outside the 2-ulp RMS gate the + correct call passes. + """ + args, ref = _v41() + tp_args = _v41_tp_args() + gen = torch.Generator(device=DEV).manual_seed(4144) + num_tokens = 1024 + data, sf, x = _rand_activation(num_tokens, V41_H, ref["g1"], gen) + ids, wts = _v41_routing(num_tokens, gen) + ref_full, ref_win, ref_tp, _ = _ref_v41(x, ids, wts, ref) + exp_full = ref_full.to(torch.bfloat16) + scal = _v41_scalars(ref) + good = _call(data, sf, args, V41_E, V41_K, topk_ids=ids, topk_weights=wts, **scal)[0] + _assert_moe_close(good, exp_full) + _, good_rms = _dev(good, exp_full) + off = V41_EP # window 1 + win1 = _window_args(args, off, V41_EP) + up_p, gt_p, up_s, gt_s, _, _ = ref["raw"] + w1_swapped, s1_swapped = _prep_fc1(gt_p, up_p, gt_s, up_s) + bad = { + "clamp limit in the true domain": ( + _call( + data, + sf, + args, + V41_E, + V41_K, + topk_ids=ids, + topk_weights=wts, + **{**scal, "gemm1_clamp_limit": torch.full((V41_E,), V41_LIMIT, device=DEV)}, + )[0], + exp_full, + ), + "FC1 stacked [gate ; up]": ( + _call( + data, + sf, + {**args, "gemm1_weights": w1_swapped, "gemm1_weights_scale": s1_swapped}, + V41_E, + V41_K, + topk_ids=ids, + topk_weights=wts, + **scal, + )[0], + exp_full, + ), + "EP window 1 at offset 0": ( + _call( + data, + sf, + win1, + V41_E, + V41_K, + local_expert_offset=0, + local_num_experts=V41_EP, + topk_ids=ids, + topk_weights=wts, + **_v41_slice(scal, off, V41_EP), + )[0], + ref_win[1].to(torch.bfloat16), + ), + "scalars shifted by one expert": ( + _call( + data, + sf, + args, + V41_E, + V41_K, + topk_ids=ids, + topk_weights=wts, + **{k: v.roll(1).contiguous() for k, v in scal.items()}, + )[0], + exp_full, + ), + "one TP4 rank left out of the sum": ( + sum( + _call(data, sf, ta, V41_E, V41_K, topk_ids=ids, topk_weights=wts, **scal)[0].float() + for ta in tp_args[:-1] + ).to(torch.bfloat16), + exp_full, + ), + } + # the correct window-1 call, for the offset control's scale + right = _call( + data, + sf, + win1, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=V41_EP, + topk_ids=ids, + topk_weights=wts, + **_v41_slice(scal, off, V41_EP), + )[0] + _assert_moe_close(right, ref_win[1].to(torch.bfloat16)) + lines = [] + for name, (y, exp) in bad.items(): + elt, rms = _dev(y, exp) + assert rms > 25.0, ( + f"{name}: only {rms:.1f} rms ulp from the correct result — not discriminated" + ) + lines.append(f"{name} {rms:.0f}") + # the swapped-FC1 mistake is also visible from the reference side + swapped_ref = _ref_v41(x, ids, wts, ref, swap_gate_up=True)[0].to(torch.bfloat16) + _, rms_swapped_ref = _dev(good, swapped_ref) + assert rms_swapped_ref > 25.0, rms_swapped_ref + dropped = int(((ids >= off) & (ids < off + V41_EP)).any(dim=1).sum()) + print( + f" test_v41_wrong_operands_are_discriminated OK (correct call {good_rms:.2f} rms ulp; " + f"wrong: {', '.join(lines)} rms ulp; kernel vs a [gate ; up] reference " + f"{rms_swapped_ref:.0f}; the offset-0 window dropped the slots of {dropped} tokens)" + ) + + def test_r1_autotuner_cache_is_inert(): """A warm autotuner profiling cache does not change the R1-geometry result. @@ -1535,6 +2068,4223 @@ def entries(): ) +def test_v41_autotuner_cache_is_inert(): + """A warm autotuner cache does not change the V4.1 results. + + The V4.1 layer presents three tuning keys of its own — the 384-expert stack + and the EP4 windows at `I = 2304` (96 local experts), and the MoE-TP4 ranks + at `I = 576`. This warms all three through `autotune()` at 8192 rows, then + re-runs every layout (full, four windows, four TP ranks) over a spread of + token counts and demands **bitwise** equality with the cold results. + Clears the cache at the end, like the R1 test. + """ + args, ref = _v41() + tp_args = _v41_tp_args() + scal = _v41_scalars(ref) + layouts = ( + [("full", 0, V41_E, args, scal)] + + [ + ( + f"ep{r}", + r * V41_EP, + V41_EP, + _window_args(args, r * V41_EP, V41_EP), + _v41_slice(scal, r * V41_EP, V41_EP), + ) + for r in range(V41_E // V41_EP) + ] + + [(f"tp{r}", 0, V41_E, ta, scal) for r, ta in enumerate(tp_args)] + ) + tokens = (1, 8, 32, 256, 4096, 8192) + gen = torch.Generator(device=DEV).manual_seed(4145) + cases = {} + for num_tokens in tokens: + data, sf, _ = _rand_activation(num_tokens, V41_H, ref["g1"], gen) + cases[num_tokens] = (data, sf) + _v41_routing(num_tokens, gen) + + def run(layout, num_tokens): + _, off, n, la, ls = layout + data, sf, ids, wts = cases[num_tokens] + return _call( + data, + sf, + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=n, + topk_ids=ids, + topk_weights=wts, + **ls, + )[0] + + tuner = AutoTuner.get() + op = "trtllm::fp4_block_scale_moe_runner" + + def entries(): + return [k for k in tuner.profiling_cache.cache if k[0] == op] + + assert not entries(), "the profiling cache already holds entries for this op" + cold = {(lay[0], t): run(lay, t) for lay in layouts for t in tokens} + assert not entries(), "an ordinary call wrote a profiling-cache entry" + with autotune(): + for lay in layouts: + if lay[0] in ("full", "ep0", "tp0"): + run(lay, tokens[-1]) + keys = entries() + unique_ids = sorted({str(k[2]) for k in keys}) + assert len(unique_ids) == 3, f"expected three tuning keys, got {unique_ids}" + tactics = {str(tuner.profiling_cache.cache[k][1]) for k in keys} + per_uid = { + u: {str(tuner.profiling_cache.cache[k][1]) for k in keys if str(k[2]) == u} + for u in unique_ids + } + assert all(v - {"-1"} for v in per_uid.values()), ( + f"a tuning key recorded only the fallback tactic, so warm == cold proves nothing there: {per_uid}" + ) + for lay in layouts: + for num_tokens in tokens: + warm = run(lay, num_tokens) + assert torch.equal(warm, cold[(lay[0], num_tokens)]), ( + f"a warm autotuner cache changed the {lay[0]} result at T={num_tokens}: " + f"{_dev(warm, cold[(lay[0], num_tokens)].float())} (elt / rms ulp)" + ) + tuner.clear_cache() + print( + f" test_v41_autotuner_cache_is_inert OK ({len(keys)} tuned entries over unique_ids " + f"{unique_ids}, {len(tactics)} distinct tactics; all {len(layouts) * len(tokens)} " + f"(layout, T) results bitwise unchanged)" + ) + + +# ── V4.1 decode under CUDA-graph capture ────────────────────────────────── +# The V4.1 target calls this op inside its decode CUDA graphs, as an EP4 window +# (tep4 / dep4: 96 local experts, I = 2304) or a MoE-TP4 rank (tp4: all 384 +# experts, I = 576), with the FC2-input global scale LAYER-level — one +# g2 = 1 / max_e w2.input_scale[e], i.e. min_e g2[e], in every slot — and its +# output buffer allocated inside the captured region. +# +# The captured row counts, per rank. `_GSET` is every CUDA-graph batch G an +# engine at max_batch_size <= 256 captures: every G 1..32 plus TRT-LLM's padded +# list at 256 (`_generate_cuda_graph_batch_sizes`: 1, 2, 4, 8, 16..128 in steps +# of 8, 192, 256); the padding-off list at 256 (1..32, 64, 128, 256) is a +# subset. A captured step is generation-only with next_n 1 (decode) or 6 +# (DSpark target-verify: 1 + 5 drafts), so a tep4 EP4 window and a tp4 rank +# route T = G and 6G of their own rows (82 counts, 1..1536). dep4 gathers the +# four ranks' rows, each padded to the same G, ahead of the router onto its EP4 +# windows: T = 4G and 24G (82 counts, 4..6144). The routed MoE chunks at 8192 +# rows (`_moe_chunks`), so no captured call is split. dep4 pads each rank's +# block to the largest rank's row count with zero rows, which the router sends +# to one set of six experts; the "pad" case models that at every dep4 count. + +_GSET = tuple(range(1, 33)) + (40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256) +_V41_RANK_TOKENS = tuple(sorted(set(_GSET) | {6 * g for g in _GSET})) # tep4 / tp4: G, 6G +_V41_DEP4_TOKENS = tuple(sorted({4 * g for g in _GSET} | {24 * g for g in _GSET})) # dep4: 4G, 24G +_V41_GATHERED_TOKENS = tuple( + n for n in _V41_DEP4_TOKENS if n not in _V41_RANK_TOKENS +) # EP4 windows only +_V41_GRAPH_TOKENS = _V41_RANK_TOKENS +_MOE_OP = "trtllm::fp4_block_scale_moe_runner" + + +def _v41_counts(base, nloc): + """The captured row counts of one rank layout: `base`, plus dep4's gathered + verify counts on an EP4 window (a TP4 rank is never gathered).""" + return base + (_V41_GATHERED_TOKENS if nloc < V41_E else ()) + + +def _moe_cache_entries(): + return [k for k in AutoTuner.get().profiling_cache.cache if k[0] == _MOE_OP] + + +def _v41_layer_g2(): + """The V4.1 layer with its FC2-input global scale made layer-level.""" + args, ref = _v41() + return args, dict(ref, g2=torch.full_like(ref["g2"], float(ref["g2"].min()))) + + +def _v41_graph_layouts(args, ref_l): + """(name, offset, local experts, operands, scalars) of the eight rank layouts.""" + scal = _v41_scalars(ref_l) + lays = [ + ( + f"ep{r}", + r * V41_EP, + V41_EP, + _window_args(args, r * V41_EP, V41_EP), + _v41_slice(scal, r * V41_EP, V41_EP), + ) + for r in range(V41_E // V41_EP) + ] + lays += [(f"tp{r}", 0, V41_E, ta, scal) for r, ta in enumerate(_v41_tp_args())] + return lays + + +def _v41_case(kind, n, off, nloc, g1, gen, ranks=4): + """One replay's inputs. `skew` puts every token on the same six local experts; + `none` routes every token outside this rank's window; `pad` is a dep member's + zero-padded gathered batch over `ranks` ranks (dep4 by default: EP4 windows, + `n` a multiple of 4; dep2 and dep8 on the EP2 and EP8 windows).""" + data, sf, x = _rand_activation(n, V41_H, g1, gen) + ids, wts = _v41_routing(n, gen) + if kind == "skew": + ids = ( + (off + torch.arange(V41_K, device=DEV, dtype=torch.int32)).expand(n, V41_K).contiguous() + ) + elif kind == "pad": + # The gathered batch: `ranks` equal rank blocks, rank 0 full, the others + # ending in 0..rows-1 zero rows, every zero row routed to the same six + # experts (spread over this window) with the same weights. + rows = n // ranks + pad = torch.zeros(n, dtype=torch.bool, device=DEV) + for b in range(1, ranks): + p = int(torch.randint(0, rows, (1,), device=DEV, generator=gen)) + pad[(b + 1) * rows - p : (b + 1) * rows] = True + six = off + torch.tensor([5, 17, 33, 50, 71, 90], device=DEV, dtype=torch.int32) % nloc + ids[pad] = six + wts[pad] = torch.tensor([0.4, 0.3, 0.3, 0.2, 0.2, 0.1], device=DEV, dtype=torch.bfloat16) + data[pad] = 0 + sf.view(n, V41_H // SV)[pad] = 0 + x[pad] = 0 + elif kind == "none": + nxt = (off + nloc) % V41_E + pick = torch.rand(n, nloc, device=DEV, generator=gen).argsort(dim=1)[:, :V41_K] + ids = (nxt + pick).to(torch.int32).contiguous() + return kind, data, sf, x, ids, wts + + +class _V41Graph: + """One captured call over static input buffers, the way a decode graph holds it. + + With `chunks` (row ranges `[a, b)` tiling `n`, as `_moe_chunks` cuts a gathered + batch) the graph holds one call per chunk, each on its rows of the static + buffers and into its own output allocated inside the captured region.""" + + def __init__(self, lay, n, first, chunks=None): + self.lay, self.n = lay, n + self.chunks = chunks or [(0, n)] + self.data = torch.empty(n, V41_H // 2, dtype=torch.uint8, device=DEV) + self.sf = torch.empty(n * V41_H // SV, dtype=torch.uint8, device=DEV) + self.ids = torch.empty(n, V41_K, dtype=torch.int32, device=DEV) + self.wts = torch.empty(n, V41_K, dtype=torch.bfloat16, device=DEV) + self._load(first) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + self._body() + torch.cuda.current_stream().wait_stream(side) + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph): + self.out = self._body() + + def _body(self): + _, off, nloc, la, ls = self.lay + outs = [] + for a, b in self.chunks: + out = torch.empty(b - a, V41_H, dtype=torch.bfloat16, device=DEV) + _call( + self.data[a:b], + self.sf[a * V41_H // SV : b * V41_H // SV], + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=nloc, + topk_ids=self.ids[a:b], + topk_weights=self.wts[a:b], + output=out, + **ls, + ) + outs.append(out) + return outs[0] if len(outs) == 1 else outs + + def _load(self, case): + _, data, sf, _, ids, wts = case + self.data.copy_(data) + self.sf.copy_(sf) + self.ids.copy_(ids) + self.wts.copy_(wts) + + def replay(self, case): + self._load(case) + self.graph.replay() + torch.cuda.synchronize() + return self.out.clone() if len(self.chunks) == 1 else torch.cat(self.out) + + def kernels(self, case): + """Names of the CUDA kernels one replay launches.""" + self._load(case) + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + self.graph.replay() + torch.cuda.synchronize() + return [ + e.name + for e in prof.events() + if e.device_type == torch.autograd.DeviceType.CUDA and e.self_device_time_total > 0 + ] + + +def _v41_eager(lay, case, chunks=None): + """The same call outside any graph, into a NaN-filled output buffer; with + `chunks`, one call per chunk, each into its own NaN-filled buffer.""" + _, off, nloc, la, ls = lay + _, data, sf, _, ids, wts = case + outs = [] + for a, b in chunks or [(0, data.shape[0])]: + out = torch.full((b - a, V41_H), float("nan"), dtype=torch.bfloat16, device=DEV) + _call( + data[a:b], + sf[a * V41_H // SV : b * V41_H // SV], + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=nloc, + topk_ids=ids[a:b], + topk_weights=wts[a:b], + output=out, + **ls, + ) + outs.append(out) + return outs[0] if len(outs) == 1 else torch.cat(outs) + + +def test_v41_cuda_graph_replay_is_bitwise_eager(): + """Captured once per (rank layout, T), replayed on routing the capture never saw. + + Cold autotuner cache (fallback tactic), the target's call shape: pre-routed, + clamp at 10, layer-level g2, `output` allocated inside the captured region. + Each graph is captured after one side-stream warm-up call and replayed on two + uniform routings, a skewed one (every token on the same six local experts) + and — EP windows — one routing nothing into the window, then on the capture's + own data again. At dep4's gathered counts an EP window also replays a + zero-padded batch (`pad`). Every replay must be **bitwise** equal to an + eager call on the same values; the replayed output is also held to the + torch reference's gates (the empty-window replay to exact zeros). + Row counts: G and 6G over `_GSET` (82 counts, 1..1536) on every layout, and + dep4's gathered 4G and 24G (82 counts, 4..6144) on the EP4 windows. The + reference is `_ref_v41_fast`, first tied bitwise to `_ref_v41`. + """ + assert not _moe_cache_entries(), "the cold-cache graph test ran with a warm cache" + args, ref_l = _v41_layer_g2() + layouts = _v41_graph_layouts(args, ref_l) + g1 = ref_l["g1"] + gen = torch.Generator(device=DEV).manual_seed(4146) + # The fast reference is the certified one, bit for bit (windows and TP slices). + for kind, n, off, nloc in ( + ("uniform", 48, 0, V41_E), + ("pad", 768, 96, V41_EP), + ("skew", 1536, 0, V41_E), + ): + c = _v41_case(kind, n, off, nloc, g1, gen) + fast, slow = _ref_v41_fast(c[3], c[4], c[5], ref_l), _ref_v41(c[3], c[4], c[5], ref_l) + for r in range(4): + assert torch.equal(fast[0][r], slow[1][r]) and torch.equal(fast[1][r], slow[2][r]), ( + kind, + n, + r, + ) + counts = sorted(set(_V41_GRAPH_TOKENS) | set(_V41_GATHERED_TOKENS)) + assert len(_V41_GRAPH_TOKENS) == 82 and len(_V41_DEP4_TOKENS) == 82 and len(counts) == 122 + assert max(_V41_GRAPH_TOKENS) == 1536 and max(counts) == 6144 + worst = (0.0, 0.0) + worst_at = [None, None] + worst_big = (0.0, 0.0) + replays = graphs = 0 + per_layout = {} + for n in counts: + lays = [lay for lay in layouts if n in _v41_counts(_V41_GRAPH_TOKENS, lay[2])] + shared = [_v41_case("uniform", n, 0, V41_E, g1, gen) for _ in range(2)] + refs = [ + _ref_v41_fast(c[3], c[4], c[5], ref_l, slices=n in _V41_GRAPH_TOKENS) for c in shared + ] + for lay in lays: + name, off, nloc = lay[:3] + r = int(name[2:]) + own = [_v41_case("skew", n, off, nloc, g1, gen)] + if nloc < V41_E: + own.append(_v41_case("none", n, off, nloc, g1, gen)) + if n in _V41_DEP4_TOKENS: + own.append(_v41_case("pad", n, off, nloc, g1, gen)) + exps = [rf[0 if nloc < V41_E else 1][r] for rf in refs] + for c in own: + rf = ( + _ref_v41_fast(c[3], c[4], c[5], ref_l, slices=nloc == V41_E) + if c[0] != "none" + else None + ) + exps.append(rf[0 if nloc < V41_E else 1][r] if rf else None) + cases = shared + own + shared[:1] + exps = exps + exps[:1] + g = _V41Graph(lay, n, cases[0]) + prev = None + for case, exp in zip(cases, exps): + y = g.replay(case) + eager = _v41_eager(lay, case) + assert torch.equal(y, eager), ( + f"{name} T={n} {case[0]}: graph replay differs from eager " + f"{_dev(y, eager.float())} (elt / rms ulp)" + ) + if exp is None: + assert torch.equal(y, torch.zeros_like(y)), ( + f"{name} T={n}: an empty window is not zero" + ) + else: + exp = exp.to(torch.bfloat16) + _assert_moe_close(y, exp) + dv = _dev(y, exp) + for i in range(2): + if dv[i] > worst[i]: + worst_at[i] = (name, n, case[0]) + worst = tuple(max(a, b) for a, b in zip(worst, dv)) + if n > max(_V41_GRAPH_TOKENS): + worst_big = tuple(max(a, b) for a, b in zip(worst_big, dv)) + # consecutive replays compute different results, so a replay that + # kept serving the previous inputs could not pass (two all-zero + # results — few tokens, none routed into this window — are exempt; + # the skew case is never zero and the empty-window case always is) + assert prev is None or not torch.equal(prev, y) or not y.any(), ( + f"{name} T={n}: two cases gave one result" + ) + prev = y + replays += 1 + graphs += 1 + per_layout[name] = per_layout.get(name, 0) + 1 + del g + del refs, shared + assert not _moe_cache_entries() + assert per_layout == { + **{f"ep{r}": 122 for r in range(4)}, + **{f"tp{r}": 82 for r in range(4)}, + }, per_layout + print( + f" test_v41_cuda_graph_replay_is_bitwise_eager OK ({graphs} (layout, T) graphs over " + f"{len(layouts)} layouts {per_layout}, {replays} replays bitwise equal to eager; " + f"worst vs reference {worst[0]:.2f} elt at {worst_at[0]} / {worst[1]:.2f} rms ulp at {worst_at[1]}, " + f"at T 1537..6144 (dep4 only) " + f"{worst_big[0]:.2f} elt / {worst_big[1]:.2f} rms ulp; layer g2 {float(ref_l['g2'][0]):.4g}; " + f"EP4 counts {counts}; TP4 counts {list(_V41_GRAPH_TOKENS)})" + ) + + +def test_v41_cuda_graph_swizzled_scales_are_caught_at_gathered_counts(): + """Control: the graph harness sees the one silent layout mistake at dep4's counts. + + dep4's gathered counts 384, 768, 1536 and 6144 (24G at G 16, 32, 64, 256) + are multiples of 128, where a 128x4-swizzled activation-scale buffer has + exactly the linear one's byte count and passes the op's size check. + Replaying a captured EP4 window (offsets 0 and 288) on the same scale bytes + in swizzled order must fail the reference gate, while the linear bytes of + the same case pass it. + """ + args, ref_l = _v41_layer_g2() + layouts = [lay for lay in _v41_graph_layouts(args, ref_l) if lay[0] in ("ep0", "ep3")] + g1 = ref_l["g1"] + gen = torch.Generator(device=DEV).manual_seed(4148) + seen = [] + for n in (384, 768, 1536, 6144): + case = _v41_case("uniform", n, 0, V41_E, g1, gen) + windows = _ref_v41_windows(case[3], case[4], case[5], ref_l) + kind, data, sf, x, ids, wts = case + swz = _swizzle_scales(sf.view(1, n, V41_H // SV)).reshape(-1).contiguous() + assert swz.numel() == sf.numel() and not torch.equal(swz, sf) + for lay in layouts: + name, off = lay[:2] + exp = windows[off // V41_EP].to(torch.bfloat16) + g = _V41Graph(lay, n, case) + _assert_moe_close(g.replay(case), exp) + bad = g.replay((kind, data, swz, x, ids, wts)) + elt, rms = _dev(bad, exp) + assert rms > 2.0 and elt > 16.0, ( + f"{name} T={n}: swizzled scales passed ({elt:.1f} / {rms:.2f} ulp)" + ) + seen.append((name, n, round(elt), round(rms))) + del g + print( + " test_v41_cuda_graph_swizzled_scales_are_caught_at_gathered_counts OK " + f"(swizzled replay vs reference, elt / rms ulp: {seen}; gates 16 / 2)" + ) + + +# The first row count of each kernel change the eager profile of +# `test_v41_row_axis_by_bucket` finds besides the tuner buckets and untuned +# tiles: the routing kernel's instantiation (5, 65, 129) and the finalize +# kernel's vector-load variant (60). +_V41_ROUTING_FIRSTS = (5, 60, 65, 129) + + +def _v41_segment(n: int, nloc: int, warm: bool): + """(lo, hi): the row counts around `n` that launch `n`'s kernels, per the eager + profile of `test_v41_row_axis_by_bucket` -- `n`'s tuner bucket (warm) or untuned + token tile (cold), cut at the routing / finalize kernel changes and, on an untuned + TP4 rank, at 3 (its FC1 config).""" + firsts = set(_V41_ROUTING_FIRSTS) + if warm: + firsts |= {1 << k for k in range(1, 14)} + else: + firsts |= { + m for m in range(2, 8193) if _v41_cold_tile(m, nloc) != _v41_cold_tile(m - 1, nloc) + } + if nloc == V41_E: + firsts.add(3) + lo = max([1] + [f for f in firsts if f <= n]) + hi = min([8193] + [f for f in firsts if f > n]) - 1 + return lo, hi + + +def test_v41_cuda_graph_tactic_is_fixed_at_capture(): + """A graph keeps the tactic it was captured with; no tactic changes the result. + + The engine autotunes during warm-up and captures its decode graphs after, so + its graphs hold tuned tactics. Here: graphs are captured cold (fallback + tactic), then the cache is warmed through `autotune()` at 8192 rows for the + two V4.1 tuning keys; the cold graphs replay unchanged — same kernels, same + bits — while graphs captured warm launch the tuned kernels (different ones + for some (layout, T), which is what makes the comparison mean something) and + are still **bitwise** equal to the cold results, and so to a warm eager call. + Clearing the cache after capture changes neither. Clears the cache at the + end. Row counts: G and 6G over `_GSET` on every layout, and dep4's gathered + 4G and 24G on the EP4 windows (the counts of the replay test). + + Per-bucket kernel choice: on EP4 window 0 and TP4 rank 0 at every captured + count, and on all four EP4 windows at dep4's counts above 1536, the kernels a + graph replays (captured cold, then warm) are exactly the kernels an eager call + launches at that count in the same tuner state, and at the lower and upper + edge of the count's bucket (warm) or untuned tile (cold) (`_v41_segment`) -- + the edges `test_v41_row_axis_by_bucket` holds to the reference. + """ + assert not _moe_cache_entries(), "the profiling cache already holds entries for this op" + args, ref_l = _v41_layer_g2() + layouts = _v41_graph_layouts(args, ref_l) + g1 = ref_l["g1"] + gen = torch.Generator(device=DEV).manual_seed(4147) + kgen = torch.Generator(device=DEV).manual_seed(4157) + eager_k = {"cold": {}, "warm": {}} + + def eager_kernels(state, lay, m): + """Kernels of one eager call at row count m on `lay`, output preallocated.""" + key = (lay[0], m) + if key not in eager_k[state]: + _, off, nloc, la, ls = lay + c = _v41_case("uniform", m, 0, V41_E, g1, kgen) + out = torch.empty(m, V41_H, dtype=torch.bfloat16, device=DEV) + + def run(): + _call( + c[1], + c[2], + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=nloc, + topk_ids=c[4], + topk_weights=c[5], + output=out, + **ls, + ) + + run() + eager_k[state][key] = _kernel_names(run) + return eager_k[state][key] + + big = max(_V41_GRAPH_TOKENS) + segs = {"cold": {}, "warm": {}} + + def check_bucket(state, lay, n, k_graph): + if not (lay[0] in ("ep0", "tp0") or n > big): + return + lo, hi = _v41_segment(n, lay[2], state == "warm") + for m in (n, lo, hi): + k = eager_kernels(state, lay, m) + assert tuple(k_graph) == k, ( + f"{state} {lay[0]} T={n}: the graph's kernels differ from an eager call at T={m} " + f"(segment {lo}..{hi}): {k_graph} vs {k}" + ) + segs[state][(lay[0], n)] = (lo, hi) + + cold = {} + for lay in layouts: + name, off, nloc = lay[:3] + for n in _v41_counts(_V41_GRAPH_TOKENS, nloc): + kinds = ("uniform", "skew") + (("none",) if nloc < V41_E else ()) + cases = [_v41_case(k, n, off, nloc, g1, gen) for k in kinds] + g = _V41Graph(lay, n, cases[0]) + k_cold = g.kernels(cases[0]) + check_bucket("cold", lay, n, k_cold) + cold[(name, n)] = (g, cases, [_v41_eager(lay, c) for c in cases], k_cold) + assert not _moe_cache_entries(), "an ordinary or captured call wrote a profiling-cache entry" + + with autotune(): + for lay in layouts: + if lay[0] in ("ep0", "tp0"): + _, off, nloc, la, ls = lay + c = _v41_case("uniform", 8192, off, nloc, g1, gen) + _call( + c[1], + c[2], + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=nloc, + topk_ids=c[4], + topk_weights=c[5], + **ls, + ) + keys = _moe_cache_entries() + + def tactic(k): + t = AutoTuner.get().profiling_cache.cache[k][1] + return tuple(t) if isinstance(t, (list, tuple)) else (t,) + + tuned = {(str(k[2]), k[3][2][0]): tactic(k) for k in keys} + uids = sorted({u for u, _ in tuned}) + assert uids == ["(6, 2304, 96, 0)", "(6, 576, 384, 0)"], uids + # every token bucket a graph batch maps to (the power of two at or below T) + for u in uids: + for n in _v41_counts(_V41_GRAPH_TOKENS, V41_EP if u == "(6, 2304, 96, 0)" else V41_E): + b = 1 << (n.bit_length() - 1) + assert tuned.get((u, b), (-1,)) != (-1,), ( + f"bucket {b} of {u} is not tuned: {tuned.get((u, b))}" + ) + + changed = [] + for lay in layouts: + name = lay[0] + for n in _v41_counts(_V41_GRAPH_TOKENS, lay[2]): + g, cases, want, k_cold = cold[(name, n)] + for c, w in zip(cases, want): + assert torch.equal(g.replay(c), w), ( + f"{name} T={n} {c[0]}: a cold graph changed after warming" + ) + assert torch.equal(_v41_eager(lay, c), w), ( + f"{name} T={n} {c[0]}: warm eager != cold eager" + ) + assert g.kernels(cases[0]) == k_cold, ( + f"{name} T={n}: a cold graph re-read the tuner at replay" + ) + del g + warm = _V41Graph(lay, n, cases[0]) + k_warm = warm.kernels(cases[0]) + check_bucket("warm", lay, n, k_warm) + if k_warm != k_cold: + changed.append((name, n)) + for c, w in zip(cases, want): + assert torch.equal(warm.replay(c), w), ( + f"{name} T={n} {c[0]}: warm graph != cold result" + ) + cold[(name, n)] = (warm, cases, want, k_warm) + assert changed, ( + "no graph captured warm launched different kernels — the comparison proved nothing" + ) + AutoTuner.get().clear_cache() + for (name, n), (warm, cases, want, k_warm) in cold.items(): + assert warm.kernels(cases[0]) == k_warm, f"{name} T={n}: a warm graph re-read the tuner" + for c, w in zip(cases, want): + assert torch.equal(warm.replay(c), w), ( + f"{name} T={n} {c[0]}: warm graph changed after clear" + ) + del cold + assert len(segs["cold"]) == len(segs["warm"]) == 122 + 82 + 3 * 10, ( + len(segs["cold"]), + len(segs["warm"]), + ) + above = sorted(n for n in _V41_GATHERED_TOKENS if n > big) + by_seg = { + state: { + f"{lo}..{hi}": [n for n in above if segs[state][("ep0", n)] == (lo, hi)] + for lo, hi in sorted({segs[state][("ep0", n)] for n in above}) + } + for state in ("cold", "warm") + } + print( + f" test_v41_cuda_graph_tactic_is_fixed_at_capture OK ({len(keys)} tuned entries over {uids}, " + f"{len(set(tuned.values()))} distinct tactics; " + f"{sum(len(_v41_counts(_V41_GRAPH_TOKENS, lay[2])) for lay in layouts)} " + f"(layout, T) graphs, cold- and warm-captured, bitwise equal on every replay; warm capture launched " + f"different kernels at {len(changed)} of them, e.g. {changed[:4]}; at T > {big}: " + f"{[c for c in changed if c[1] > big]}; per-bucket kernel check at {len(segs['cold'])} (layout, T) " + f"per state: graph kernels == eager kernels at T and at its segment's edges; dep4 counts {above} by " + f"segment: cold {by_seg['cold']}, warm {by_seg['warm']})" + ) + + +# ── V4.1 eager row axis, bucket by bucket ────────────────────────────────── +# The op picks its kernels by row count in two places, so its row axis is +# certified by bucket rather than by a prefix argument: +# - the autotuner's token bucket (`FP4BlockScaleMoERunner.get_dynamic_tensor_specs`, +# `_torch/custom_ops/trtllm_gen_custom_ops.py`): the largest power of two at or +# below T, capped at 8192, so the buckets are [2^k, 2^(k+1) - 1] for k < 13 and +# {8192}. A warm cache holds one tactic per bucket, tuned at its lower edge; +# - with no tuned tactic, the C++ runner (`FP4BlockScaleMoeRunner::run`, +# `thop/moe/trtllmGen/fp4BlockScaleMoe.cpp`) picks the token tile as the next +# power of two of the integer part of T * top_k / local_num_experts, clamped to +# 8..256 — so the tile edges fall at T = 16 * (tile + 1) for an EP4 window and +# 64 * (tile + 1) for a TP4 rank — and takes the first valid config of that tile. +# Around the GEMMs, kernels also change with T independently of both: the +# pre-routed routing step runs a single-block kernel up to 4 rows and a +# single-cluster kernel above (`RoutingFromTopKIds.cu`), whose instantiation +# changes again above 64 and 128 rows, and the finalize kernel turns into its +# vector-load variant from 60 rows. Profiled at every T from 1 to 8192 on one +# EP4 window and one TP4 rank, cold and warm, the launched kernels change only +# at the bucket edges, the untuned-tile edges and these (plus 2 / 3, where the +# untuned TP4 rank's FC1 config changes); both sides of each are run below. +# Every row count the target passes is at most 8192 (`_moe` chunks the gathered +# rows by 8192). + +_V41_DUMPED_TOKENS = (48, 56, 192, 1018, 1935, 2048, 3552, 3560) # EP4 (TP4: 48, 1018, 1935, 2048) +# Every row count the batch-256 / long-context engine dumps show on the EP4 windows +# (tep4 graph and eager at max_batch_size 256, L 1M at batch 8, L 128K at batch 32) +# and the TP4 ranks (tp4 graph at 256), captured and eager. +_V41_DUMPED_TOKENS_B256 = ( + 1, + 2, + 4, + 5, + 6, + 7, + 8, + 12, + 16, + 18, + 24, + 29, + 32, + 34, + 48, + 60, + 96, + 108, + 128, + 134, + 144, + 168, + 192, + 228, + 240, + 288, + 330, + 336, + 384, + 432, + 462, + 480, + 528, + 546, + 576, + 610, + 624, + 672, + 720, + 768, + 780, + 852, + 1152, + 1172, + 1178, + 1536, + 2252, + 2255, + 2259, + 2778, + 2784, + 3870, + 4208, + 4245, + 4388, + 4394, + 4591, + 4868, + 8116, + 8124, + 8128, + 8130, + 8137, + 8141, + 8148, + 8192, +) +_V41_KERNEL_EDGES = ( + 2, + 3, + 4, + 5, + 59, + 60, + 64, + 65, + 128, + 129, +) # routing / finalize / TP4 cold FC1 config + + +def _v41_bucket_tokens(): + """Lower edge, one interior point and upper edge of every autotuner bucket.""" + out, b = set(), 1 + while b <= 8192: + hi = min(2 * b - 1, 8192) + out |= {b, (b + hi) // 2, hi} + b *= 2 + return out + + +def _v41_cold_tile(n: int, nloc: int) -> int: + """The untuned runner's token tile (the C++ runner's rule, int truncation and all).""" + v = int(float(n * V41_K) / nloc) + p = 0 if v <= 0 else 1 << (v - 1).bit_length() + return min(max(p, 8), 256) + + +def _v41_tile_edges(nloc: int) -> set: + """The last and first row count of every untuned-tile change up to 8192.""" + out = set() + for n in range(2, 8193): + if _v41_cold_tile(n, nloc) != _v41_cold_tile(n - 1, nloc): + out |= {n - 1, n} + return out + + +def _kernel_names(fn): + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + return tuple( + e.name + for e in prof.events() + if e.device_type == torch.autograd.DeviceType.CUDA and e.self_device_time_total > 0 + ) + + +def test_v41_row_axis_by_bucket(): + """Every tuner bucket and untuned tile up to 8192, at the V4.1 layouts, cold and warm. + + For each EP4 window (96 local experts, I = 2304) and each TP4 rank (all 384, + I = 576), with the target's call (pre-routed, clamp 10, layer-level g2): the + lower edge, an interior point and the upper edge of every autotuner bucket, + both sides of every untuned-tile edge of that layout and of every other + observed kernel change, and every row count the runtime-value dump shows. Each call is held to the torch reference; + rows need not be bitwise equal across row counts. Then the cache is warmed + through `autotune()` at 8192 rows (one tactic per bucket, tuned at its lower + edge) and every call is re-run on the same inputs against the reference — + the warm tactic at a bucket's interior and upper edge is the one engines + run. Clears the cache at the end. + """ + assert not _moe_cache_entries(), "the profiling cache already holds entries for this op" + args, ref_l = _v41_layer_g2() + layouts = _v41_graph_layouts(args, ref_l) + g1 = ref_l["g1"] + gen = torch.Generator(device=DEV).manual_seed(4149) + common = ( + _v41_bucket_tokens() + | set(_V41_KERNEL_EDGES) + | set(_V41_DUMPED_TOKENS) + | set(_V41_DUMPED_TOKENS_B256) + ) + counts = { + V41_EP: sorted(common | _v41_tile_edges(V41_EP)), + V41_E: sorted(common | _v41_tile_edges(V41_E)), + } + assert {143, 144, 2063, 2064} <= set(counts[V41_EP]) and {575, 576, 4159, 4160} <= set( + counts[V41_E] + ) + all_counts = sorted(set(counts[V41_EP]) | set(counts[V41_E])) + + def run(lay, case): + _, off, nloc, la, ls = lay + _, data, sf, _, ids, wts = case + return _call( + data, + sf, + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=nloc, + topk_ids=ids, + topk_weights=wts, + **ls, + )[0] + + cases, cold, kern = {}, {}, {} + worst = {"ep": (0.0, 0.0), "tp": (0.0, 0.0)} + worst_at = {"ep": [None, None], "tp": [None, None]} + for n in all_counts: + case = _v41_case("uniform", n, 0, V41_E, g1, gen) + _, windows, slices, _ = _ref_v41(case[3], case[4], case[5], ref_l) + exps = {} + for lay in layouts: + name, off, nloc = lay[:3] + if n not in counts[nloc]: + continue + exp = (windows if nloc < V41_E else slices)[int(name[2:])].to(torch.bfloat16) + y = run(lay, case) + _assert_moe_close(y, exp) + dv = _dev(y, exp) + for i in range(2): + if dv[i] > worst[name[:2]][i]: + worst_at[name[:2]][i] = (name, n) + worst[name[:2]] = tuple(max(a, b) for a, b in zip(worst[name[:2]], dv)) + cold[(name, n)] = y + exps[name] = exp + if name in ("ep0", "tp0"): + kern[(name, n)] = _kernel_names(lambda: run(lay, case)) + cases[n] = (case, exps) + assert not _moe_cache_entries(), "an ordinary call wrote a profiling-cache entry" + + def runs(name): + """Consecutive row counts launching the same kernels, as (first, last) ranges.""" + out = [] + for n in counts[V41_EP if name == "ep0" else V41_E]: + if out and kern[(name, n)] == out[-1][2]: + out[-1][1] = n + else: + out.append([n, n, kern[(name, n)]]) + return [(a, b) for a, b, _ in out] + + cold_runs = {name: runs(name) for name in ("ep0", "tp0")} + + with autotune(): + for lay in layouts: + if lay[0] in ("ep0", "tp0"): + c = _v41_case("uniform", 8192, lay[1], lay[2], g1, gen) + run(lay, c) + uids = sorted({str(k[2]) for k in _moe_cache_entries()}) + assert uids == ["(6, 2304, 96, 0)", "(6, 576, 384, 0)"], uids + same = total = 0 + worst_warm = (0.0, 0.0) + for n in all_counts: + case, exps = cases[n] + for lay in layouts: + name = lay[0] + if name not in exps: + continue + y = run(lay, case) + _assert_moe_close(y, exps[name]) + worst_warm = tuple(max(a, b) for a, b in zip(worst_warm, _dev(y, exps[name]))) + same += int(torch.equal(y, cold[(name, n)])) + total += 1 + if name in ("ep0", "tp0"): + kern[(name, n)] = _kernel_names(lambda: run(lay, case)) + warm_runs = {name: runs(name) for name in ("ep0", "tp0")} + AutoTuner.get().clear_cache() + print( + f" test_v41_row_axis_by_bucket OK (EP4 windows at {len(counts[V41_EP])} row counts, TP4 ranks at " + f"{len(counts[V41_E])}, 1..8192; cold worst EP {worst['ep'][0]:.2f} / {worst['ep'][1]:.2f}, TP " + f"{worst['tp'][0]:.2f} / {worst['tp'][1]:.2f} elt / rms ulp (at {worst_at}); warm worst {worst_warm[0]:.2f} / " + f"{worst_warm[1]:.2f}; warm bitwise equal to cold on {same} of {total}; kernel runs cold {cold_runs}; " + f"warm {warm_runs})" + ) + + +# ── dep4's zero-padding skew, eager, at every chunk size ─────────────────── +# dep4 pads every rank's rows to the step's largest per-rank count with zero +# rows before the all-gather (`_moe` in the unified target), so the gathered +# batch is four equal blocks, rank r's real rows first and its padding after, +# and `_moe_chunks` cuts it into calls of 8192 rows and one remainder: every +# eager row count is a multiple of 4 in 4..8192. A zero row has zero router +# logits, so the router sends every one of them to the six bias-top experts +# with weight 1.5 / 6 = 0.25, and `fp4_quantize` turns it into zero codes and +# zero scale bytes. A chunk can therefore hold up to 8191 padding rows (a rank +# with one real token behind a full rank) or none but padding (the tail of a +# short rank's block), all on six experts, which in one window can be all six. + +_DEP4_SIX_SPREAD = (5, 17, 101, 200, 290, 377) # 2 / 1 / 1 / 2 of the six per EP4 window + +# Every row count the dep4 accuracy-gate preview dumped for the EP4 windows +# (spool/logs/acc-preview-dep4.log, lines of exactly one 5-tuple; garbled lines +# dropped). All are multiples of 4. +_DEP4_DUMPED_TOKENS = ( + 4, + 8, + 16, + 32, + 64, + 96, + 128, + 236, + 256, + 348, + 356, + 396, + 468, + 512, + 544, + 568, + 580, + 1256, + 1580, + 1968, + 2016, + 2320, + 2344, + 2408, + 2440, + 2500, + 2512, + 2520, + 2528, + 2644, + 2648, + 2652, + 2664, + 2696, + 2700, + 2708, + 2716, + 2720, + 2732, + 2736, + 2740, + 2748, + 2768, + 2772, + 2792, + 2812, + 2820, + 2824, + 2828, + 2860, + 2864, + 2876, + 2884, + 2888, + 2900, + 2908, + 2928, + 2932, + 2940, + 2968, + 2976, + 3000, + 3004, + 3024, + 3028, + 3052, + 3056, + 3072, + 3076, + 3080, + 3088, + 3092, + 3104, + 3112, + 3116, + 3120, + 3132, + 3136, + 3140, + 3144, + 3148, + 3152, + 3156, + 3160, + 3164, + 3168, + 3176, + 3180, + 3196, + 3204, + 3216, + 3224, + 3228, + 3232, + 3236, + 3240, + 3244, + 3248, + 3260, + 3264, + 3268, + 3276, + 3296, + 3304, + 3308, + 3316, + 3320, + 3328, + 3340, + 3348, + 3352, + 3368, + 3372, + 3376, + 3380, + 3384, + 3392, + 3400, + 3404, + 3412, + 3416, + 3420, + 3424, + 3428, + 3432, + 3436, + 3440, + 3444, + 3448, + 3456, + 3460, + 3464, + 3468, + 3476, + 3488, + 3496, + 3500, + 3504, + 3512, + 3516, + 3524, + 3528, + 3532, + 3536, + 3540, + 3544, + 3548, + 3552, + 3556, + 3560, + 3564, + 3568, + 3572, + 3580, + 3592, + 3596, + 3600, + 3604, + 3608, + 3616, + 3620, + 3624, + 3628, + 3632, + 3636, + 3640, + 3644, + 3648, + 3652, + 3656, + 3660, + 3664, + 3668, + 3672, + 3676, + 3680, + 3688, + 3692, + 3696, + 3700, + 3704, + 3708, + 3712, + 3716, + 3724, + 3728, + 3736, + 3740, + 3748, + 3752, + 3756, + 3760, + 3764, + 3768, + 3772, + 3776, + 3780, + 3784, + 3792, + 3796, + 3800, + 3804, + 3808, + 3812, + 3816, + 3820, + 3824, + 3828, + 3832, + 3836, + 3840, + 3844, + 3848, + 3852, + 3856, + 3868, + 3872, + 3880, + 3884, + 3888, + 3892, + 3896, + 3900, + 3904, + 3912, + 3924, + 3928, + 3932, + 3936, + 3940, + 3944, + 3952, + 3956, + 3960, + 3968, + 3976, + 3980, + 3984, + 3988, + 3992, + 3996, + 4004, + 4012, + 4016, + 4020, + 4028, + 4032, + 4044, + 4056, + 4068, + 4072, + 4076, + 4080, + 4084, + 4088, + 4092, + 4100, + 4108, + 4112, + 4116, + 4120, + 4124, + 4132, + 4136, + 4140, + 4144, + 4152, + 4156, + 4168, + 4176, + 4180, + 4184, + 4196, + 4200, + 4204, + 4208, + 4224, + 4228, + 4236, + 4248, + 4252, + 4260, + 4268, + 4276, + 4280, + 4284, + 4296, + 4308, + 4316, + 4324, + 4344, + 4352, + 4356, + 4372, + 4380, + 4384, + 4392, + 4396, + 4404, + 4440, + 4444, + 4448, + 4452, + 4460, + 4464, + 4480, + 4492, + 4496, + 4500, + 4504, + 4520, + 4524, + 4528, + 4540, + 4548, + 4568, + 4584, + 4600, + 4604, + 4612, + 4628, + 4644, + 4664, + 4668, + 4672, + 4680, + 4688, + 4692, + 4696, + 4704, + 4720, + 4724, + 4740, + 4748, + 4752, + 4764, + 4804, + 4808, + 4812, + 4832, + 4836, + 4852, + 4864, + 4916, + 4924, + 4928, + 4940, + 4968, + 4992, + 5028, + 5072, + 5168, + 5216, + 5328, + 5372, + 5496, + 5732, + 5816, + 5876, + 5968, + 5976, + 6008, + 6152, + 6164, + 6168, + 6232, + 6260, + 6272, + 6324, + 6368, + 6384, + 6416, + 6420, + 6456, + 6504, + 6524, + 6540, + 6544, + 6580, + 6624, + 6628, + 6636, + 6668, + 6708, + 6744, + 6764, + 6768, + 6772, + 6792, + 6832, + 6852, + 6860, + 6868, + 6896, + 6924, + 6928, + 6932, + 6940, + 6948, + 6956, + 6960, + 6968, + 6984, + 7000, + 7012, + 7016, + 7056, + 7060, + 7088, + 7092, + 7096, + 7108, + 7120, + 7128, + 7148, + 7152, + 7156, + 7176, + 7184, + 7240, + 7248, + 7284, + 7292, + 7296, + 7316, + 7328, + 7352, + 7360, + 7368, + 7400, + 7420, + 7432, + 7464, + 7468, + 7492, + 7544, + 7576, + 7596, + 7612, + 7648, + 7660, + 7772, + 7776, + 7804, + 7876, + 7908, + 7928, + 7932, + 7936, + 7944, + 7996, + 8008, + 8012, + 8020, + 8056, + 8136, + 8140, + 8152, + 8172, + 8192, +) +# Every row count the batch-256 dep4 engine dumped (graph and eager; 99-logs/cg/cg-b256-dep4). +_DEP4_DUMPED_TOKENS_B256 = ( + 4, + 8, + 24, + 28, + 48, + 64, + 96, + 192, + 384, + 576, + 768, + 804, + 844, + 960, + 1152, + 1344, + 1536, + 1728, + 1920, + 2112, + 2304, + 2388, + 2496, + 2688, + 2880, + 3072, + 4608, + 5328, + 6144, + 8168, + 8192, +) + + +def _dep4_six(w): + """The six experts every padding row is routed to: all inside EP4 window ``w``, + or (``w`` None) spread over the four windows.""" + if w is None: + return _DEP4_SIX_SPREAD + return tuple(w * V41_EP + e for e in (5, 17, 33, 50, 71, 90)) + + +def _dep4_eager_counts(): + """dep4's eager row counts (multiples of 4) at every tuner bucket's lower edge, one + interior point and upper edge, and on both sides of every untuned-tile and other + kernel edge of an EP4 window.""" + out, b = set(), 4 + while b <= 8192: + vals = list(range(b, min(2 * b - 1, 8192) + 1, 4)) + out |= {vals[0], vals[len(vals) // 2], vals[-1]} + b *= 2 + for e in _v41_tile_edges(V41_EP) | set(_V41_KERNEL_EDGES): + out |= {v for v in (4 * (e // 4), 4 * -(-e // 4)) if 4 <= v <= 8192} + return sorted(out) + + +def _dep4_pad_mask(n, cpu): + """Padding rows of one `_moe_chunks` chunk of `n` rows cut from a random dep4 gathered + batch: rows = (8192 j + n) / 4 per rank (the chunk is the last one; for n = 8192 any + chunk of 8192 (j + 1) rows), one rank full (the step's largest), the rank whose block + overlaps the chunk most holding ONE real token, the other two a random count.""" + j = int(torch.randint(0, 4, (1,), generator=cpu)) + if n == 8192: + rows, start = 2048 * (j + 1), 8192 * int(torch.randint(0, j + 1, (1,), generator=cpu)) + else: + rows, start = (8192 * j + n) // 4, 8192 * j + ov = [max(0, min(start + n, (r + 1) * rows) - max(start, r * rows)) for r in range(4)] + big = max(range(4), key=lambda r: (ov[r], float(torch.rand(1, generator=cpu)))) + rest = [r for r in range(4) if r != big] + full = min(rest, key=lambda r: (ov[r], float(torch.rand(1, generator=cpu)))) + counts = [int(torch.randint(1, rows + 1, (1,), generator=cpu)) for _ in range(4)] + counts[full], counts[big] = rows, 1 + pad = torch.zeros(4 * rows, dtype=torch.bool) + for r, c in enumerate(counts): + pad[r * rows + c : (r + 1) * rows] = True + return pad[start : start + n].to(DEV), (rows, counts, start) + + +def _dep4_case(kind, n, w, g1, seed): + """One eager call's inputs (deterministic in ``seed``). ``dep4``: a chunk of a padded + gathered batch (`_dep4_pad_mask`); ``one_real``: one real row then n - 1 padding rows; + ``all_pad``: n padding rows; ``skew6``: n real rows all routed to the six experts (the + per-expert count at its maximum, n, with nonzero data). Padding rows are zero codes and + zero scale bytes, routed to ``_dep4_six(w)`` with weight 0.25 each.""" + gen = torch.Generator(device=DEV).manual_seed(seed) + cpu = torch.Generator().manual_seed(seed) + data, sf, x = _rand_activation(n, V41_H, g1, gen) + ids, wts = _v41_routing(n, gen) + six = torch.tensor(_dep4_six(w), device=DEV, dtype=torch.int32) + layout = None + if kind == "skew6": + pad = torch.zeros(n, dtype=torch.bool, device=DEV) + ids = six.expand(n, V41_K).contiguous() + elif kind == "dep4": + pad, layout = _dep4_pad_mask(n, cpu) + else: + pad = torch.ones(n, dtype=torch.bool, device=DEV) + pad[0] = kind == "all_pad" + ids[pad] = six + wts[pad] = 0.25 + data[pad] = 0 + sf.view(n, V41_H // SV)[pad] = 0 + x[pad] = 0 + return (kind, data, sf, x, ids, wts), pad, layout + + +_V41_DQ = None + + +def _v41_dequant_cache(): + """The V4.1 layer's dequantized weights (code x block scale), cached in bf16 -- + exact: an e2m1 code times a power-of-two e4m3 scale has at most two significant + bits -- so the eager sweep's reference does not re-dequantize 1152 matrices per call.""" + global _V41_DQ + if _V41_DQ is None: + _, ref = _v41() + _V41_DQ = { + k: torch.stack([_dequant(c[e], s[e]).to(torch.bfloat16) for e in range(V41_E)]) + for k, (c, s) in ((k, ref[k]) for k in ("up", "gate", "down")) + } + return _V41_DQ + + +def _ref_v41_fast(x, ids, wts, ref, slices=True): + """`_ref_v41`'s four EP4 window shares and (with `slices`) its four MoE-TP slice + shares, computed the same way (same per-expert row order, same fp32 ops) from the + cached dequantized weights; tied bitwise to `_ref_v41` in the tests that use it.""" + dq = _v41_dequant_cache() + flat = ids.flatten().long() + order = torch.argsort(flat, stable=True) + counts = torch.bincount(flat, minlength=V41_E).tolist() + tok_all, slot_all = order // V41_K, order % V41_K + g2 = ref["g2"].tolist() + windows = [ + torch.zeros(x.shape[0], V41_H, dtype=torch.float32, device=DEV) + for _ in range(V41_E // V41_EP) + ] + parts = [torch.zeros_like(windows[0]) for _ in range(V41_I // V41_TP_I)] if slices else None + start = 0 + for e, c in enumerate(counts): + if c == 0: + continue + tok, slot = tok_all[start : start + c], slot_all[start : start + c] + start += c + xe = x[tok] + up = xe @ (dq["up"][e].float() / ref["gw1"][e]).t() + gate = xe @ (dq["gate"][e].float() / ref["gw1"][e]).t() + gate = torch.clamp(gate, max=V41_LIMIT) + up = torch.clamp(up, min=-V41_LIMIT, max=V41_LIMIT) + act = up * gate * torch.sigmoid(gate) + act = _q_nvfp4(act, g2[e]) / g2[e] + wd = dq["down"][e].float() / ref["gw2"][e] + w = wts[tok, slot].float().unsqueeze(1) + windows[e // V41_EP].index_add_(0, tok, (act @ wd.t()) * w) + if slices: + for s in range(len(parts)): + a, b = s * V41_TP_I, (s + 1) * V41_TP_I + parts[s].index_add_(0, tok, (act[:, a:b] @ wd[:, a:b].t()) * w) + return windows, parts + + +def _ref_v41_windows(x, ids, wts, ref): + """`_ref_v41`'s four EP4 window shares (`_ref_v41_fast` without the TP slices).""" + return _ref_v41_fast(x, ids, wts, ref, slices=False)[0] + + +def test_v41_dep4_padding_skew_eager_by_bucket(): + """dep4's zero-padding skew on the EP4 windows, eager, at every chunk size the target passes. + + The target's call (pre-routed, clamp 10, layer-level g2, `output=`) on each of the four + 96-local windows (I = 2304). Row counts: every autotuner bucket's lower edge, an interior + point and upper edge in dep4's multiples of 4, and both sides of every untuned-tile and + other kernel edge (`_dep4_eager_counts`). At each count, each run on all four windows: + a random dep4 chunk (`_dep4_pad_mask`) with the six padding experts spread over the + windows and, for every window, with all six inside it; one real row behind n - 1 + padding rows; n padding rows; and, for every window, n real rows all on its six experts + (up to 8192 rows per expert). Each call is held to the torch reference; every row the + reference leaves at zero (padding rows, rows with no slot in the window) must be + exactly zero. Run cold (the untuned tile rule), then with the cache warmed by + `autotune()` at 8192 rows (one tuned tactic per bucket, the engine's state): every + case again, plus a dep4 chunk at every row count the dep4 preview and the + batch-256 dep4 engine dumped. The warm + calls write no cache entry (the tactic stays the bucket's), and the kernels launched + at a row count are the same for uniform, skewed and all-padding routing, cold and warm. + Control: a reference that loses the last quarter of one expert's rows (a kernel whose + per-expert capacity assumed balanced routing) fails the gate. Clears the cache at the end. + """ + assert not _moe_cache_entries(), "the profiling cache already holds entries for this op" + args, ref_l = _v41_layer_g2() + layouts = [lay for lay in _v41_graph_layouts(args, ref_l) if lay[2] == V41_EP] + assert [lay[1] for lay in layouts] == [0, 96, 192, 288] + g1 = ref_l["g1"] + counts = _dep4_eager_counts() + assert {4, 8188, 8192, 140, 144, 2060, 2064} <= set(counts) and all(n % 4 == 0 for n in counts) + + # The fast reference is the certified one, bit for bit. + for kind, n, w, seed in ( + ("dep4", 1024, 1, 4150), + ("skew6", 256, 2, 4151), + ("dep4", 8192, None, 4152), + ): + case, _, _ = _dep4_case(kind, n, w, g1, seed) + fast = _ref_v41_windows(case[3], case[4], case[5], ref_l) + slow = _ref_v41(case[3], case[4], case[5], ref_l)[1] + for r in range(4): + assert torch.equal(fast[r], slow[r]), (kind, n, r) + + def plan(n): + """(kind, window of the six, seed) of every case at row count n.""" + s = 1_000_003 * n + cases = [("dep4", None, s)] + [("dep4", w, s + 1 + w) for w in range(4)] + cases += [("one_real", (n // 4) % 4, s + 5), ("all_pad", (n // 4 + 1) % 4, s + 6)] + return cases + [("skew6", w, s + 7 + w) for w in range(4)] + + def run(lay, case): + _, off, nloc, la, ls = lay + _, data, sf, _, ids, wts = case + out = torch.full((data.shape[0], V41_H), float("nan"), dtype=torch.bfloat16, device=DEV) + _call( + data, + sf, + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=nloc, + topk_ids=ids, + topk_weights=wts, + output=out, + **ls, + ) + return out + + stats = { + "calls": 0, + "zero_rows": 0, + "pad_rows": 0, + "max_expert_rows": 0, + "worst": (0.0, 0.0), + "worst_skew": (0.0, 0.0), + "worst_pad": (0.0, 0.0), + "worst_dump": (0.0, 0.0), + } + kern = {} + + def check(tag, n, kind, w, seed, profile_it): + case, pad, _ = _dep4_case(kind, n, w, g1, seed) + refs = _ref_v41_windows(case[3], case[4], case[5], ref_l) + per_expert = torch.bincount(case[4].flatten().long(), minlength=V41_E) + stats["max_expert_rows"] = max(stats["max_expert_rows"], int(per_expert.max())) + stats["pad_rows"] += int(pad.sum()) + for lay, rf in zip(layouts, refs): + name = lay[0] + y = run(lay, case) + exp = rf.to(torch.bfloat16) + zero = (exp == 0).all(dim=1) + assert pad.logical_not().logical_or(zero).all(), ( + f"{tag} {name} T={n} {kind}: a padding row is nonzero" + ) + assert (y[zero] == 0).all(), ( + f"{tag} {name} T={n} {kind} six@{w}: a row the reference leaves zero is not" + ) + stats["zero_rows"] += int(zero.sum()) + if not zero.all(): + _assert_moe_close(y, exp) + dv = _dev(y[~zero], exp[~zero]) + keys = ["worst", "worst_skew" if kind == "skew6" else "worst_pad"] + ( + ["worst_dump"] if tag == "dump" else [] + ) + for key in keys: + stats[key] = tuple(max(a, b) for a, b in zip(stats[key], dv)) + stats["calls"] += 1 + if profile_it and name == f"ep{w}": + kern[(tag, n, kind)] = _kernel_names(lambda: run(lay, case)) + return case + + uniform_kernels = {} + for tag in ("cold", "warm"): + if tag == "warm": + with autotune(): + _, off, nloc, la, ls = layouts[0] + c = _v41_case( + "uniform", 8192, off, nloc, g1, torch.Generator(device=DEV).manual_seed(4153) + ) + _call( + c[1], + c[2], + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=nloc, + topk_ids=c[4], + topk_weights=c[5], + **ls, + ) + uids = sorted({str(k[2]) for k in _moe_cache_entries()}) + assert uids == ["(6, 2304, 96, 0)"], uids + tuned = len(_moe_cache_entries()) + for n in counts: + u = _v41_case( + "uniform", n, 0, V41_EP, g1, torch.Generator(device=DEV).manual_seed(4154 + n) + ) + uniform_kernels[(tag, n)] = _kernel_names(lambda: run(layouts[0], u)) + for kind, w, seed in plan(n): + check( + tag, + n, + kind, + w, + seed, + profile_it=kind == "all_pad" or (kind == "skew6" and w == 0), + ) + for kind in ("skew6", "all_pad"): + if (tag, n, kind) in kern: + assert kern[(tag, n, kind)] == uniform_kernels[(tag, n)], ( + f"{tag} T={n} {kind}: the routing changed the launched kernels" + ) + if tag == "warm": + for i, n in enumerate(sorted(set(_DEP4_DUMPED_TOKENS) | set(_DEP4_DUMPED_TOKENS_B256))): + check("dump", n, "dep4", None if i % 5 == 4 else i % 4, 2_000_003 * n, False) + assert len(_moe_cache_entries()) == tuned, "a skewed call re-tuned a bucket" + assert (tag == "warm") == bool(_moe_cache_entries()), ( + "an ordinary call wrote a profiling-cache entry" + ) + + # Control: a kernel whose per-expert capacity assumed balanced routing would lose rows of + # the overloaded experts; a reference missing the last quarter of one expert's rows is + # what it would return, and the gate must reject it. + ctrl = [] + for n in (256, 2048, 8192): + case, _, _ = _dep4_case("skew6", n, 0, g1, 4155 + n) + y = run(layouts[0], case) + lost = case[5].clone() + lost[3 * n // 4 :, 5] = 0 # those rows' slot on the sixth expert contributes nothing + bad = _ref_v41_windows(case[3], case[4], lost, ref_l)[0].to(torch.bfloat16) + ctrl.append((n, *(round(v, 1) for v in _dev(y, bad)))) + try: + _assert_moe_close(y, bad) + except AssertionError: + continue + raise AssertionError( + f"T={n}: the gate accepts a result missing a quarter of one expert's rows" + ) + AutoTuner.get().clear_cache() + global _V41_DQ + _V41_DQ = None + torch.cuda.empty_cache() + print( + f" test_v41_dep4_padding_skew_eager_by_bucket OK ({len(counts)} dep4 row counts {counts[:6]}..{counts[-3:]} " + f"x 11 cases x 4 windows, cold and warm ({tuned} tuned entries), plus " + f"{len(set(_DEP4_DUMPED_TOKENS) | set(_DEP4_DUMPED_TOKENS_B256))} dumped counts " + f"warm: {stats['calls']} calls, {stats['pad_rows']} padding rows, max rows on one expert " + f"{stats['max_expert_rows']}, {stats['zero_rows']} rows exactly zero; worst elt / rms ulp all " + f"{stats['worst'][0]:.2f} / {stats['worst'][1]:.2f}, all-on-six {stats['worst_skew'][0]:.2f} / " + f"{stats['worst_skew'][1]:.2f}, padded chunks {stats['worst_pad'][0]:.2f} / {stats['worst_pad'][1]:.2f}, " + f"dumped {stats['worst_dump'][0]:.2f} / {stats['worst_dump'][1]:.2f}; kernels at a row count are " + f"routing-independent at all {len(counts)} counts cold and warm; lost-rows control (T, elt, rms) {ctrl})" + ) + + +# ── V4.1 rank forms of the tp 2 / 8 members: EP2, EP8 windows, MoE-TP2 ranks ─ +# ("E6" in the names below). Besides tep4 / dep4 / tp4 (above), the unified +# target admits tep8 and dep8 +# (tp 8 over two trays) and, with two pipeline stages, pp2 x {tep2, dep2, tp2, +# tep4, dep4, tp4} (`_parallel` in `models/deepseek_v41/routing.py`). A pipeline +# stage runs its layers' MoE exactly as the same member without stages does, so +# the rank forms this op sees that the sections above do not cover are: +# - EP2: 192 local experts at offsets 0 / 192, I = 2304 (tep2, dep2); +# - EP8: 48 local experts at offsets 0, 48, ..., 336, I = 2304 (tep8, dep8); +# - MoE-TP2: all 384 experts, I = 1152, rank r holding intermediate columns +# [1152r, 1152r + 1152) of every expert (tp2). +# Under capture every rank routes G and 6G of its own rows over `_GSET`; dep2's +# EP2 windows route the gathered 2G and 12G (to 3072) and dep8's EP8 windows the +# gathered 8G and 48G (to 12288). The target cuts a routed batch into calls of +# at most 8192 rows (`_moe_chunks`), so dep8's 48G at G 192 and 256 (9216 and +# 12288 rows) is two calls captured in one graph: 8192 + 1024 and 8192 + 4096. +# Eager, every call is at most 8192 rows, as for the members above; dep2 and +# dep8 pad and gather like dep4, so their eager chunks are multiples of 2 and 8. +# +# Kernel choice, profiled at every T 1..8192 on EP2 window 0, EP8 window 0 and +# TP2 rank 0, cold and after `autotune()` (a probe; the tests below profile their +# sweep points and assert the same edges): the launched kernels change only at +# the routing / finalize edges of the forms above (4/5, 59/60, 64/65, 128/129), +# the untuned token tile's edges (the C++ runner's rule, `_v41_cold_tile`) or, +# warm, the tuner's power-of-two bucket edges, plus 1/2 on an untuned TP2 rank +# (its FC1 config; 2/3 on a TP4 rank). + +_E6 = { + "ep2": (192, V41_I, 2), + "ep8": (48, V41_I, 8), + "tp2": (V41_E, 1152, 2), +} # local experts, I, ranks +_MOE_MAX_T = 8192 +_E6_GATHERED = { + "ep2": tuple(sorted({2 * g for g in _GSET} | {12 * g for g in _GSET})), # dep2: 2G, 12G + "ep8": tuple(sorted({8 * g for g in _GSET} | {48 * g for g in _GSET})), # dep8: 8G, 48G + "tp2": (), # a MoE-TP rank is never gathered +} +_E6_UIDS = {"ep2": "(6, 2304, 192, 0)", "ep8": "(6, 2304, 48, 0)", "tp2": "(6, 1152, 384, 0)"} +# Every runner row count the E6 members' engine dumps show (tep8 boot and GSM8K, dep8 boot and +# GSM8K, pp2 x tp2 and pp2 x tp4 GSM8K, per rank, lines of exactly one well-formed +# tuple): the EP8 windows' 1140 counts, the TP2 ranks' 631 and the TP4 ranks' 646, +# 1437 distinct, 1..8192. +_E6_DUMPED_TOKENS = ( + 1, + 2, + 6, + 7, + 8, + 12, + 16, + 22, + 23, + 24, + 32, + 38, + 40, + 48, + 54, + 56, + 64, + 80, + 96, + 104, + 120, + 128, + 129, + 130, + 140, + 144, + 152, + 160, + 192, + 232, + 240, + 248, + 304, + 308, + 312, + 314, + 336, + 344, + 384, + 400, + 408, + 416, + 424, + 432, + 436, + 464, + 472, + 496, + 504, + 512, + 546, + 552, + 558, + 600, + 616, + 624, + 640, + 680, + 735, + 767, + 768, + 772, + 776, + 783, + 785, + 795, + 798, + 802, + 805, + 807, + 808, + 809, + 812, + 815, + 816, + 817, + 818, + 820, + 821, + 823, + 826, + 829, + 830, + 832, + 835, + 836, + 838, + 839, + 840, + 842, + 847, + 848, + 849, + 851, + 852, + 853, + 854, + 856, + 860, + 861, + 862, + 865, + 866, + 867, + 869, + 870, + 871, + 872, + 873, + 874, + 875, + 876, + 877, + 878, + 880, + 882, + 883, + 884, + 885, + 886, + 887, + 888, + 889, + 890, + 892, + 893, + 896, + 897, + 898, + 899, + 900, + 904, + 905, + 906, + 907, + 908, + 909, + 910, + 911, + 912, + 913, + 914, + 915, + 916, + 917, + 919, + 920, + 921, + 922, + 923, + 924, + 925, + 926, + 927, + 928, + 929, + 931, + 932, + 933, + 934, + 935, + 937, + 938, + 939, + 941, + 942, + 943, + 944, + 945, + 946, + 947, + 948, + 949, + 950, + 951, + 952, + 953, + 954, + 955, + 956, + 957, + 958, + 959, + 960, + 961, + 962, + 963, + 964, + 965, + 966, + 967, + 968, + 969, + 970, + 971, + 972, + 973, + 974, + 975, + 976, + 977, + 978, + 979, + 980, + 981, + 982, + 983, + 984, + 985, + 986, + 987, + 988, + 989, + 990, + 991, + 992, + 993, + 994, + 995, + 996, + 997, + 998, + 1000, + 1001, + 1002, + 1003, + 1004, + 1005, + 1006, + 1008, + 1009, + 1010, + 1011, + 1012, + 1013, + 1014, + 1015, + 1016, + 1017, + 1018, + 1019, + 1020, + 1021, + 1022, + 1023, + 1024, + 1025, + 1026, + 1027, + 1028, + 1030, + 1031, + 1032, + 1033, + 1034, + 1035, + 1036, + 1037, + 1038, + 1039, + 1040, + 1041, + 1042, + 1043, + 1044, + 1045, + 1046, + 1047, + 1048, + 1049, + 1050, + 1052, + 1053, + 1054, + 1055, + 1056, + 1057, + 1058, + 1059, + 1060, + 1061, + 1062, + 1063, + 1064, + 1065, + 1066, + 1067, + 1068, + 1069, + 1070, + 1071, + 1072, + 1073, + 1074, + 1075, + 1076, + 1077, + 1078, + 1079, + 1080, + 1081, + 1082, + 1083, + 1084, + 1085, + 1086, + 1087, + 1088, + 1089, + 1090, + 1091, + 1092, + 1093, + 1094, + 1095, + 1096, + 1097, + 1098, + 1099, + 1100, + 1101, + 1102, + 1103, + 1104, + 1105, + 1106, + 1107, + 1108, + 1109, + 1110, + 1111, + 1112, + 1113, + 1114, + 1115, + 1116, + 1117, + 1118, + 1119, + 1120, + 1121, + 1122, + 1124, + 1125, + 1126, + 1127, + 1128, + 1129, + 1130, + 1131, + 1132, + 1133, + 1135, + 1136, + 1137, + 1138, + 1139, + 1140, + 1141, + 1142, + 1143, + 1144, + 1145, + 1146, + 1147, + 1148, + 1149, + 1150, + 1151, + 1152, + 1153, + 1154, + 1155, + 1156, + 1158, + 1159, + 1160, + 1161, + 1163, + 1165, + 1166, + 1167, + 1168, + 1169, + 1170, + 1172, + 1173, + 1174, + 1175, + 1176, + 1177, + 1178, + 1179, + 1180, + 1182, + 1183, + 1184, + 1185, + 1186, + 1187, + 1189, + 1190, + 1191, + 1193, + 1194, + 1197, + 1198, + 1199, + 1200, + 1201, + 1202, + 1204, + 1205, + 1206, + 1207, + 1209, + 1210, + 1211, + 1212, + 1214, + 1217, + 1220, + 1222, + 1224, + 1225, + 1226, + 1228, + 1229, + 1230, + 1231, + 1232, + 1234, + 1235, + 1236, + 1241, + 1243, + 1244, + 1248, + 1250, + 1251, + 1252, + 1253, + 1254, + 1256, + 1258, + 1260, + 1262, + 1265, + 1266, + 1267, + 1268, + 1270, + 1271, + 1272, + 1273, + 1275, + 1276, + 1277, + 1278, + 1279, + 1280, + 1281, + 1283, + 1284, + 1285, + 1286, + 1287, + 1288, + 1290, + 1292, + 1295, + 1297, + 1298, + 1301, + 1304, + 1305, + 1306, + 1308, + 1312, + 1314, + 1315, + 1316, + 1322, + 1323, + 1325, + 1326, + 1327, + 1329, + 1331, + 1332, + 1335, + 1336, + 1340, + 1356, + 1357, + 1358, + 1360, + 1363, + 1364, + 1368, + 1384, + 1386, + 1387, + 1388, + 1390, + 1397, + 1400, + 1403, + 1421, + 1422, + 1423, + 1432, + 1447, + 1455, + 1459, + 1465, + 1481, + 1484, + 1493, + 1506, + 1510, + 1516, + 1528, + 1530, + 1536, + 1552, + 1554, + 1556, + 1560, + 1568, + 1577, + 1584, + 1589, + 1592, + 1599, + 1601, + 1603, + 1604, + 1607, + 1609, + 1613, + 1616, + 1619, + 1620, + 1622, + 1627, + 1628, + 1630, + 1632, + 1634, + 1638, + 1639, + 1640, + 1642, + 1644, + 1646, + 1651, + 1652, + 1653, + 1656, + 1657, + 1661, + 1663, + 1665, + 1667, + 1670, + 1671, + 1672, + 1673, + 1675, + 1677, + 1679, + 1680, + 1686, + 1687, + 1688, + 1691, + 1692, + 1693, + 1695, + 1697, + 1703, + 1704, + 1705, + 1706, + 1707, + 1708, + 1710, + 1711, + 1713, + 1715, + 1716, + 1717, + 1718, + 1719, + 1720, + 1722, + 1723, + 1724, + 1726, + 1727, + 1728, + 1729, + 1730, + 1731, + 1740, + 1741, + 1742, + 1744, + 1746, + 1747, + 1748, + 1749, + 1750, + 1751, + 1752, + 1753, + 1754, + 1755, + 1757, + 1759, + 1760, + 1761, + 1762, + 1764, + 1765, + 1766, + 1767, + 1768, + 1769, + 1770, + 1771, + 1772, + 1774, + 1775, + 1776, + 1777, + 1779, + 1781, + 1783, + 1784, + 1786, + 1788, + 1790, + 1793, + 1794, + 1795, + 1796, + 1798, + 1799, + 1800, + 1802, + 1803, + 1804, + 1805, + 1806, + 1807, + 1808, + 1810, + 1811, + 1812, + 1813, + 1814, + 1816, + 1817, + 1818, + 1820, + 1821, + 1822, + 1823, + 1824, + 1825, + 1826, + 1827, + 1828, + 1829, + 1830, + 1831, + 1832, + 1833, + 1835, + 1836, + 1837, + 1839, + 1840, + 1841, + 1842, + 1843, + 1844, + 1845, + 1846, + 1847, + 1848, + 1850, + 1851, + 1852, + 1853, + 1854, + 1855, + 1860, + 1862, + 1863, + 1864, + 1865, + 1866, + 1867, + 1868, + 1869, + 1870, + 1871, + 1872, + 1873, + 1874, + 1876, + 1877, + 1878, + 1879, + 1880, + 1882, + 1883, + 1884, + 1885, + 1886, + 1889, + 1890, + 1891, + 1892, + 1893, + 1894, + 1896, + 1897, + 1899, + 1902, + 1904, + 1905, + 1906, + 1907, + 1908, + 1910, + 1911, + 1914, + 1915, + 1916, + 1917, + 1918, + 1919, + 1920, + 1922, + 1923, + 1924, + 1925, + 1926, + 1927, + 1929, + 1931, + 1932, + 1933, + 1934, + 1936, + 1937, + 1938, + 1939, + 1940, + 1942, + 1943, + 1944, + 1945, + 1946, + 1948, + 1949, + 1950, + 1951, + 1952, + 1953, + 1954, + 1955, + 1956, + 1958, + 1960, + 1961, + 1962, + 1964, + 1965, + 1966, + 1969, + 1970, + 1971, + 1972, + 1974, + 1975, + 1976, + 1977, + 1978, + 1979, + 1980, + 1981, + 1982, + 1983, + 1984, + 1985, + 1986, + 1987, + 1988, + 1989, + 1990, + 1992, + 1993, + 1994, + 1996, + 1997, + 1998, + 1999, + 2000, + 2001, + 2002, + 2004, + 2005, + 2007, + 2008, + 2009, + 2010, + 2012, + 2013, + 2014, + 2017, + 2019, + 2020, + 2021, + 2024, + 2025, + 2026, + 2027, + 2031, + 2032, + 2035, + 2036, + 2037, + 2038, + 2039, + 2040, + 2041, + 2042, + 2043, + 2045, + 2047, + 2048, + 2049, + 2051, + 2052, + 2053, + 2054, + 2055, + 2056, + 2057, + 2059, + 2061, + 2062, + 2063, + 2064, + 2069, + 2070, + 2071, + 2072, + 2075, + 2077, + 2081, + 2082, + 2088, + 2089, + 2090, + 2091, + 2093, + 2095, + 2098, + 2099, + 2100, + 2101, + 2102, + 2104, + 2105, + 2106, + 2107, + 2108, + 2111, + 2114, + 2115, + 2118, + 2119, + 2120, + 2121, + 2123, + 2125, + 2130, + 2131, + 2132, + 2136, + 2137, + 2139, + 2140, + 2146, + 2148, + 2149, + 2152, + 2154, + 2155, + 2156, + 2159, + 2160, + 2162, + 2163, + 2164, + 2166, + 2167, + 2168, + 2169, + 2173, + 2175, + 2177, + 2178, + 2180, + 2184, + 2187, + 2188, + 2189, + 2190, + 2191, + 2193, + 2195, + 2198, + 2203, + 2206, + 2209, + 2213, + 2216, + 2217, + 2219, + 2222, + 2224, + 2230, + 2231, + 2234, + 2239, + 2242, + 2246, + 2247, + 2252, + 2255, + 2257, + 2259, + 2267, + 2272, + 2273, + 2280, + 2282, + 2284, + 2291, + 2292, + 2298, + 2301, + 2312, + 2313, + 2320, + 2322, + 2335, + 2336, + 2344, + 2348, + 2349, + 2351, + 2354, + 2356, + 2359, + 2362, + 2371, + 2373, + 2377, + 2380, + 2382, + 2388, + 2390, + 2395, + 2410, + 2413, + 2416, + 2424, + 2427, + 2432, + 2440, + 2444, + 2458, + 2460, + 2463, + 2471, + 2480, + 2483, + 2487, + 2490, + 2492, + 2495, + 2498, + 2501, + 2506, + 2507, + 2509, + 2513, + 2516, + 2520, + 2527, + 2528, + 2529, + 2538, + 2540, + 2541, + 2543, + 2544, + 2545, + 2549, + 2553, + 2560, + 2562, + 2566, + 2568, + 2576, + 2578, + 2579, + 2581, + 2584, + 2586, + 2589, + 2595, + 2597, + 2598, + 2601, + 2602, + 2603, + 2604, + 2607, + 2611, + 2613, + 2615, + 2616, + 2620, + 2621, + 2623, + 2626, + 2629, + 2630, + 2641, + 2642, + 2644, + 2646, + 2653, + 2654, + 2656, + 2659, + 2660, + 2662, + 2665, + 2670, + 2673, + 2676, + 2679, + 2686, + 2688, + 2689, + 2691, + 2694, + 2695, + 2697, + 2698, + 2699, + 2701, + 2703, + 2704, + 2715, + 2716, + 2720, + 2724, + 2725, + 2726, + 2727, + 2730, + 2734, + 2736, + 2737, + 2738, + 2743, + 2745, + 2749, + 2751, + 2755, + 2758, + 2762, + 2763, + 2769, + 2770, + 2771, + 2773, + 2774, + 2784, + 2785, + 2789, + 2790, + 2792, + 2795, + 2796, + 2797, + 2799, + 2800, + 2803, + 2805, + 2806, + 2808, + 2811, + 2813, + 2814, + 2818, + 2824, + 2826, + 2838, + 2840, + 2845, + 2846, + 2847, + 2851, + 2852, + 2855, + 2858, + 2861, + 2863, + 2865, + 2866, + 2871, + 2872, + 2875, + 2879, + 2880, + 2884, + 2887, + 2888, + 2891, + 2893, + 2894, + 2900, + 2906, + 2908, + 2913, + 2917, + 2918, + 2925, + 2931, + 2937, + 2938, + 2943, + 2948, + 2951, + 2954, + 2958, + 2962, + 2963, + 2966, + 2967, + 2974, + 2976, + 2977, + 2982, + 2987, + 2988, + 2990, + 3015, + 3017, + 3022, + 3024, + 3025, + 3032, + 3036, + 3040, + 3041, + 3043, + 3044, + 3048, + 3055, + 3060, + 3064, + 3067, + 3072, + 3084, + 3086, + 3092, + 3094, + 3100, + 3102, + 3116, + 3117, + 3120, + 3122, + 3130, + 3140, + 3142, + 3143, + 3161, + 3181, + 3182, + 3195, + 3205, + 3214, + 3221, + 3242, + 3257, + 3271, + 3278, + 3280, + 3332, + 3333, + 3342, + 3352, + 3353, + 3358, + 3383, + 3384, + 3408, + 3409, + 3422, + 3428, + 3449, + 3451, + 3458, + 3491, + 3494, + 3499, + 3502, + 3510, + 3528, + 3537, + 3542, + 3543, + 3560, + 3569, + 3581, + 3591, + 3613, + 3616, + 3638, + 3648, + 3656, + 3664, + 3671, + 3672, + 3683, + 3695, + 3704, + 3706, + 3717, + 3719, + 3723, + 3725, + 3736, + 3814, + 3820, + 3826, + 3827, + 3834, + 3840, + 3858, + 3862, + 3893, + 3920, + 3928, + 3947, + 3983, + 3986, + 3990, + 4017, + 4046, + 4057, + 4093, + 4128, + 4169, + 4178, + 4185, + 4222, + 4239, + 4246, + 4261, + 4276, + 4285, + 4288, + 4298, + 4299, + 4378, + 4386, + 4500, + 4526, + 4581, + 4587, + 4589, + 4664, + 4671, + 4752, + 4788, + 4792, + 4825, + 4920, + 4952, + 4968, + 5040, + 5192, + 5232, + 5248, + 5272, + 5312, + 5360, + 5378, + 5384, + 5440, + 5512, + 5600, + 5648, + 5656, + 5696, + 5704, + 5744, + 5771, + 5784, + 5800, + 5912, + 5928, + 6032, + 6048, + 6056, + 6104, + 6112, + 6144, + 6152, + 6216, + 6240, + 6320, + 6368, + 6384, + 6392, + 6432, + 6456, + 6464, + 6472, + 6488, + 6504, + 6520, + 6528, + 6552, + 6560, + 6568, + 6608, + 6640, + 6648, + 6704, + 6744, + 6768, + 6800, + 6808, + 6824, + 6864, + 6880, + 6896, + 6904, + 6928, + 6936, + 6976, + 7000, + 7016, + 7040, + 7048, + 7064, + 7088, + 7120, + 7128, + 7160, + 7176, + 7192, + 7200, + 7248, + 7288, + 7304, + 7328, + 7344, + 7416, + 7440, + 7456, + 7512, + 7568, + 7576, + 7592, + 7632, + 7640, + 7648, + 7664, + 7672, + 7680, + 7696, + 7704, + 7712, + 7760, + 7824, + 7848, + 7872, + 7896, + 7904, + 7912, + 7944, + 7952, + 7968, + 7992, + 8000, + 8016, + 8024, + 8040, + 8048, + 8080, + 8088, + 8098, + 8104, + 8112, + 8120, + 8127, + 8128, + 8136, + 8144, + 8154, + 8158, + 8159, + 8165, + 8168, + 8176, + 8177, + 8179, + 8181, + 8184, + 8192, +) + + +def _moe_chunks(total: int): + """Row ranges `[a, b)` of the target's cut of a routed batch into calls of at + most 8192 rows (`_moe_chunks` in the unified target's `modeling.py`).""" + out, a = [], 0 + while a < total: + b = min(total, a + _MOE_MAX_T) + out.append((a, b)) + a = b + return out + + +def _e6_counts(fam): + """Captured row counts of one E6 rank form: G and 6G, plus the dep member's + gathered counts on an EP window.""" + return tuple(sorted(set(_V41_RANK_TOKENS) | set(_E6_GATHERED[fam]))) + + +def _e6_layouts(fam, args, ref_l, ranks=None): + """(name, offset, local experts, operands, scalars) of an E6 family's ranks.""" + nloc, inter, n_ranks = _E6[fam] + scal = _v41_scalars(ref_l) + ranks = range(n_ranks) if ranks is None else ranks + if fam.startswith("ep"): + return [ + ( + f"{fam}.{r}", + r * nloc, + nloc, + _window_args(args, r * nloc, nloc), + _v41_slice(scal, r * nloc, nloc), + ) + for r in ranks + ] + tp = _v41_tp_args(inter) + return [(f"{fam}.{r}", 0, V41_E, tp[r], scal) for r in ranks] + + +def _e6_part(fam): + """The `_ref_v41_parts` key of an E6 family's rank shares.""" + nloc, inter, _ = _E6[fam] + return ("ep", nloc) if fam.startswith("ep") else ("tp", inter) + + +def _v41_kernel_firsts(nloc: int, inter: int, warm: bool) -> set: + """First row count of every kernel change of a rank form (`nloc` local experts, + `inter` intermediate columns) up to 8192, as the eager profile at every T found: + the routing / finalize edges, then the tuner's bucket edges (warm) or the untuned + token tile's edges and, on an untuned MoE-TP rank, its FC1 config change (at 3 on a + TP4 rank, at 2 on a TP2 rank).""" + firsts = set(_V41_ROUTING_FIRSTS) + if warm: + return firsts | {1 << k for k in range(1, 14)} + firsts |= {m for m in range(2, 8193) if _v41_cold_tile(m, nloc) != _v41_cold_tile(m - 1, nloc)} + return firsts | {(V41_E, V41_TP_I): {3}, (V41_E, 1152): {2}}.get((nloc, inter), set()) + + +def _e6_segment(n: int, nloc: int, inter: int, warm: bool): + """(lo, hi): the row counts around `n` that launch `n`'s kernels.""" + firsts = _v41_kernel_firsts(nloc, inter, warm) + lo = max([1] + [f for f in firsts if f <= n]) + hi = min([8193] + [f for f in firsts if f > n]) - 1 + return lo, hi + + +def _ref_v41_parts(x, ids, wts, ref, eps=(), tps=()): + """`_ref_v41_fast` over any rank partitions at once: for each width in `eps` the + expert-parallel windows of that many local experts (384: the whole layer), for each + width in `tps` the MoE-TP slices of that many intermediate columns. Same per-expert + row order and fp32 ops as `_ref_v41`, to which it is tied bitwise at 384 / 96 / 576 + in `test_v41_e6_rank_forms_geometry`. Returns {("ep" | "tp", width): [shares]}.""" + dq = _v41_dequant_cache() + flat = ids.flatten().long() + order = torch.argsort(flat, stable=True) + counts = torch.bincount(flat, minlength=V41_E).tolist() + tok_all, slot_all = order // V41_K, order % V41_K + g2 = ref["g2"].tolist() + + def zeros(k): + return [torch.zeros(x.shape[0], V41_H, dtype=torch.float32, device=DEV) for _ in range(k)] + + out = {("ep", w): zeros(V41_E // w) for w in eps} + out.update({("tp", w): zeros(V41_I // w) for w in tps}) + start = 0 + for e, c in enumerate(counts): + if c == 0: + continue + tok, slot = tok_all[start : start + c], slot_all[start : start + c] + start += c + xe = x[tok] + up = xe @ (dq["up"][e].float() / ref["gw1"][e]).t() + gate = xe @ (dq["gate"][e].float() / ref["gw1"][e]).t() + gate = torch.clamp(gate, max=V41_LIMIT) + up = torch.clamp(up, min=-V41_LIMIT, max=V41_LIMIT) + act = up * gate * torch.sigmoid(gate) + act = _q_nvfp4(act, g2[e]) / g2[e] + wd = dq["down"][e].float() / ref["gw2"][e] + w = wts[tok, slot].float().unsqueeze(1) + if eps: + y = (act @ wd.t()) * w + for ew in eps: + out[("ep", ew)][e // ew].index_add_(0, tok, y) + for tw in tps: + for s in range(V41_I // tw): + a, b = s * tw, (s + 1) * tw + out[("tp", tw)][s].index_add_(0, tok, (act[:, a:b] @ wd[:, a:b].t()) * w) + return out + + +def _worst(acc, key, dv): + acc[key] = tuple(max(a, b) for a, b in zip(acc.get(key, (0.0, 0.0)), dv)) + + +def test_v41_e6_rank_forms_geometry(): + """The E6 rank forms at the V4.1 geometry, with per-expert global scales. + + EP2 (two 192-expert windows), EP8 (eight 48-expert windows) and MoE-TP2 (two + I = 1152 ranks), at every T of `_V41_TOKENS` (1..8192): `num_experts` 384 on + every call, pre-routed, the clamp at 10 / alpha1, every expert with its own + gw1, gw2 and g2. Each window matches a reference masking every slot outside it + and each TP2 rank its share of the reference; tokens with no slot in a window + come back bitwise zero; the bf16 outputs of a member's ranks summed match the + full-layer reference. The reference's new partitions are tied to the certified + one first: at 384 / 96 / 576 `_ref_v41_parts` is `_ref_v41` bit for bit, and its + EP8 / EP2 windows and TP2 slices sum (fp32) to the EP4 windows and TP4 slices + within 1e-5 of the row magnitude. Controls: EP8 window 1 issued at offset 0, + and a TP2 sum without rank 1, fail the gate. + """ + args, ref = _v41() + gen = torch.Generator(device=DEV).manual_seed(4160) + g1 = ref["g1"] + fams = {fam: _e6_layouts(fam, args, ref) for fam in _E6} + worst, zero_rows, ctrl, tie_worst = {}, {}, [], 0.0 + for num_tokens in _V41_TOKENS: + data, sf, x = _rand_activation(num_tokens, V41_H, g1, gen) + ids, wts = _v41_routing(num_tokens, gen) + parts = _ref_v41_parts(x, ids, wts, ref, eps=(V41_E, 192, V41_EP, 48), tps=(V41_TP_I, 1152)) + full = parts[("ep", V41_E)][0] + if num_tokens == 1024: + slow_full, slow_win, slow_tp, _ = _ref_v41(x, ids, wts, ref) + assert torch.equal(full, slow_full) + for r in range(4): + assert torch.equal(parts[("ep", V41_EP)][r], slow_win[r]), r + assert torch.equal(parts[("tp", V41_TP_I)][r], slow_tp[r]), r + row = full.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + pairs = [(("ep", 48), ("ep", V41_EP), w) for w in range(4)] + pairs += [(("ep", V41_EP), ("ep", 192), w) for w in range(2)] + pairs += [(("tp", V41_TP_I), ("tp", 1152), w) for w in range(2)] + for fine, coarse, w in pairs: + err = ( + ((parts[fine][2 * w] + parts[fine][2 * w + 1] - parts[coarse][w]).abs() / row) + .max() + .item() + ) + assert err < 1e-5, f"{fine} pair {w} vs {coarse} at T={num_tokens}: {err:.2e}" + tie_worst = max(tie_worst, err) + exp_full = full.to(torch.bfloat16) + for fam, lays in fams.items(): + nloc = _E6[fam][0] + outs = [] + for r, (name, off, n_loc, la, ls) in enumerate(lays): + y = _call( + data, + sf, + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=n_loc, + topk_ids=ids, + topk_weights=wts, + **ls, + )[0] + exp = parts[_e6_part(fam)][r].to(torch.bfloat16) + _assert_moe_close(y, exp) + _worst(worst, fam, _dev(y, exp)) + if fam.startswith("ep"): + outside = ~((ids >= off) & (ids < off + nloc)).any(dim=1) + assert torch.equal(y[outside], torch.zeros_like(y[outside])), ( + f"{name} wrote rows routed entirely elsewhere at T={num_tokens}" + ) + zero_rows[fam] = max(zero_rows.get(fam, 0), int(outside.sum())) + outs.append(y) + total = sum(o.float() for o in outs).to(torch.bfloat16) + _assert_moe_close(total, exp_full) + _worst(worst, fam + " sum", _dev(total, exp_full)) + if num_tokens == 1024 and fam in ("ep8", "tp2"): + if fam == "ep8": + _, _, n_loc, la, ls = lays[1] + bad = _call( + data, + sf, + la, + V41_E, + V41_K, + local_expert_offset=0, + local_num_experts=n_loc, + topk_ids=ids, + topk_weights=wts, + **ls, + )[0] + exp, what = parts[("ep", 48)][1].to(torch.bfloat16), "EP8 window 1 at offset 0" + else: + bad, exp, what = outs[0], exp_full, "TP2 sum without rank 1" + ctrl.append((what, *(round(v, 1) for v in _dev(bad, exp)))) + try: + _assert_moe_close(bad, exp) + except AssertionError: + continue + raise AssertionError(f"the gate accepts {what}") + fmt = {k: f"{v[0]:.2f} / {v[1]:.2f}" for k, v in sorted(worst.items())} + print( + f" test_v41_e6_rank_forms_geometry OK (worst elt / rms ulp: {fmt}; bitwise-zero rows per window " + f"<= {zero_rows}; new partitions vs certified ones <= {tie_worst:.1e} of the row; controls (elt, rms): {ctrl})" + ) + + +def test_v41_e6_row_axis_by_bucket(): + """Every tuner bucket and untuned tile up to 8192 at the E6 rank forms, cold and warm, + and every dumped row count warm. + + On every EP2 window, EP8 window and TP2 rank, with the target's call (pre-routed, + clamp 10, layer-level g2): the lower edge, an interior point and the upper edge of + every autotuner bucket, and both sides of every untuned-tile edge of the form and of + every other kernel change, each against the torch reference on fresh inputs. The + kernels launched on EP2 / EP8 window 0 and TP2 rank 0 change between two neighbouring + counts of the sweep only across a row count `_v41_kernel_firsts` names, and cold at + every one of them. Then the cache is warmed through `autotune()` at 8192 rows on each + form and on TP4 rank 0, and every count runs again. Warm, every row count of + `_E6_DUMPED_TOKENS` also runs on every E6 rank and on the four TP4 ranks, on the first + n rows of one 8192-row batch against the same rows of its reference. Clears the cache. + """ + assert not _moe_cache_entries(), "the profiling cache already holds entries for this op" + args, ref_l = _v41_layer_g2() + fams = {fam: _e6_layouts(fam, args, ref_l) for fam in _E6} + tp4 = [lay for lay in _v41_graph_layouts(args, ref_l) if lay[0].startswith("tp")] + g1 = ref_l["g1"] + gen = torch.Generator(device=DEV).manual_seed(4162) + base = _v41_bucket_tokens() | set(_V41_KERNEL_EDGES) + counts = {fam: sorted(base | _v41_tile_edges(_E6[fam][0])) for fam in _E6} + assert {287, 288, 4127, 4128} <= set(counts["ep2"]) and {71, 72, 1031, 1032} <= set( + counts["ep8"] + ) + assert {1, 2, 575, 576, 4159, 4160} <= set(counts["tp2"]) + all_counts = sorted(set().union(*map(set, counts.values()))) + + def run(lay, data, sf, ids, wts): + _, off, nloc, la, ls = lay + return _call( + data, + sf, + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=nloc, + topk_ids=ids, + topk_weights=wts, + **ls, + )[0] + + cases, cold, kern, worst = {}, {}, {}, {} + for n in all_counts: + _, data, sf, x, ids, wts = _v41_case("uniform", n, 0, V41_E, g1, gen) + parts = _ref_v41_parts(x, ids, wts, ref_l, eps=(192, 48), tps=(1152,)) + exps = {} + for fam, lays in fams.items(): + if n not in counts[fam]: + continue + for r, lay in enumerate(lays): + exp = parts[_e6_part(fam)][r].to(torch.bfloat16) + y = run(lay, data, sf, ids, wts) + _assert_moe_close(y, exp) + _worst(worst, ("cold", fam), _dev(y, exp)) + cold[(lay[0], n)] = y + exps[lay[0]] = exp + if r == 0: + kern[("cold", fam, n)] = _kernel_names(lambda: run(lay, data, sf, ids, wts)) + cases[n] = ((data, sf, ids, wts), exps) + del parts, x + assert not _moe_cache_entries(), "an ordinary call wrote a profiling-cache entry" + + with autotune(): + for lay in [lays[0] for lays in fams.values()] + [tp4[0]]: + c = _v41_case("uniform", 8192, lay[1], lay[2], g1, gen) + run(lay, c[1], c[2], c[4], c[5]) + uids = sorted({str(k[2]) for k in _moe_cache_entries()}) + assert uids == sorted([*_E6_UIDS.values(), "(6, 576, 384, 0)"]), uids + same = total = 0 + for n in all_counts: + (data, sf, ids, wts), exps = cases[n] + for fam, lays in fams.items(): + if n not in counts[fam]: + continue + for r, lay in enumerate(lays): + y = run(lay, data, sf, ids, wts) + _assert_moe_close(y, exps[lay[0]]) + _worst(worst, ("warm", fam), _dev(y, exps[lay[0]])) + same += int(torch.equal(y, cold[(lay[0], n)])) + total += 1 + if r == 0: + kern[("warm", fam, n)] = _kernel_names(lambda: run(lay, data, sf, ids, wts)) + del cases, cold + + runs, changes = {}, {} + for state in ("cold", "warm"): + for fam in _E6: + nloc, inter, _ = _E6[fam] + firsts = _v41_kernel_firsts(nloc, inter, state == "warm") + cs = counts[fam] + seen, out = set(), [] + for a, b in zip(cs, cs[1:]): + if kern[(state, fam, a)] != kern[(state, fam, b)]: + assert b == a + 1 and b in firsts, ( + f"{state} {fam}: the kernels change between T={a} and T={b}, not at a known edge" + ) + seen.add(b) + if state == "cold": + want = {f for f in firsts if f in cs and f - 1 in cs} + assert seen == want, ( + f"cold {fam}: expected kernel changes at {sorted(want - seen)} not seen" + ) + for n in cs: + if out and kern[(state, fam, n)] == out[-1][2]: + out[-1][1] = n + else: + out.append([n, n, kern[(state, fam, n)]]) + runs[(state, fam)] = [(a, b) for a, b, _ in out] + changes[(state, fam)] = sorted(seen) + + # Every dumped row count, warm, on the first n rows of one 8192-row batch. + _, data, sf, x, ids, wts = _v41_case("uniform", 8192, 0, V41_E, g1, gen) + parts = _ref_v41_parts(x, ids, wts, ref_l, eps=(192, 48), tps=(1152, V41_TP_I)) + exp_all = {key: [p.to(torch.bfloat16) for p in v] for key, v in parts.items()} + del parts, x + forms = dict(fams, tp4=tp4) + dumped_calls = 0 + for n in _E6_DUMPED_TOKENS: + for fam, lays in forms.items(): + key = ("tp", V41_TP_I) if fam == "tp4" else _e6_part(fam) + for r, lay in enumerate(lays): + y = run(lay, data[:n], sf[: n * V41_H // SV], ids[:n], wts[:n]) + exp = exp_all[key][r][:n] + _assert_moe_close(y, exp) + _worst(worst, ("dump", fam), _dev(y, exp)) + dumped_calls += 1 + entries = len(_moe_cache_entries()) + AutoTuner.get().clear_cache() + fmt = {f"{k[0]} {k[1]}": f"{v[0]:.2f} / {v[1]:.2f}" for k, v in sorted(worst.items())} + print( + f" test_v41_e6_row_axis_by_bucket OK (row counts per form: " + f"{ {fam: len(c) for fam, c in counts.items()} }, every rank, 1..8192; worst elt / rms ulp {fmt}; " + f"warm bitwise equal to cold on {same} of {total}; {entries} tuned entries over {uids}; " + f"{len(_E6_DUMPED_TOKENS)} dumped counts x {sum(len(v) for v in forms.values())} ranks = {dumped_calls} " + f"warm calls; kernel changes at {changes}; kernel runs {runs})" + ) + + +def test_v41_e6_cuda_graph_replay_is_bitwise_eager(): + """Captured once per (E6 rank, T), replayed on data the capture never saw, cold tuner. + + As `test_v41_cuda_graph_replay_is_bitwise_eager`, at the E6 forms: every EP2 window + at G and 6G and dep2's gathered 2G and 12G (115 counts, 1..3072), every EP8 window + at G and 6G and dep8's gathered 8G and 48G (130 counts, 1..12288), and both TP2 + ranks at G and 6G (82 counts, 1..1536). At 9216 and 12288 the graph holds the + target's two calls (`_moe_chunks`: 8192 + 1024 and 8192 + 4096), each on its rows + of the static buffers and into its own output. Each graph replays two uniform + routings, a skewed one (every token on the same six local experts), on an EP window + one routing nothing into it and, at the dep member's gathered counts, its + zero-padded batch over 2 or 8 rank blocks (`pad`), then the capture's own data. + Every replay is bitwise equal to the eager calls on the same values (one per + chunk), and within the torch reference's gates (the empty window exactly zero); + consecutive replays differ, so a stale replay could not pass. + """ + assert not _moe_cache_entries(), "the cold-cache graph test ran with a warm cache" + args, ref_l = _v41_layer_g2() + fams = {fam: _e6_layouts(fam, args, ref_l) for fam in _E6} + g1 = ref_l["g1"] + gen = torch.Generator(device=DEV).manual_seed(4163) + fam_counts = {fam: _e6_counts(fam) for fam in _E6} + assert [len(fam_counts[f]) for f in ("ep2", "ep8", "tp2")] == [115, 130, 82] + assert [max(fam_counts[f]) for f in ("ep2", "ep8", "tp2")] == [3072, 12288, 1536] + counts = sorted(set().union(*map(set, fam_counts.values()))) + worst, worst_at, worst_split = {}, {}, (0.0, 0.0) + replays = graphs = 0 + per_layout, split = {}, [] + for n in counts: + chunks = _moe_chunks(n) + need = [fam for fam in _E6 if n in fam_counts[fam]] + eps = tuple(_E6[f][0] for f in need if f.startswith("ep")) + tps = tuple(_E6[f][1] for f in need if f.startswith("tp")) + shared = [_v41_case("uniform", n, 0, V41_E, g1, gen) for _ in range(2)] + refs = [_ref_v41_parts(c[3], c[4], c[5], ref_l, eps, tps) for c in shared] + for fam in need: + nloc, inter, ranks = _E6[fam] + key = _e6_part(fam) + for r, lay in enumerate(fams[fam]): + name, off = lay[:2] + own = [_v41_case("skew", n, off, nloc, g1, gen)] + if fam.startswith("ep"): + own.append(_v41_case("none", n, off, nloc, g1, gen)) + if n in _E6_GATHERED[fam]: + own.append(_v41_case("pad", n, off, nloc, g1, gen, ranks=ranks)) + exps = [rf[key][r] for rf in refs] + for c in own: + if c[0] == "none": + exps.append(None) + else: + sub = ((nloc,), ()) if fam.startswith("ep") else ((), (inter,)) + exps.append(_ref_v41_parts(c[3], c[4], c[5], ref_l, *sub)[key][r]) + cases = shared + own + shared[:1] + exps = exps + exps[:1] + g = _V41Graph(lay, n, cases[0], chunks) + prev = None + for case, exp in zip(cases, exps): + y = g.replay(case) + eager = _v41_eager(lay, case, chunks) + assert torch.equal(y, eager), ( + f"{name} T={n} {case[0]}: graph replay differs from eager " + f"{_dev(y, eager.float())} (elt / rms ulp)" + ) + if exp is None: + assert torch.equal(y, torch.zeros_like(y)), ( + f"{name} T={n}: an empty window is not zero" + ) + else: + exp = exp.to(torch.bfloat16) + _assert_moe_close(y, exp) + dv = _dev(y, exp) + for i in range(2): + if dv[i] > worst.get(fam, (0.0, 0.0))[i]: + worst_at[(fam, i)] = (name, n, case[0]) + _worst(worst, fam, dv) + if len(chunks) > 1: + worst_split = tuple(max(a, b) for a, b in zip(worst_split, dv)) + assert prev is None or not torch.equal(prev, y) or not y.any(), ( + f"{name} T={n}: two cases gave one result" + ) + prev = y + replays += 1 + graphs += 1 + per_layout[name] = per_layout.get(name, 0) + 1 + if len(chunks) > 1: + split.append((name, n)) + del g + del refs, shared + assert not _moe_cache_entries() + assert per_layout == { + **{f"ep2.{r}": 115 for r in range(2)}, + **{f"ep8.{r}": 130 for r in range(8)}, + **{f"tp2.{r}": 82 for r in range(2)}, + }, per_layout + assert len(split) == 16, split # 9216 and 12288 on each of the eight EP8 windows + fmt = {k: f"{v[0]:.2f} / {v[1]:.2f}" for k, v in worst.items()} + print( + f" test_v41_e6_cuda_graph_replay_is_bitwise_eager OK ({graphs} (rank, T) graphs {per_layout}, " + f"{replays} replays bitwise equal to eager; worst vs reference elt / rms ulp {fmt} (at {worst_at}); " + f"two-call graphs {len(split)} at 9216 / 12288: {worst_split[0]:.2f} / {worst_split[1]:.2f}; " + f"counts EP2 {list(fam_counts['ep2'])}; EP8 {list(fam_counts['ep8'])}; TP2 {list(fam_counts['tp2'])})" + ) + + +def test_v41_e6_cuda_graph_swizzled_scales_are_caught(): + """Control: the graph harness sees the swizzled activation-scale mistake at the E6 + gathered counts, the two-call graph included. + + At dep2's 3072 (12G at G 256) on EP2 windows 0 and 1, and at dep8's 6144 and 12288 + (48G at G 128 and 256; 12288 as 8192 + 4096) on EP8 windows 0 and 7, every call's + row count is a multiple of 128, where a 128x4-swizzled scale buffer has the linear + one's byte count. Replaying on each call's scale bytes in swizzled order must fail + the reference gate, while the linear bytes pass. + """ + args, ref_l = _v41_layer_g2() + g1 = ref_l["g1"] + gen = torch.Generator(device=DEV).manual_seed(4167) + seen = [] + for fam, ranks, n in (("ep2", (0, 1), 3072), ("ep8", (0, 7), 6144), ("ep8", (0, 7), 12288)): + nloc = _E6[fam][0] + chunks = _moe_chunks(n) + case = _v41_case("uniform", n, 0, V41_E, g1, gen) + kind, data, sf, x, ids, wts = case + windows = _ref_v41_parts(x, ids, wts, ref_l, eps=(nloc,))[("ep", nloc)] + swz = torch.cat( + [ + _swizzle_scales( + sf[a * V41_H // SV : b * V41_H // SV].view(1, b - a, V41_H // SV) + ).reshape(-1) + for a, b in chunks + ] + ).contiguous() + assert swz.numel() == sf.numel() and not torch.equal(swz, sf) + for lay in _e6_layouts(fam, args, ref_l, ranks): + name, off = lay[:2] + exp = windows[off // nloc].to(torch.bfloat16) + g = _V41Graph(lay, n, case, chunks) + _assert_moe_close(g.replay(case), exp) + bad = g.replay((kind, data, swz, x, ids, wts)) + elt, rms = _dev(bad, exp) + assert rms > 2.0 and elt > 16.0, ( + f"{name} T={n}: swizzled scales passed ({elt:.1f} / {rms:.2f} ulp)" + ) + seen.append((name, n, round(elt), round(rms))) + del g + print( + " test_v41_e6_cuda_graph_swizzled_scales_are_caught OK " + f"(swizzled replay vs reference, elt / rms ulp: {seen}; gates 16 / 2)" + ) + + +def test_v41_e6_cuda_graph_tactic_is_fixed_at_capture(): + """At the E6 forms: a graph keeps the tactic it was captured with; no tactic seen + changes the result. + + As `test_v41_cuda_graph_tactic_is_fixed_at_capture`, one group of ranks at a time + (EP2 windows 0-1, EP8 windows 0-3, EP8 windows 4-7, TP2 ranks 0-1), each from a + cold cache. The group's graphs at every captured count (`_e6_counts`, the two-call + graphs at 9216 and 12288 included) are captured cold; the cache is warmed through + `autotune()` at 8192 rows on the group's first rank (every token bucket a captured + call maps to must then hold a non-fallback tactic); the cold graphs replay + unchanged, kernels and bits; graphs captured warm replay bitwise equal to the cold + results and to a warm eager call; clearing the cache changes neither. + + Per-bucket kernel choice: a graph launches exactly the kernels of eager calls at + its calls' row counts in the same tuner state, and each of those equals the + kernels of eager calls at the lower and upper edge of the count's bucket (warm) or + untuned tile (cold) (`_e6_segment`), checked on the group's first rank at every + count and on every rank at the counts above 1536. + """ + assert not _moe_cache_entries(), "the profiling cache already holds entries for this op" + args, ref_l = _v41_layer_g2() + g1 = ref_l["g1"] + gen = torch.Generator(device=DEV).manual_seed(4164) + kgen = torch.Generator(device=DEV).manual_seed(4165) + + def tactic(k): + t = AutoTuner.get().profiling_cache.cache[k][1] + return tuple(t) if isinstance(t, (list, tuple)) else (t,) + + groups = (("ep2", (0, 1)), ("ep8", (0, 1, 2, 3)), ("ep8", (4, 5, 6, 7)), ("tp2", (0, 1))) + summary, changed, checked, graphs = [], [], {"cold": 0, "warm": 0}, 0 + for fam, ranks in groups: + assert not _moe_cache_entries() + nloc, inter, _ = _E6[fam] + lays = _e6_layouts(fam, args, ref_l, ranks) + counts = _e6_counts(fam) + eager_k = {"cold": {}, "warm": {}} + + def eager_kernels(state, lay, m): + key = (lay[0], m) + if key not in eager_k[state]: + _, off, nl, la, ls = lay + c = _v41_case("uniform", m, 0, V41_E, g1, kgen) + out = torch.empty(m, V41_H, dtype=torch.bfloat16, device=DEV) + + def run(): + _call( + c[1], + c[2], + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=nl, + topk_ids=c[4], + topk_weights=c[5], + output=out, + **ls, + ) + + run() + eager_k[state][key] = _kernel_names(run) + return eager_k[state][key] + + segs = {"cold": {}, "warm": {}} + + def check_bucket(state, lay, n, k_graph): + if not (lay[0] == lays[0][0] or n > 1536): + return + sizes = [b - a for a, b in _moe_chunks(n)] + want = sum((eager_kernels(state, lay, m) for m in sizes), ()) + assert tuple(k_graph) == want, ( + f"{state} {lay[0]} T={n}: the graph's kernels differ from eager calls at {sizes}: {k_graph} vs {want}" + ) + for m in sizes: + lo, hi = _e6_segment(m, nloc, inter, state == "warm") + for e in (lo, hi): + assert eager_kernels(state, lay, e) == eager_kernels(state, lay, m), ( + f"{state} {lay[0]} T={m}: eager kernels differ at the segment edge {e} ({lo}..{hi})" + ) + if lay[0] == lays[0][0]: + segs[state].setdefault(f"{lo}..{hi}", set()).add(m) + checked[state] += 1 + + held = {} + for lay in lays: + name, off = lay[:2] + for n in counts: + chunks = _moe_chunks(n) + kinds = ("uniform", "skew") + (("none",) if fam.startswith("ep") else ()) + cases = [ + c[:3] + (None,) + c[4:] + for c in (_v41_case(k, n, off, nloc, g1, gen) for k in kinds) + ] + g = _V41Graph(lay, n, cases[0], chunks) + k_cold = g.kernels(cases[0]) + check_bucket("cold", lay, n, k_cold) + held[(name, n)] = (g, cases, [_v41_eager(lay, c, chunks) for c in cases], k_cold) + assert not _moe_cache_entries(), ( + "an ordinary or captured call wrote a profiling-cache entry" + ) + + with autotune(): + _, off, nl, la, ls = lays[0] + c = _v41_case("uniform", 8192, off, nl, g1, gen) + _call( + c[1], + c[2], + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=nl, + topk_ids=c[4], + topk_weights=c[5], + **ls, + ) + keys = _moe_cache_entries() + tuned = {(str(k[2]), k[3][2][0]): tactic(k) for k in keys} + assert sorted({u for u, _ in tuned}) == [_E6_UIDS[fam]], tuned + for n in counts: + for a, b in _moe_chunks(n): + bkt = 1 << ((b - a).bit_length() - 1) + assert tuned.get((_E6_UIDS[fam], bkt), (-1,)) != (-1,), ( + f"bucket {bkt} of {fam} is not tuned" + ) + + grp_changed = [] + for lay in lays: + name = lay[0] + for n in counts: + chunks = _moe_chunks(n) + g, cases, want, k_cold = held[(name, n)] + for c, w in zip(cases, want): + assert torch.equal(g.replay(c), w), ( + f"{name} T={n} {c[0]}: a cold graph changed after warming" + ) + assert torch.equal(_v41_eager(lay, c, chunks), w), ( + f"{name} T={n} {c[0]}: warm eager != cold" + ) + assert g.kernels(cases[0]) == k_cold, ( + f"{name} T={n}: a cold graph re-read the tuner at replay" + ) + del g + warm = _V41Graph(lay, n, cases[0], chunks) + k_warm = warm.kernels(cases[0]) + check_bucket("warm", lay, n, k_warm) + if k_warm != k_cold: + grp_changed.append((name, n)) + for c, w in zip(cases, want): + assert torch.equal(warm.replay(c), w), ( + f"{name} T={n} {c[0]}: warm graph != cold result" + ) + held[(name, n)] = (warm, cases, want, k_warm) + graphs += 1 + AutoTuner.get().clear_cache() + for (name, n), (warm, cases, want, k_warm) in held.items(): + assert warm.kernels(cases[0]) == k_warm, f"{name} T={n}: a warm graph re-read the tuner" + for c, w in zip(cases, want): + assert torch.equal(warm.replay(c), w), ( + f"{name} T={n} {c[0]}: warm graph changed after clear" + ) + del held + torch.cuda.empty_cache() + changed += grp_changed + summary.append( + ( + fam, + ranks, + len(keys), + len(set(tuned.values())), + len(grp_changed), + { + s: { + k: sorted(v) + for k, v in sorted( + segs[s].items(), key=lambda kv: int(kv[0].split("..")[0]) + ) + } + for s in ("cold", "warm") + }, + ) + ) + assert changed, ( + "no graph captured warm launched different kernels — the comparison proved nothing" + ) + print( + f" test_v41_e6_cuda_graph_tactic_is_fixed_at_capture OK ({graphs} (rank, T) graphs, cold- and " + f"warm-captured, bitwise equal on every replay; warm capture launched different kernels at " + f"{len(changed)} of them, e.g. {changed[:4]}; per-bucket kernel check at {checked} graphs; per group " + f"(family, ranks, tuned entries, distinct tactics, changed, window-0 call sizes by segment): {summary})" + ) + + +# dep2's and dep8's zero-padding skew. Like dep4, both pad every rank's rows to the +# step's largest per-rank count with zero rows before the all-gather, and +# `_moe_chunks` cuts the gathered batch into calls of 8192 rows and one remainder: +# dep2's eager row counts are multiples of 2 in 2..8192, dep8's multiples of 8 in +# 8..8192. Every padding row goes to the router's six bias-top experts with weight +# 0.25; on an EP8 window of 48 experts all six can be local. + +_DEP_SIX_SPREAD = {192: (5, 17, 101, 200, 290, 377), 48: (5, 60, 110, 170, 250, 377)} + + +def _dep_six(w, nloc): + """The six experts every padding row is routed to: all inside window ``w`` of + ``nloc`` experts, or (``w`` None) spread over the windows.""" + if w is None: + return _DEP_SIX_SPREAD[nloc] + return tuple(w * nloc + e % nloc for e in (5, 17, 33, 50, 71, 90)) + + +def _dep_eager_counts(ranks, nloc): + """A dep member's eager row counts (multiples of ``ranks``) at every tuner bucket's + lower edge, one interior point and upper edge, and on both sides of every + untuned-tile and other kernel edge of its ``nloc``-expert window.""" + out, b = set(), ranks + while b <= 8192: + vals = list(range(b, min(2 * b - 1, 8192) + 1, ranks)) + out |= {vals[0], vals[len(vals) // 2], vals[-1]} + b *= 2 + for e in _v41_tile_edges(nloc) | set(_V41_KERNEL_EDGES): + out |= {v for v in (ranks * (e // ranks), ranks * -(-e // ranks)) if ranks <= v <= 8192} + return sorted(out) + + +def _dep_pad_mask(n, cpu, ranks): + """Padding rows of one `_moe_chunks` chunk of `n` rows cut from a random gathered + batch of ``ranks`` rank blocks of `rows` rows (at most 16384 per rank): the chunk + index j is random below 2 * ranks, the batch holds 8192 (j + 1) rows when n is 8192 + (the chunk any of them) and 8192 j + n otherwise (the chunk the last); one rank is + full (the step's largest), the rank whose block overlaps the chunk most holds ONE + real token, the others a random count.""" + j = int(torch.randint(0, 2 * ranks, (1,), generator=cpu)) + if n == 8192: + rows, start = ( + 8192 * (j + 1) // ranks, + 8192 * int(torch.randint(0, j + 1, (1,), generator=cpu)), + ) + else: + rows, start = (8192 * j + n) // ranks, 8192 * j + ov = [max(0, min(start + n, (r + 1) * rows) - max(start, r * rows)) for r in range(ranks)] + big = max(range(ranks), key=lambda r: (ov[r], float(torch.rand(1, generator=cpu)))) + rest = [r for r in range(ranks) if r != big] + full = min(rest, key=lambda r: (ov[r], float(torch.rand(1, generator=cpu)))) + counts = [int(torch.randint(1, rows + 1, (1,), generator=cpu)) for _ in range(ranks)] + counts[full], counts[big] = rows, 1 + pad = torch.zeros(ranks * rows, dtype=torch.bool) + for r, c in enumerate(counts): + pad[r * rows + c : (r + 1) * rows] = True + return pad[start : start + n].to(DEV) + + +def _dep_case(kind, n, w, g1, seed, ranks, nloc): + """As `_dep4_case`, for a dep member of ``ranks`` ranks on ``nloc``-expert windows + (kinds ``dep``, ``one_real``, ``all_pad``, ``skew6``).""" + gen = torch.Generator(device=DEV).manual_seed(seed) + cpu = torch.Generator().manual_seed(seed) + data, sf, x = _rand_activation(n, V41_H, g1, gen) + ids, wts = _v41_routing(n, gen) + six = torch.tensor(_dep_six(w, nloc), device=DEV, dtype=torch.int32) + if kind == "skew6": + pad = torch.zeros(n, dtype=torch.bool, device=DEV) + ids = six.expand(n, V41_K).contiguous() + elif kind == "dep": + pad = _dep_pad_mask(n, cpu, ranks) + else: + pad = torch.ones(n, dtype=torch.bool, device=DEV) + pad[0] = kind == "all_pad" + ids[pad] = six + wts[pad] = 0.25 + data[pad] = 0 + sf.view(n, V41_H // SV)[pad] = 0 + x[pad] = 0 + return (kind, data, sf, x, ids, wts), pad + + +def _dep_padding_skew(fam): + """dep2's (``ep2``) or dep8's (``ep8``) zero-padding skew on every window of the + member, eager, cold then warm; see the two tests below.""" + assert not _moe_cache_entries(), "the profiling cache already holds entries for this op" + nloc, _, ranks = _E6[fam] + args, ref_l = _v41_layer_g2() + layouts = _e6_layouts(fam, args, ref_l) + g1 = ref_l["g1"] + counts = _dep_eager_counts(ranks, nloc) + assert all(n % ranks == 0 for n in counts) and counts[0] == ranks and counts[-1] == 8192 + reach = {e for e in _v41_kernel_firsts(nloc, V41_I, False) if e % ranks == 0} + assert all(e in counts and e - ranks in counts for e in reach), sorted(reach - set(counts)) + + def plan(n): + """(kind, window of the six, seed) of every case at row count n.""" + s = 1_000_003 * n + 7 * ranks + cases = [("dep", None, s)] + [("dep", w, s + 1 + w) for w in range(ranks)] + cases += [ + ("one_real", (n // ranks) % ranks, s + ranks + 1), + ("all_pad", (n // ranks + 1) % ranks, s + ranks + 2), + ] + return cases + [("skew6", w, s + ranks + 3 + w) for w in range(ranks)] + + def run(lay, case): + _, off, nl, la, ls = lay + _, data, sf, _, ids, wts = case + out = torch.full((data.shape[0], V41_H), float("nan"), dtype=torch.bfloat16, device=DEV) + _call( + data, + sf, + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=nl, + topk_ids=ids, + topk_weights=wts, + output=out, + **ls, + ) + return out + + stats = {"calls": 0, "zero_rows": 0, "pad_rows": 0, "max_expert_rows": 0} + worst = {} + kern = {} + + def check(tag, n, kind, w, seed, profile_it): + case, pad = _dep_case(kind, n, w, g1, seed, ranks, nloc) + refs = _ref_v41_parts(case[3], case[4], case[5], ref_l, eps=(nloc,))[("ep", nloc)] + per_expert = torch.bincount(case[4].flatten().long(), minlength=V41_E) + stats["max_expert_rows"] = max(stats["max_expert_rows"], int(per_expert.max())) + stats["pad_rows"] += int(pad.sum()) + for lay, rf in zip(layouts, refs): + name = lay[0] + y = run(lay, case) + exp = rf.to(torch.bfloat16) + zero = (exp == 0).all(dim=1) + assert pad.logical_not().logical_or(zero).all(), ( + f"{tag} {name} T={n} {kind}: a padding row is nonzero" + ) + assert (y[zero] == 0).all(), ( + f"{tag} {name} T={n} {kind} six@{w}: a row the reference leaves zero is not" + ) + stats["zero_rows"] += int(zero.sum()) + if not zero.all(): + _assert_moe_close(y, exp) + dv = _dev(y[~zero], exp[~zero]) + _worst(worst, "all", dv) + _worst(worst, "all-on-six" if kind == "skew6" else "padded", dv) + if tag == "dump": + _worst(worst, "dumped", dv) + stats["calls"] += 1 + if profile_it and name == f"{fam}.{w}": + kern[(tag, n, kind)] = _kernel_names(lambda: run(lay, case)) + + uniform_kernels = {} + tuned = 0 + dumped = [n for n in _E6_DUMPED_TOKENS if n % ranks == 0] if fam == "ep8" else [] + for tag in ("cold", "warm"): + if tag == "warm": + with autotune(): + _, off, nl, la, ls = layouts[0] + c = _v41_case( + "uniform", 8192, off, nl, g1, torch.Generator(device=DEV).manual_seed(4168) + ) + _call( + c[1], + c[2], + la, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=nl, + topk_ids=c[4], + topk_weights=c[5], + **ls, + ) + uids = sorted({str(k[2]) for k in _moe_cache_entries()}) + assert uids == [_E6_UIDS[fam]], uids + tuned = len(_moe_cache_entries()) + for n in counts: + u = _v41_case( + "uniform", n, 0, nloc, g1, torch.Generator(device=DEV).manual_seed(4169 + n) + ) + uniform_kernels[(tag, n)] = _kernel_names(lambda: run(layouts[0], u)) + for kind, w, seed in plan(n): + check( + tag, + n, + kind, + w, + seed, + profile_it=kind == "all_pad" or (kind == "skew6" and w == 0), + ) + for kind in ("skew6", "all_pad"): + if (tag, n, kind) in kern: + assert kern[(tag, n, kind)] == uniform_kernels[(tag, n)], ( + f"{tag} T={n} {kind}: the routing changed the launched kernels" + ) + if tag == "warm": + for i, n in enumerate(dumped): + check( + "dump", + n, + "dep", + None if i % (ranks + 1) == ranks else i % (ranks + 1), + 2_000_003 * n, + False, + ) + assert len(_moe_cache_entries()) == tuned, "a skewed call re-tuned a bucket" + assert (tag == "warm") == bool(_moe_cache_entries()), ( + "an ordinary call wrote a profiling-cache entry" + ) + + # Control: a result missing the last quarter of one expert's rows (a kernel whose + # per-expert capacity assumed balanced routing) must fail the gate. + ctrl = [] + for n in (256, 2048, 8192): + case, _ = _dep_case("skew6", n, 0, g1, 4170 + n, ranks, nloc) + y = run(layouts[0], case) + lost = case[5].clone() + lost[3 * n // 4 :, 5] = 0 + bad = _ref_v41_parts(case[3], case[4], lost, ref_l, eps=(nloc,))[("ep", nloc)][0].to( + torch.bfloat16 + ) + ctrl.append((n, *(round(v, 1) for v in _dev(y, bad)))) + try: + _assert_moe_close(y, bad) + except AssertionError: + continue + raise AssertionError( + f"T={n}: the gate accepts a result missing a quarter of one expert's rows" + ) + AutoTuner.get().clear_cache() + fmt = {k: f"{v[0]:.2f} / {v[1]:.2f}" for k, v in worst.items()} + return ( + f"{len(counts)} row counts {counts[:6]}..{counts[-3:]} x {len(plan(8))} cases x {ranks} windows, cold and " + f"warm ({tuned} tuned entries), plus {len(dumped)} dumped counts warm: {stats['calls']} calls, " + f"{stats['pad_rows']} padding rows, max rows on one expert {stats['max_expert_rows']}, " + f"{stats['zero_rows']} rows exactly zero; worst elt / rms ulp {fmt}; kernels at a row count are " + f"routing-independent at all {len(counts)} counts cold and warm; lost-rows control (T, elt, rms) {ctrl}" + ) + + +def test_v41_e6_dep2_padding_skew_eager_by_bucket(): + """dep2's zero-padding skew on both EP2 windows, eager, at every chunk size the target passes. + + As `test_v41_dep4_padding_skew_eager_by_bucket` with two rank blocks: row counts in + multiples of 2 at every tuner bucket's lower edge, an interior point and upper edge, + and on both sides of every untuned-tile and other kernel edge of a 192-expert window + (`_dep_eager_counts`). At each count, each case runs on both windows with `output=` + NaN-filled: a random chunk of a padded gathered batch (`_dep_pad_mask`) with the six + padding experts spread over the windows and with all six inside each window in turn; + one real row behind n - 1 padding rows; n padding rows; and, for each window, n real + rows all on its six experts. Held to the torch reference; every row the reference + leaves at zero must be exactly zero. Cold, then warm (`autotune()` at 8192 rows): no + warm call writes a cache entry, and the kernels at a row count are the same for + uniform, skewed and all-padding routing. Control: a reference that loses the last + quarter of one expert's rows fails the gate. Clears the cache at the end. + """ + print(f" test_v41_e6_dep2_padding_skew_eager_by_bucket OK ({_dep_padding_skew('ep2')})") + + +def test_v41_e6_dep8_padding_skew_eager_by_bucket(): + """dep8's zero-padding skew on all eight EP8 windows, eager, at every chunk size the target passes. + + As `test_v41_e6_dep2_padding_skew_eager_by_bucket` with eight rank blocks and 48-expert + windows (row counts in multiples of 8; every case on all eight windows; all six + padding experts inside each window in turn, and n real rows on each window's six), + and the warm pass adds a padded chunk at every multiple of 8 among + `_E6_DUMPED_TOKENS`. + """ + print(f" test_v41_e6_dep8_padding_skew_eager_by_bucket OK ({_dep_padding_skew('ep8')})") + + def test_do_finalize_and_output_buffer(): """Both output modes, and what `do_finalize=False` actually returns.""" num_experts, hidden, inter, top_k, num_tokens = 8, 256, 128, 2, 16 @@ -1849,3 +6599,164 @@ def test_strided_inputs_are_silently_wrong(): "a strided hidden_states no longer changes the result — re-check the guard" ) print(" test_strided_inputs_are_silently_wrong OK") + + +def _certify_v41_cell(cell) -> None: + """One DeepSeek-V4.1 cell: the target's call on one rank, against the entry's + reference, which is first tied to `_ref_v41_parts`. + + The layer is `_v41()`'s -- per-expert `gw1`, `gw2`, `g2`, the clamp at 10 + passed as `10 / output1_scale_gate_scalar` -- or, for the layer-level cells, + `_v41_layer_g2()`'s, the form the target passes. With `output` the kernel + writes into a NaN-filled `output=` buffer, as the target's call does, and + returns the empty list element; otherwise it returns the result. + + An EP window's entry reference is `_ref_v41_parts`' share of that window bit + for bit (same per-expert row order, same fp32 ops), and a token with no slot + in the window must come back bitwise zero. A MoE-TP rank's entry reference + runs its FC1 GEMM over the rank's own columns while `_ref_v41_parts` slices + the full-width one; the two GEMMs may round an fp32 element differently and + flip one e2m1 code of the intermediate, so that tie takes the entry's gate + rather than bitwise equality. + """ + spec = cell.spec + assert (spec["hidden"], spec["num_experts"], spec["top_k"]) == (V41_H, V41_E, V41_K), spec + assert spec["swiglu_limit"] == V41_LIMIT, spec + args, ref = _v41() if spec["g2"] == "expert" else _v41_layer_g2() + off, n, inter = spec["local_expert_offset"], spec["local_num_experts"], spec["intermediate"] + tokens = spec["tokens"] + gen = torch.Generator(device=DEV).manual_seed(zlib.crc32(cell.why.encode())) + + data, sf, x = _rand_activation(tokens, V41_H, ref["g1"], gen) + ids, wts = _v41_routing(tokens, gen) + scal = _v41_scalars(ref) + if spec["tp_rank"] is None: + kargs = args if n == V41_E else _window_args(args, off, n) + kscal = _v41_slice(scal, off, n) + ops = _window_ref(ref, off, n) + share = _ref_v41_parts(x, ids, wts, ref, eps=(n,))[("ep", n)][off // n] + else: + r = spec["tp_rank"] + a, b = r * inter, (r + 1) * inter + kargs = _v41_tp_args(inter)[r] + kscal = scal + ops = dict( + up=(ref["up"][0][:, a:b], ref["up"][1][:, a:b]), + gate=(ref["gate"][0][:, a:b], ref["gate"][1][:, a:b]), + down=(ref["down"][0][:, :, a:b], ref["down"][1][:, :, a // SV : b // SV]), + ) + share = _ref_v41_parts(x, ids, wts, ref, tps=(inter,))[("tp", inter)][r] + assert kargs["intermediate_size"] == inter, (kargs["intermediate_size"], inter) + win = slice(off, off + n) + + # 1. The entry's reference against this file's, on identical inputs. + entry_ref = moe.reference( + x, + ids, + wts, + ops["up"], + ops["gate"], + ops["down"], + ref["g2"][win], + local_expert_offset=off, + local_num_experts=n, + fc1_weight_global_scale=ref["gw1"][win], + fc2_weight_global_scale=ref["gw2"][win], + swiglu_limit=spec["swiglu_limit"], + ) + if spec["tp_rank"] is None: + torch.testing.assert_close(entry_ref, share, rtol=0.0, atol=0.0) + else: + moe.compare(entry_ref.to(torch.bfloat16), share.to(torch.bfloat16)) + + # 2. The entry's reference against the kernel, through the entry's gate. + kw = {} + if spec["output"]: + kw["output"] = torch.full((tokens, V41_H), float("nan"), dtype=torch.bfloat16, device=DEV) + res = _call( + data, + sf, + kargs, + V41_E, + V41_K, + local_expert_offset=off, + local_num_experts=n, + topk_ids=ids, + topk_weights=wts, + **kw, + **kscal, + ) + assert len(res) == 1, [tuple(t.shape) for t in res] + if spec["output"]: + assert res[0].numel() == 0, tuple(res[0].shape) + out = kw.get("output", res[0]) + assert out.shape == (tokens, V41_H), out.shape + if spec["tp_rank"] is None: + outside = ~((ids >= off) & (ids < off + n)).any(dim=1) + assert torch.equal(out[outside], torch.zeros_like(out[outside])), ( + "a token with no slot in the window did not come back zero" + ) + moe.compare(out, entry_ref.to(torch.bfloat16)) + + +@pytest.mark.parametrize("cell", moe.CELLS, ids=[c.why[:44] for c in moe.CELLS]) +def test_certified_cells(cell): + """Drive every configuration the entry claims, through the entry's own + reference and its own gate. + + The cases above test the op; this one tests the *claim*. A cell list that + nothing executes reads as coverage and is not. + + Each cell cross-checks `moe.reference` against this file's `_ref_moe` + before either is used to gate the kernel. Two references disagreeing is a + different failure from the kernel moving, and collapsing them into one + assert loses which happened. + """ + spec = cell.spec + if spec.get("model") == "deepseek-v4.1": + _certify_v41_cell(cell) + return + args, ref = _dsr1() + off, n = spec["local_expert_offset"], spec["local_num_experts"] + num_experts, hidden, top_k = spec["num_experts"], spec["hidden"], spec["top_k"] + gen = torch.Generator(device=DEV).manual_seed(abs(hash(cell.why)) % 2**31) + g1 = 137.0 + + data, sf, x = _rand_activation(spec["tokens"], hidden, g1, gen) + ids, wts = _routing(spec["tokens"], num_experts, top_k, gen) + g2 = _calibrate_g2(x, ids, ref) + s1, sg, s2 = _scalars(num_experts, g1, g2) + + win_args = args if n == num_experts else _window_args(args, off, n) + win_ref = ref if n == num_experts else _window_ref(ref, off, n) + + # 1. The entry's reference against this file's, on identical inputs. + entry_ref = moe.reference( + x, + ids, + wts, + win_ref["up"], + win_ref["gate"], + win_ref["down"], + g2, + local_expert_offset=off, + local_num_experts=n, + ) + local_ref = _ref_moe(x, ids, wts, win_ref, g2, offset=off, num_local=n) + torch.testing.assert_close(entry_ref, local_ref, rtol=0.0, atol=0.0) + + # 2. The entry's reference against the kernel, through the entry's gate. + y = _call( + data, + sf, + win_args, + num_experts, + top_k, + local_expert_offset=off, + local_num_experts=n, + topk_ids=ids, + topk_weights=wts, + **_window_scalars(s1, sg, s2, off, n), + )[0] + assert y.shape == (spec["tokens"], hidden), y.shape + moe.compare(y, entry_ref.to(torch.bfloat16)) diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py index 71795816de8b..4eef29b98ff1 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py +++ b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py @@ -2,16 +2,25 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the mxe4m3_mxe2m1_block_scale_moe_runner catalog entry.""" +import contextlib import math +import zlib +import pytest import torch +from torch.profiler import ProfilerActivity, profile, record_function from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.mxe4m3_mxe2m1_block_scale_moe_runner import ( mxe4m3_mxe2m1_block_scale_moe_runner as moe, ) +from tensorrt_llm._torch.autotuner import AutoTuner, autotune, fallback_tactic + +__extra_import_path__ = [".."] +from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "mxe4m3_mxe2m1_block_scale_moe_runner requires a CUDA device" + DEV = "cuda" # The reference GEMMs must be true fp32; TF32 would leave the reference with # 10 mantissa bits, coarser than the bf16 output it is meant to bound. @@ -71,15 +80,15 @@ def _exp_base(k: int) -> int: return 127 + round(0.5 * math.log2(0.01057 / k)) -def _rand_mxfp4(e: int, n: int, k: int, gen: torch.Generator): +def _rand_mxfp4(e: int, n: int, k: int, gen: torch.Generator, exp_shift: int = 0): """Random mxfp4 expert stack. Returns `(packed [E, N, K/2] uint8, scales [E, N, K/32] uint8, codes [E, N, K] uint8)`. Codes and e8m0 exponents are drawn directly, so the fp32 value of every weight is exact — the reference never has to model a - quantizer. + quantizer. `exp_shift` scales the whole stack by `2^exp_shift`. """ - base = _exp_base(k) + base = _exp_base(k) + exp_shift codes = torch.randint(0, 16, (e, n, k), dtype=torch.uint8, device=DEV, generator=gen) exps = torch.randint( base, base + 6, (e, n, k // SV), dtype=torch.uint8, device=DEV, generator=gen @@ -203,27 +212,44 @@ def _prep_fc2(dn_p, dn_s, dn_b, h2_pad, i_pad): ) -def _build(num_experts: int, hidden: int, inter: int, seed: int): - """Build one MoE layer: kernel-ready tensors plus the reference operands.""" +def _build( + num_experts: int, + hidden: int, + inter: int, + seed: int, + fc1_exp_shift: int = 0, + with_bias: bool = True, +): + """Build one MoE layer: kernel-ready tensors plus the reference operands. + + `fc1_exp_shift` scales both FC1 halves by `2^fc1_exp_shift` (so a clamp + limit bites more often); `with_bias=False` passes both biases as `None` + and gives the reference zero biases. + """ gen = torch.Generator(device=DEV).manual_seed(seed) i_pad = _pad_up(inter, WEIGHT_ALIGN) h1_pad = _pad_up(hidden, HIDDEN_ALIGN) h2_pad = _pad_up(hidden, WEIGHT_ALIGN) - up_p, up_s, up_c = _rand_mxfp4(num_experts, inter, hidden, gen) - gt_p, gt_s, gt_c = _rand_mxfp4(num_experts, inter, hidden, gen) + up_p, up_s, up_c = _rand_mxfp4(num_experts, inter, hidden, gen, fc1_exp_shift) + gt_p, gt_s, gt_c = _rand_mxfp4(num_experts, inter, hidden, gen, fc1_exp_shift) dn_p, dn_s, dn_c = _rand_mxfp4(num_experts, hidden, inter, gen) - up_b = torch.randn(num_experts, inter, device=DEV, generator=gen) - gt_b = torch.randn(num_experts, inter, device=DEV, generator=gen) - dn_b = torch.randn(num_experts, hidden, device=DEV, generator=gen) * 0.05 + if with_bias: + up_b = torch.randn(num_experts, inter, device=DEV, generator=gen) + gt_b = torch.randn(num_experts, inter, device=DEV, generator=gen) + dn_b = torch.randn(num_experts, hidden, device=DEV, generator=gen) * 0.05 + else: + up_b = torch.zeros(num_experts, inter, device=DEV) + gt_b = torch.zeros(num_experts, inter, device=DEV) + dn_b = torch.zeros(num_experts, hidden, device=DEV) w1, s1, b1 = _prep_fc1(up_p, gt_p, up_s, gt_s, up_b, gt_b, i_pad, h1_pad) w2, s2, b2 = _prep_fc2(dn_p, dn_s, dn_b, h2_pad, i_pad) args = dict( gemm1_weights=w1, gemm1_weights_scale=s1, - gemm1_bias=b1, + gemm1_bias=b1 if with_bias else None, gemm2_weights=w2, gemm2_weights_scale=s2, - gemm2_bias=b2, + gemm2_bias=b2 if with_bias else None, intermediate_size=i_pad, valid_hidden_size=hidden, valid_intermediate_size=inter, @@ -255,34 +281,37 @@ def _call(data, sf, args, num_experts, top_k, **kw): kw.setdefault("valid_hidden_size", args["valid_hidden_size"]) kw.setdefault("valid_intermediate_size", args["valid_intermediate_size"]) kw.setdefault("intermediate_size", args["intermediate_size"]) - return moe( - kw.pop("routing_logits"), - kw.pop("routing_bias"), - data, - sf, - args["gemm1_weights"], - args["gemm1_weights_scale"], - args["gemm1_bias"], - kw.pop("gemm1_alpha"), - kw.pop("gemm1_beta"), - kw.pop("gemm1_clamp_limit"), - args["gemm2_weights"], - args["gemm2_weights_scale"], - args["gemm2_bias"], - num_experts, - top_k, - kw.pop("n_group"), - kw.pop("topk_group"), - kw.pop("intermediate_size"), - kw.pop("valid_hidden_size"), - kw.pop("valid_intermediate_size"), - kw.pop("local_expert_offset"), - kw.pop("local_num_experts"), - kw.pop("routed_scaling_factor"), - kw.pop("routing_method_type"), - kw.pop("act_type"), - **kw, - ) + # _call is the single call site: run the guard over every call it makes. + # The negative cases reach the raw op directly when they want it unguarded. + with validating(moe): + return moe( + kw.pop("routing_logits"), + kw.pop("routing_bias"), + data, + sf, + args["gemm1_weights"], + args["gemm1_weights_scale"], + args["gemm1_bias"], + kw.pop("gemm1_alpha"), + kw.pop("gemm1_beta"), + kw.pop("gemm1_clamp_limit"), + args["gemm2_weights"], + args["gemm2_weights_scale"], + args["gemm2_bias"], + num_experts, + top_k, + kw.pop("n_group"), + kw.pop("topk_group"), + kw.pop("intermediate_size"), + kw.pop("valid_hidden_size"), + kw.pop("valid_intermediate_size"), + kw.pop("local_expert_offset"), + kw.pop("local_num_experts"), + kw.pop("routed_scaling_factor"), + kw.pop("routing_method_type"), + kw.pop("act_type"), + **kw, + ) # ── reference ───────────────────────────────────────────────────────────── @@ -356,6 +385,9 @@ def _ref_moe( num_local: int | None = None, swap_gate_up: bool = False, quantize_intermediate: bool = True, + recipe: str | None = None, + weights_before_down: bool = False, + stats: dict | None = None, ): """Native-torch MoE over dequantized mxfp4 weights, fp32 throughout. @@ -363,6 +395,12 @@ def _ref_moe( `[offset, offset + num_local)`. `scales[t, j]` multiplies slot `j`'s expert output; nothing is renormalized. The FC1 activation is requantized to MXFP8 before FC2, which is what the kernel does. + + `weights_before_down=True` is the DeepSeek-V4.1 checkpoint reference's + order instead (`Expert.forward` in the checkpoint's `inference/model.py`): + the routing weight multiplies the activation, which is rounded to bf16 and + then MXFP8-quantized as FC2's input. `stats`, when given, accumulates how + many FC1 elements the clamp limit changed. """ num_tokens, hidden = x_valid.shape up_c, up_s, up_b = ref["up"] @@ -383,15 +421,22 @@ def _ref_moe( gate = xe @ _dequant(gt_c[local_e], gt_s[local_e]).t() + gt_b[local_e] if limit is not None: lim = float(limit[local_e]) + if stats is not None: + stats["gate"] = stats.get("gate", 0) + int((gate > lim).sum()) + stats["up"] = stats.get("up", 0) + int((up.abs() > lim).sum()) + stats["n"] = stats.get("n", 0) + gate.numel() gate = gate.clamp(max=lim) up = up.clamp(-lim, lim) a = 1.0 if alpha is None else float(alpha[local_e]) b = 0.0 if beta is None else float(beta[local_e]) act = (up + b) * gate * torch.sigmoid(a * gate) + w = scales[tok, slot].float().unsqueeze(1) + if weights_before_down: + act = (act * w).to(torch.bfloat16).float() if quantize_intermediate: - act = _q_intermediate(act) + act = _q_intermediate(act, recipe) y = act @ _dequant(dn_c[local_e], dn_s[local_e]).t() + dn_b[local_e] - out.index_add_(0, tok, y * scales[tok, slot].float().unsqueeze(1)) + out.index_add_(0, tok, y if weights_before_down else y * w) return out @@ -415,8 +460,9 @@ def _assert_moe_close(y: torch.Tensor, ref: torch.Tensor) -> None: one output ulp of a two-GEMM chain, and per-element `rtol` is meaningless where cancellation drives `|ref|` to zero. So the element gate is 8 ulp of the row's largest magnitude and the aggregate gate is 4 ulp of relative - RMS. Worst values measured over every configuration covered here: 2.0 ulp - element-wise (128 experts, 8192 tokens, H=I=2880) and 0.87 ulp RMS. Both + RMS. Worst values measured over every configuration covered here: 2.2 ulp + element-wise (a skewed-routing capture replay at the DSpark drafter + geometry; 2.0 at 128 experts, 8192 tokens, H=I=2880) and 0.87 ulp RMS. Both gates were sized against deliberately wrong computations when they were set: a reference that skips the FC1-output requantization lands at 19.6 / 10.6 ulp, and a gate/up-swapped or unshuffled operand at 100+ ulp. @@ -710,6 +756,306 @@ def test_intermediate_is_mxfp8_quantized(): print(" test_intermediate_is_mxfp8_quantized OK") +def test_dspark_intermediate_is_mxfp8_quantized(): + """DeepSeek-V4.1 DSpark drafter geometry, plain clamped SwiGLU: the FC1 + epilogue's MXFP8 intermediate, read out bit-exactly. + + H = 5120, I = 2304 is not square, so FC2 is a *partial* identity + (`down[i, i] = 1` for `i < I`): output columns `[0, I)` are the kernel's own + intermediate and `[I, H)` must be zero. alpha/beta are left at `None` and + the clamp limit is 10 -- the V4.1 `Expert` (gate <= 10, up in [-10, 10], + `silu(gate) * up`). Two FC1 scales: one where the clamp bites on ~10% of + gate and ~20% of up elements, one where the saturating-scale branch is hit + by ~19% of blocks. + """ + hidden, inter, num_experts, num_tokens = 5120, 2304, 2, 16 + for seed, shift in ((81, 1), (83, 0)): + gen = torch.Generator(device=DEV).manual_seed(seed) + up_p, up_s, up_c = _rand_mxfp4(num_experts, inter, hidden, gen, shift) + gt_p, gt_s, gt_c = _rand_mxfp4(num_experts, inter, hidden, gen, shift) + dn_c = torch.zeros(num_experts, hidden, inter, dtype=torch.uint8, device=DEV) + dn_c[:, torch.arange(inter), torch.arange(inter)] = 2 # e2m1 code 2 = 1.0 + dn_p = (dn_c[..., 0::2] | (dn_c[..., 1::2] << 4)).contiguous() + dn_s = torch.full((num_experts, hidden, inter // SV), 127, dtype=torch.uint8, device=DEV) + zero1 = torch.zeros(num_experts, inter, device=DEV) + zero2 = torch.zeros(num_experts, hidden, device=DEV) + w1, s1, _ = _prep_fc1(up_p, gt_p, up_s, gt_s, zero1, zero1, inter, hidden) + w2, s2, _ = _prep_fc2(dn_p, dn_s, zero2, hidden, inter) + args = dict( + gemm1_weights=w1, + gemm1_weights_scale=s1, + gemm1_bias=None, + gemm2_weights=w2, + gemm2_weights_scale=s2, + gemm2_bias=None, + intermediate_size=inter, + valid_hidden_size=hidden, + valid_intermediate_size=inter, + ) + data, sf, xv = _rand_mxfp8(num_tokens, hidden, hidden, gen) + ids = torch.zeros(num_tokens, 1, dtype=torch.int32, device=DEV) + wts = torch.ones(num_tokens, 1, dtype=torch.bfloat16, device=DEV) + limit = torch.full((num_experts,), 10.0, dtype=torch.float32, device=DEV) + out = _call( + data, sf, args, num_experts, 1, topk_ids=ids, topk_weights=wts, gemm1_clamp_limit=limit + ).float() + assert torch.count_nonzero(out[:, inter:]) == 0, "columns past I must be zero" + up = xv @ _dequant(up_c[0], up_s[0]).t() + gate = xv @ _dequant(gt_c[0], gt_s[0]).t() + act = up.clamp(-10.0, 10.0) * torch.nn.functional.silu(gate.clamp(max=10.0)) + recipe = _scale_recipe() + assert torch.equal(out[:, :inter], _q_intermediate(act, recipe)), ( + f"the FC1 epilogue's MXFP8 requantization at H=5120/I=2304 is not the {recipe!r} " + "block scale; see test_intermediate_is_mxfp8_quantized" + ) + other = _OCP_SCALE if recipe == _ROUND_UP_SCALE else _ROUND_UP_SCALE + assert not torch.equal(out[:, :inter], _q_intermediate(act, other)), ( + f"the {other!r} scale fits too -- this case does not separate the recipes" + ) + assert not torch.equal(out[:, :inter], act.to(torch.bfloat16).float()) + if shift == 1: + assert (gate > 10.0).float().mean() > 0.05 and ( + up.abs() > 10.0 + ).float().mean() > 0.05, "the clamp limit does not bite in this case" + print(" test_dspark_intermediate_is_mxfp8_quantized OK") + + +def _gate_rejects(y: torch.Tensor, ref: torch.Tensor) -> bool: + try: + _assert_moe_close(y, ref) + except AssertionError: + return True + return False + + +def test_dspark_drafter_surface(): + """DeepSeek-V4.1 DSpark drafter MoE: H=5120, I=2304, 128 experts all local + (offset 0), top-3, clamp limit 10 per expert, no alpha/beta, no biases, + pre-routed bf16 weights; decode blocks T = 5..160 and prefill rows to 8192. + + Deliberately wrong references (no clamp, the gpt-oss clamp 7, gate/up + swapped, the gpt-oss alpha/beta, no intermediate requantization) must all + land outside the gate, and the checkpoint's own order -- routing weight + multiplied into the activation before FC2's MXFP8 input quantization -- + is measured to be a different computation from this kernel's. + """ + num_experts, hidden, inter, top_k = 128, 5120, 2304, 3 + args, ref = _build(num_experts, hidden, inter, seed=300, fc1_exp_shift=1, with_bias=False) + gen = ref["gen"] + assert ref["h1_pad"] == hidden and args["intermediate_size"] == inter, ( + "5120/2304 need no padding" + ) + limit = torch.full((num_experts,), 10.0, dtype=torch.float32, device=DEV) + for num_tokens in (1, 5, 40, 85, 160, 1024, 8192): + data, sf, xv = _rand_mxfp8(num_tokens, hidden, hidden, gen) + ids, _ = _routing(num_tokens, num_experts, top_k, gen) + wts = (torch.rand(num_tokens, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + out = _call( + data, + sf, + args, + num_experts, + top_k, + topk_ids=ids, + topk_weights=wts, + gemm1_clamp_limit=limit, + ) + assert out.shape == (num_tokens, hidden), out.shape + stats = {} + exp = _ref_moe(xv, ids, wts, ref, None, None, limit, stats=stats).to(torch.bfloat16) + _assert_moe_close(out, exp) + elt, rms = _dev(out, exp) + print( + f" T={num_tokens}: {elt:.2f} ulp elt, {rms:.2f} ulp RMS; clamp changed " + f"{stats['gate'] / stats['n']:.3f} of gate, {stats['up'] / stats['n']:.3f} of up" + ) + if num_tokens >= 40: + assert stats["gate"] / stats["n"] > 0.02 and stats["up"] / stats["n"] > 0.02, stats + if num_tokens != 160: + continue + wrong = { + "no clamp": _ref_moe(xv, ids, wts, ref, None, None, None), + "clamp 7": _ref_moe(xv, ids, wts, ref, None, None, limit * 0.7), + "gate/up swapped": _ref_moe(xv, ids, wts, ref, None, None, limit, swap_gate_up=True), + "gpt-oss alpha/beta": _ref_moe( + xv, + ids, + wts, + ref, + torch.full((num_experts,), 1.702, device=DEV), + torch.ones(num_experts, device=DEV), + limit, + ), + "no intermediate requant": _ref_moe( + xv, ids, wts, ref, None, None, limit, quantize_intermediate=False + ), + } + for name, bad in wrong.items(): + elt, rms = _dev(out, bad.to(torch.bfloat16)) + print(f" WRONG {name}: {elt:.2f} ulp elt, {rms:.2f} ulp RMS") + assert _gate_rejects(out, bad.to(torch.bfloat16)), ( + f"the {name} reference passed the gate" + ) + # the checkpoint's order: weight * act -> bf16 -> MXFP8 -> FC2. The + # kernel multiplies after FC2 instead; the two differ by a whole + # requantization's worth of rounding (measured ~10 ulp RMS). + ckpt = _ref_moe(xv, ids, wts, ref, None, None, limit, weights_before_down=True) + elt, rms = _dev(out, ckpt.to(torch.bfloat16)) + print(f" checkpoint order: {elt:.2f} ulp elt, {rms:.2f} ulp RMS") + assert rms > 6.0, f"checkpoint-order reference unexpectedly close: {rms:.2f} ulp RMS" + + # fed by the certified quantizer: at H = 5120 alignment 512 adds no padding + for num_tokens in (40, 160): + xb = torch.randn(num_tokens, hidden, device=DEV, dtype=torch.bfloat16, generator=gen) * 2 + data, sf = torch.ops.trtllm.mxfp8_quantize(xb, False, 512) + assert data.shape == (num_tokens, hidden) and sf.shape == (num_tokens * hidden // SV,) + ids, _ = _routing(num_tokens, num_experts, top_k, gen) + wts = (torch.rand(num_tokens, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + out = _call( + data, + sf, + args, + num_experts, + top_k, + topk_ids=ids, + topk_weights=wts, + gemm1_clamp_limit=limit, + ) + exp = _ref_moe(_dequant_mxfp8(data, sf), ids, wts, ref, None, None, limit) + _assert_moe_close(out, exp.to(torch.bfloat16)) + elt, rms = _dev(out, exp.to(torch.bfloat16)) + print(f" mxfp8_quantize-fed T={num_tokens}: {elt:.2f} ulp elt, {rms:.2f} ulp RMS") + + # CUDA-graph capture at a decode size, replayed over new inputs written + # into the captured buffers, equals an eager call on those inputs bitwise + num_tokens = 40 + data, sf, _ = _rand_mxfp8(num_tokens, hidden, hidden, gen) + ids, _ = _routing(num_tokens, num_experts, top_k, gen) + wts = (torch.rand(num_tokens, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + kw = dict(topk_ids=ids, topk_weights=wts, gemm1_clamp_limit=limit) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + _call(data, sf, args, num_experts, top_k, **kw) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = _call(data, sf, args, num_experts, top_k, **kw) + data2, sf2, _ = _rand_mxfp8(num_tokens, hidden, hidden, gen) + ids2, _ = _routing(num_tokens, num_experts, top_k, gen) + wts2 = (torch.rand(num_tokens, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + data.copy_(data2) + sf.copy_(sf2) + ids.copy_(ids2) + wts.copy_(wts2) + graph.replay() + torch.cuda.synchronize() + eager = _call( + data2, + sf2, + args, + num_experts, + top_k, + topk_ids=ids2, + topk_weights=wts2, + gemm1_clamp_limit=limit, + ) + assert torch.equal(captured, eager), "graph replay over new inputs differs from eager" + print(" test_dspark_drafter_surface OK") + + +# The DSpark drafter's MoE runs on its 5-token draft block: T = 5G for the +# drafter batch G = 1..32, inside the target's decode CUDA graphs (the capture +# set 5 * {1..8, 16, 24, 32} is a subset) and eagerly in mixed steps. +# Replicated on every rank, so no dep4 gather reaches it. +_DSPARK_BLOCK_TOKENS = tuple(range(5, 161, 5)) + + +def _skewed_routing(num_tokens, num_experts, top_k, gen): + """Every token routed to the same top_k experts (first, middle and last), in a per-token order: + the most unbalanced routing a captured launch can meet.""" + hot = torch.tensor( + [0, num_experts // 2, num_experts - 1][:top_k], dtype=torch.int32, device=DEV + ) + ids = torch.stack( + [hot[torch.randperm(top_k, device=DEV, generator=gen)] for _ in range(num_tokens)] + ) + wts = (torch.rand(num_tokens, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + return ids.contiguous(), wts + + +def test_dspark_every_block_size_eager_and_capture(): + """DSpark drafter MoE at every T = 5..160 step 5: eager against the reference, then captured. + + Per T: an eager call on uniform random routing matches the kernel-order + reference. The call is then warmed up on a side stream and captured with + the op's autotuner cold, so the kernel configuration the graph holds is the + one chosen for this T at capture. Two replays follow, each after writing + new activations and new routing into the captured buffers: new uniform + routing, then every token on the same three experts (first, middle, last: + the most unbalanced load). Each replay is bitwise equal to an eager call + on the same inputs, and the skewed one also matches the reference. + """ + num_experts, hidden, inter, top_k = 128, 5120, 2304, 3 + args, ref = _build(num_experts, hidden, inter, seed=400, fc1_exp_shift=1, with_bias=False) + gen = ref["gen"] + limit = torch.full((num_experts,), 10.0, dtype=torch.float32, device=DEV) + worst = [0.0, 0.0] + for num_tokens in _DSPARK_BLOCK_TOKENS: + data, sf, xv = _rand_mxfp8(num_tokens, hidden, hidden, gen) + ids, _ = _routing(num_tokens, num_experts, top_k, gen) + wts = (torch.rand(num_tokens, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + kw = dict(topk_ids=ids, topk_weights=wts, gemm1_clamp_limit=limit) + out = _call(data, sf, args, num_experts, top_k, **kw) + exp = _ref_moe(xv, ids, wts, ref, None, None, limit).to(torch.bfloat16) + _assert_moe_close(out, exp) + elt, rms = _dev(out, exp) + worst = [max(worst[0], elt), max(worst[1], rms)] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + _call(data, sf, args, num_experts, top_k, **kw) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = _call(data, sf, args, num_experts, top_k, **kw) + for skewed in (False, True): + d2, s2, x2 = _rand_mxfp8(num_tokens, hidden, hidden, gen) + if skewed: + i2, w2 = _skewed_routing(num_tokens, num_experts, top_k, gen) + else: + i2, _ = _routing(num_tokens, num_experts, top_k, gen) + w2 = (torch.rand(num_tokens, top_k, device=DEV, generator=gen) * 1.5).to( + torch.bfloat16 + ) + data.copy_(d2) + sf.copy_(s2) + ids.copy_(i2) + wts.copy_(w2) + captured.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + eager = _call( + d2, + s2, + args, + num_experts, + top_k, + topk_ids=i2, + topk_weights=w2, + gemm1_clamp_limit=limit, + ) + assert torch.equal(captured, eager), (num_tokens, "skewed" if skewed else "uniform") + if skewed: + exp2 = _ref_moe(x2, i2, w2, ref, None, None, limit).to(torch.bfloat16) + _assert_moe_close(captured, exp2) + del graph, captured + print( + f" test_dspark_every_block_size_eager_and_capture OK: worst {worst[0]:.2f} ulp elt, " + f"{worst[1]:.2f} ulp RMS over T 5..160" + ) + + def test_activation_variants(): """alpha/beta/clamp_limit are independently optional and per-expert.""" num_experts, hidden, inter, top_k, num_tokens = 8, 1024, 512, 3, 24 @@ -1554,3 +1900,792 @@ def test_determinism_and_purity(): assert torch.equal(data, data_snap) and torch.equal(sf, sf_snap) assert torch.equal(ids, ids_snap) and torch.equal(wts, wts_snap) print(" test_determinism_and_purity OK") + + +# ── the DSpark drafter at max_batch_size 256 ───────────────────────────── +# +# The drafter runs this op once per stage on its N = 5G block rows, G the +# number of drafting requests on the rank (at most max_batch_size = 256), so +# N <= 1280 and N is always a multiple of 5. Under CUDA graphs G is a captured +# batch size: every G in Gset, the integers 1..32 plus TRT-LLM's padded list +# for max_batch_size 256. The target calls the op inside the autotuner's +# `fallback_tactic()` scope, in which `AutoTuner.choose_one` +# returns the untuned fallback without reading the profiling cache, so every +# rank takes the same tactic; the tests below call it exactly that way and +# record the tactic the AutoTuner hands the op on every call. + +_OP = "trtllm::mxe4m3_mxe2m1_block_scale_moe_runner" +_DSPARK_GSET = tuple(range(1, 33)) + (40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256) +_DSPARK_CAPTURE_N = tuple(5 * g for g in _DSPARK_GSET) +_DSPARK_N_MAX = 5 * 256 +# Every drafter-MoE row count the runtime-value dumps of the max_batch_size +# 256 engines recorded (padded CUDA-graph, eager, tp4 and dep4 legs) plus the +# long-context legs. +_DSPARK_DUMPED_N = ( + 5, 10, 15, 20, 40, 50, 80, 90, 110, 115, 120, 125, 140, 160, 190, 200, 240, 275, 280, 320, + 360, 370, 385, 400, 440, 455, 480, 520, 560, 580, 585, 600, 640, 650, 665, 710, 715, 720, + 960, 1280, +) # fmt: skip + + +def _drafter_layer(seed): + """The drafter's routed experts: 128 local, H 5120, I 2304, clamp 10 per expert, no biases.""" + args, ref = _build(128, 5120, 2304, seed=seed, fc1_exp_shift=1, with_bias=False) + limit = torch.full((128,), 10.0, dtype=torch.float32, device=DEV) + return args, ref, limit + + +def _pinned(data, sf, args, num_experts, top_k, **kw): + """The call as the target's drafter makes it: inside `fallback_tactic()`.""" + with fallback_tactic(): + return _call(data, sf, args, num_experts, top_k, **kw) + + +@contextlib.contextmanager +def _tactics(): + """Record `(tactic, tuning mode, capturing)` for every AutoTuner.choose_one call this op makes.""" + tuner = AutoTuner.get() + real = tuner.choose_one + rec = [] + + def spy(custom_op, runners, tuning_config, inputs, **kw): + out = real(custom_op, runners, tuning_config, inputs, **kw) + if custom_op == _OP: + rec.append((out[1], tuner.is_tuning_mode, torch.cuda.is_current_stream_capturing())) + return out + + tuner.choose_one = spy + try: + yield rec + finally: + del tuner.choose_one + + +def _op_cache_entries(): + return len(AutoTuner.get().profiling_cache.get_specific_custom_op(_OP)) + + +def _kernel_map(run, rows) -> dict: + """{N: the device activity one call at N launched, in launch order}. + + Each call runs in its own `record_function` range, bracketed by two + synchronizations, so the device events that start inside the range are + exactly the ones that call launched (the range's own device-side + annotation is excluded by name). + """ + names = {} + for i in range(0, len(rows), 256): + chunk = rows[i : i + 256] + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for n in chunk: + with record_function(f"N={n}"): + run(n) + torch.cuda.synchronize() + evs = prof.events() + ranges = [ + (int(e.name[2:]), e.time_range.start, e.time_range.end) + for e in evs + if e.name.startswith("N=") and e.device_type == torch.autograd.DeviceType.CPU + ] + device = sorted( + (e.time_range.start, e.name) + for e in evs + if e.device_type == torch.autograd.DeviceType.CUDA and not e.name.startswith("N=") + ) + for n, lo, hi in ranges: + names[n] = tuple(nm for t, nm in device if lo <= t <= hi) + del prof + return names + + +def _runs(names, rows): + """Maximal runs `[first N, last N, activity]` of consecutive N that launch the same device activity.""" + runs = [] + for n in rows: + assert names.get(n), f"no device activity attributed to the call at N={n}" + if runs and runs[-1][2] == names[n]: + runs[-1][1] = n + else: + runs.append([n, n, names[n]]) + return runs + + +def test_dspark_eager_every_row_count_to_1280(): + """Every drafter row count N = 1..1280, called as the target calls it, against the reference. + + One 1280-row activation and uniform random top-3 routing; the call at N + reads their first N rows and writes the first N rows of a NaN-poisoned + buffer (the target passes `output=`). Every N is held to the reference's + first N rows at the entry's gates, and a second call on the same inputs + must be bitwise equal to the first (every rank of the replicated drafter + computes the same rows). Every call is pinned as in the target and must + take the untuned fallback tactic. + + The cold tactic picks its kernels by N (the tokens-per-expert tile and the + default GEMM configuration, plus the routing kernels), so the test then + profiles every call 1..1280, derives the kernel map (maximal runs of + consecutive N launching the same device activity), and checks the first, + middle and last N of every run again with every token routed to the same + three experts (the most unbalanced load), against the reference (every + switch point at its edges). + """ + num_experts, hidden, top_k = 128, 5120, 3 + args, ref, limit = _drafter_layer(seed=500) + gen = ref["gen"] + n_max, sfw = _DSPARK_N_MAX, hidden // SV + data, sf, xv = _rand_mxfp8(n_max, hidden, hidden, gen) + ids, _ = _routing(n_max, num_experts, top_k, gen) + wts = (torch.rand(n_max, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + exp = _ref_moe(xv, ids, wts, ref, None, None, limit).to(torch.bfloat16) + buf = torch.empty(n_max, hidden, dtype=torch.bfloat16, device=DEV) + + def run(n, out=None, i=None, w=None): + out = buf[:n] if out is None else out + _pinned( + data[:n], + sf[: n * sfw], + args, + num_experts, + top_k, + topk_ids=ids[:n] if i is None else i, + topk_weights=wts[:n] if w is None else w, + gemm1_clamp_limit=limit, + output=out, + ) + return out + + assert _op_cache_entries() == 0, ( + "a cached tuned entry would move the pinned call off the fallback" + ) + with _tactics() as rec: + full = run( + n_max, torch.full((n_max, hidden), float("nan"), dtype=torch.bfloat16, device=DEV) + ) + _assert_moe_close(full, exp) + worst, off_prefix = [0.0, 0.0], [] + for n in range(1, n_max + 1): + buf[:n].fill_(float("nan")) + out = run(n) + _assert_moe_close(out, exp[:n]) + elt, rms = _dev(out, exp[:n]) + worst = [max(worst[0], elt), max(worst[1], rms)] + again = run(n, torch.full((n, hidden), float("nan"), dtype=torch.bfloat16, device=DEV)) + assert torch.equal(again, out), f"two identical calls at N={n} differ" + if not torch.equal(out, full[:n]): + off_prefix.append(n) + calls = len(rec) + rows = list(range(1, n_max + 1)) + names = _kernel_map(run, rows) + runs = _runs(names, rows) + edges = sorted({m for lo, hi, _ in runs for m in (lo, (lo + hi) // 2, hi)}) + edge_worst = [0.0, 0.0] + for n in edges: + i2, w2 = _skewed_routing(n, num_experts, top_k, gen) + out = run( + n, torch.full((n, hidden), float("nan"), dtype=torch.bfloat16, device=DEV), i2, w2 + ) + exp2 = _ref_moe(xv[:n], i2, w2, ref, None, None, limit).to(torch.bfloat16) + _assert_moe_close(out, exp2) + elt, rms = _dev(out, exp2) + edge_worst = [max(edge_worst[0], elt), max(edge_worst[1], rms)] + assert calls == 1 + 2 * n_max, calls + assert {r[0] for r in rec} == {-1}, sorted({str(r[0]) for r in rec}) + assert _op_cache_entries() == 0 + assert set(_DSPARK_DUMPED_N) <= set(rows) + # The cold path's kernel map as measured on sm_103: at 2 the FC1/FC2 + # configurations change (N = 1 runs FC1 with a static scheduler and FC2 + # split-K 2); the routing kernel changes at 5 (block -> dynamic block), + # 17, 65 and 129 (cluster kernel sizes 256 / 512 / 1024); the finalize + # kernel at 60 (vectorized loads); and the tokens-per-expert tile + # (clamp(next_pow2(floor(3N / 128)), 8, 256)) goes 8 -> 16 at 384 and + # 16 -> 32 at 726. The contract lists these; a move means both are stale. + assert [lo for lo, _, _ in runs] == [1, 2, 5, 17, 60, 65, 129, 384, 726], [ + lo for lo, _, _ in runs + ] + # Rows do not depend on N: every N >= 2 is bitwise the first N rows of + # the N = 1280 call; N = 1 differs only through its split-K FC2. + assert off_prefix == [1], off_prefix[:20] + kinds = sorted({nm for _, _, nm in runs}) + print( + f" test_dspark_eager_every_row_count_to_1280: {n_max} row counts in the gates, worst " + f"{worst[0]:.2f} ulp elt / {worst[1]:.2f} ulp RMS; repeat calls bitwise equal; {len(rec)} calls " + f"all on tactic -1; not bitwise equal to the first N rows of the N={n_max} call at " + f"{len(off_prefix)} N (first {off_prefix[:12]})" + ) + print( + f" kernel map: {len(runs)} runs of {len(kinds)} distinct sequences over N 1..{n_max}; switch " + f"points (first N of each run) {[lo for lo, _, _ in runs]}; {len(edges)} run edges/middles on " + f"skewed routing: worst {edge_worst[0]:.2f} ulp elt / {edge_worst[1]:.2f} ulp RMS" + ) + prev = () + for lo, hi, nm in runs: + changed = [x for x in nm if x not in prev] + print( + f" N {lo}..{hi}: {len(nm)} kernels; new vs previous run: {' || '.join(x[:400] for x in changed)}" + ) + prev = nm + + +def test_dspark_capture_every_5G_over_gset(): + """Capture at every drafter block size N = 5G for G in Gset (46 values, 5..1280), as the target runs it. + + Per N, in the engine's order: an eager call at that shape (the warm-up, + checked against the reference), a side-stream warm-up, then the capture + of the pinned call writing into an `output` allocated inside the capture + (as the target's `_moe` does). Every other capture is made while an + outer `autotune()` is open (the engine's tuning pass also runs the + drafter). Two replays follow, each after writing new activations and new + routing into the captured buffers -- uniform random routing, then every + token on the same three experts -- with the output poisoned first. Each + replay must be bitwise equal to an eager pinned call on the same inputs, + inside the gates against the reference, and differ from the capture-time + output on more than 90% of elements (a replay that reran on the old + inputs is visible). Every call, captured ones included, must get the + fallback tactic, and the op's AutoTuner cache must stay empty. + """ + num_experts, hidden, top_k = 128, 5120, 3 + args, ref, limit = _drafter_layer(seed=600) + gen = ref["gen"] + assert len(_DSPARK_CAPTURE_N) == 46 and max(_DSPARK_CAPTURE_N) == _DSPARK_N_MAX + assert _op_cache_entries() == 0 + worst, stale_min = [0.0, 0.0], 1.0 + with _tactics() as rec: + for i, n in enumerate(_DSPARK_CAPTURE_N): + data, sf, xv = _rand_mxfp8(n, hidden, hidden, gen) + ids, _ = _routing(n, num_experts, top_k, gen) + wts = (torch.rand(n, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + kw = dict(topk_ids=ids, topk_weights=wts, gemm1_clamp_limit=limit) + warm = torch.full((n, hidden), float("nan"), dtype=torch.bfloat16, device=DEV) + _pinned(data, sf, args, num_experts, top_k, output=warm, **kw) + _assert_moe_close( + warm, _ref_moe(xv, ids, wts, ref, None, None, limit).to(torch.bfloat16) + ) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + side = torch.empty(n, hidden, dtype=torch.bfloat16, device=DEV) + _pinned(data, sf, args, num_experts, top_k, output=side, **kw) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with autotune() if i % 2 else contextlib.nullcontext(): + with torch.cuda.graph(graph): + captured = torch.empty(n, hidden, dtype=torch.bfloat16, device=DEV) + _pinned(data, sf, args, num_experts, top_k, output=captured, **kw) + for skewed in (False, True): + d2, s2, x2 = _rand_mxfp8(n, hidden, hidden, gen) + if skewed: + i2, w2 = _skewed_routing(n, num_experts, top_k, gen) + else: + i2, _ = _routing(n, num_experts, top_k, gen) + w2 = (torch.rand(n, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + data.copy_(d2) + sf.copy_(s2) + ids.copy_(i2) + wts.copy_(w2) + captured.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + eager = torch.full((n, hidden), float("nan"), dtype=torch.bfloat16, device=DEV) + _pinned( + d2, s2, args, num_experts, top_k, + topk_ids=i2, topk_weights=w2, gemm1_clamp_limit=limit, output=eager, + ) # fmt: skip + assert torch.equal(captured, eager), (n, "skewed" if skewed else "uniform") + exp2 = _ref_moe(x2, i2, w2, ref, None, None, limit).to(torch.bfloat16) + _assert_moe_close(captured, exp2) + elt, rms = _dev(captured, exp2) + worst = [max(worst[0], elt), max(worst[1], rms)] + frac = (captured != warm).float().mean().item() + stale_min = min(stale_min, frac) + assert frac > 0.9, (n, skewed, frac) + del graph, captured, side + captured_calls = [r for r in rec if r[2]] + assert len(captured_calls) == len(_DSPARK_CAPTURE_N), len(captured_calls) + assert len(rec) == 5 * len(_DSPARK_CAPTURE_N), len(rec) + assert {r[0] for r in rec} == {-1}, sorted({str(r[0]) for r in rec}) + assert sum(r[1] for r in rec) == len(_DSPARK_CAPTURE_N) // 2, ( + "the outer autotune() was not open at capture" + ) + assert _op_cache_entries() == 0 + print( + f" test_dspark_capture_every_5G_over_gset OK: {len(_DSPARK_CAPTURE_N)} captures " + f"({len(_DSPARK_CAPTURE_N) // 2} inside an open autotune()), {2 * len(_DSPARK_CAPTURE_N)} replays bitwise = " + f"eager and in the gates (worst {worst[0]:.2f} ulp elt / {worst[1]:.2f} ulp RMS), replay vs " + f"capture-time output differing on >= {stale_min:.3f} of elements; {len(rec)} calls " + f"({len(captured_calls)} under capture, {sum(r[1] for r in rec)} with the outer tuning mode on) all on " + f"tactic -1; N {list(_DSPARK_CAPTURE_N)}" + ) + + +def test_pin_holds_the_cold_tactic_inside_a_tuning_pass(): + """The target's pin inside the engine's tuning pass, with a tuned entry cached, and a control. + + (1) A pinned call made inside an open `autotune()` -- the engine's warm-up + tuning pass, which runs the drafter -- takes the fallback tactic although + tuning mode is on at the call, leaves the outer tuning mode on afterwards, + adds no cache entry, and is bitwise equal to a call made outside any + tuning context, at every dumped row count, every captured N and every + cold-tile edge. + (2) Control: the same call UNPINNED inside `autotune()` is tuned -- the + recorder sees a tactic other than the fallback and the cache gains an + entry -- so the recorder in (1) would have seen a tuned call. + (3) With that entry cached, the pinned call still takes the fallback and + the cold bits, while `autotune(tune_mode=False)` alone (which only stops + tuning) takes the cached tactic. + """ + num_experts, hidden, top_k = 128, 5120, 3 + args, ref, limit = _drafter_layer(seed=700) + gen = ref["gen"] + tuner = AutoTuner.get() + rows = sorted(set(_DSPARK_DUMPED_N) | set(_DSPARK_CAPTURE_N) | {1, 383, 384, 725, 726}) + n_max, sfw = max(rows), hidden // SV + data, sf, _ = _rand_mxfp8(n_max, hidden, hidden, gen) + ids, _ = _routing(n_max, num_experts, top_k, gen) + wts = (torch.rand(n_max, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + + def call(n, pin=True): + out = torch.full((n, hidden), float("nan"), dtype=torch.bfloat16, device=DEV) + (_pinned if pin else _call)( + data[:n], sf[: n * sfw], args, num_experts, top_k, + topk_ids=ids[:n], topk_weights=wts[:n], gemm1_clamp_limit=limit, output=out, + ) # fmt: skip + return out + + assert _op_cache_entries() == 0 + cold = {n: call(n) for n in rows} + cache = tuner.profiling_cache.cache + try: + with _tactics() as rec: + restored, diff = [], [] + with autotune(): + for n in rows: + if not torch.equal(call(n), cold[n]): + diff.append(n) + restored.append(tuner.is_tuning_mode) + assert diff == [], f"pinned inside autotune() != cold at N {diff}" + assert all(restored), "the pin did not restore the outer tuning mode" + assert ( + len(rec) == len(rows) and {r[0] for r in rec} == {-1} and all(r[1] for r in rec) + ), rec + assert _op_cache_entries() == 0 + pinned_calls = len(rec) + + rec.clear() + with autotune(skip_dynamic_tuning_buckets=True): + call(5, pin=False) + tuned = rec[-1] + assert tuned[1], "the unpinned call inside autotune() did not run in tuning mode" + assert tuned[0] != -1, "the control call was not tuned" + assert _op_cache_entries() >= 1 + + rec.clear() + assert torch.equal(call(5), cold[5]) + assert [r[0] for r in rec] == [-1], (rec, tuned) + rec.clear() + with autotune(tune_mode=False): + call(5, pin=False) + untuned_only = rec[-1][0] + assert untuned_only == tuned[0], (rec, tuned) + print( + f" test_pin_holds_the_cold_tactic_inside_a_tuning_pass OK: {pinned_calls} pinned calls " + f"inside autotune() at N {rows[0]}..{rows[-1]} ({len(rows)} row counts) on tactic -1, bitwise = " + f"cold, cache empty; control: unpinned call tuned to {tuned[0]}; with that entry cached the pinned " + f"call takes -1 and the cold bits, autotune(tune_mode=False) alone takes {untuned_only}" + ) + finally: + for key in [k for k in cache if k[0] == _OP]: + del cache[key] + assert _op_cache_entries() == 0 + with _tactics() as rec: + assert torch.equal(call(5), cold[5]) + assert [r[0] for r in rec] == [-1], rec + + +def test_tactics_change_output_bits(): + """Why the drafter pins the cold tactic: the op's tunable tactics do not all give the same bits. + + At one row count per cold tile band and at the drafter's largest block + (N 5, 160, 385, 1280), every tactic the AutoTuner may choose from + (`get_valid_tactics` of the op's TunableRunner) is run on the same inputs + and compared bitwise with the pinned (cold) call, and against the + reference at the entry's gates. A rank that tunes on its own can therefore + land in a different bit class from its neighbour; the pinned call cannot. + """ + from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import MxE4m3MxE2m1BlockScaleMoERunner + + num_experts, hidden, inter, top_k = 128, 5120, 2304, 3 + args, ref, limit = _drafter_layer(seed=800) + gen = ref["gen"] + runner = MxE4m3MxE2m1BlockScaleMoERunner( + num_experts, top_k, None, None, inter, hidden, inter, 0, num_experts, None, 1, 0 + ) + summary = [] + for n in (5, 160, 385, 1280): + data, sf, xv = _rand_mxfp8(n, hidden, hidden, gen) + ids, _ = _routing(n, num_experts, top_k, gen) + wts = (torch.rand(n, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + exp = _ref_moe(xv, ids, wts, ref, None, None, limit).to(torch.bfloat16) + cold = torch.full((n, hidden), float("nan"), dtype=torch.bfloat16, device=DEV) + _pinned( + data, + sf, + args, + num_experts, + top_k, + topk_ids=ids, + topk_weights=wts, + gemm1_clamp_limit=limit, + output=cold, + ) + inputs = [ + None, None, data, sf, + args["gemm1_weights"], args["gemm1_weights_scale"], None, None, None, limit, + args["gemm2_weights"], args["gemm2_weights_scale"], None, wts, ids, + ] # fmt: skip + tactics = runner.get_valid_tactics(inputs, None) + differ, outside, worst = [], [], [0.0, 0.0] + for tac in tactics: + out = torch.full((n, hidden), float("nan"), dtype=torch.bfloat16, device=DEV) + runner.forward(inputs, tactic=list(tac), output=out) + if not torch.equal(out, cold): + differ.append(tuple(tac)) + if _gate_rejects(out, exp): + outside.append(tuple(tac)) + elt, rms = _dev(out, exp) + worst = [max(worst[0], elt), max(worst[1], rms)] + summary.append((n, len(tactics), len(differ), len(outside), worst, differ[:4], outside[:4])) + for n, total, nd, no, w, d4, o4 in summary: + print( + f" N={n}: {total} valid tactics, {nd} not bitwise equal to the cold call (first {d4}), " + f"{no} outside the gates (first {o4}); worst {w[0]:.2f} ulp elt / {w[1]:.2f} ulp RMS" + ) + assert all(total > 0 for _, total, *_ in summary) + assert all(no == 0 for _, _, _, no, *_ in summary), "a valid tactic lands outside the gates" + assert any(nd > 0 for _, _, nd, *_ in summary), ( + "every tactic gave the cold bits: this case separates nothing" + ) + + +def test_pin_ignores_a_cached_tactic(): + """The target's pin takes the fallback even when the process holds a tuned entry for this op. + + The drafter never tunes this op, but a process can still hold tuned + entries for it: `TLLM_AUTOTUNER_CACHE_PATH` makes the engine's warm-up + load an autotuner cache file (`autotune(cache_path=...)`), which another + engine may have tuned, rank by rank. The test writes such a file the + engine's way. It tunes the op unpinned at one row count per cold tile band + and at the drafter maximum (N 5, 160, 385, 1280), sets each entry to a + valid tactic whose bits differ from the fallback's where one exists (none + does at 1280, where the first valid tactic is used), saves the cache, + drops the entries, and loads the file through `autotune(cache_path=...)`. + + Control: the cache is then hit -- `autotune(tune_mode=False)` alone takes + the cached tactic and, at N 5, 160 and 385, gives different bits. The + target's pin must take tactic -1 and give the uncached call's bits at every + N: eagerly, inside an open `autotune()`, and under CUDA-graph capture + replayed on new inputs. + """ + import tempfile + + from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import MxE4m3MxE2m1BlockScaleMoERunner + + def _tac(t): + return list(t) if isinstance(t, (list, tuple)) else t + + num_experts, hidden, inter, top_k = 128, 5120, 2304, 3 + args, ref, limit = _drafter_layer(seed=900) + gen = ref["gen"] + tuner = AutoTuner.get() + cache = tuner.profiling_cache.cache + rows = (5, 160, 385, 1280) + + def case(n): + data, sf, _ = _rand_mxfp8(n, hidden, hidden, gen) + ids, _ = _routing(n, num_experts, top_k, gen) + wts = (torch.rand(n, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + return data, sf, ids, wts + + def call(fn, c, out=None): + data, sf, ids, wts = c + out = ( + torch.full((data.shape[0], hidden), float("nan"), dtype=torch.bfloat16, device=DEV) + if out is None + else out + ) + fn( + data, + sf, + args, + num_experts, + top_k, + topk_ids=ids, + topk_weights=wts, + gemm1_clamp_limit=limit, + output=out, + ) + return out + + def old_pin(*a, **kw): + with autotune(tune_mode=False): + return _call(*a, **kw) + + assert _op_cache_entries() == 0 + cases = {n: case(n) for n in rows} + fresh = {n: case(n) for n in rows} + cold = {n: call(_pinned, cases[n]) for n in rows} + cold_fresh = {n: call(_pinned, fresh[n]) for n in rows} + runner = MxE4m3MxE2m1BlockScaleMoERunner( + num_experts, top_k, None, None, inter, hidden, inter, 0, num_experts, None, 1, 0 + ) + chosen, differs = {}, {} + try: + for n in rows: + data, sf, ids, wts = cases[n] + inputs = [ + None, None, data, sf, + args["gemm1_weights"], args["gemm1_weights_scale"], None, None, None, limit, + args["gemm2_weights"], args["gemm2_weights_scale"], None, wts, ids, + ] # fmt: skip + tactics = [list(t) for t in runner.get_valid_tactics(inputs, None)] + pick = None + for tac in tactics: + out = torch.full((n, hidden), float("nan"), dtype=torch.bfloat16, device=DEV) + runner.forward(inputs, tactic=tac, output=out) + if not torch.equal(out, cold[n]): + pick = tac + break + differs[n] = pick is not None + chosen[n] = pick if pick is not None else tactics[0] + before = set(k for k in cache if k[0] == _OP) + with autotune(skip_dynamic_tuning_buckets=True): + call(_call, cases[n]) + new = set(k for k in cache if k[0] == _OP) - before + assert new, f"the unpinned call at N={n} inside autotune() wrote no cache entry" + for key in new: + runner_id, _, min_time = cache[key] + cache[key] = (runner_id, chosen[n], min_time) + assert [differs[n] for n in rows] == [True, True, True, False], differs + written = {k: _tac(v[1]) for k, v in cache.items() if k[0] == _OP} + with tempfile.TemporaryDirectory() as tmp: + path = f"{tmp}/autotuner_cache.json" + tuner.profiling_cache.save_cache(path, tuner.mapping.rank) + for key in [k for k in cache if k[0] == _OP]: + del cache[key] + assert _op_cache_entries() == 0 + with autotune(tune_mode=False, cache_path=path): + pass + cache = tuner.profiling_cache.cache + loaded = {k: _tac(v[1]) for k, v in cache.items() if k[0] == _OP} + assert loaded == written, (loaded, written) + + with _tactics() as rec: + for n in rows: + rec.clear() + out = call(old_pin, cases[n]) + assert [_tac(r[0]) for r in rec] == [chosen[n]], (n, rec) + assert torch.equal(out, cold[n]) != differs[n], n + print( + f" control: autotune(tune_mode=False) takes the loaded tactics {[chosen[n] for n in rows]}" + ) + + with _tactics() as rec: + for n in rows: + assert torch.equal(call(_pinned, cases[n]), cold[n]), ( + f"pinned call != uncached call at N={n}" + ) + with autotune(): + assert torch.equal(call(_pinned, cases[n]), cold[n]), ( + f"pinned in autotune() != uncached at N={n}" + ) + for i, n in enumerate(rows): + c = case(n) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + call(_pinned, c) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with autotune() if i % 2 else contextlib.nullcontext(): + with torch.cuda.graph(graph): + captured = call( + _pinned, c, torch.empty(n, hidden, dtype=torch.bfloat16, device=DEV) + ) + for dst, src in zip(c, fresh[n]): + dst.copy_(src) + captured.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, cold_fresh[n]), ( + f"replay with a cached tactic present != uncached at N={n}" + ) + del graph, captured + assert {str(r[0]) for r in rec} == {"-1"}, sorted({str(r[0]) for r in rec}) + assert sum(r[2] for r in rec) == len(rows) + print( + f" test_pin_ignores_a_cached_tactic OK: with tactics {[chosen[n] for n in rows]} loaded for " + f"N {list(rows)}, {len(rec)} pinned calls ({len(rows)} under capture) took -1 and the uncached bits" + ) + finally: + for key in [k for k in tuner.profiling_cache.cache if k[0] == _OP]: + del tuner.profiling_cache.cache[key] + assert _op_cache_entries() == 0 + + +def _window(args, ref, offset: int, num_local: int): + """The `(args, ref)` pair narrowed to one expert-parallel window.""" + sub = dict(args) + for key in ( + "gemm1_weights", + "gemm1_weights_scale", + "gemm1_bias", + "gemm2_weights", + "gemm2_weights_scale", + "gemm2_bias", + ): + sub[key] = args[key][offset : offset + num_local].contiguous() + sub_ref = dict(ref) + for role in ("up", "gate", "down"): + c, s, b = ref[role] + sub_ref[role] = ( + c[offset : offset + num_local], + s[offset : offset + num_local], + b[offset : offset + num_local], + ) + return sub, sub_ref + + +_DSPARK_LAYER = None + + +def _dspark_layer(): + """The drafter layer of `test_dspark_drafter_surface`, built once for the cells.""" + global _DSPARK_LAYER + if _DSPARK_LAYER is None: + _DSPARK_LAYER = _drafter_layer(seed=300) + return _DSPARK_LAYER + + +def _certify_dspark_cell(cell) -> None: + """One DeepSeek-V4.1 drafter cell, called the way the drafter calls it. + + Pre-routed with bf16 weights in `[0, 1.5)`, the clamp at 10 per expert, no + alpha / beta, no biases. A pinned cell runs inside `fallback_tactic()` and + writes a NaN-filled `output=` buffer, getting back the empty tensor; the + other form returns a fresh tensor. The entry's reference is tied bitwise to + this file's `_ref_moe` before it gates the kernel. + """ + spec = cell.spec + num_experts, hidden, inter, top_k = ( + spec["num_experts"], + spec["hidden"], + spec["intermediate"], + spec["top_k"], + ) + assert (num_experts, hidden, inter, top_k) == (128, 5120, 2304, 3), spec + assert spec["local_expert_offset"] == 0 and spec["local_num_experts"] == num_experts, spec + assert not spec["routed_by_kernel"], spec + args, ref, limit = _dspark_layer() + assert ref["h1_pad"] == hidden and args["intermediate_size"] == inter, "no padding at 5120/2304" + assert torch.equal(limit, torch.full_like(limit, spec["swiglu_limit"])), limit + tokens = spec["tokens"] + gen = torch.Generator(device=DEV).manual_seed(zlib.crc32(cell.why.encode())) + data, sf, xv = _rand_mxfp8(tokens, hidden, hidden, gen) + ids, _ = _routing(tokens, num_experts, top_k, gen) + wts = (torch.rand(tokens, top_k, device=DEV, generator=gen) * 1.5).to(torch.bfloat16) + + # 1. The entry's reference against this file's, on identical inputs. + entry_ref = moe.reference(xv, ids, wts, ref["up"], ref["gate"], ref["down"], clamp_limit=limit) + local_ref = _ref_moe(xv, ids, wts, ref, None, None, limit) + torch.testing.assert_close(entry_ref, local_ref, rtol=0.0, atol=0.0) + + # 2. The entry's reference against the kernel, through the entry's gate. + kw = dict(topk_ids=ids, topk_weights=wts, gemm1_clamp_limit=limit) + if spec["output"]: + out = torch.full((tokens, hidden), float("nan"), dtype=torch.bfloat16, device=DEV) + kw["output"] = out + run = _pinned if spec["pinned"] else _call + y = run(data, sf, args, num_experts, top_k, **kw) + if spec["output"]: + assert y.numel() == 0, tuple(y.shape) + y = out + assert y.shape == (tokens, hidden), y.shape + moe.compare(y, entry_ref.to(torch.bfloat16)) + + +@pytest.mark.parametrize("cell", moe.CELLS, ids=[c.why[:44] for c in moe.CELLS]) +def test_certified_cells(cell): + """Drive every configuration the entry claims, through the entry's own + reference and its own gate. + + The cases above test the op; this one tests the *claim*. A cell list that + nothing executes reads as coverage and is not. + + Each cell cross-checks `moe.reference` against this file's `_ref_moe` + before either is used to gate the kernel. Two references disagreeing is a + different failure from the kernel moving, and collapsing them into one + assert loses which happened. That cross-check earns its keep here: both + references model the FC1 epilogue's architecture-specific block-scale + recipe, and a disagreement there is bit-exact rather than a tolerance. + + `routed_by_kernel` is honoured rather than approximated. The gpt-oss target + hands the call logits and lets it route, so those cells do too; the expected + routing is recomputed here the way Renormalize defines it (top-k, then + softmax over the kept logits), which is what makes the reference comparable + at all. + """ + spec = cell.spec + if spec.get("model") == "deepseek-v4.1": + _certify_dspark_cell(cell) + return + args, ref = _gpt_oss() + off, n = spec["local_expert_offset"], spec["local_num_experts"] + num_experts, hidden, top_k = spec["num_experts"], spec["hidden"], spec["top_k"] + gen = torch.Generator(device=DEV).manual_seed(abs(hash(cell.why)) % 2**31) + + data, sf, xv = _rand_mxfp8(spec["tokens"], hidden, ref["h1_pad"], gen) + alpha, beta, limit = _swiglu_params(n) + sub_args, sub_ref = (args, ref) if n == num_experts else _window(args, ref, off, n) + + if spec["routed_by_kernel"]: + logits = _logits(spec["tokens"], num_experts, gen) + top_v, top_i = torch.topk(logits.float(), top_k, dim=-1) + ids, wts = top_i.to(torch.int32), torch.softmax(top_v, dim=-1) + entry = dict(routing_logits=logits, routing_method_type=1) # Renormalize + else: + ids, wts = _routing(spec["tokens"], num_experts, top_k, gen) + entry = dict(topk_ids=ids, topk_weights=wts) + + # 1. The entry's reference against this file's, on identical inputs. + entry_ref = moe.reference( + xv, + ids, + wts, + sub_ref["up"], + sub_ref["gate"], + sub_ref["down"], + alpha=alpha, + beta=beta, + clamp_limit=limit, + local_expert_offset=off, + local_num_experts=n, + ) + local_ref = _ref_moe(xv, ids, wts, sub_ref, alpha, beta, limit, offset=off, num_local=n) + torch.testing.assert_close(entry_ref, local_ref, rtol=0.0, atol=0.0) + + # 2. The entry's reference against the kernel, through the entry's gate. + y = _call( + data, + sf, + sub_args, + num_experts, + top_k, + local_expert_offset=off, + local_num_experts=n, + gemm1_alpha=alpha, + gemm1_beta=beta, + gemm1_clamp_limit=limit, + **entry, + ) + assert y.shape == (spec["tokens"], hidden), y.shape + moe.compare(y, entry_ref.to(torch.bfloat16)) diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py deleted file mode 100644 index 3bb045cbea66..000000000000 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py +++ /dev/null @@ -1,502 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the noaux_tc_op catalog entry.""" - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.noaux_tc_op import noaux_tc_op - -assert torch.cuda.is_available(), "noaux_tc_op requires a CUDA device" - - -def _kernel_sigmoid(x: torch.Tensor) -> torch.Tensor: - """The sigmoid the kernel actually evaluates: 0.5 * tanh(x/2) + 0.5, fp32. - - Algebraically identical to 1/(1+exp(-x)) but numerically different in the - tails: it saturates to exactly 1.0 at x >= ~17 and to exactly 0.0 at - x <= ~-18.5. Pinned by test_sigmoid_is_the_tanh_form. - """ - return 0.5 * torch.tanh(0.5 * x.float()) + 0.5 - - -def _ref( - router_logits: torch.Tensor, - bias: torch.Tensor, - n_group: int, - topk_group: int, - topk: int, - routed_scaling_factor: float, - sigmoid=_kernel_sigmoid, -) -> tuple[torch.Tensor, torch.Tensor]: - """fp32 reference built from native torch ops only. - - A *stable* descending sort keeps equal selection scores in ascending expert - order, which is the kernel's observed tie-break. - """ - num_experts = router_logits.shape[-1] - scores = sigmoid(router_logits) - choice = scores + bias.float() - if n_group > 1: - grouped = choice.view(-1, n_group, num_experts // n_group) - group_score = torch.topk(grouped, k=2, dim=-1).values.sum(-1) - keep = torch.sort(group_score, dim=-1, descending=True, stable=True).indices[:, :topk_group] - mask = torch.zeros_like(group_score).scatter_(-1, keep, 1.0) - mask = mask.unsqueeze(-1).expand_as(grouped).reshape(choice.shape) - choice = torch.where(mask.bool(), choice, torch.tensor(float("-inf"), device=choice.device)) - ids = torch.sort(choice, dim=-1, descending=True, stable=True).indices[:, :topk] - weights = torch.gather(scores, 1, ids) - # The scores are summed in fp32, but the division and the scaling are - # evaluated in fp64 (pinned by test_routed_scaling_factor: an all-fp32 - # normalization disagrees on a third of the elements at - # routed_scaling_factor = 2.5, an fp64 one is bit-exact). - total = weights.sum(-1, keepdim=True).double() - weights = weights.double() / (total + 1e-20) * routed_scaling_factor - return weights.to(router_logits.dtype), ids.to(torch.int32) - - -def _check( - router_logits: torch.Tensor, - bias: torch.Tensor, - n_group: int, - topk_group: int, - topk: int, - routed_scaling_factor: float, -) -> tuple[torch.Tensor, torch.Tensor]: - weights, ids = noaux_tc_op( - router_logits, bias, n_group, topk_group, topk, routed_scaling_factor - ) - ref_weights, ref_ids = _ref( - router_logits, bias, n_group, topk_group, topk, routed_scaling_factor - ) - num_tokens = router_logits.shape[0] - assert weights.shape == (num_tokens, topk), weights.shape - assert ids.shape == (num_tokens, topk), ids.shape - assert weights.dtype == router_logits.dtype, weights.dtype - assert ids.dtype == torch.int32, ids.dtype - assert weights.is_contiguous() and ids.is_contiguous() - assert weights.device == router_logits.device - assert ids.device == router_logits.device - torch.testing.assert_close(ids, ref_ids) - # Dual gate. The fp32 reference is not an approximation of this kernel, it - # is the kernel — the only observed disagreement is the final rounding of - # the renormalized weight, which can step by one last-bit unit because the - # kernel adds up the topk scores in a different order than torch.sum - # (measured: 1 element in 49152 at 8192x72 bf16). So the tolerance is - # exactly one ulp (rtol = the dtype's eps is the largest relative size one - # ulp can have; atol = 0), and on top of that almost every element must be - # bit-identical. - torch.testing.assert_close( - weights, - ref_weights, - rtol=torch.finfo(router_logits.dtype).eps, - atol=0.0, - ) - mismatched = int((weights != ref_weights).sum()) - budget = 1 + weights.numel() // 10000 - assert mismatched <= budget, ( - f"{mismatched} of {weights.numel()} weights are not bit-exact (budget " - f"{budget}): dtype={router_logits.dtype} shape={tuple(router_logits.shape)} " - f"n_group={n_group} topk={topk}" - ) - return weights, ids - - -def test_deepseek_v3_lite_config() -> None: - # 72 routed experts, top-6, n_group=1, topk_group=1, scaling 2.0, bf16 - # router logits and a bf16 [72] correction bias. Decode-like through - # prefill-like token counts. - torch.manual_seed(0) - bias = torch.randn(72, dtype=torch.bfloat16, device="cuda") * 0.1 - for num_tokens in [1, 2, 7, 64, 512, 2048, 4096, 8192]: - logits = torch.randn(num_tokens, 72, dtype=torch.bfloat16, device="cuda") - weights, _ = _check(logits, bias, 1, 1, 6, 2.0) - # renormalized then scaled: every row sums to routed_scaling_factor - torch.testing.assert_close( - weights.float().sum(-1), - torch.full((num_tokens,), 2.0, device="cuda"), - rtol=8e-3, # 6 bf16 addends, each rounded to 8 mantissa bits - atol=0.0, - ) - - -def test_matches_hf_deepseek_sigmoid_reference() -> None: - # The HF DeepSeek-V3 gate written with the textbook sigmoid. Over the - # logit range a real router produces (|logit| <= 8) the kernel's tanh-form - # sigmoid and torch.sigmoid agree to well inside dtype tolerance. - torch.manual_seed(1) - for dtype in [torch.bfloat16, torch.float32]: - for num_tokens in [1, 256, 2048]: - logits = (torch.randn(num_tokens, 72, device="cuda") * 2.5).to(dtype) - bias = (torch.randn(72, device="cuda") * 0.1).to(dtype) - weights, ids = noaux_tc_op(logits, bias, 1, 1, 6, 2.0) - ref_weights, ref_ids = _ref( - logits, bias, 1, 1, 6, 2.0, sigmoid=lambda x: torch.sigmoid(x.float()) - ) - torch.testing.assert_close(ids, ref_ids) - torch.testing.assert_close(weights, ref_weights) - - -def test_selection_uses_bias_but_weights_do_not() -> None: - # The trap this entry exists to pin: the bias enters selection only. - # A bias large enough to reorder the experts must move the ids, and the - # returned weights must come from the *unbiased* sigmoid. - torch.manual_seed(2) - logits = torch.randn(512, 72, dtype=torch.float32, device="cuda") - bias = torch.randn(72, dtype=torch.float32, device="cuda") - weights, ids = noaux_tc_op(logits, bias, 1, 1, 6, 2.0) - - unbiased_ids = torch.sort( - _kernel_sigmoid(logits), dim=-1, descending=True, stable=True - ).indices[:, :6] - assert not torch.equal(ids, unbiased_ids.to(torch.int32)), ( - "the bias did not affect selection — test input too weak" - ) - - scores = _kernel_sigmoid(logits) - gathered = torch.gather(scores, 1, ids.long()) - expect = gathered / (gathered.sum(-1, keepdim=True) + 1e-20) * 2.0 - torch.testing.assert_close(weights, expect) - - # negative control: gathering the *biased* score instead is a different, - # plausible-looking answer that the kernel does not produce - biased = torch.gather(scores + bias, 1, ids.long()) - wrong = biased / (biased.sum(-1, keepdim=True) + 1e-20) * 2.0 - assert (weights - wrong).abs().max().item() > 1e-2, ( - "biased-weight reference is indistinguishable — test input too weak" - ) - - -def test_slot_order_is_descending_by_biased_score() -> None: - torch.manual_seed(3) - logits = torch.randn(256, 72, dtype=torch.bfloat16, device="cuda") - bias = torch.randn(72, dtype=torch.bfloat16, device="cuda") - _, ids = noaux_tc_op(logits, bias, 1, 1, 6, 2.0) - choice = _kernel_sigmoid(logits) + bias.float() - selected = torch.gather(choice, 1, ids.long()) - assert bool((selected[:, :-1] >= selected[:, 1:]).all()), ( - "slots are not ordered by decreasing selection score" - ) - # the same row is *not* sorted by the unbiased score, i.e. the ordering - # key really is the biased one - unbiased = torch.gather(_kernel_sigmoid(logits), 1, ids.long()) - assert not bool((unbiased[:, :-1] >= unbiased[:, 1:]).all()) - - -def test_ties_break_toward_lower_expert_index() -> None: - # Opposite to torch.topk, which emits the larger index first among equals. - logits = torch.zeros(4, 16, dtype=torch.bfloat16, device="cuda") - bias = torch.zeros(16, dtype=torch.bfloat16, device="cuda") - logits[0, :] = 1.0 - logits[1, 3] = 5.0 - logits[1, 7] = 5.0 - logits[2, ::2] = 2.0 - logits[3, 0] = 1.0 - logits[3, 15] = 1.0 - _, ids = noaux_tc_op(logits, bias, 1, 1, 4, 2.0) - expected = torch.tensor( - [[0, 1, 2, 3], [3, 7, 0, 1], [0, 2, 4, 6], [0, 15, 1, 2]], - dtype=torch.int32, - device="cuda", - ) - torch.testing.assert_close(ids, expected) - _check(logits, bias, 1, 1, 4, 2.0) - - # tie-saturated inputs: ties from the logits, and ties created by a coarse - # bias on top of continuous logits - torch.manual_seed(4) - for num_tokens, num_experts, topk in [(64, 128, 8), (512, 72, 6), (256, 16, 4)]: - coarse = torch.randint(0, 3, (num_tokens, num_experts), device="cuda").to(torch.bfloat16) - zero = torch.zeros(num_experts, dtype=torch.bfloat16, device="cuda") - _check(coarse, zero, 1, 1, topk, 2.0) - smooth = torch.randn(num_tokens, num_experts, dtype=torch.bfloat16, device="cuda") - coarse_bias = torch.randint(0, 2, (num_experts,), device="cuda").to(torch.bfloat16) - _check(smooth, coarse_bias, 1, 1, topk, 2.0) - - -def test_dtype_matrix() -> None: - # weights come back in the *router_logits* dtype; the bias dtype is - # independent, except that fp16 logits reject a bf16 bias. - torch.manual_seed(5) - base = torch.randn(128, 72, device="cuda") - bias_base = torch.randn(72, device="cuda") * 0.1 - dtypes = [torch.bfloat16, torch.float16, torch.float32] - for logits_dtype in dtypes: - for bias_dtype in dtypes: - if logits_dtype is torch.float16 and bias_dtype is torch.bfloat16: - continue # rejected; covered by test_op_rejects_unsupported_domains - logits = base.to(logits_dtype) - bias = bias_base.to(bias_dtype) - weights, _ = _check(logits, bias, 1, 1, 6, 2.0) - assert weights.dtype == logits_dtype, ( - logits_dtype, - bias_dtype, - weights.dtype, - ) - - -def test_routed_scaling_factor() -> None: - torch.manual_seed(6) - logits = torch.randn(256, 72, dtype=torch.float32, device="cuda") - bias = torch.randn(72, dtype=torch.float32, device="cuda") * 0.1 - for scaling in [1.0, 2.0, 2.5, 3.0, 0.5, 0.0, -1.0, 1000.0]: - weights, _ = _check(logits, bias, 1, 1, 6, scaling) - torch.testing.assert_close(weights.sum(-1), torch.full((256,), scaling, device="cuda")) - # negative control for the fp64 normalization: an all-fp32 reference - # is a different answer for scaling factors that are not a power of two - scores = _kernel_sigmoid(logits) - _, ids = noaux_tc_op(logits, bias, 1, 1, 6, scaling) - gathered = torch.gather(scores, 1, ids.long()) - fp32_form = gathered / (gathered.sum(-1, keepdim=True) + 1e-20) * scaling - differs = int((weights != fp32_form).sum()) - if scaling in (2.5, 3.0, 1000.0): - assert differs > 0.1 * weights.numel(), ( - f"fp32 and fp64 normalization are indistinguishable at {scaling}" - ) - - -def test_sigmoid_is_the_tanh_form() -> None: - # Pins the kernel's sigmoid out to |logits| ~ 40, where the tanh form has - # long saturated. _check's reference is the tanh form; the exp form is - # ruled out separately in test_saturation_floor_and_ceiling. - torch.manual_seed(7) - for dtype in [torch.bfloat16, torch.float16, torch.float32]: - for scale in [1.0, 4.0, 10.0, 20.0, 40.0]: - logits = (torch.randn(256, 72, device="cuda") * scale).to(dtype) - bias = (torch.randn(72, device="cuda") * 0.1).to(dtype) - weights, _ = _check(logits, bias, 1, 1, 6, 2.0) - assert bool(torch.isfinite(weights.float()).all()) - - -def test_saturation_floor_and_ceiling() -> None: - # The tail behaviour that follows from the tanh form. - bias = torch.zeros(16, dtype=torch.float32, device="cuda") - high = torch.full((2, 16), 20.0, dtype=torch.float32, device="cuda") - weights, _ = noaux_tc_op(high, bias, 1, 1, 4, 2.0) - # every selected score saturates to exactly 1.0 -> equal weights - assert torch.equal(weights, torch.full_like(weights, 0.5)), weights - - low = torch.full((2, 16), -20.0, dtype=torch.float32, device="cuda") - weights, _ = noaux_tc_op(low, bias, 1, 1, 4, 2.0) - # every selected score saturates to exactly 0.0; the guard term in the - # denominator turns 0/0 into exactly 0 instead of NaN - assert torch.equal(weights, torch.zeros_like(weights)), weights - - # Negative control that separates the two sigmoid forms: with every - # selected logit inside the tanh form's lossy band, a 1/(1+exp(-x)) - # reference lands three orders of magnitude outside the one-ulp gate - # _check applies (fp32 eps * 0.33 ~ 4e-8). - torch.manual_seed(8) - band = (torch.rand(128, 72, device="cuda") * 12.0 - 20.0).to(torch.float32) - zero = torch.zeros(72, dtype=torch.float32, device="cuda") - weights, _ = _check(band, zero, 1, 1, 6, 2.0) - exp_weights, _ = _ref(band, zero, 1, 1, 6, 2.0, sigmoid=lambda x: torch.sigmoid(x.float())) - assert (weights - exp_weights).abs().max().item() > 1e-5, ( - "exp-form and tanh-form references are indistinguishable here" - ) - - -def test_grouped_routing() -> None: - # n_group > 1: group score = sum of the two largest biased scores in the - # group; the topk_group best groups survive, the rest are masked out. - torch.manual_seed(9) - for num_experts, n_group, topk_group, topk in [ - (256, 8, 4, 8), # DeepSeek-V3 full - (128, 4, 2, 6), - (72, 8, 2, 6), - (72, 4, 2, 6), - (64, 8, 4, 8), - ]: - for num_tokens in [1, 64, 1024]: - logits = torch.randn(num_tokens, num_experts, dtype=torch.bfloat16, device="cuda") - bias = torch.randn(num_experts, dtype=torch.bfloat16, device="cuda") * 0.1 - _check(logits, bias, n_group, topk_group, topk, 2.0) - # tie-saturated grouping - coarse = torch.randint(0, 3, (64, num_experts), device="cuda").to(torch.bfloat16) - zero = torch.zeros(num_experts, dtype=torch.bfloat16, device="cuda") - _check(coarse, zero, n_group, topk_group, topk, 2.0) - - # keeping every group is the same as not grouping at all - logits = torch.randn(64, 256, dtype=torch.bfloat16, device="cuda") - bias = torch.randn(256, dtype=torch.bfloat16, device="cuda") * 0.1 - grouped = noaux_tc_op(logits, bias, 8, 8, 8, 2.0) - flat = noaux_tc_op(logits, bias, 1, 1, 8, 2.0) - assert torch.equal(grouped[0], flat[0]) and torch.equal(grouped[1], flat[1]) - - -def test_topk_group_ignored_when_n_group_is_one() -> None: - torch.manual_seed(10) - logits = torch.randn(64, 72, dtype=torch.bfloat16, device="cuda") - bias = torch.randn(72, dtype=torch.bfloat16, device="cuda") * 0.1 - reference = noaux_tc_op(logits, bias, 1, 1, 6, 2.0) - for topk_group in [0, 1, 2, 5, 32, 100]: - weights, ids = noaux_tc_op(logits, bias, 1, topk_group, 6, 2.0) - assert torch.equal(weights, reference[0]) and torch.equal(ids, reference[1]), ( - f"topk_group={topk_group} changed the result at n_group=1" - ) - - -def test_num_experts_and_topk_sweep() -> None: - torch.manual_seed(11) - for num_experts in [1, 2, 7, 32, 64, 72, 100, 128, 256, 257, 512, 1024]: - logits = torch.randn(8, num_experts, dtype=torch.bfloat16, device="cuda") - bias = torch.randn(num_experts, dtype=torch.bfloat16, device="cuda") - _check(logits, bias, 1, 1, min(6, num_experts), 2.0) - logits = torch.randn(64, 64, dtype=torch.bfloat16, device="cuda") - bias = torch.randn(64, dtype=torch.bfloat16, device="cuda") - for topk in [1, 2, 3, 6, 8, 16, 31, 32]: - _check(logits, bias, 1, 1, topk, 2.0) - # topk == num_experts returns the full ranking - small = torch.randn(16, 8, dtype=torch.bfloat16, device="cuda") - small_bias = torch.randn(8, dtype=torch.bfloat16, device="cuda") - _check(small, small_bias, 1, 1, 8, 2.0) - - -def test_degenerate_shapes() -> None: - logits = torch.randn(8, 72, dtype=torch.bfloat16, device="cuda") - bias = torch.randn(72, dtype=torch.bfloat16, device="cuda") - weights, ids = noaux_tc_op(logits, bias, 1, 1, 0, 2.0) - assert weights.shape == (8, 0) and ids.shape == (8, 0) - empty = torch.randn(0, 72, dtype=torch.bfloat16, device="cuda") - weights, ids = noaux_tc_op(empty, bias, 1, 1, 6, 2.0) - assert weights.shape == (0, 6) and ids.shape == (0, 6) - assert weights.dtype == torch.bfloat16 and ids.dtype == torch.int32 - - # Empty inputs must leave subsequent CUDA work usable. - assert torch.ones(1, device="cuda").item() == 1 - - -def test_input_not_mutated_and_deterministic() -> None: - torch.manual_seed(12) - logits = torch.randn(4096, 72, dtype=torch.bfloat16, device="cuda") - bias = torch.randn(72, dtype=torch.bfloat16, device="cuda") - logits_snapshot, bias_snapshot = logits.clone(), bias.clone() - weights_a, ids_a = noaux_tc_op(logits, bias, 1, 1, 6, 2.0) - weights_b, ids_b = noaux_tc_op(logits, bias, 1, 1, 6, 2.0) - assert torch.equal(logits, logits_snapshot), "router_logits was mutated" - assert torch.equal(bias, bias_snapshot), "bias was mutated" - assert torch.equal(weights_a, weights_b) and torch.equal(ids_a, ids_b), ( - "two identical calls disagreed" - ) - assert weights_a.data_ptr() != weights_b.data_ptr() - assert weights_a.data_ptr() != logits.data_ptr() - - -def test_op_rejects_unsupported_domains() -> None: - logits = torch.randn(8, 72, dtype=torch.bfloat16, device="cuda") - bias = torch.randn(72, dtype=torch.bfloat16, device="cuda") - big = torch.randn(8, 256, dtype=torch.bfloat16, device="cuda") - big_bias = torch.randn(256, dtype=torch.bfloat16, device="cuda") - cases = [ - ("1D router_logits", lambda: noaux_tc_op(logits[0], bias, 1, 1, 6, 2.0)), - ( - "3D router_logits", - lambda: noaux_tc_op(logits.unsqueeze(0), bias, 1, 1, 6, 2.0), - ), - ( - "fp64 router_logits", - lambda: noaux_tc_op(logits.double(), bias, 1, 1, 6, 2.0), - ), - ("int32 router_logits", lambda: noaux_tc_op(logits.int(), bias, 1, 1, 6, 2.0)), - ("fp64 bias", lambda: noaux_tc_op(logits, bias.double(), 1, 1, 6, 2.0)), - ("int32 bias", lambda: noaux_tc_op(logits, bias.int(), 1, 1, 6, 2.0)), - ( - "fp16 logits + bf16 bias", - lambda: noaux_tc_op(logits.half(), bias, 1, 1, 6, 2.0), - ), - ("2D bias", lambda: noaux_tc_op(logits, bias.unsqueeze(0), 1, 1, 6, 2.0)), - ("short bias", lambda: noaux_tc_op(logits, bias[:71], 1, 1, 6, 2.0)), - ("cpu tensors", lambda: noaux_tc_op(logits.cpu(), bias.cpu(), 1, 1, 6, 2.0)), - ("bias on cpu", lambda: noaux_tc_op(logits, bias.cpu(), 1, 1, 6, 2.0)), - ("negative topk", lambda: noaux_tc_op(logits, bias, 1, 1, -1, 2.0)), - ("topk 33", lambda: noaux_tc_op(big, big_bias, 1, 1, 33, 2.0)), - ( - "num_experts 1025", - lambda: noaux_tc_op( - torch.randn(8, 1025, dtype=torch.bfloat16, device="cuda"), - torch.randn(1025, dtype=torch.bfloat16, device="cuda"), - 1, - 1, - 6, - 2.0, - ), - ), - ( - "n_group does not divide num_experts", - lambda: noaux_tc_op(logits, bias, 5, 1, 6, 2.0), - ), - ( - "n_group 33", - lambda: noaux_tc_op( - torch.randn(8, 132, dtype=torch.bfloat16, device="cuda"), - torch.randn(132, dtype=torch.bfloat16, device="cuda"), - 33, - 1, - 6, - 2.0, - ), - ), - ( - "grouped topk_group > n_group", - lambda: noaux_tc_op(big, big_bias, 8, 9, 8, 2.0), - ), - ("grouped topk > 8", lambda: noaux_tc_op(big, big_bias, 8, 4, 9, 2.0)), - ( - "grouped experts_per_group > 32", - lambda: noaux_tc_op(big, big_bias, 4, 2, 8, 2.0), - ), - ( - "grouped num_experts > 256", - lambda: noaux_tc_op( - torch.randn(8, 512, dtype=torch.bfloat16, device="cuda"), - torch.randn(512, dtype=torch.bfloat16, device="cuda"), - 8, - 4, - 8, - 2.0, - ), - ), - ] - for tag, call in cases: - try: - call() - except (RuntimeError, ValueError, NotImplementedError): - continue - raise AssertionError(f"op accepted an unsupported domain: {tag}") - - -def test_wrapper_rejects_silently_wrong_inputs() -> None: - torch.manual_seed(13) - bias = torch.randn(72, dtype=torch.bfloat16, device="cuda") - wide = torch.randn(8, 144, dtype=torch.bfloat16, device="cuda") - logits = torch.randn(8, 72, dtype=torch.bfloat16, device="cuda") - wide_bias = torch.randn(144, dtype=torch.bfloat16, device="cuda") - - # A column slice is not routed on its own elements: the kernel reads the - # first 8*72 elements of the underlying storage as a dense [8, 72] buffer. - sliced = wide[:, :72] - assert not sliced.is_contiguous() - reinterpreted = wide.reshape(-1)[: 8 * 72].reshape(8, 72).contiguous() - raw = torch.ops.trtllm.noaux_tc_op(sliced, bias, 1, 1, 6, 2.0) - dense = torch.ops.trtllm.noaux_tc_op(reinterpreted, bias, 1, 1, 6, 2.0) - correct = torch.ops.trtllm.noaux_tc_op(sliced.contiguous(), bias, 1, 1, 6, 2.0) - assert torch.equal(raw[1], dense[1]), "strided read is not a dense reinterpretation" - assert not torch.equal(raw[1], correct[1]), "strided view happened to be correct" - - bad_inputs = [ - ("strided router_logits", lambda: noaux_tc_op(sliced, bias, 1, 1, 6, 2.0)), - ( - "transposed router_logits", - lambda: noaux_tc_op(wide.t()[:72, :], bias, 1, 1, 6, 2.0), - ), - ("strided bias", lambda: noaux_tc_op(logits, wide_bias[::2], 1, 1, 6, 2.0)), - ("topk > num_experts", lambda: noaux_tc_op(logits, bias, 1, 1, 73, 2.0)), - ] - for tag, call in bad_inputs: - try: - call() - except AssertionError: - continue - raise AssertionError(f"wrapper accepted a silently-wrong input: {tag}") - - # the same data made contiguous routes correctly - _check(sliced.contiguous(), bias, 1, 1, 6, 2.0) diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_fused_add_rmsnorm.py b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_fused_add_rmsnorm.py index 4b5858a26502..e1de4353c7a0 100644 --- a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_fused_add_rmsnorm.py +++ b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_fused_add_rmsnorm.py @@ -1,69 +1,96 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the flashinfer_fused_add_rmsnorm catalog entry.""" +"""GPU test for the flashinfer_fused_add_rmsnorm catalog entry. +The entry owns what is certified: `CELLS` is the list, `reference` and +`compare` are the gate, `is_valid` is the guard. This file owns only what a +cell cannot carry -- turning a spec into real tensors -- plus the inputs the +guard exists to refuse, which are claims about inputs no cell describes. +""" + +import pytest import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( - flashinfer_fused_add_rmsnorm, + flashinfer_fused_add_rmsnorm as op, ) +__extra_import_path__ = [".."] +from _validating import validating # noqa: E402 — needs the path declared above + assert torch.cuda.is_available(), "flashinfer_fused_add_rmsnorm requires a CUDA device" -def _ref( - x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float -) -> tuple[torch.Tensor, torch.Tensor]: - """fp32-accumulated reference for the fused op. +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + shape, dtype = spec["shape"], spec["dtype"] + x = torch.randn(shape, generator=g, device="cuda").to(dtype) + residual = torch.randn(shape, generator=g, device="cuda").to(dtype) + weight = torch.randn(shape[-1], generator=g, device="cuda").to(dtype) + return x, residual, weight, spec["eps"] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:40] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Every certified configuration, kernel against reference. + + Driven inside `validating`, so each cell also proves the guard admits the + inputs a shipped target actually passes. + """ + x, residual, weight, eps = _build(cell.spec, seed=abs(hash(cell.why)) % 2**31) + ref_x, ref_residual = op.reference(x, residual, weight, eps) + with validating(op): + out_x, out_residual = op(x, residual, weight, eps) + assert out_x is x and out_residual is residual, "the op writes in place" + op.compare(residual, ref_residual) + op.compare(x, ref_x) + + +def _one(**over): + g = torch.Generator(device="cuda").manual_seed(7) + kw = dict(rows=4, hidden=2880, dtype=torch.bfloat16) + kw.update(over) + shape = (kw["rows"], kw["hidden"]) + return ( + torch.randn(shape, generator=g, device="cuda").to(kw["dtype"]), + torch.randn(shape, generator=g, device="cuda").to(kw["dtype"]), + torch.randn(shape[-1], generator=g, device="cuda").to(kw["dtype"]), + ) + + +def test_guard_refuses_fp8() -> None: + """Accepted, and silently plausible: the case a caller cannot detect.""" + x, residual, weight = _one(dtype=torch.float8_e4m3fn) + with pytest.raises(AssertionError, match="dequantize"): + with validating(op): + op(x, residual, weight, 1e-6) + + +def test_guard_refuses_a_strided_normalized_dim() -> None: + """`stride(-1) != 1` reads the wrong elements rather than failing.""" + buf = torch.randn(8, 2880, 2, dtype=torch.bfloat16, device="cuda") + x, residual = buf[..., 0], buf[..., 1] + assert x.stride(-1) != 1 + weight = torch.randn(2880, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(x, residual, weight, 1e-6) + + +def test_guard_refuses_aliased_buffers() -> None: + """The op writes both tensors; one buffer passed twice yields neither.""" + x, _, weight = _one() + with pytest.raises(AssertionError, match="distinct buffers"): + with validating(op): + op(x, x, weight, 1e-6) + + +def test_the_guard_is_not_on_the_hot_path() -> None: + """Outside `validating`, the wrapper goes straight to the kernel. - h = fp32(x) + fp32(residual); the norm reads the fp32 h, not the - rounded residual output, matching the kernel. + Stated as a test because it is the property that lets the guard be as + strict as it likes: a served engine never runs it. """ - h = x.float() + residual.float() - normed = h * torch.rsqrt(h.pow(2).mean(dim=-1, keepdim=True) + eps) - x_out = (normed * weight.float()).to(x.dtype) - residual_out = h.to(residual.dtype) - return x_out, residual_out - - -def _check(x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float) -> None: - ref_x, ref_res = _ref(x, residual, weight, eps) - flashinfer_fused_add_rmsnorm(x, residual, weight, eps) - torch.testing.assert_close(residual, ref_res) - torch.testing.assert_close(x, ref_x) - - -def test_bf16_2d() -> None: - torch.manual_seed(0) - # decode-like (few tokens) and prefill-like (many tokens) shapes - for num_tokens, hidden in [(1, 4096), (4, 5120), (2048, 4096)]: - x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device="cuda") - r = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device="cuda") - w = torch.randn(hidden, dtype=torch.bfloat16, device="cuda") - _check(x, r, w, 1e-6) - - -def test_bf16_unaligned_hidden() -> None: - # hidden sizes not divisible by the 128-bit vector width - torch.manual_seed(1) - for hidden in [111, 1152]: - x = torch.randn(16, hidden, dtype=torch.bfloat16, device="cuda") - r = torch.randn(16, hidden, dtype=torch.bfloat16, device="cuda") - w = torch.randn(hidden, dtype=torch.bfloat16, device="cuda") - _check(x, r, w, 1e-6) - - -def test_bf16_strided_rows() -> None: - # last-dim contiguous slices of wider buffers (row stride != hidden); - # row stride 8192 is divisible by the kernel vector size - torch.manual_seed(2) - x_buf = torch.randn(8, 8192, dtype=torch.bfloat16, device="cuda") - r_buf = torch.randn(8, 8192, dtype=torch.bfloat16, device="cuda") - x, r = x_buf[:, :4096], r_buf[:, :4096] - assert not x.is_contiguous() and x.stride(-1) == 1 - w = torch.randn(4096, dtype=torch.bfloat16, device="cuda") - # snapshot the untouched right halves before the in-place call - right_before = torch.cat([x_buf[:, 4096:], r_buf[:, 4096:]]).clone() - _check(x, r, w, 1e-6) - # mutation must land in the parent buffers' left halves only - assert torch.equal(torch.cat([x_buf[:, 4096:], r_buf[:, 4096:]]), right_before) + x, residual, weight = _one(dtype=torch.float8_e4m3fn) + # fp8 is what the guard refuses; unguarded, the kernel answers anyway. + op(x, residual, weight, 1e-6) diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py index 8492fdc49eeb..9d49f7821592 100644 --- a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py +++ b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py @@ -1,60 +1,979 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the flashinfer_rmsnorm catalog entry.""" +"""GPU test for the flashinfer_rmsnorm catalog entry. +The entry owns what is certified: `CELLS` is the list, `reference` and +`compare` are the gate, `is_valid` is the guard. This file owns only what a +cell cannot carry -- turning a spec into real tensors -- plus the inputs the +guard exists to refuse. +""" + +import pytest import torch +from flashinfer.norm import rmsnorm as _flashinfer_norm_rmsnorm +from flashinfer.norm.kernels import rmsnorm as _cute_rmsnorm_kernels +from torch.profiler import ProfilerActivity, profile from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( - flashinfer_rmsnorm, + flashinfer_rmsnorm as op, ) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +__extra_import_path__ = [".."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "flashinfer_rmsnorm requires a CUDA device" +#: The DeepSeek-V4.1 vision tower's residual stream carries massive activations +#: in these channels (|x| ~ 2000). +_VIT_MASSIVE = (53, 147, 308) + + +def _v41_row_scales(rows: int) -> torch.Tensor: + """Row i's magnitude class, i % 4: unit; 1e-10, where mean(x^2) is about eps 1e-20 and eps + decides the result; 1e-11, where mean(x^2) is far below eps; all-zero.""" + classes = torch.tensor([1.0, 1e-10, 1e-11, 0.0], device="cuda") + return classes[torch.arange(rows, device="cuda") % 4] + + +def _vit_rows(x: torch.Tensor, g: torch.Generator) -> torch.Tensor: + """The vision tower's row classes, row i taking i % 4, in place on fp32 `x`: unit; unit with + the massive channels at +-2000; 1e-3 scale, where mean(x^2) is about eps 1e-6; all-zero.""" + cls = torch.arange(x.shape[0], device="cuda") % 4 + x[cls == 2] *= 1e-3 + x[cls == 3] = 0.0 + sign = torch.where( + torch.rand(x.shape[0], len(_VIT_MASSIVE), device="cuda", generator=g) < 0.5, -1.0, 1.0 + ) + massive = x[:, list(_VIT_MASSIVE)] + x[:, list(_VIT_MASSIVE)] = torch.where((cls == 1)[:, None], 2000.0 * sign, massive) + return x + -def _ref_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: - """fp32-accumulated reference: x / sqrt(mean(x^2, -1) + eps) * weight.""" - xf = x.float() - normed = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) - return (normed * weight.float()).to(x.dtype) +def _build(spec, seed): + """Tensors for one cell; the optional keys default to a contiguous randn input. + `row_stride` / `col_offset`: `x` is the column view `parent[:, col_offset : col_offset + width]` + of a `[rows, row_stride]` parent, as the DeepSeek-V4.1 target slices its qkv and ring buffers. + `data`: the rows' magnitude pattern -- "eps_classes" (`_v41_row_scales`) or "vit" (`_vit_rows`, + with the weight drawn near 1 as the vision tower's is). + """ + g = torch.Generator(device="cuda").manual_seed(seed) + shape, dtype = spec["shape"], spec["dtype"] + stride, offset, data = spec.get("row_stride"), spec.get("col_offset", 0), spec.get("data") + x = torch.randn(shape if stride is None else (shape[0], stride), generator=g, device="cuda") + if data == "eps_classes": + x = x * _v41_row_scales(shape[0])[:, None] + elif data == "vit": + x = _vit_rows(x, g) + x = x.to(dtype) + if stride is not None: + x = x[:, offset : offset + shape[-1]] + weight = torch.randn(shape[-1], generator=g, device="cuda") + if data == "vit": + weight = 1.0 + 0.5 * weight + return x, weight.to(dtype), spec["eps"] -def _check(x: torch.Tensor, weight: torch.Tensor, eps: float) -> None: - out = flashinfer_rmsnorm(x, weight, eps) - ref = _ref_rmsnorm(x, weight, eps) + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + x, weight, eps = _build(cell.spec, seed=abs(hash(cell.why)) % 2**31) + out = op(x, weight, eps) assert out.shape == x.shape and out.dtype == x.dtype - torch.testing.assert_close(out, ref) - - -def test_bf16_2d() -> None: - torch.manual_seed(0) - # decode-like (few tokens) and prefill-like (many tokens) shapes - for num_tokens, hidden in [(1, 4096), (4, 5120), (2048, 4096)]: - x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device="cuda") - w = torch.randn(hidden, dtype=torch.bfloat16, device="cuda") - _check(x, w, 1e-6) - - -def test_bf16_3d_qk_norm_shape() -> None: - torch.manual_seed(1) - x = torch.randn(16, 32, 128, dtype=torch.bfloat16, device="cuda") - w = torch.randn(128, dtype=torch.bfloat16, device="cuda") - _check(x, w, 1e-5) - - -def test_bf16_unaligned_hidden() -> None: - # hidden sizes not divisible by the 128-bit vector width - torch.manual_seed(2) - for hidden in [111, 1152]: - x = torch.randn(16, hidden, dtype=torch.bfloat16, device="cuda") - w = torch.randn(hidden, dtype=torch.bfloat16, device="cuda") - _check(x, w, 1e-6) - - -def test_bf16_strided_rows() -> None: - # last-dim contiguous slice of a wider buffer (row stride != hidden) - torch.manual_seed(3) - buf = torch.randn(8, 8192, dtype=torch.bfloat16, device="cuda") - x = buf[:, :4096] - assert not x.is_contiguous() and x.stride(-1) == 1 - w = torch.randn(4096, dtype=torch.bfloat16, device="cuda") - _check(x, w, 1e-6) + assert out.data_ptr() != x.data_ptr(), "the op must return a new tensor" + if cell.spec.get("data") is not None: + # The magnitude-class cells: zero rows come back exactly 0, and nothing overflows. + assert torch.isfinite(out).all() and (out[3::4] == 0).all() + op.compare(out, op.reference(x, weight, eps)) + + +def test_guard_refuses_fp8_before_the_kernel_answers() -> None: + """The hazard `is_valid` exists for: accepted, and silently plausible. + + Driving the op directly shows why the guard is not cosmetic -- it returns a + tensor rather than raising, so a caller who forgot to dequantize gets a + number back. + """ + x = torch.randn(4, 2880, device="cuda").to(torch.float8_e4m3fn) + weight = torch.randn(2880, device="cuda").to(torch.float8_e4m3fn) + unguarded = torch.ops.trtllm.flashinfer_rmsnorm(x, weight, 1e-6) + assert unguarded.dtype is torch.float8_e4m3fn, "the hazard this guard covers is gone" + + with pytest.raises(AssertionError, match="dequantize"): + with validating(op): + op(x, weight, 1e-6) + + +def test_guard_refuses_a_strided_normalized_dim() -> None: + """`x.stride(-1) != 1` reads the wrong elements rather than failing.""" + buf = torch.randn(8, 2880, 2, dtype=torch.bfloat16, device="cuda") + x = buf[..., 0] + assert x.stride(-1) != 1 + weight = torch.randn(2880, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(x, weight, 1e-6) + + +# --------------------------------------------------------------------------- +# DeepSeek-V4.1. The core target's call sites, all bf16 at eps = 1e-20 (the +# checkpoint's rms_norm_eps): +# width 5120: attn_norm / ffn_norm / final norm on contiguous [T, 5120] +# width 1280: q_norm on the column view qkv[:, :1280] of a [T, 1792] buffer +# width 512: kv_norm on the column view qkv[:, 1280:] (row stride 1792, +# element offset 1280), comp_norm on contiguous [n, 512] +# width 128: index-K norm on contiguous [n, 128] +# and the DSpark drafter's ring write: kv3 = wkv3(main_x) is a [M, 1536] +# buffer and each of the three stages normalizes its own width-512 column view +# of it, split(kv3, [512] * 3, -1): row stride 1536, element offsets 0, 512 +# and 1024. Rows: the engine's per-rank token counts, 1..16384 (the largest +# per-rank max_num_tokens an envelope engine runs). The cells hold the two +# ends; the row sweep below covers every count between, and the capture tests +# the engines' whole capture set. +_V41_EPS = 1e-20 +_V41_ROWS = (1, 2, 3, 7, 24, 32, 323, 1005, 4617, 8192, 16384) +_V41_GRAPH_ROWS = (1, 2, 4, 8, 16, 24, 32) +_RING_PARENT = 1536 +_RING_OFFSET = {"ring0": 0, "ring1": 512, "ring2": 1024} + + +def _v41_parent(rows: int, layout: str) -> torch.Tensor: + """The buffer the target's input is a view of. + + [rows, 1792] for the two column views of the qkv projection, [rows, 1536] + for the ring views, [rows, width] otherwise. Rows cycle through the four + magnitude classes of `_v41_row_scales`. + """ + if layout in ("q", "kv"): + width = 1792 + elif layout in _RING_OFFSET: + width = _RING_PARENT + else: + width = {"hidden": 5120, "comp": 512, "idx": 128}[layout] + x = torch.randn(rows, width, device="cuda") * _v41_row_scales(rows)[:, None] + return x.to(torch.bfloat16) + + +def _v41_view(parent: torch.Tensor, layout: str, rows: int) -> torch.Tensor: + """The first `rows` rows of `parent`, in the layout the target hands the op.""" + if layout == "q": + return parent[:rows, :1280] + if layout == "kv": + return parent[:rows, 1280:] + if layout in _RING_OFFSET: + off = _RING_OFFSET[layout] + return parent[:rows, off : off + 512] + return parent[:rows] + + +def _v41_input(rows: int, layout: str) -> torch.Tensor: + """A [rows, width] input in the target's layout (see `_v41_parent`).""" + return _v41_view(_v41_parent(rows, layout), layout, rows) + + +def test_v41_eps_1e20_widths_and_layouts() -> None: + """Fresh inputs at sampled row counts; eps 1e-20 is honoured, not clamped.""" + torch.manual_seed(30) + for layout in ("hidden", "q", "kv", "comp", "idx"): + for rows in _V41_ROWS: + x = _v41_input(rows, layout) + if layout in ("q", "kv"): + # A single row reports contiguous (torch ignores a size-1 + # dim's stride) and takes the kernel's compact path; every + # other row count takes the strided path. + assert x.stride() == (1792, 1) and x.is_contiguous() == (rows == 1) + w = torch.randn(x.shape[-1], dtype=torch.bfloat16, device="cuda") + out = op(x, w, _V41_EPS) + assert out.shape == x.shape and out.dtype == x.dtype + assert torch.isfinite(out).all() + op.compare(out, op.reference(x, w, _V41_EPS)) + if rows >= 4: + # Zero rows come back zero; eps-dominated rows are O(1), not + # O(x): the same gate rejects an eps of 1e-6 on those rows. + assert (out[3::4] == 0).all() + tiny = slice(1, None, 4) + wrong = op.reference(x, w, 1e-6) + assert not torch.allclose( + out[tiny].float(), wrong[tiny].float(), rtol=1.6e-2, atol=1e-5 + ) + + +def test_v41_eps_1e20_cuda_graph() -> None: + """Each core layout captured alone after a side-stream warm-up, replayed on two fresh inputs.""" + torch.manual_seed(31) + for layout in ("hidden", "q", "kv", "comp", "idx"): + width = {"hidden": 5120, "q": 1280, "kv": 512, "comp": 512, "idx": 128}[layout] + w = torch.randn(width, dtype=torch.bfloat16, device="cuda") + for rows in _V41_GRAPH_ROWS: + x = _v41_input(rows, layout) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(x, w, _V41_EPS) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op(x, w, _V41_EPS) + for _ in range(2): + x.copy_(_v41_input(rows, layout)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + op.compare(out, op.reference(x, w, _V41_EPS)) + + +# ─── Every row count 1..16384 at every layout, and the capture union ───────── +# +# The row axis of a per-row op is certified as a range: a run at N_max = 16384 +# against the reference, every row count r = 1..N_max bitwise equal to the +# first r rows of that run, and a probe that the row count compiles nothing +# new. On the certified path (FlashInfer's CuTe DSL RMSNormKernel) one CTA +# normalizes several rows -- 2 at width 5120, 4 at 1280 and 512, 8 at 128, +# read from the kernel object on sm_103 -- and a predicate on the row count +# masks the rows at or past r in the last CTA. So r reaches the kernel twice: +# as the grid size and as the bound of that mask. Compiled kernels are looked +# up by dtype, width, weight bias, PDL, SM version and whether the rows are +# compact; the row count is a symbolic 64-bit integer. +# +# The column views are rows of one parent, sliced exactly as the target +# slices it. + +_N_MAX = 16384 +_CORE_LAYOUTS = ("hidden", "q", "kv", "comp", "idx") +_LAYOUTS = _CORE_LAYOUTS + tuple(_RING_OFFSET) +_COLUMN_VIEWS = ("q", "kv") + tuple(_RING_OFFSET) +_WIDTH = { + "hidden": 5120, + "q": 1280, + "kv": 512, + "comp": 512, + "idx": 128, + "ring0": 512, + "ring1": 512, + "ring2": 512, +} +# The capture union per rank. G ranges over every CUDA-graph batch size an +# envelope engine captures: every G 1..32, plus TRT-LLM's padded default list +# for cuda_graph_config {enable_padding: true, max_batch_size: 256} +# (40..128 step 8, 192, 256); the padding-off list at 256 (1..32, 64, 128, +# 256) is a subset. A captured step is generation-only, so the rows are G for +# a decode step, 6G for the DSpark target-verify step (1 + 5 draft tokens per +# request), 3G for that step's ratio-2 compressed rows (each verify request +# reserves ceil(6 / 2) = 3 of them, so the compressed-latent [n, 512] and +# index-K [n, 128] norms run at n = 3G), and 5G for the drafter's draft block. +# These norms run on the rank's own tokens in every member, never on a +# dep4-gathered set. Every core layout is captured at every count of the +# union, including the counts only some of the layouts reach. +_BATCHES = tuple(range(1, 33)) + tuple(range(40, 129, 8)) + (192, 256) +_CAPTURE_ROWS = tuple(sorted({m * b for b in _BATCHES for m in (1, 3, 5, 6)})) +assert len(_BATCHES) == 46 and len(_CAPTURE_ROWS) == 126 and max(_CAPTURE_ROWS) == 1536 +# The profiler probe's row counts: 1..9, the capture union, the edges of the +# 2-, 4- and 8-row CTAs around 16, 64, 128, 4096, 8192 and 12288, two prefill +# counts from the sampled rows, and the top of the range. +_PROBE_ROWS = tuple( + sorted( + set(range(1, 10)) + | set(_CAPTURE_ROWS) + | {15, 17, 63, 64, 65, 127, 128, 129, 323, 1005, 4095, 4096, 4097, 4617, 8191, 8192, 8193} + | {12287, 12288, 12289, 16383, 16384} + ) +) + + +def _compiled_variants() -> int: + """How many RMSNormKernel variants the certified path has compiled in this process.""" + return _cute_rmsnorm_kernels._get_compiled_rmsnorm_kernel.cache_info().misses + + +def _rmsnorm_into( + x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor, eps: float = _V41_EPS +) -> None: + """The call the op's body makes (`flashinfer.norm.rmsnorm`, same PDL setting), into a caller-owned `out`. + + Used only to see what the kernel writes past the last row: the op returns + a fresh tensor, so the memory after its last row is not the caller's to + inspect. + """ + _flashinfer_norm_rmsnorm(x, weight, eps, out=out, enable_pdl=get_env_enable_pdl()) + + +def _ordered_bits(t: torch.Tensor) -> torch.Tensor: + b = t.view(torch.int16).int() + return torch.where(b < 0, -(b & 0x7FFF), b) + + +def _error_stats(out: torch.Tensor, ref: torch.Tensor): + """(max |out - ref|, max distance in bf16 ulps, fraction of elements bit-identical).""" + diff = (out.float() - ref.float()).abs().max().item() + ulps = int((_ordered_bits(out) - _ordered_bits(ref)).abs().max().item()) + exact = (out.view(torch.int16) == ref.view(torch.int16)).float().mean().item() + return diff, ulps, exact + + +@pytest.mark.parametrize("layout", _LAYOUTS) +def test_v41_every_row_count_is_the_same_per_row_function(layout: str) -> None: + """Every r = 1..16384: the op on the first r rows returns exactly the first r rows of the 16384-row result. + + The 16384-row result is held to the fp32 reference, so every row count is + covered. Each output is poisoned with NaN before it is freed, so a row the + kernel skipped cannot pick up the previous call's correct bits from a + recycled allocator block. At every r the same kernel also runs into a + NaN-filled caller buffer one row longer than r, and the row past the last + must stay NaN. The sweep compiles no new kernel variant, and the parent + buffer (all its columns, for the views) is left bitwise untouched. + """ + torch.manual_seed(40 + _LAYOUTS.index(layout)) + width = _WIDTH[layout] + parent = _v41_parent(_N_MAX, layout) + snapshot = parent.clone() + w = torch.randn(width, dtype=torch.bfloat16, device="cuda") + x = _v41_view(parent, layout, _N_MAX) + full = op(x, w, _V41_EPS) + ref = op.reference(x, w, _V41_EPS) + op.compare(full, ref) + assert torch.isfinite(full).all() and (full[3::4] == 0).all() + if layout in _RING_OFFSET: + # The DSpark ring view gives bitwise what the op gives on a contiguous + # copy of the same rows (the `comp` layout's kernel). + assert full.is_contiguous() + assert torch.equal( + full.view(torch.int16), + op(x.contiguous(), w, _V41_EPS).view(torch.int16), + ) + diff, ulps, exact = _error_stats(full, ref) + print( + f"\n[{layout}] {_N_MAX} rows vs reference: max abs diff {diff:.6g}, " + f"max {ulps} bf16 ulp, bit-exact {exact:.6f}" + ) + + # A one-row view reports contiguous, so for the column views the sweep + # reaches two variants; compile both before counting. + op(_v41_view(parent, layout, 1), w, _V41_EPS) + variants = _compiled_variants() + + full_bits = full.view(torch.int16) + buf = torch.empty(_N_MAX + 1, width, dtype=torch.bfloat16, device="cuda") + mismatch = torch.zeros(_N_MAX + 1, dtype=torch.bool, device="cuda") + mismatch_into = torch.zeros(_N_MAX + 1, dtype=torch.bool, device="cuda") + guard = torch.zeros(_N_MAX + 1, dtype=torch.int64, device="cuda") + for r in range(1, _N_MAX + 1): + xr = _v41_view(parent, layout, r) + out = op(xr, w, _V41_EPS) + assert out.shape == (r, width) and out.dtype == torch.bfloat16 + mismatch[r] = (out.view(torch.int16) != full_bits[:r]).any() + out.fill_(float("nan")) + del out + buf[: r + 1].fill_(float("nan")) + _rmsnorm_into(xr, w, buf[:r]) + mismatch_into[r] = (buf[:r].view(torch.int16) != full_bits[:r]).any() + guard[r] = (~buf[r].isnan()).sum() + torch.cuda.synchronize() + assert _compiled_variants() == variants, "the row sweep compiled a new kernel variant" + bad = mismatch.nonzero().flatten().tolist() + bad_into = mismatch_into.nonzero().flatten().tolist() + over = guard.nonzero().flatten().tolist() + assert not bad, ( + f"{len(bad)} row counts differ from the {_N_MAX}-row result's rows, first {bad[:10]}" + ) + assert not bad_into, ( + f"{len(bad_into)} row counts differ into a caller buffer, first {bad_into[:10]}" + ) + assert not over, f"{len(over)} row counts wrote the row past the last, first {over[:10]}" + assert torch.equal(parent.view(torch.int16), snapshot.view(torch.int16)), ( + "the op wrote its input" + ) + print( + f"[{layout}] r = 1..{_N_MAX}: every prefix bitwise equal, row past the last untouched, " + f"no compile during the sweep ({variants} variants compiled in this process)" + ) + + +def test_v41_row_sweep_harness_sees_a_row_shift_and_a_tail_write() -> None: + """Controls for the sweep: its bitwise check rejects a read one row off, and its guard sees a write at row r.""" + torch.manual_seed(47) + for layout in _LAYOUTS: + width = _WIDTH[layout] + parent = _v41_parent(66, layout) + w = torch.randn(width, dtype=torch.bfloat16, device="cuda") + full = op(_v41_view(parent, layout, 64), w, _V41_EPS) + for r in (1, 2, 7, 63, 64): + shifted = op(_v41_view(parent[1:], layout, r), w, _V41_EPS) + n_bad = int((shifted.view(torch.int16) != full.view(torch.int16)[:r]).sum().item()) + print(f"[control row shift, {layout}] r={r}: {n_bad} / {r * width} elements differ") + assert n_bad > r * width // 2 + buf = torch.full((66, width), float("nan"), dtype=torch.bfloat16, device="cuda") + _rmsnorm_into(_v41_view(parent, layout, 64), w, buf[:64]) + assert buf[64].isnan().all() + _rmsnorm_into(_v41_view(parent, layout, 65), w, buf[:65]) # one row too many + written = int((~buf[64].isnan()).sum().item()) + print( + f"[control tail write, {layout}] a 65-row call leaves {written} / {width} non-NaN in row 64" + ) + assert written == width + + +@pytest.mark.parametrize("layout", _LAYOUTS) +def test_v41_one_compiled_variant_per_layout(layout: str) -> None: + """Probe: each call launches one kernel, and every row count >= 2 launches the same one. + + At one row a column view reports contiguous (torch ignores a size-1 dim's + stride) and launches the compact variant, the one a contiguous [1, width] + input launches; the contiguous layouts launch the same kernel at one row + as at every other count. No call in the probe compiles a new variant. + """ + torch.manual_seed(48) + width = _WIDTH[layout] + parent = _v41_parent(_N_MAX, layout) + w = torch.randn(width, dtype=torch.bfloat16, device="cuda") + + def launched(x: torch.Tensor) -> str: + with profile(activities=[ProfilerActivity.CUDA]) as prof: + op(x, w, _V41_EPS) + torch.cuda.synchronize() + names = [e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA] + assert len(names) == 1, (tuple(x.shape), names) + return names[0] + + for r in (1, 2): + op(_v41_view(parent, layout, r), w, _V41_EPS) + variants = _compiled_variants() + by_rows = {r: launched(_v41_view(parent, layout, r)) for r in _PROBE_ROWS} + compact_one_row = launched(torch.randn(1, width, device="cuda").to(torch.bfloat16)) + assert _compiled_variants() == variants, "the probe compiled a new kernel variant" + names = set(by_rows.values()) + assert all("RMSNormKernel" in n for n in names), names + many = {by_rows[r] for r in _PROBE_ROWS if r >= 2} + assert len(many) == 1, many + assert by_rows[1] == compact_one_row + if layout in _COLUMN_VIEWS: + assert by_rows[1] not in many + else: + assert by_rows[1] in many + print(f"\n[{layout}] {len(_PROBE_ROWS)} row counts, kernels launched:") + for n in sorted(names): + print(f" {n} (rows {[r for r in _PROBE_ROWS if by_rows[r] == n][:3]}...)") + + +def test_v41_capture_union_replay_is_bitwise_eager() -> None: + """The five calls captured at every row count of the capture union, replayed on new data. + + As the engine does it: an eager call at each shape before its capture, + one graph per row count holding the five calls, all captured into one + shared memory pool, largest first, then replayed in an interleaved order + over three rounds. Before each replay new data is written into the + captured buffers (the whole [rows, 1792] qkv buffer for the two column + views) and every output is poisoned with NaN. Each replayed output must + equal an eager call on the same data bitwise, pass the reference gate, and + differ from the previous round's output. A control round then writes the + new data where the graph does not read it: the replay matches the eager + result on the stale buffer and not the one on the new data, so the bitwise + check would catch a replay that read stale input. + """ + torch.manual_seed(49) + weights = { + layout: torch.randn(_WIDTH[layout], dtype=torch.bfloat16, device="cuda") + for layout in _CORE_LAYOUTS + } + source = {"hidden": "hidden", "q": "qkv", "kv": "qkv", "comp": "comp", "idx": "idx"} + + def fresh(rows: int) -> dict: + return { + "hidden": _v41_parent(rows, "hidden"), + "qkv": _v41_parent(rows, "q"), + "comp": _v41_parent(rows, "comp"), + "idx": _v41_parent(rows, "idx"), + } + + pool = torch.cuda.graph_pool_handle() + graphs = {} + for rows in sorted(_CAPTURE_ROWS, reverse=True): + bufs = fresh(rows) + views = {layout: _v41_view(bufs[source[layout]], layout, rows) for layout in _CORE_LAYOUTS} + for layout in _CORE_LAYOUTS: # the engine's warm-up: one eager call at this shape + op(views[layout], weights[layout], _V41_EPS) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + outs = { + layout: op(views[layout], weights[layout], _V41_EPS) for layout in _CORE_LAYOUTS + } + graphs[rows] = (graph, bufs, views, outs) + torch.cuda.synchronize() + + order = _CAPTURE_ROWS[::2] + _CAPTURE_ROWS[1::2][::-1] + previous = {} + worst = {layout: (0.0, 0, 1.0) for layout in _CORE_LAYOUTS} + for rnd in range(3): + for rows in order: + graph, bufs, views, outs = graphs[rows] + for key, new in fresh(rows).items(): + bufs[key].copy_(new) + for out in outs.values(): + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + for layout in _CORE_LAYOUTS: + out, x, w = outs[layout], views[layout], weights[layout] + assert out.shape == (rows, _WIDTH[layout]) + eager = op(x, w, _V41_EPS) + assert torch.equal(out.view(torch.int16), eager.view(torch.int16)), ( + layout, + rows, + rnd, + ) + ref = op.reference(x, w, _V41_EPS) + op.compare(out, ref) + if (layout, rows) in previous: + assert not torch.equal( + out.view(torch.int16), previous[layout, rows].view(torch.int16) + ) + previous[layout, rows] = out.clone() + d, u, e = _error_stats(out, ref) + wd, wu, we = worst[layout] + worst[layout] = (max(wd, d), max(wu, u), min(we, e)) + + for rows in order: + graph, bufs, views, outs = graphs[rows] + elsewhere = fresh(rows) + for out in outs.values(): + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + for layout in _CORE_LAYOUTS: + w = weights[layout] + stale = op(views[layout], w, _V41_EPS) + new = op(_v41_view(elsewhere[source[layout]], layout, rows), w, _V41_EPS) + assert torch.equal(outs[layout].view(torch.int16), stale.view(torch.int16)), ( + layout, + rows, + ) + assert not torch.equal(outs[layout].view(torch.int16), new.view(torch.int16)), ( + layout, + rows, + ) + + print(f"\n{len(_CAPTURE_ROWS)} graphs, rows {_CAPTURE_ROWS}, 3 rounds; replay vs reference:") + for layout, (d, u, e) in worst.items(): + print(f" [{layout}] max abs diff {d:.6g}, max {u} bf16 ulp, min bit-exact {e:.6f}") + + +# ─── The DSpark drafter's ring write: three width-512 views of [M, 1536] ───── +# +# `_ring_write` computes kv3 = wkv3(main_x), a fresh [M, 1536] bf16 buffer, and +# normalizes each stage's column view split(kv3, [512] * 3, -1)[s] with that +# stage's kv_norm at eps 1e-20, out of place. M is n <= 128 at prefill seeding +# (eager), G (the bonus row) and 5G (the multi-accept back-fill), G <= +# max_batch_size 256; a draft step is captured at the engine's graph batch +# sizes and runs eager at any other G. The row sweep above covers every M +# 1..16384 at the three views (layouts ring0/ring1/ring2); this section +# captures at every G 1..256 (a superset of the captured sizes), 5G, and the +# capture union up to 1280, and adds the stride/offset generalization. +_DSPARK_G = tuple(range(1, 257)) +_DSPARK_RING_ROWS = tuple( + sorted( + set(_DSPARK_G) + | {5 * g for g in _DSPARK_G} + | {r for r in _CAPTURE_ROWS if r <= 5 * max(_DSPARK_G)} + ) +) +assert len(_DSPARK_RING_ROWS) == 473 and max(_DSPARK_RING_ROWS) == 1280 + + +def test_v41_dspark_ring_views_captured_at_every_draft_batch() -> None: + """The drafter's three kv_norm calls captured at M = G and 5G for every G 1..256, replayed on new data. + + One graph per M holds the three calls on the three column views of one + [M, 1536] buffer, as `_ring_write` issues them. As the engine does it: an + eager call at each shape before its capture, all graphs in one shared + memory pool, largest first, replays interleaved over three rounds with new + data in the whole [M, 1536] buffer and NaN-poisoned outputs. Every replay + equals an eager call on the same views and an eager call on a contiguous + copy of them bitwise, passes the reference gate, and differs from the + previous round's. A control round writes the new data where the graph does + not read it and the replay then matches only the stale result. + """ + torch.manual_seed(60) + rings = tuple(_RING_OFFSET) + weights = {layout: torch.randn(512, dtype=torch.bfloat16, device="cuda") for layout in rings} + pool = torch.cuda.graph_pool_handle() + graphs = {} + for rows in sorted(_DSPARK_RING_ROWS, reverse=True): + kv3 = _v41_parent(rows, "ring0") + views = {layout: _v41_view(kv3, layout, rows) for layout in rings} + for layout in rings: + op(views[layout], weights[layout], _V41_EPS) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + outs = {layout: op(views[layout], weights[layout], _V41_EPS) for layout in rings} + for out in outs.values(): + assert out.shape == (rows, 512) + # Contiguous at M >= 2; at M = 1 it keeps the view's row stride. + assert out.stride() == ((512, 1) if rows >= 2 else (_RING_PARENT, 1)) + graphs[rows] = (graph, kv3, views, outs) + torch.cuda.synchronize() + + order = _DSPARK_RING_ROWS[::2] + _DSPARK_RING_ROWS[1::2][::-1] + previous = {} + worst = (0.0, 0, 1.0) + for rnd in range(3): + for rows in order: + graph, kv3, views, outs = graphs[rows] + kv3.copy_(_v41_parent(rows, "ring0")) + for out in outs.values(): + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + for layout in rings: + out, x, w = outs[layout], views[layout], weights[layout] + bits = out.view(torch.int16) + assert torch.equal(bits, op(x, w, _V41_EPS).view(torch.int16)), ( + layout, + rows, + rnd, + ) + assert torch.equal(bits, op(x.contiguous(), w, _V41_EPS).view(torch.int16)) + ref = op.reference(x, w, _V41_EPS) + op.compare(out, ref) + if (layout, rows) in previous: + assert not torch.equal(bits, previous[layout, rows]), (layout, rows, rnd) + previous[layout, rows] = bits.clone() + d, u, e = _error_stats(out, ref) + worst = (max(worst[0], d), max(worst[1], u), min(worst[2], e)) + + for rows in order: + graph, kv3, views, outs = graphs[rows] + elsewhere = _v41_parent(rows, "ring0") + for out in outs.values(): + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + for layout in rings: + w = weights[layout] + stale = op(views[layout], w, _V41_EPS) + new = op(_v41_view(elsewhere, layout, rows), w, _V41_EPS) + assert torch.equal(outs[layout].view(torch.int16), stale.view(torch.int16)), ( + layout, + rows, + ) + assert not torch.equal(outs[layout].view(torch.int16), new.view(torch.int16)), ( + layout, + rows, + ) + print( + f"\n[dspark ring] {len(_DSPARK_RING_ROWS)} graphs x 3 views, M {_DSPARK_RING_ROWS}, 3 rounds; " + f"replay vs reference: max abs diff {worst[0]:.6g}, max {worst[1]} bf16 ulp, min bit-exact {worst[2]:.6f}" + ) + + +# The stride/offset generalization: every row stride S that is a multiple of 8 +# elements (16 bytes) from 512 to 2048, every column offset o that is a multiple +# of 8 with o + 512 <= S, at the CTA-edge row counts below (4 rows per CTA at +# width 512) up to the drafter's 160. A one-row view reports contiguous and +# takes the compact kernel, which needs a 128-byte-aligned start: at one row +# only the offsets that are multiples of 64 elements run (the parents here are +# fresh allocations, so 128-byte aligned); the others raise ValueError. +_GEN_STRIDES = tuple(range(512, 2049, 8)) +_GEN_ROWS = (2, 3, 4, 5, 8, 9, 64, 65, 127, 128, 129, 159, 160) +_GEN_N = 160 + + +def test_v41_dspark_any_16_byte_aligned_stride_and_offset() -> None: + """Width 512, every 16-byte-aligned (row stride, column offset) pair in 512..2048: bitwise the contiguous copy. + + Per pair, the view's first r rows (r in `_GEN_ROWS`, and r = 1 where the + offset is 128-byte aligned) give bitwise the first r rows of the op on a + contiguous copy of the 160-row view; that contiguous result passes the + reference gate, zero rows come back 0, the parent buffer is not written, + and no kernel variant compiles (the row stride is a symbolic argument of + the one strided kernel). A one-row view at an offset that is not a + multiple of 64 elements raises ValueError. Mismatches accumulate on the + device; one sync at the end. + """ + torch.manual_seed(61) + w = torch.randn(512, dtype=torch.bfloat16, device="cuda") + op(_v41_view(_v41_parent(2, "ring1"), "ring1", 2), w, _V41_EPS) + op(_v41_input(2, "comp"), w, _V41_EPS) + op(_v41_input(1, "comp"), w, _V41_EPS) + variants = _compiled_variants() + scales = _v41_row_scales(_GEN_N)[:, None] + bad_bits = torch.zeros((), dtype=torch.int64, device="cuda") + bad_gate = torch.zeros((), dtype=torch.int64, device="cuda") + bad_zero = torch.zeros((), dtype=torch.int64, device="cuda") + bad_input = torch.zeros((), dtype=torch.int64, device="cuda") + pairs = one_row_ran = one_row_raised = 0 + for stride in _GEN_STRIDES: + parent = (torch.randn(_GEN_N, stride, device="cuda") * scales).to(torch.bfloat16) + snapshot = parent.clone() + for off in range(0, stride - 512 + 1, 8): + view = parent[:, off : off + 512] + contig = view.contiguous() + full = op(contig, w, _V41_EPS) + ref = op.reference(contig, w, _V41_EPS) + bad_gate += ( + (full.float() - ref.float()).abs() > 1e-5 + 1.6e-2 * ref.float().abs() + ).sum() + bad_zero += (full[3::4] != 0).sum() + full_bits = full.view(torch.int16) + for r in _GEN_ROWS: + out = op(view[:r], w, _V41_EPS) + bad_bits += (out.view(torch.int16) != full_bits[:r]).sum() + if off % 64 == 0: + out = op(view[:1], w, _V41_EPS) + bad_bits += (out.view(torch.int16) != full_bits[:1]).sum() + one_row_ran += 1 + else: + with pytest.raises(ValueError, match="Misaligned Tensor data"): + op(view[:1], w, _V41_EPS) + one_row_raised += 1 + pairs += 1 + bad_input += (parent.view(torch.int16) != snapshot.view(torch.int16)).sum() + torch.cuda.synchronize() + print( + f"\n[dspark generalization] {len(_GEN_STRIDES)} strides 512..2048 step 8, {pairs} (stride, offset) pairs, " + f"rows {_GEN_ROWS}: {int(bad_bits)} elements differ from the contiguous copy, {int(bad_gate)} outside " + f"the gate, {int(bad_zero)} nonzero zero-row elements, {int(bad_input)} input elements written, " + f"{_compiled_variants() - variants} variants compiled; one row: {one_row_ran} pairs ran (offset % 64 == 0), " + f"{one_row_raised} raised ValueError" + ) + assert pairs == sum((s - 512) // 8 + 1 for s in _GEN_STRIDES) == one_row_ran + one_row_raised + assert int(bad_bits) == 0 and int(bad_gate) == 0 and int(bad_zero) == 0 and int(bad_input) == 0 + assert _compiled_variants() == variants, "a stride or offset compiled a new kernel variant" + + +@pytest.mark.parametrize( + "stride, off, rows", + [ + (1536, 4, 8), + (1536, 1, 8), + (1536, 516, 8), + (1540, 0, 8), + (1538, 512, 8), + (520, 2, 8), + (1536, 8, 1), + (1536, 520, 1), + (1792, 1272, 1), + ], +) +def test_v41_dspark_misaligned_stride_or_offset_raises(stride: int, off: int, rows: int) -> None: + """Misaligned views raise ValueError; nothing runs. + + At two or more rows, a row stride or column offset that is not a multiple of + 8 elements (16 bytes). At one row, where the compact kernel runs, a start + that is not 128-byte aligned. + """ + parent = torch.randn(rows, stride, device="cuda").to(torch.bfloat16) + w = torch.randn(512, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="Misaligned Tensor data|Invalid mX.strides"): + op(parent[:, off : off + 512], w, _V41_EPS) + torch.cuda.synchronize() + + +# ─── The DeepSeek-V4.1 vision tower: width 1024, eps 1e-6, rows 1..51111 ───── +# +# The ViT normalizes its +# residual stream with this op at width 1024 and eps 1e-6 -- norm1 and norm2 +# of each of the 32 blocks and the final norm -- on a contiguous [P, 1024] +# bf16 tensor, P = every patch of one request (at most 17037 over the 900 +# MMMU val documents). The row range is certified up to 51111 = 3 x 17037 +# rows. The residual stream carries massive activations (channels 53, 147 +# and 308 reach |x| ~ 2000 in the golden), so rows cycle through: unit; +# unit with those three channels at +-2000; 1e-3 scale, where mean(x^2) ~ +# eps and eps sets the result; all-zero. +_VIT_WIDTH = 1024 +_VIT_EPS = 1e-6 +_VIT_N_MAX = 51111 +# The prototype's whole-tower CUDA graph ran at the golden cases' patch +# counts; 17037 is the MMMU maximum. +_VIT_GRAPH_ROWS = (1517, 4680, 5325, 8772, 17037) +_VIT_PROBE_ROWS = ( + 1, + 2, + 3, + 4, + 5, + 7, + 8, + 9, + 15, + 16, + 17, + 63, + 64, + 65, + 127, + 128, + 129, + 1517, + 4680, + 5325, + 8191, + 8192, + 8193, + 8772, + 16383, + 16384, + 16385, + 17037, + 32767, + 32768, + 32769, + 51110, + 51111, +) + + +def _vit_input(rows: int, seed: int) -> torch.Tensor: + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(rows, _VIT_WIDTH, device="cuda", generator=g) + return _vit_rows(x, g).to(torch.bfloat16) + + +def _vit_weight(seed: int) -> torch.Tensor: + g = torch.Generator(device="cuda").manual_seed(seed) + return (1.0 + 0.5 * torch.randn(_VIT_WIDTH, device="cuda", generator=g)).to(torch.bfloat16) + + +def test_v41_vit_every_row_count_is_the_same_per_row_function() -> None: + """Width 1024, eps 1e-6: every r = 1..51111 is bitwise the first r rows of the 51111-row call. + + The 51111-row result is held to the fp32 reference. Each output is + poisoned with NaN before it is freed; at every r the same kernel also runs + into a NaN-filled caller buffer one row longer than r, whose last row must + stay NaN; the sweep compiles no new kernel variant; the input is left + bitwise untouched. + """ + x = _vit_input(_VIT_N_MAX, 80) + snapshot = x.clone() + w = _vit_weight(81) + full = op(x, w, _VIT_EPS) + ref = op.reference(x, w, _VIT_EPS) + op.compare(full, ref) + assert torch.isfinite(full).all() and (full[3::4] == 0).all() + diff, ulps, exact = _error_stats(full, ref) + rpb = _cute_rmsnorm_kernels.RMSNormKernel( + _cute_rmsnorm_kernels.get_cutlass_dtype("bfloat16"), _VIT_WIDTH, 0.0, sm_version=103 + ).rows_per_block + print( + f"\n[vit] {_VIT_N_MAX} rows vs reference: max abs diff {diff:.6g}, max {ulps} bf16 ulp, " + f"bit-exact {exact:.6f}; rows per CTA {rpb}" + ) + variants = _compiled_variants() + full_bits = full.view(torch.int16) + buf = torch.empty(_VIT_N_MAX + 1, _VIT_WIDTH, dtype=torch.bfloat16, device="cuda") + mismatch = torch.zeros(_VIT_N_MAX + 1, dtype=torch.bool, device="cuda") + mismatch_into = torch.zeros(_VIT_N_MAX + 1, dtype=torch.bool, device="cuda") + guard = torch.zeros(_VIT_N_MAX + 1, dtype=torch.int64, device="cuda") + for r in range(1, _VIT_N_MAX + 1): + out = op(x[:r], w, _VIT_EPS) + mismatch[r] = (out.view(torch.int16) != full_bits[:r]).any() + out.fill_(float("nan")) + del out + buf[: r + 1].fill_(float("nan")) + _rmsnorm_into(x[:r], w, buf[:r], _VIT_EPS) + mismatch_into[r] = (buf[:r].view(torch.int16) != full_bits[:r]).any() + guard[r] = (~buf[r].isnan()).sum() + torch.cuda.synchronize() + assert _compiled_variants() == variants, "the row sweep compiled a new kernel variant" + bad = mismatch.nonzero().flatten().tolist() + bad_into = mismatch_into.nonzero().flatten().tolist() + over = guard.nonzero().flatten().tolist() + assert not bad, ( + f"{len(bad)} row counts differ from the {_VIT_N_MAX}-row result's rows, first {bad[:10]}" + ) + assert not bad_into, ( + f"{len(bad_into)} row counts differ into a caller buffer, first {bad_into[:10]}" + ) + assert not over, f"{len(over)} row counts wrote the row past the last, first {over[:10]}" + assert torch.equal(x.view(torch.int16), snapshot.view(torch.int16)), "the op wrote its input" + print( + f"[vit] r = 1..{_VIT_N_MAX}: every prefix bitwise equal, row past the last untouched, no compile" + ) + + +def test_v41_vit_eps_and_row_controls() -> None: + """eps 1e-6 is honoured, and a one-row shift is rejected. + + The core's 1e-20 and a 1e-5 are rejected on the eps-scale rows. + """ + x = _vit_input(4096, 82) + w = _vit_weight(83) + out = op(x, w, _VIT_EPS) + op.compare(out, op.reference(x, w, _VIT_EPS)) + tiny = slice(2, None, 4) + for wrong_eps in (1e-20, 1e-5): + wrong = op.reference(x, w, wrong_eps) + bad = _bf16_outside(out[tiny], wrong[tiny]) + print( + f"\n[vit control eps {wrong_eps:g}] {bad:.4f} of the eps-scale rows' elements outside the bf16 gate" + ) + assert bad > 0.5 + for r in (1, 2, 7, 64, 4095): + shifted = op(x[1 : r + 1], w, _VIT_EPS) + n_bad = int((shifted.view(torch.int16) != out[:r].view(torch.int16)).sum()) + print(f"[vit control row shift] r={r}: {n_bad} / {r * _VIT_WIDTH} elements differ") + assert n_bad > r * _VIT_WIDTH // 2 + + +def _bf16_outside(a: torch.Tensor, b: torch.Tensor) -> float: + """Fraction of elements outside torch's default bf16 gate (rtol 1.6e-2, atol 1e-5).""" + af, bf = a.float(), b.float() + return float(((af - bf).abs() > 1e-5 + 1.6e-2 * bf.abs()).float().mean()) + + +def test_v41_vit_one_compiled_variant() -> None: + """Every probed row count 1..51111 launches the same single RMSNormKernel, and nothing compiles.""" + x = _vit_input(_VIT_N_MAX, 84) + w = _vit_weight(85) + op(x[:2], w, _VIT_EPS) + variants = _compiled_variants() + names = set() + for r in _VIT_PROBE_ROWS: + with profile(activities=[ProfilerActivity.CUDA]) as prof: + op(x[:r], w, _VIT_EPS) + torch.cuda.synchronize() + launched = [ + e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA + ] + assert len(launched) == 1, (r, launched) + names.add(launched[0]) + assert _compiled_variants() == variants, "the probe compiled a new kernel variant" + assert len(names) == 1 and "RMSNormKernel" in next(iter(names)), names + print(f"\n[vit] {len(_VIT_PROBE_ROWS)} row counts, one kernel: {next(iter(names))}") + + +def test_v41_vit_cuda_graph_at_the_prototype_sizes() -> None: + """Captured at the prototype's graph sizes into one pool, after an eager call. + + Replays on new data are bitwise eager. + """ + w = _vit_weight(86) + pool = torch.cuda.graph_pool_handle() + graphs = {} + for rows in sorted(_VIT_GRAPH_ROWS, reverse=True): + x = _vit_input(rows, 87) + op(x, w, _VIT_EPS) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + out = op(x, w, _VIT_EPS) + graphs[rows] = (graph, x, out) + previous = {} + for rnd in range(2): + for rows in _VIT_GRAPH_ROWS: + graph, x, out = graphs[rows] + x.copy_(_vit_input(rows, 88 + rnd)) + out.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), op(x, w, _VIT_EPS).view(torch.int16)), rows + op.compare(out, op.reference(x, w, _VIT_EPS)) + if rows in previous: + assert not torch.equal(out.view(torch.int16), previous[rows].view(torch.int16)), ( + rows + ) + previous[rows] = out.clone() + print( + f"\n[vit] captured at {_VIT_GRAPH_ROWS}: every replay bitwise equal to eager and in the gate" + ) diff --git a/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py index 1e9d8c7b4fab..ad9477d5bfed 100644 --- a/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py +++ b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py @@ -1,118 +1,243 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the fp4_quantize catalog entry (NVFP4: sf_vec_size=16, e4m3 scales).""" +"""GPU test for the fp4_quantize catalog entry. +NVFP4 (sf_vec_size=16, e4m3 scales, both scale layouts) and MXFP4 (sf_vec_size=32, UE8M0 +scales, linear layout: the DeepSeek-V4.1 index K / index Q quantizer). The DeepSeek-V4.1 +section below keeps the sweeps behind the V4.1 cells: every row count, CUDA-graph capture of +the target's calls, and the MoE input's chunk views. +""" + +import pytest import torch from torch.profiler import ProfilerActivity, profile import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( - fp4_quantize, + fp4_quantize as op, ) -assert torch.cuda.is_available(), "fp4_quantize requires a CUDA device" +__extra_import_path__ = [".."] +from _validating import validating # noqa: E402 — needs the path declared above -DEV = torch.device("cuda") -VEC = 16 -E4M3_MAX = 448.0 -E2M1_MAX = 6.0 +assert torch.cuda.is_available(), "fp4_quantize requires a CUDA device" -# e2m1 code -> value. code = (exponent << 1) | mantissa, bit 3 is the sign. -E2M1_VALUES = torch.tensor( - [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32, device=DEV -) -# Midpoints between consecutive e2m1 magnitudes: a value landing exactly here is -# a rounding tie. `_TIE_UP[i]` is True when rounding *up* at midpoint i yields the -# even code (round-to-nearest-even keeps it) and False when it yields the odd one -# (round-to-nearest-even goes down instead). -E2M1_MIDPOINTS = torch.tensor( - [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], dtype=torch.float32, device=DEV +_E2M1_VALUES = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32, device="cuda" ) -_TIE_UP = torch.tensor([False, True, False, True, False, True, False], dtype=torch.bool, device=DEV) - -# The kernel builds its per-block output scale through two `rcp.approx.ftz.f32` -# reciprocals (~2^-23 relative error each), so a value sitting within a few fp32 -# ulps of an e2m1 midpoint can round to either neighbour. Everything further away -# than this window must match the reference bit for bit. The largest deviation -# observed on this machine was 7.9e-8 relative, ~12x inside this window. -TIE_WINDOW = 2.0**-20 - - -def _e2m1_codes(v: torch.Tensor) -> torch.Tensor: - """fp32 -> 4-bit e2m1 codes (bit 3 = sign), round-to-nearest-even, saturating to +-6.""" - a = v.abs().unsqueeze(-1) - gt = (a > E2M1_MIDPOINTS).sum(-1) # strict: ties provisionally round down - tie_up = ((a == E2M1_MIDPOINTS) & _TIE_UP).any(-1) # ties whose even side is up - code = (gt + tie_up).to(torch.uint8) - neg = (v < 0) | ((v == 0) & torch.signbit(v)) - return code | (neg.to(torch.uint8) << 3) - - -def _ref(x: torch.Tensor, gs: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Independent native-torch reference: (packed data [M, K/2], sf bytes [M, K/16], scaled). - - Per 16-element block along the last dim (rows = product of the leading dims): - `vecmax = max|x|`; the block scale is `e4m3(gs * vecmax / 6)` (6 = e2m1 max, - e4m3 conversion saturating at 448); the data is the fp32 input times - `gs / scale`, rounded to e2m1. `scaled` is that pre-rounding product, returned - so the caller can identify rounding ties. + + +def _static_g1(calib): + """A modelopt-style static activation scale, fixed from a calibration batch. + + modelopt stores `input_scale = amax / (448 * 6)`; the op takes its reciprocal. + Calibrating on a *different* batch than the one quantized is what serving does: + `g` is then not `448*6/amax` of the tensor in hand. """ - k = x.shape[-1] - m = x.numel() // k - xr = x.reshape(m, k) - vecmax = xr.abs().reshape(m, k // VEC, VEC).amax(-1).float() - # torch's fp32 -> e4m3 cast emits NaN above 464 while the kernel's cast - # saturates; clamping first makes the two agree (448 < v <= 464 rounds to - # 448 either way). - sf = (gs * (vecmax / E2M1_MAX)).clamp(max=E4M3_MAX).to(torch.float8_e4m3fn) - out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) - scaled = (xr.float().reshape(m, k // VEC, VEC) * out_scale.unsqueeze(-1)).reshape(m, k) - codes = _e2m1_codes(scaled) - return codes[:, 0::2] | (codes[:, 1::2] << 4), sf.view(torch.uint8), scaled - - -def _unpack(packed: torch.Tensor) -> torch.Tensor: - """[M, K/2] packed bytes -> [M, K] e2m1 codes (element 2i in the low nibble).""" - m, kh = packed.shape - codes = torch.empty(m, kh * 2, dtype=torch.uint8, device=packed.device) + input_scale = calib.abs().max().float() / (448.0 * 6.0) + return (1.0 / input_scale).reshape(1).contiguous() + + +def _build(spec, seed): + """Cell spec -> the op's five arguments. + + Optional keys, absent from the deepseek-r1 cells: `ue8m0` (default False); + `gs` -- "amax" (default: 448*6/amax of the input), "static" (`_static_g1` of a + separate [4096, K] batch, DeepSeek-V4.1's routed-expert scale), or None (MXFP4); + `block_mag` -- each 32-block of the input times its own 2^j, j in -20..20. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(spec["shape"], generator=g, device="cuda") + if spec.get("block_mag", False): + j = torch.randint(-20, 21, (x.numel() // 32, 1), generator=g, device="cuda") + x = (x.view(-1, 32) * torch.exp2(j.float())).view(spec["shape"]) + x = x.to(spec["dtype"]) + kind = spec.get("gs", "amax") + if kind == "amax": + # The canonical NVFP4 activation global scale, 448*6/amax. + gs = (448.0 * 6.0 / x.abs().max().float()).reshape(1) + elif kind == "static": + calib = torch.randn(4096, x.shape[-1], generator=g, device="cuda").to(spec["dtype"]) + gs = _static_g1(calib) + else: + assert kind is None, kind + gs = None + return x, gs, spec["vec"], spec.get("ue8m0", False), spec["swizzled"] + + +def _unpack(packed): + """[M, K/2] packed bytes -> [M, K] codes (element 2i in the low nibble).""" + m, half = packed.shape + codes = torch.empty(m, half * 2, dtype=torch.uint8, device=packed.device) codes[:, 0::2] = packed & 0xF codes[:, 1::2] = packed >> 4 return codes -def _dequant(packed: torch.Tensor, sf: torch.Tensor, gs: float) -> torch.Tensor: - """data * scale / global_scale -- the value an NVFP4 consumer reconstructs. +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Codes bit-exact away from the rounding ties the entry declines to predict. - `sf` must be the linear-layout scale buffer of the same call. + `near_tie` is the entry's own statement about where the kernel's approximate + reciprocals decide the direction. Everything outside it has to match; the + count inside it is asserted small, so the window cannot quietly grow into a + licence for the kernel to be wrong. + + MXFP4 cells have no tie window: the scale is a power of two, so `x * 2^-e` is + exact and a midpoint stays exactly a midpoint. Their data is bit-exact, ties + included. """ + x, gs, vec, ue8m0, swizzled = _build(cell.spec, seed=abs(hash(cell.why)) % 2**31) + with validating(op): + packed, sf = op(x, gs, vec, ue8m0, swizzled) + ref_packed, ref_sf = op.reference(x, gs, vec, ue8m0, swizzled) + + assert torch.equal(sf, ref_sf), "block scales are exact; a mismatch is a defect" + if ue8m0: + assert packed.shape == (*x.shape[:-1], x.shape[-1] // 2) + assert torch.equal(packed, ref_packed), "MXFP4 data is bit-exact, ties included" + return + + tie = op.near_tie(op.scaled(x, gs, vec)) + codes, ref_codes = _unpack(packed), _unpack(ref_packed) + assert torch.equal(codes[~tie], ref_codes[~tie]), "codes differ away from a rounding tie" + # Measured across these cells: 0.13% to 0.24%. Higher than a random landing + # would give by some three orders of magnitude, because most of these are + # *exact* midpoint hits rather than near misses -- a bf16 input times an + # e4m3 scale lands on a grid coarse enough to hit midpoints often. The bound + # is 4x the worst measured, tight enough that a window quietly widening into + # a licence for the kernel to be wrong would trip it. The deepseek-v4.1 cells' + # static scale, calibrated on another batch, gave 0.05% to 0.29% over the + # sweeps in the DeepSeek-V4.1 section below. + assert tie.sum().item() <= 0.01 * tie.numel(), ( + f"{tie.sum().item()} of {tie.numel()} elements sit in the tie window; " + "a window this wide is no longer a statement about rounding" + ) + + +def test_the_quantization_round_trips() -> None: + """Independent of the reference: dequantizing has to return the input. + + The cells compare two implementations of one arithmetic, which would both be + wrong together if the arithmetic were. This bounds the error against the + input instead -- e2m1 carries one mantissa bit, so at worst half a step of + the block's own scale. + """ + x = torch.randn(64, 7168, dtype=torch.bfloat16, device="cuda") + gs = (448.0 * 6.0 / x.abs().max().float()).reshape(1) + packed, sf = op(x, gs, 16, False, False) + codes = _unpack(packed) - v = E2M1_VALUES[(codes & 7).long()] - v = torch.where((codes & 8).bool(), -v, v) - rows, k = codes.shape - scale = sf.view(rows, k // VEC).view(torch.float8_e4m3fn).float() - return v * scale.repeat_interleave(VEC, dim=-1) / gs + value = _E2M1_VALUES[(codes & 7).long()] + value = torch.where((codes & 8).bool(), -value, value) + scale = sf.view(64, 7168 // 16).view(torch.float8_e4m3fn).float() + deq = value * scale.repeat_interleave(16, dim=-1) / gs.item() + + block_amax = x.float().view(64, 7168 // 16, 16).abs().amax(-1, keepdim=True) + err = (deq - x.float()).view(64, 7168 // 16, 16).abs().amax(-1, keepdim=True) + assert torch.all(err <= block_amax * 0.5 + 1e-6) + + +def test_guard_refuses_a_per_token_global_scale() -> None: + """Element 0 is applied to every row and the rest are ignored.""" + x = torch.randn(8, 128, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="exactly one element"): + with validating(op): + op(x, torch.ones(8, device="cuda"), 16, False, True) + + +def test_guard_refuses_a_partial_block() -> None: + x = torch.randn(8, 120, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="whole number of 16-element"): + with validating(op): + op(x, torch.ones(1, device="cuda"), 16, False, True) + + +def test_guard_refuses_a_non_contiguous_input() -> None: + x = torch.randn(8, 256, dtype=torch.bfloat16, device="cuda")[:, ::2] + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(x, torch.ones(1, device="cuda"), 16, False, True) + + +# ─── DeepSeek-V4.1 ──────────────────────────────────────────────────────────── +# The V4.1 target quantizes, per step: index K `[n, 128]` and index Q `[n, 32, 128]` +# in MXFP4 mode (`op(x, None, 32, True, False)`, the checkpoint's `fp4_act_quant(x, +# 32)`), and the routed-expert input `[n, 5120]` in NVFP4 mode, linear layout, with the +# layer's static FC1 activation scale `g1 = 1 / input_scale` as a `[1]` fp32 device +# tensor. Decode steps run inside CUDA graphs. The cells cover the ends of each range; +# the tests below sweep it. + +DEV = torch.device("cuda") +NV_VEC = 16 +MX_VEC = 32 +INDEX_HEAD_DIM = 128 +INDEX_N_HEADS = 32 +V41_HIDDEN = 5120 +# The reference floors the block max at 6 * 2^-126, so its smallest scale is 2^-126 +# (UE8M0 byte 0x01). The op has no floor: a block whose max is at or below 3 * 2^-126 +# (so that max * fp32(1/6) <= 2^-127) gets scale byte 0x00 and is written as signed +# zeros. Every block with a larger max matches the reference bit for bit. +MX_REF_AMAX_FLOOR = 6 * 2.0**-126 +MX_OP_ZERO_AMAX = 3 * 2.0**-126 +_INV6_F32 = torch.tensor(1.0 / 6.0, dtype=torch.float32, device=DEV) +# e2m1 midpoints, and the even code each one rounds to under round-to-nearest-even +_MX_MIDS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_MX_MID_CODES = (0, 2, 2, 4, 4, 6, 6) +_MX_GRID = (0.5, 1.0, 1.5, 2.0, 3.0, 4.0) + + +def _mx(x): + return op(x, None, MX_VEC, True, False) -def _near_tie(scaled: torch.Tensor) -> torch.Tensor: - """Elements whose pre-rounding value sits within TIE_WINDOW (relative) of an e2m1 midpoint.""" - a = scaled.abs().unsqueeze(-1) - return ((a - E2M1_MIDPOINTS).abs() <= TIE_WINDOW * a).any(-1) +def _mx_ref(x): + return op.reference(x, None, MX_VEC, True, False) -def _assert_data( - got: torch.Tensor, ref: torch.Tensor, scaled: torch.Tensor -) -> tuple[int, int, int]: - """Bit-exact away from e2m1 rounding ties; one adjacent code at a tie. +def _mx_wrong(x, ties="even", scale="ceil"): + """Deliberately wrong MXFP4 variants, for the controls: (packed [M, K/2], sf [M, K/32]). - Returns `(near_tie_count, mismatch_count, total)`. The two counts are - different quantities: the first is how many elements sit inside the tie - window at all (the check's blind spot), the second how many of those the - kernel actually resolved the other way. + `ties="zero"` rounds e2m1 ties toward zero instead of to even; `scale="floor"` takes + 2^floor instead of 2^ceil for the block scale. With the defaults it is the reference. + """ + k = x.shape[-1] + m = x.numel() // k + xf = x.reshape(m, k // MX_VEC, MX_VEC).float() + bits = (xf.abs().amax(-1).clamp_min(MX_REF_AMAX_FLOOR) * _INV6_F32).view(torch.int32) + e = ((bits >> 23) & 0xFF) - 127 + if scale == "ceil": + e = e + ((bits & 0x7FFFFF) != 0).to(torch.int32) + s = ((e + 127) << 23).view(torch.float32) + a = (xf / s.unsqueeze(-1)).clamp(-6.0, 6.0).abs() + + def rnd(v): + return torch.round(v) if ties == "even" else torch.ceil(v - 0.5) + + r = torch.where(a < 2, rnd(a * 2) / 2, torch.where(a < 4, rnd(a), rnd(a / 2) * 2)) + code = torch.where(r < 2, r * 2, torch.where(r < 4, r + 2, r / 2 + 4)).to(torch.uint8) + codes = (code | (torch.signbit(xf).to(torch.uint8) << 3)).reshape(m, k) + return codes[:, 0::2] | (codes[:, 1::2] << 4), (e + 127).to(torch.uint8) + + +def _block_mag_rows(n, gen=None): + """[n, 128] bf16 with a different power-of-two magnitude per 32-block (2^-20..2^20).""" + mag = torch.exp2(torch.randint(-20, 21, (n, 4, 1), device=DEV, generator=gen).float()) + x = torch.randn(n, 4, MX_VEC, device=DEV, generator=gen) * mag + return x.to(torch.bfloat16).view(n, INDEX_HEAD_DIM) + + +def _assert_data(got, ref, scaled): + """NVFP4 data: bit-exact away from e2m1 rounding ties; one adjacent code at a tie. + + Returns `(near_tie_count, mismatch_count, total)`. The two counts are different + quantities: the first is how many elements sit inside the tie window at all (the + check's blind spot), the second how many of those the kernel resolved the other way. """ g = _unpack(got) r = _unpack(ref) - near_tie = _near_tie(scaled) - torch.testing.assert_close(g[~near_tie], r[~near_tie]) + near_tie = op.near_tie(scaled) + assert torch.equal(g[~near_tie], r[~near_tie]), "codes differ away from a rounding tie" diff = (g != r) & near_tie if diff.any(): # at a tie the kernel may pick either neighbour, never anything else @@ -122,535 +247,641 @@ def _assert_data( return int(near_tie.sum()), int(diff.sum()), diff.numel() -def _swizzle_index(rows: int, cols: int) -> torch.Tensor: - """[rows, cols] scale coordinates -> flat offsets in the 128x4-swizzled buffer.""" - padded_cols = _pad_up(cols, 4) - r = torch.arange(rows, device=DEV).view(-1, 1) - c = torch.arange(cols, device=DEV).view(1, -1) - return ( - (c % 4) - + (c // 4) * (4 * 128) - + (r % 32) * 16 - + ((r % 128) // 32) * 4 - + (r // 128) * (128 * padded_cols) - ) - - -def _pad_up(x: int, m: int) -> int: - return (x + m - 1) // m * m - - -def _global_scale(x: torch.Tensor) -> tuple[torch.Tensor, float]: - """The canonical NVFP4 activation global scale, 448*6/amax, as a [1] fp32 tensor.""" - gs = (E4M3_MAX * E2M1_MAX / x.abs().max().float()).reshape(1) - return gs, gs.item() - - -def test_linear_layout_bf16() -> None: - """Linear scale order, DeepSeek-V3-Lite widths, decode- through prefill-sized rows. - - Rows straddle 1024 and widths straddle a multiple of 512 -- the two - conditions a TMA high-throughput variant would key on. This build has no - such variant (see test_r1_one_kernel_at_every_width); the coverage is kept - because a future one would. - """ - torch.manual_seed(0) - worst = 0.0 - for t in (1, 2, 7, 64, 1023, 1024, 2048): - for k in (2560, 3072, 12288): - if t * k > 16 * 1024 * 1024: # keeps the reference's [M, K, 7] temporaries bounded - continue - x = torch.randn(t, k, dtype=torch.bfloat16, device=DEV) - gs, gsf = _global_scale(x) - data, sf = fp4_quantize(x, gs, VEC, False, False) - assert data.shape == (t, k // 2) and data.dtype == torch.uint8 - assert data.is_contiguous() - assert sf.shape == (t * k // VEC,) and sf.dtype == torch.uint8 - ref_data, ref_sf, scaled = _ref(x, gsf) - torch.testing.assert_close(sf.view(t, k // VEC), ref_sf) - _, mismatch, total = _assert_data(data, ref_data, scaled) - worst = max(worst, mismatch / total) - # the kernel resolves a near-tie the other way from the exact-fp32 reference - # on at most a fraction of a percent of elements; how many is set by the - # global scale rather than by the shape - assert worst < 0.005, worst - - -def test_swizzled_layout_bf16() -> None: - """Swizzled scales carry the same bytes at the 128x4 offsets; the data is layout-independent.""" - torch.manual_seed(1) - for t, k in ( - (1, 2560), - (7, 2560), - (129, 2560), - (1024, 2560), - (200, 3072), - (3, 112), - ): - x = torch.randn(t, k, dtype=torch.bfloat16, device=DEV) - gs, gsf = _global_scale(x) - cols = k // VEC - data_sw, sf_sw = fp4_quantize(x, gs, VEC, False, True) - data_li, sf_li = fp4_quantize(x, gs, VEC, False, False) - assert sf_sw.shape == (_pad_up(t, 128) * _pad_up(cols, 4),) - assert torch.equal(data_sw, data_li), "data must not depend on the scale layout" - ref_data, ref_sf, scaled = _ref(x, gsf) - torch.testing.assert_close(sf_li.view(t, cols), ref_sf) - _assert_data(data_sw, ref_data, scaled) - idx = _swizzle_index(t, cols).reshape(-1) - torch.testing.assert_close(sf_sw[idx].view(t, cols), ref_sf) - # every offset not addressed by a real (row, col) pair is row/column padding - rest = torch.ones_like(sf_sw, dtype=torch.bool) - rest[idx] = False - assert (sf_sw[rest] == 0).all(), "swizzled padding must be zero" - - -def test_3d_input_collapses_leading_dims() -> None: - """data keeps the leading dims; the scale buffer is laid out for their product.""" - torch.manual_seed(3) - x = torch.randn(2, 5, 2560, dtype=torch.bfloat16, device=DEV) - gs, gsf = _global_scale(x) - data, sf = fp4_quantize(x, gs, VEC, False, False) - assert data.shape == (2, 5, 1280) - assert sf.shape == (10 * 160,) - ref_data, ref_sf, scaled = _ref(x, gsf) - torch.testing.assert_close(sf.view(10, 160), ref_sf) - _assert_data(data.reshape(10, 1280), ref_data, scaled) - data_sw, sf_sw = fp4_quantize(x, gs, VEC, False, True) - assert sf_sw.shape == (128 * 160,) - assert torch.equal(data_sw, data) - - -def test_e2m1_ties_round_to_nearest_even() -> None: - """global_scale=1 with a block max of 6 makes the output scale exactly 1, exposing the tie rule.""" - gs = torch.ones(1, dtype=torch.float32, device=DEV) - mids = [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0] - # the even-code neighbour of each midpoint - expect = [0.0, 1.0, 1.0, 2.0, 2.0, 4.0, 4.0] - x = torch.zeros(2, 16, dtype=torch.bfloat16, device=DEV) - x[0, 0] = E2M1_MAX - x[1, 0] = -E2M1_MAX - for i, mid in enumerate(mids): - x[0, 1 + i] = mid - x[1, 1 + i] = -mid - data, sf = fp4_quantize(x, gs, VEC, False, False) - assert sf.tolist() == [0x38, 0x38], sf.tolist() # e4m3 0x38 == 1.0 - deq = _dequant(data, sf, 1.0) - for i, want in enumerate(expect): - assert deq[0, 1 + i].item() == want, (mids[i], deq[0, 1 + i].item()) - assert deq[1, 1 + i].item() == -want, (mids[i], deq[1, 1 + i].item()) - assert deq[0, 0].item() == E2M1_MAX and deq[1, 0].item() == -E2M1_MAX - - -def test_dequant_error_bound() -> None: - """data * sf / global_scale reconstructs the input within the nvfp4 rounding bound.""" - torch.manual_seed(4) - x = torch.randn(1024, 2560, dtype=torch.bfloat16, device=DEV) * 3.0 - gs, gsf = _global_scale(x) - data, sf = fp4_quantize(x, gs, VEC, False, False) - deq = _dequant(data, sf, gsf) - scale = sf.view(1024, 160).view(torch.float8_e4m3fn).float() - scale = scale.repeat_interleave(VEC, dim=-1) - # In the block-scaled domain s = x*gs/sf the e2m1 grid has spacing 0.5 below - # 2, 1.0 on [2,4), 2.0 on [4,6] and saturates at 6, so the rounding error is - # at most max(0.25, |s|/4). Dividing back by gs/sf gives the bound below; - # |s|/4 / (gs/sf) is exactly |x|/4. - bound = torch.maximum(0.25 * scale / gsf, x.float().abs() / 4) - err = (deq - x.float()).abs() - assert (err <= bound).all(), (err - bound).max().item() - # hard gate: the quantization is the reference, not merely inside the bound - ref_data, ref_sf, scaled = _ref(x, gsf) - torch.testing.assert_close(sf.view(1024, 160), ref_sf) - _, mismatch, total = _assert_data(data, ref_data, scaled) - assert mismatch / total < 0.005, mismatch / total - # a block whose scale rounds down lets its max element clip to 6 -- that is - # the format, and it is inside the bound above - assert (scaled.abs() > E2M1_MAX).any() - - -def test_scale_saturation() -> None: - """A global_scale calibrated on a smaller amax saturates the scale and clips every lane.""" - x = torch.zeros(1, 16, dtype=torch.bfloat16, device=DEV) - x[0, :4] = torch.tensor([1.0, -2.0, 0.5, 3.0]) - gs = torch.tensor([E4M3_MAX * E2M1_MAX / 0.001], dtype=torch.float32, device=DEV) - data, sf = fp4_quantize(x, gs, VEC, False, False) - assert sf.tolist() == [0x7E] # e4m3 0x7E == 448, the finite max - # every nonzero lane clips to +-6 regardless of its magnitude - codes = _unpack(data) - assert codes[0, :4].tolist() == [7, 15, 7, 7] - assert (codes[0, 4:] == 0).all() - deq = _dequant(data, sf, gs.item()) - clipped = torch.full((4,), E2M1_MAX * E4M3_MAX / gs.item(), device=DEV) - clipped[1] = -clipped[1] - torch.testing.assert_close(deq[0, :4], clipped) - assert (deq[0, 4:] == 0.0).all() - - -def test_scale_underflow() -> None: - """A global_scale so small that the block scale rounds to 0 zeroes the whole block.""" - x = torch.zeros(1, 16, dtype=torch.bfloat16, device=DEV) - x[0, 0] = 1.0 - x[0, 1] = -0.5 - # e4m3's smallest subnormal is 2^-9; gs*vecmax/6 = 2^-11 rounds to zero - gs = torch.tensor([E2M1_MAX * 2.0**-11], dtype=torch.float32, device=DEV) - data, sf = fp4_quantize(x, gs, VEC, False, False) - assert sf.tolist() == [0x00] - # the reciprocal of the zero scale is +inf: every lane saturates to +6, - # including the exact zeros (0 * inf = NaN, which the e2m1 cast saturates) - codes = _unpack(data) - assert codes[0, 0].item() == 7 and codes[0, 1].item() == 15 - assert (codes[0, 2:] == 7).all() - # with a zero scale the consumer reads the block back as all zeros - assert (_dequant(data, sf, gs.item()) == 0.0).all() - # one binade higher the scale is the e4m3 minimum subnormal and the block is exact - gs2 = torch.tensor([E2M1_MAX * 2.0**-9], dtype=torch.float32, device=DEV) - data2, sf2 = fp4_quantize(x, gs2, VEC, False, False) - assert sf2.tolist() == [0x01] - deq2 = _dequant(data2, sf2, gs2.item()) - assert deq2[0, 0].item() == 1.0 and deq2[0, 1].item() == -0.5 - - -def test_degenerate_global_scale() -> None: - """A zero or negative global_scale is accepted and silently unusable.""" - x = torch.zeros(1, 16, dtype=torch.bfloat16, device=DEV) - x[0, :4] = torch.tensor([1.0, -2.0, 0.5, 3.0]) - # g = 0: the block scale is 0 and every lane saturates, so a consumer - # dividing by g reads NaN - data, sf = fp4_quantize(x, torch.zeros(1, device=DEV), VEC, False, False) - assert sf.tolist() == [0x00] - assert (_unpack(data)[0] == 7).all() - # g < 0: the scale byte carries e4m3's sign bit, which every consumer reads - # as an unsigned UE4M3 magnitude; the data bytes are those of |g| - data_n, sf_n = fp4_quantize(x, -torch.ones(1, device=DEV), VEC, False, False) - data_p, sf_p = fp4_quantize(x, torch.ones(1, device=DEV), VEC, False, False) - assert sf_n.tolist() == [0xB0] and sf_p.tolist() == [0x30] # -0.5 and +0.5 - assert torch.equal(data_n, data_p) - - -def test_reciprocal_global_scale_is_silently_destructive() -> None: - """A modelopt checkpoint stores amax/(448*6); passing it unreciprocated is accepted. - - The mistake costs a factor of `(448*6/amax)^2`, which for normed activations - puts almost every block under the e4m3 scale floor -- but not all of them, so - the result is neither an error nor an obviously dead tensor. - """ - torch.manual_seed(12) - x = torch.randn(64, 7168, dtype=torch.bfloat16, device=DEV) - amax = x.abs().max().float() - stored = (amax / (E4M3_MAX * E2M1_MAX)).reshape(1) # what is on disk - data, sf = fp4_quantize(x, stored, VEC, False, False) - zero_scales = (sf == 0).float().mean().item() - assert zero_scales > 0.9, zero_scales - assert zero_scales < 1.0, zero_scales - deq = _dequant(data, sf, stored.item()) - # the blocks that survive saturate: reconstructed magnitudes above the true - # amax, next to blocks that read back as exact zeros - assert deq.abs().max().item() > amax.item() - assert (deq == 0).any() - - -def test_zero_block_and_signed_zero() -> None: - """An all-zero block emits a zero scale and zero data; -0.0 keeps its sign bit.""" - z = torch.zeros(1, 32, dtype=torch.bfloat16, device=DEV) - gs = torch.tensor([100.0], dtype=torch.float32, device=DEV) - data, sf = fp4_quantize(z, gs, VEC, False, False) - assert sf.tolist() == [0x00, 0x00] - assert (data == 0).all() - y = torch.zeros(1, 16, dtype=torch.bfloat16, device=DEV) - y[0, 0] = 6.0 - y[0, 1] = -0.0 - y[0, 2] = 0.0 - codes = _unpack(fp4_quantize(y, torch.ones(1, device=DEV), VEC, False, False)[0]) - assert codes[0, 1].item() == 8, "negative zero must keep its sign bit" - assert codes[0, 2].item() == 0 - - -def test_nonfinite_inputs() -> None: - """Non-finite lanes are accepted silently; both outcomes are destructive.""" - gs = torch.ones(1, dtype=torch.float32, device=DEV) - # +inf drives the block max to inf, so the scale saturates at 448 and every - # finite lane is scaled by 1/448 -- small lanes round to zero. - x = torch.zeros(1, 16, dtype=torch.bfloat16, device=DEV) - x[0, 0] = float("inf") - x[0, 1] = 1.0 - data, sf = fp4_quantize(x, gs, VEC, False, False) - assert sf.tolist() == [0x7E] - codes = _unpack(data) - assert codes[0, 0].item() == 7, "inf saturates to +6" - assert (codes[0, 1:] == 0).all() - # NaN is excluded from the block max, so the scale comes from the finite - # lanes; the NaN lane itself saturates to +6 and stays confined. - y = torch.zeros(1, 16, dtype=torch.bfloat16, device=DEV) - y[0, 0] = float("nan") - y[0, 1] = 1.0 - data, sf = fp4_quantize(y, gs, VEC, False, False) - assert sf.tolist() == [0x23] # e4m3 nearest to 1/6 - codes = _unpack(data) - assert codes[0, 0].item() == 7 - assert codes[0, 1].item() == 7 # 1.0 * (1/0.171875) = 5.82 -> 6 - assert (codes[0, 2:] == 0).all() - - -def test_zero_rows() -> None: - """An empty batch is accepted, not rejected: both outputs come back empty.""" - gs = torch.ones(1, dtype=torch.float32, device=DEV) - x = torch.randn(0, 2560, dtype=torch.bfloat16, device=DEV) - for swizzled in (False, True): - data, sf = fp4_quantize(x, gs, VEC, False, swizzled) - assert data.shape == (0, 1280) and data.dtype == torch.uint8 - assert sf.shape == (0,) and sf.dtype == torch.uint8 - - # Empty inputs must leave subsequent CUDA work usable. - assert torch.ones(1, device="cuda").item() == 1 - - -def test_input_not_mutated_and_deterministic() -> None: - torch.manual_seed(5) - x = torch.randn(64, 2560, dtype=torch.bfloat16, device=DEV) - gs, _ = _global_scale(x) - before = x.clone() - a = fp4_quantize(x, gs, VEC, False, True) - b = fp4_quantize(x, gs, VEC, False, True) - assert torch.equal(x, before) - assert torch.equal(a[0], b[0]) and torch.equal(a[1], b[1]) - - -def test_multi_element_global_scale() -> None: - """The op silently uses global_scale[0] for every row; the wrapper rejects that shape.""" - torch.manual_seed(6) - x = torch.randn(8, 256, dtype=torch.bfloat16, device=DEV) - per_row = torch.tensor( - [1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0], dtype=torch.float32, device=DEV - ) - # raw op: accepted, and identical to broadcasting the first element - got = torch.ops.trtllm.fp4_quantize(x, per_row, VEC, False, False)[1] - want = torch.ops.trtllm.fp4_quantize(x, per_row[:1], VEC, False, False)[1] - assert torch.equal(got, want), "multi-element global_scale is not applied per row" - try: - fp4_quantize(x, per_row, VEC, False, False) - except AssertionError: - pass - else: - raise AssertionError("wrapper must reject a multi-element global_scale") - - -# DeepSeek-R1 activation widths: model hidden (dense-MLP / shared-expert / MoE -# input), dense-MLP intermediate, shared-expert intermediate. All three are -# multiples of 512, the width half of the condition a TMA high-throughput -# variant would key on if this build had one. -R1_WIDTHS = (7168, 18432, 2048) -# The receipt's existing row range plus max_num_tokens = 8192, which is also the -# MoE path's chunk bound. -R1_ROWS = (1, 2, 7, 64, 1023, 1024, 2048, 8192) -# Row-chunk budget for the reference: its [rows, K, 7] midpoint comparison is -# the memory ceiling, and the reference is row-independent so chunking is exact. REF_CHUNK_ELEMS = 16 * 1024 * 1024 -def _assert_call_matches_ref( - x: torch.Tensor, data: torch.Tensor, sf_linear: torch.Tensor, gsf: float -) -> tuple[int, int, int]: - """Row-chunked bit-exact check of one linear-layout call. +def _assert_call_matches_ref(x, data, sf_linear, g1): + """Row-chunked check of one NVFP4 linear-layout call against `op.reference`. - Returns `(near_tie_count, mismatch_count, total)` summed over the chunks. + Scale bytes exact, data as `_assert_data`. Returns the summed counts. """ rows, k = x.shape - cols = k // VEC + cols = k // NV_VEC step = max(1, REF_CHUNK_ELEMS // k) - n_tie = n_mis = n_all = 0 + n_tie = n_miss = n_all = 0 for lo in range(0, rows, step): hi = min(rows, lo + step) - ref_data, ref_sf, scaled = _ref(x[lo:hi], gsf) - torch.testing.assert_close(sf_linear.view(rows, cols)[lo:hi], ref_sf) - tie, mismatch, total = _assert_data(data[lo:hi], ref_data, scaled) - n_tie += tie - n_mis += mismatch - n_all += total - return n_tie, n_mis, n_all - - -def test_r1_widths_both_layouts_bf16() -> None: - """DeepSeek-R1 activation widths in both scale layouts, decode- through 8192-row prefill. - - 7168 is the model hidden, taken swizzled by the dense NVFP4 GEMM and linear - by the trtllm-gen MoE runner; 18432 is the dense-MLP intermediate and 2048 - the shared-expert intermediate, both taken swizzled. Every width is a - multiple of 16 with `K/16` already a multiple of 4, so the swizzled buffer - carries row padding but no column padding. + ref_data, ref_sf = op.reference(x[lo:hi], g1, NV_VEC, False, False) + assert torch.equal(sf_linear.view(rows, cols)[lo:hi], ref_sf.view(hi - lo, cols)) + tie, miss, total = _assert_data(data[lo:hi], ref_data, op.scaled(x[lo:hi], g1, NV_VEC)) + n_tie, n_miss, n_all = n_tie + tie, n_miss + miss, n_all + total + return n_tie, n_miss, n_all + + +def _launched_kernels(call): + """Names of the CUDA kernels one warm call launches.""" + call() + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + call() + torch.cuda.synchronize() + return [ + e.name + for e in prof.events() + if e.device_type == torch.autograd.DeviceType.CUDA and e.self_device_time_total > 0 + ] + + +def test_v41_mxfp4_every_bf16_block_amax() -> None: + """Every positive finite bf16 value as a 32-block max: the UE8M0 rule at every boundary. + + Subnormals included. The max lane rotates through all 32 positions and alternates + sign, and the other 31 lanes are random in [-max, max]. Five all-zero blocks with + different sign patterns close the set. """ - worst_tie = worst_mis = 0.0 - seed = 100 - for k in R1_WIDTHS: - cols = k // VEC - for t in R1_ROWS: - torch.manual_seed(seed) - seed += 1 - x = torch.randn(t, k, dtype=torch.bfloat16, device=DEV) - gs, gsf = _global_scale(x) - data_li, sf_li = fp4_quantize(x, gs, VEC, False, False) - data_sw, sf_sw = fp4_quantize(x, gs, VEC, False, True) - assert data_li.shape == (t, k // 2) and data_li.dtype == torch.uint8 - assert data_li.is_contiguous() and sf_li.is_contiguous() - assert sf_li.shape == (t * cols,) and sf_li.dtype == torch.uint8 - assert sf_sw.shape == (_pad_up(t, 128) * _pad_up(cols, 4),) - # the same bytes out of two separate calls: layout-independent data - assert torch.equal(data_sw, data_li), "data must not depend on the layout" - tie, mismatch, total = _assert_call_matches_ref(x, data_li, sf_li, gsf) - worst_tie = max(worst_tie, tie / total) - worst_mis = max(worst_mis, mismatch / total) - idx = _swizzle_index(t, cols).reshape(-1) - torch.testing.assert_close(sf_sw[idx], sf_li) - rest = torch.ones_like(sf_sw, dtype=torch.bool) - rest[idx] = False - assert (sf_sw[rest] == 0).all(), "swizzled padding must be zero" - del x, data_li, sf_li, data_sw, sf_sw, idx, rest - # Ceilings on the check's blind spot, not correctness gates -- correctness is - # the bit-exactness above. Both quantities are set by the global scale rather - # than by the shape (see the next test), so they are ceilings with room, not - # fits: the worst of these 24 shapes was 0.60 % near-tie / 0.14 % mismatch. - assert worst_mis < 0.005, worst_mis - assert worst_tie < 0.02, worst_tie - - -def test_r1_one_kernel_at_every_width() -> None: - """One call launches exactly one CUDA kernel, and it is the same one at every shape. - - Rows (7 vs 8192) and width (multiple of 512 or not) do not select a - different kernel in this build: there is no TMA high-throughput variant on - this op's path to switch to. + torch.manual_seed(20) + pos_bits = torch.arange(1, 0x7F80, dtype=torch.int32).to(torch.int16) # 0x0001..0x7F7F + amax = pos_bits.view(torch.bfloat16).float().to(DEV) + n = amax.numel() + idx = torch.arange(n, device=DEV) + u = torch.rand(n, MX_VEC, device=DEV) * 2 - 1 + blocks = (u * amax[:, None]).to(torch.bfloat16).float() + blocks[idx, idx % MX_VEC] = torch.where(idx % 2 == 0, amax, -amax) + neg0 = -torch.zeros(MX_VEC, device=DEV) + pos0 = torch.zeros(MX_VEC, device=DEV) + lane = torch.arange(MX_VEC, device=DEV) + zeros = torch.stack( + [ + pos0, + neg0, + torch.where(lane % 2 == 0, neg0, pos0), + torch.where(lane % 2 == 1, neg0, pos0), + torch.where(lane < 16, neg0, pos0), + ] + ) + x = torch.cat([blocks, zeros]).to(torch.bfloat16).view(-1, INDEX_HEAD_DIM) # 8161 rows + blk_amax = x.float().view(-1, MX_VEC).abs().amax(-1) + assert torch.equal(blk_amax[:n], amax) + + with validating(op): + data, sf = _mx(x) + rows = x.shape[0] + assert data.shape == (rows, INDEX_HEAD_DIM // 2) and sf.shape == (rows * 4,) + ref_data, ref_sf = _mx_ref(x) + got_c = _unpack(data).view(-1, MX_VEC) + ref_c = _unpack(ref_data).view(-1, MX_VEC) + + # bit-exact wherever the block max is above 3 * 2^-126 + live = blk_amax > MX_OP_ZERO_AMAX + assert torch.equal(sf[live], ref_sf[live]) + assert torch.equal(got_c[live], ref_c[live]) + # the bf16 range reaches every UE8M0 scale from 2^-126 to 2^126 + assert int(sf[live].min()) == 1 and int(sf[live].max()) == 253 + # each boundary: max = 6 * 2^e is scale 2^e exactly, the next bf16 up is 2^(e+1) + e = torch.arange(-126, 126, device=DEV) + at = (6 * torch.exp2(e.double())).to(torch.bfloat16).view(torch.int16).long() - 1 + assert torch.equal(sf[at], (e + 127).to(torch.uint8)) + assert torch.equal(sf[at + 1], (e + 128).to(torch.uint8)) + + # At or below 3 * 2^-126 the op writes scale byte 0x00 where the reference's floor + # gives 0x01, and the block comes out as signed zeros: bf16 bit patterns + # 0x0001..0x0140, plus the five all-zero blocks. + dead = ~live + assert int(dead.sum()) == 0x140 + 5, int(dead.sum()) + sign_only = torch.signbit(x.float().view(-1, MX_VEC)).to(torch.uint8) << 3 + assert (sf[dead] == 0).all() and (ref_sf[dead] == 1).all() + assert torch.equal(got_c[dead], sign_only[dead]) + # The reference also rounds such a block to signed zeros while max <= 2^-128. + # Above that it keeps nonzero codes, so there the two value sets differ. + tiny = dead & (blk_amax <= 2.0**-128) + assert torch.equal(ref_c[tiny], sign_only[tiny]) + lost = dead & (blk_amax > 2.0**-128) + assert int(lost.sum()) == 0x140 - 0x20, int(lost.sum()) + assert ((ref_c[lost] & 7) != 0).any(-1).all() + + # Deliberately wrong control: 2^floor instead of 2^ceil. The op disagrees with it + # on every live block whose (floored) max / 6 is not a power of two: 31940 of the + # 32319. The 379 it agrees on are the 252 maxima 1.5 * 2^j and the 128 maxima + # floored to 6 * 2^-126 (one of them in both sets). + _, fl_sf = _mx_wrong(x, scale="floor") + ratio = (blk_amax.clamp_min(MX_REF_AMAX_FLOOR) * _INV6_F32).view(torch.int32) + pow2 = (ratio & 0x7FFFFF) == 0 + moved = fl_sf.view(-1) != sf + assert torch.equal(moved[live], ~pow2[live]) + assert int(moved[live].sum()) == 31940, int(moved[live].sum()) + + +def test_v41_mxfp4_e2m1_midpoints_every_scale() -> None: + """Every e2m1 midpoint and grid point, both signs, both nibbles, at every UE8M0 scale. + + For each exponent e in [-126, 126] one block holds max 6 * 2^e (3.5 * 2^126 at the + top, where 6 * 2^126 is not a finite bf16), which fixes the scale at exactly 2^e. + The block also holds +-{midpoints, grid points} * 2^e, and a copy rolled by one + lane puts every value in the other nibble. + """ + blocks, exps = [], [] + for e in range(-126, 127): + top = 6.0 if e <= 125 else 3.5 + mags = [top] + [v for v in _MX_MIDS + _MX_GRID if v <= top] + lanes = [v for m in mags for v in (m, -m)] + lanes += [0.0] * (MX_VEC - len(lanes)) + for block in (lanes, lanes[-1:] + lanes[:-1]): + blocks.append([v * 2.0**e for v in block]) + exps.append(e) + # pad to a whole number of 128-wide rows with the e = 0 pair + blocks += blocks[252:254] + exps += exps[252:254] + assert exps[252] == 0 and len(blocks) % 4 == 0 + xs = torch.tensor(blocks, dtype=torch.float64) + x = xs.to(torch.bfloat16) + assert torch.equal(x.double(), xs), "midpoint inputs must be exact in bf16" + x = x.to(DEV).view(-1, INDEX_HEAD_DIM) + exps = torch.tensor(exps, device=DEV) + + with validating(op): + data, sf = _mx(x) + ref_data, ref_sf = _mx_ref(x) + assert torch.equal(sf, (exps + 127).to(torch.uint8)) + assert torch.equal(sf, ref_sf) + assert torch.equal(data, ref_data) + + got = _unpack(data).view(-1, MX_VEC) + scaled = x.double().view(-1, MX_VEC) / torch.exp2(exps.double()).unsqueeze(-1) # exact + for mid, code in zip(_MX_MIDS, _MX_MID_CODES): + at = scaled.abs() == mid + assert int(at.sum()) >= 4 * 252, (mid, int(at.sum())) + want = code | (torch.signbit(scaled[at]).to(torch.uint8) << 3) + assert torch.equal(got[at], want), mid + + # Deliberately wrong control: ties toward zero. It differs from the op on exactly + # the midpoints that round-to-nearest-even takes away from zero (0.75, 1.75, 3.5), + # and nowhere else. + tz_data, _ = _mx_wrong(x, ties="zero") + tz = _unpack(tz_data).view(-1, MX_VEC) + away = (scaled.abs() == 0.75) | (scaled.abs() == 1.75) | (scaled.abs() == 3.5) + assert torch.equal(tz != got, away) + + +MX_ROWS = (1, 2, 3, 7, 31, 32, 33, 127, 128, 129, 1023, 1024, 4096, 8192, 65536, 262143, 262144) + + +def _mx_check_call(x): + """One MXFP4 call: bit-exact vs the reference, in the bytes indexer_k_cache_scatter takes.""" + with validating(op): + data, sf = _mx(x) + rows = x.numel() // INDEX_HEAD_DIM + assert data.shape == (*x.shape[:-1], INDEX_HEAD_DIM // 2) and data.dtype == torch.uint8 + assert sf.shape == (rows * 4,) and sf.dtype == torch.uint8 + assert data.is_contiguous() and sf.is_contiguous() + ref_data, ref_sf = _mx_ref(x) + assert torch.equal(data, ref_data) + assert torch.equal(sf, ref_sf) + # The scatter takes k_data = data as is and k_scale = sf.view(int32).view(T, 1). + # That is a free, 4-byte-aligned view. Word t holds block g's exponent in its byte g + # (little-endian). + word = sf.view(torch.int32).view(rows, 1) + assert word.data_ptr() == sf.data_ptr() and sf.data_ptr() % 4 == 0 + r = ref_sf.view(rows, 4).to(torch.int64) + want = r[:, 0] | (r[:, 1] << 8) | (r[:, 2] << 16) | (r[:, 3] << 24) + assert torch.equal(word.view(-1).to(torch.int64) & 0xFFFFFFFF, want) + + +def test_v41_mxfp4_index_kq_rows_bit_exact() -> None: + """Index K rows 1..262144, index Q as [8192, 32, 128], and random bf16 bit patterns.""" + for t in MX_ROWS: + torch.manual_seed(1000 + t) + _mx_check_call(_block_mag_rows(t)) + # index Q as the model forms it: 8192 tokens x 32 heads = 262144 rows, rank 3 + torch.manual_seed(2000) + q = torch.randn(8192, INDEX_N_HEADS, INDEX_HEAD_DIM, device=DEV).to(torch.bfloat16) + _mx_check_call(q) + # Random finite bf16 bit patterns: blocks spanning the whole exponent range, + # subnormal lanes included. + torch.manual_seed(2001) + bits = torch.randint(-32768, 32768, (65536, INDEX_HEAD_DIM), dtype=torch.int32, device=DEV) + xb = bits.to(torch.int16).view(torch.bfloat16) + xb = torch.where(torch.isfinite(xb), xb, torch.zeros_like(xb)) + assert (xb.float().view(-1, MX_VEC).abs().amax(-1) > MX_OP_ZERO_AMAX).all() + _mx_check_call(xb) + + +def test_v41_mxfp4_ignores_global_scale_one_kernel() -> None: + """MXFP4 mode never reads global_scale, and one call is one 32-block UE8M0 kernel launch.""" + torch.manual_seed(21) + x = torch.randn(64, INDEX_HEAD_DIM, dtype=torch.bfloat16, device=DEV) + base = _mx(x) + for g in (1.0, 2.0, 1e-3, 448.0 * 6.0 / x.abs().max().item()): + gs = torch.tensor([g], dtype=torch.float32, device=DEV) + got = op(x, gs, MX_VEC, True, False) + assert torch.equal(got[0], base[0]) and torch.equal(got[1], base[1]), g + for t in (1, 262144): + xx = torch.randn(t, INDEX_HEAD_DIM, dtype=torch.bfloat16, device=DEV) + launched = _launched_kernels(lambda: _mx(xx)) + assert len(launched) == 1, (t, launched) + assert "quantize_with_block_size" in launched[0], launched[0] + assert "__nv_bfloat16, 32, true" in launched[0], launched[0] + + +# Every decode row count 1..32, dep4's gathered MoE rows (4 ranks x 16 / 24 / 32), and +# prefill chunks up to the MoE chunk size 8192. +V41_ROWS = tuple(range(1, 33)) + (64, 96, 128, 1023, 1024, 2048, 4096, 8192) +V41_GRAPH_ROWS = tuple(range(1, 33)) + (64, 96, 128, 1024, 8192) + + +def test_v41_moe_input_k5120_every_row_count() -> None: + """NVFP4 at the V4.1 hidden width, linear and swizzled, rows 1 through 8192. + + `g` is a static `1 / input_scale` calibrated on a separate batch (the serving + regime), so it is not the canonical `448*6/amax` of the tensor quantized; the row + count and the global scale are the only things that change between calls. The + swizzled buffer is checked against the reference's swizzle (offsets and zero + padding), and its data against the linear call's. + """ + cols = V41_HIDDEN // NV_VEC # 320, a multiple of 4: row padding only when swizzled + torch.manual_seed(5120) + g1 = _static_g1(torch.randn(4096, V41_HIDDEN, device=DEV).to(torch.bfloat16)) + n_tie = n_miss = n_all = 0 + for t in V41_ROWS: + x = torch.randn(t, V41_HIDDEN, dtype=torch.bfloat16, device=DEV) + with validating(op): + data, sf = op(x, g1, NV_VEC, False, False) + assert data.shape == (t, V41_HIDDEN // 2) and sf.shape == (t * cols,) + tie, miss, total = _assert_call_matches_ref(x, data, sf, g1) + n_tie, n_miss, n_all = n_tie + tie, n_miss + miss, n_all + total + with validating(op): + data_sw, sf_sw = op(x, g1, NV_VEC, False, True) + assert torch.equal(data_sw, data), "data must not depend on the layout" + assert sf_sw.shape == (-(-t // 128) * 128 * cols,) + assert torch.equal(sf_sw, op.reference(x, g1, NV_VEC, False, True)[1]) + # Ceilings on the check's blind spot; correctness is the bit-exactness above. + # Measured over these 88 M elements: 0.20 % near-ties, 0 disagreements. + assert n_miss / n_all < 0.005 and n_tie / n_all < 0.02, (n_tie, n_miss, n_all) + + +def test_v41_cuda_graph_replay_is_bitwise_eager() -> None: + """The V4.1 target's three calls captured in one CUDA graph, replayed on new data. + + Per row count: index K `[n, 128]` and index Q `[n, 32, 128]` (MXFP4, linear) and + the MoE input `[n, 5120]` (NVFP4, linear, static `[1]` global scale) are captured + once, after one warm-up call on a side stream (what torch.cuda.graph callers do). + Each replay then quantizes whatever the static input buffers hold: three different + data sets, each bitwise equal to an eager call on the same values and matched + against the reference (MXFP4 bit-exact, NVFP4 bit-exact outside the near-tie + window). The global scale is read from device memory at replay: scaling it in + place changes the replayed bytes exactly as it changes an eager call's. + """ + torch.manual_seed(4141) + g1 = _static_g1(torch.randn(4096, V41_HIDDEN, device=DEV).to(torch.bfloat16)) + for n in V41_GRAPH_ROWS: + k_in = torch.empty(n, INDEX_HEAD_DIM, dtype=torch.bfloat16, device=DEV) + q_in = torch.empty(n, INDEX_N_HEADS, INDEX_HEAD_DIM, dtype=torch.bfloat16, device=DEV) + h_in = torch.empty(n, V41_HIDDEN, dtype=torch.bfloat16, device=DEV) + + def fill(seed: int) -> None: + gen = torch.Generator(device=DEV).manual_seed(seed) + k_in.copy_(_block_mag_rows(n, gen)) + q_in.copy_(torch.randn(n, INDEX_N_HEADS, INDEX_HEAD_DIM, device=DEV, generator=gen)) + h_in.copy_(torch.randn(n, V41_HIDDEN, device=DEV, generator=gen)) + + def body(): + kd, ks = _mx(k_in) + qd, qs = _mx(q_in) + hd, hs = op(h_in, g1, NV_VEC, False, False) + return kd, ks, qd, qs, hd, hs + + fill(10 * n) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + body() + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = body() + for seed in (10 * n, 10 * n + 1, 10 * n + 2): + fill(seed) + graph.replay() + torch.cuda.synchronize() + eager = body() + for got, want in zip(outs, eager): + assert torch.equal(got, want), (n, seed) + kd, ks, qd, qs, hd, hs = outs + ref_d, ref_s = _mx_ref(k_in) + assert torch.equal(kd, ref_d) and torch.equal(ks, ref_s), (n, seed) + ref_d, ref_s = _mx_ref(q_in) + assert torch.equal(qd, ref_d) and torch.equal(qs, ref_s), (n, seed) + _assert_call_matches_ref(h_in, hd, hs, g1) + old = g1.clone() + before = outs[5].clone() + g1.mul_(1.37) + graph.replay() + torch.cuda.synchronize() + eager = body() + g1.copy_(old) + assert torch.equal(outs[4], eager[4]) and torch.equal(outs[5], eager[5]), n + assert not torch.equal(outs[5], before), ( + "the in-place scale change did not reach the replay" + ) + del graph, outs + + +# ─── CUDA-graph capture over Gset: the target's whole step, DSpark off and on ── +# Graph batches: Gset, every G an engine of the envelope captures. That is every integer +# 1..32 (padding off at max_batch_size up to 32, and the head of the padded list) plus +# TRT-LLM's padded default list at max_batch_size 256 (1, 2, 4, 8, 16..128 step 8, 192, +# 256): 46 values. The padding-off list at 256 (1..32, 64, 128, 256) is a subset. A +# captured step is generation-only with next_n = 1 (DSpark off) or 6 (DSpark on: 1 + 5 +# drafts), so it holds T = next_n * G token rows, and the target quantizes: +# - index K of the ratio-2 KV sources at n = G * ceil(next_n / 2) +# (`num_total_compressed_tokens[2]`): G or 3G; +# - index K of the ratio-1 KV source at n = T: G or 6G; +# - index Q `[T, 32, 128]` in every index layer: G or 6G; +# - the MoE input `[T, 5120]` on tep4 / tp4: G or 6G; on dep4 the rows the 4-rank +# all-gather hands the router, `[4T, 5120]`: 4G or 24G, at most 6144 <= 8192, so +# one chunk. +# All five calls of one step are captured together, as in the target. +_GSET = tuple(range(1, 33)) + (40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256) +assert len(_GSET) == 46 +DEP4_RANKS = 4 + + +def _step_rows(next_n, G): + """Row counts of one captured step's five calls: index K ratio 2, index K ratio 1, + index Q, MoE input (tep4 / tp4), MoE input (dep4, gathered).""" + t = next_n * G + return G * -(-next_n // 2), t, t, t, DEP4_RANKS * t + + +@pytest.mark.parametrize("next_n", (1, 6), ids=("next_n1", "next_n6")) +def test_v41_cuda_graph_step_over_gset(next_n) -> None: + """The five calls of one captured step at every graph batch G of Gset, replayed on new data. + + Captured the way the engine does it: one graph per G, all in one shared memory pool, + largest G first, each after an eager call at its shapes on a side stream. Then three + interleaved replay rounds over every G, each on new data written in place: every + output bitwise equal to an eager call on the same values, MXFP4 bit-exact against + the reference, NVFP4 bit-exact outside the near-tie window. Two controls show the + harness is not blind: every replay rewrote every data buffer (none equals the + previous round's bytes), and a reference built on the inputs rolled by one row + matches none of the op's rows. The global scale is read from device memory at + replay, for both MoE calls of every G. """ - torch.manual_seed(8) - for t, k, dtype in ( - (7, 7168, torch.bfloat16), - (1024, 7168, torch.bfloat16), - (8192, 7168, torch.bfloat16), - (8192, 18432, torch.bfloat16), - (8192, 2048, torch.bfloat16), - (1024, 112, torch.bfloat16), - (8192, 7168, torch.float16), - ): - x = torch.randn(t, k, dtype=dtype, device=DEV) - gs, _ = _global_scale(x) - for swizzled in (False, True): - fp4_quantize(x, gs, VEC, False, swizzled) # warm up the launch path + torch.manual_seed(4242 + next_n) + g1 = _static_g1(torch.randn(4096, V41_HIDDEN, device=DEV).to(torch.bfloat16)) + pool = torch.cuda.graph_pool_handle() + + def alloc(G): + k2, k1, q, h, hg = _step_rows(next_n, G) + bf = torch.bfloat16 + return ( + torch.empty(k2, INDEX_HEAD_DIM, dtype=bf, device=DEV), + torch.empty(k1, INDEX_HEAD_DIM, dtype=bf, device=DEV), + torch.empty(q, INDEX_N_HEADS, INDEX_HEAD_DIM, dtype=bf, device=DEV), + torch.empty(h, V41_HIDDEN, dtype=bf, device=DEV), + torch.empty(hg, V41_HIDDEN, dtype=bf, device=DEV), + ) + + def fill(ins, seed): + gen = torch.Generator(device=DEV).manual_seed(seed) + for k_in in ins[:2]: + k_in.copy_(_block_mag_rows(k_in.shape[0], gen)) + for x in ins[2:]: + x.copy_(torch.randn(x.shape, device=DEV, generator=gen)) + + def body(ins): + outs = [] + for x in ins[:3]: + outs += _mx(x) + for x in ins[3:]: + outs += op(x, g1, NV_VEC, False, False) + return tuple(outs) + + steps = [] + for G in sorted(_GSET, reverse=True): + ins = alloc(G) + fill(ins, 1000 * G) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + body(ins) # the eager call at this step's shapes, before its capture + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + outs = body(ins) + steps.append((G, ins, graph, outs)) + steps.reverse() + + n_tie = n_miss = n_all = 0 + prev = {} + for rnd in range(1, 4): + for G, ins, graph, outs in steps: + fill(ins, 1000 * G + rnd) + graph.replay() torch.cuda.synchronize() - with profile(activities=[ProfilerActivity.CUDA]) as prof: - fp4_quantize(x, gs, VEC, False, swizzled) - torch.cuda.synchronize() - launched = [ - e.name - for e in prof.events() - if e.device_type == torch.autograd.DeviceType.CUDA and e.self_device_time_total > 0 - ] - assert len(launched) == 1, (t, k, dtype, swizzled, launched) - assert "quantize_with_block_size" in launched[0], launched[0] - assert "tma" not in launched[0].lower(), launched[0] - del x - - -def test_r1_ties_track_the_global_scale_not_the_shape() -> None: - """The near-tie population is a property of the global scale, not of M or K. - - One sample is reshaped to all three R1 widths. Scale blocks are 16 - *contiguous* elements, so the row width never moves a block boundary and the - only thing a reshape can change is `amax`; holding that fixed, all three - widths must agree on the exact count. Nudging `global_scale` off the - canonical `448*6/amax` then removes the population altogether, because a - near-tie needs `sf` to come out exactly `g*vecmax/6` -- which makes - `out_scale` exactly `6/vecmax` and puts `6*x/vecmax` on an e2m1 midpoint. + eager = body(ins) + for got, want in zip(outs, eager): + assert torch.equal(got, want), (next_n, G, rnd) + if G in prev: + for i, old_data in enumerate(prev[G]): + assert not torch.equal(outs[2 * i], old_data), ("replay left stale data", G, i) + prev[G] = tuple(o.clone() for o in outs[::2]) + for j, x in enumerate(ins[:3]): + ref_d, ref_s = _mx_ref(x) + assert torch.equal(outs[2 * j], ref_d) and torch.equal(outs[2 * j + 1], ref_s) + for j, x in enumerate(ins[3:]): + d, s = outs[6 + 2 * j], outs[7 + 2 * j] + tie, miss, total = _assert_call_matches_ref(x, d, s, g1) + n_tie, n_miss, n_all = n_tie + tie, n_miss + miss, n_all + total + # Row-addressing control: a reference on the inputs rolled by one row matches none of + # the op's rows (the checks above can see a row slip). + for G, ins, graph, outs in steps: + for j, x in enumerate(ins[:3]): + rows = x.numel() // INDEX_HEAD_DIM + if rows == 1: + continue + ref_d, _ = _mx_ref(torch.roll(x.view(rows, INDEX_HEAD_DIM), 1, 0)) + same = (outs[2 * j].view(rows, -1) == ref_d).all(-1) + assert int(same.sum()) == 0, ("row-shift control matched", G, j, int(same.sum())) + for j, x in enumerate(ins[3:]): + if x.shape[0] == 1: + continue + ref_d, _ = op.reference(torch.roll(x, 1, 0), g1, NV_VEC, False, False) + same = (outs[6 + 2 * j] == ref_d).all(-1) + assert int(same.sum()) == 0, ("row-shift control matched", G, j, int(same.sum())) + # The global scale is read from device memory at replay, for both MoE calls. + old = g1.clone() + g1.mul_(1.37) + for G, ins, graph, outs in steps: + before = (outs[7].clone(), outs[9].clone()) + graph.replay() + torch.cuda.synchronize() + eager = body(ins) + for i in range(6, 10): + assert torch.equal(outs[i], eager[i]), (G, i) + assert not torch.equal(outs[7], before[0]) and not torch.equal(outs[9], before[1]), ( + "the in-place scale change did not reach the replay", + G, + ) + g1.copy_(old) + # Measured: next_n 1, 152,371,200 NVFP4 elements, 81,261 near a tie, 12,532 resolved + # the other way; next_n 6, 914,227,200 / 484,652 / 74,032. + assert n_miss / n_all < 0.005 and n_tie / n_all < 0.02, (n_tie, n_miss, n_all) + del steps, prev + + +# ─── The row axis, 1..N_max, for the target's three call forms ───────────────── +# Rows only size the launch grid of this op, so the row axis is certified as a range: +# (a) one call at N_max against the reference, (b) every row count r in 1..N_max bitwise +# equal to the first r rows of that call (data bytes and linear scale bytes), and (c) the +# same single compiled kernel launched at every probed row count. N_max per form: +# - index K `[n, 128]`: n = T at the ratio-1 KV source, at most max_num_tokens +# (16384); certified to 32768. +# - index Q `[T, 32, 128]`: T in every index layer, including the re-index layers of +# the decoder bounded-replay pass, whose rows (the finishing requests' 128-row +# tails plus the generation rows) reach 128 x 256 = 32768 where the pass is not +# sliced to max_num_tokens; certified to 32768. +# - MoE input `[n, 5120]`: the target splits the MoE rows into chunks of at most 8192, +# so no call sees more; certified to 8192. The chunks are row views at nonzero +# offsets: test_v41_moe_input_chunk_views_of_gathered_rows. +V41_NMAX_MX = 32768 +V41_NMAX_MOE = 8192 +# At K = 128 the launcher's grid is min(rows, SMs x 128) = 19456 rows on 152 SMs; at +# K = 5120 it is min(rows, SMs x 4) = 608. Index Q presents 32 rows per token, so its grid +# saturates above 608 tokens; index K's above 19456. +_KERNEL_PROBE_ROWS = ( + 1, + 2, + 7, + 31, + 128, + 129, + 600, + 607, + 608, + 609, + 1024, + 4097, + 8191, + 8192, + 16383, + 16384, + 16385, + 19455, + 19456, + 19457, + 32767, + 32768, +) + + +def test_v41_rows_1_to_nmax_prefix_bitwise() -> None: + """Index K `[r, 128]` and index Q `[r, 32, 128]` (MXFP4) at every r in 1..32768, and + the MoE input `[r, 5120]` (NVFP4, linear, static `[1]` g1) at every r in 1..8192. + + Each prefix call reads `x[:r]` of one N_max-row input (a contiguous view, as the + target's inputs are) and must reproduce the N_max call's first r rows bit for bit: + `data[:r]` and the first `r * scale_bytes_per_row` linear scale bytes. The N_max + call itself is checked against the reference. """ - # 129024 = lcm(7168, 18432, 2048), so one buffer reshapes to all three - n = 129024 * 128 - torch.manual_seed(11) - flat = torch.randn(n, dtype=torch.bfloat16, device=DEV) - gs, gsf = _global_scale(flat) - counts = [] - for k in R1_WIDTHS: - x = flat.view(n // k, k) - data, sf = fp4_quantize(x, gs, VEC, False, False) - tie, mismatch, total = _assert_call_matches_ref(x, data, sf, gsf) - assert total == n - counts.append((tie, mismatch)) - del data, sf - assert counts[0] == counts[1] == counts[2], counts - assert counts[0][0] > 0, counts - # a global scale 0.1 % off the canonical one: no block keeps an exact sf, so - # the kernel is bit-exact against the reference with no carve-out at all - x = flat.view(n // 7168, 7168) - off = gs * 1.001 - data, sf = fp4_quantize(x, off, VEC, False, False) - tie, mismatch, _ = _assert_call_matches_ref(x, data, sf, off.item()) - assert (tie, mismatch) == (0, 0), (tie, mismatch) - # a power-of-two multiple instead scales sf exactly and leaves out_scale - # unchanged, so the population survives: it is g's mantissa that matters - pow2 = gs * 2.0 - data, sf = fp4_quantize(x, pow2, VEC, False, False) - tie, _, _ = _assert_call_matches_ref(x, data, sf, pow2.item()) - assert tie > 0, tie - - -def test_r1_dequant_error_bound() -> None: - """The nvfp4 reconstruction bound holds at the widest R1 width. - - Evaluated in float64: every operand (e2m1 value, e4m3 scale, bf16 input, - fp32 global scale) is exact there, so the comparison is the mathematical - bound rather than a tolerance. In fp32 -- what a consumer actually computes - -- an element sitting exactly on the bound lands up to half an fp32 ulp - outside it; measured 3.0e-8 on this shape. + torch.manual_seed(820) + n = V41_NMAX_MX + k = _block_mag_rows(n) + qmag = torch.exp2(torch.randint(-20, 21, (n, INDEX_N_HEADS, 4, 1), device=DEV).float()) + q = (torch.randn(n, INDEX_N_HEADS, 4, MX_VEC, device=DEV) * qmag).to(torch.bfloat16) + q = q.view(n, INDEX_N_HEADS, INDEX_HEAD_DIM) + del qmag + h = torch.randn(V41_NMAX_MOE, V41_HIDDEN, device=DEV).to(torch.bfloat16) + g1 = _static_g1(torch.randn(4096, V41_HIDDEN, device=DEV).to(torch.bfloat16)) + + def nv(x): + return op(x, g1, NV_VEC, False, False) + + # (name, N_max input, call, linear scale bytes per token row) + forms = ( + ("index K", k, _mx, INDEX_HEAD_DIM // MX_VEC), + ("index Q", q, _mx, INDEX_N_HEADS * INDEX_HEAD_DIM // MX_VEC), + ("MoE input", h, nv, V41_HIDDEN // NV_VEC), + ) + for name, x, call, sf_row in forms: + n = x.shape[0] + # (a) the N_max call against the reference + with validating(op): + fd, fs = call(x) + assert fs.numel() == n * sf_row + if call is _mx: + ref_d, ref_s = _mx_ref(x) + assert torch.equal(fd, ref_d) and torch.equal(fs, ref_s), name + del ref_d, ref_s + else: + tie, miss, total = _assert_call_matches_ref(x, fd, fs, g1) + assert miss / total < 0.005 and tie / total < 0.02, (tie, miss, total) + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + times = [] + for _ in range(5): + start.record() + call(x) + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + # Every r is run only because the N_max call is this cheap; typical best-of-5 is + # about 40 us (index K), 260 us (index Q) and 55 us (MoE input). + assert min(times) < 1.0, (name, times) + # (b) every prefix, one launch and one device-side compare each (no host sync) + bad = torch.zeros(n, dtype=torch.bool, device=DEV) + for r in range(1, n + 1): + d, s = call(x[:r]) + bad[r - 1] = (d != fd[:r]).any() | (s != fs[: r * sf_row]).any() + n_bad = int(bad.sum()) + assert n_bad == 0, (name, n_bad, bad.nonzero()[:8].view(-1).tolist()) + # Control: the compare sees a one-row slip at the edges and in the middle. + for r in (1, 2, n // 2, n - 1): + d, _ = call(x[:r]) + assert (d != fd[1 : r + 1]).any(), (name, r) + # (c) one compiled variant: the same single kernel at every probed row count + names = set() + for r in (r for r in _KERNEL_PROBE_ROWS if r <= n): + launched = _launched_kernels(lambda: call(x[:r])) + assert len(launched) == 1 and "quantize_with_block_size" in launched[0], ( + name, + r, + launched, + ) + names.add(launched[0]) + assert len(names) == 1, (name, names) + del fd, fs, bad + + +# ─── The MoE input as the target hands it over: 8192-row chunks at nonzero offsets ─ +# The target splits the MoE rows into chunks of at most 8192, and each chunk reaches +# this op as a contiguous row view of the whole `[rows, 5120]` input, at a storage +# offset of a multiple of 8192 rows. The rows are this rank's T (tep4 / tp4, up to +# 16384; the decoder pass, up to 32768) or, on dep4, the 4-rank all-gather's 4 x the +# per-rank row count (up to 65536). Totals: 65536 (the dep4 maximum), and the gathered +# counts 65512, 41804, 41764 and 38096 (4 x 16378 / 10451 / 10441 / 9524, eager steps of +# a batch-256 dep4 engine), whose last chunks of 8168, 844, 804 and 5328 rows the same +# engine shows as router row counts. +_MOE_CHUNK = 8192 +_GATHERED_TOTALS = (65536, 65512, 41804, 41764, 38096) + + +def test_v41_moe_input_chunk_views_of_gathered_rows() -> None: + """Every 8192-row chunk view of a 65536-row MoE input, and the observed last chunks. + + Each chunk is passed as the row view the target's split makes (contiguous, at a + nonzero storage offset) and must give the same bytes as the same rows at offset 0 + (a contiguous copy), and match the NVFP4 reference (bit-exact outside the near-tie + window) with the layer's static `[1]` g1, linear layout. """ - torch.manual_seed(10) - x = torch.randn(1024, 18432, dtype=torch.bfloat16, device=DEV) * 3.0 - gs, gsf = _global_scale(x) - data, sf = fp4_quantize(x, gs, VEC, False, False) - codes = _unpack(data) - v = E2M1_VALUES[(codes & 7).long()].double() - v = torch.where((codes & 8).bool(), -v, v) - scale = sf.view(1024, 1152).view(torch.float8_e4m3fn).double() - scale = scale.repeat_interleave(VEC, dim=-1) - # same bound as test_dequant_error_bound: the e2m1 subnormal half-step - # carried back through the block scale, or the half-step of the coarsest - # e2m1 binade (which also covers a block max clipping to 6) - bound = torch.maximum(0.25 * scale / gsf, x.double().abs() / 4) - err = (v * scale / gsf - x.double()).abs() - assert (err <= bound).all(), (err - bound).max().item() - - -def test_rejects_unsupported() -> None: - """Domains the contract declares unsupported must be rejected, not silently wrong.""" - x = torch.randn(8, 256, dtype=torch.bfloat16, device=DEV) - gs = torch.ones(1, dtype=torch.float32, device=DEV) - cases = { - "no global_scale on the nvfp4 path": lambda: fp4_quantize(x, None, 16, False, True), - "sf_vec_size 32 without ue8m0": lambda: fp4_quantize(x, gs, 32, False, True), - "sf_vec_size 16 with ue8m0": lambda: fp4_quantize(x, gs, 16, True, True), - "sf_vec_size 8": lambda: fp4_quantize(x, gs, 8, False, True), - "sf_vec_size 0": lambda: fp4_quantize(x, gs, 0, False, True), - "k not a multiple of 16": lambda: fp4_quantize( - torch.randn(4, 24, dtype=torch.bfloat16, device=DEV), gs, 16, False, True - ), - "fp32 input": lambda: fp4_quantize(x.float(), gs, 16, False, True), - "1-D input": lambda: fp4_quantize(x.reshape(-1), gs, 16, False, True), - "fp16 global_scale": lambda: fp4_quantize(x, gs.half(), 16, False, True), - "cpu global_scale": lambda: fp4_quantize(x, gs.cpu(), 16, False, True), - "cpu input": lambda: fp4_quantize(x.cpu(), gs.cpu(), 16, False, True), - "non-contiguous column slice": lambda: fp4_quantize( - torch.randn(8, 512, dtype=torch.bfloat16, device=DEV)[:, :256], - gs, - 16, - False, - True, - ), - "non-contiguous transpose": lambda: fp4_quantize( - torch.randn(256, 8, dtype=torch.bfloat16, device=DEV).t(), - gs, - 16, - False, - True, - ), - } - for name, fn in cases.items(): - try: - fn() - except (RuntimeError, NotImplementedError): - continue - raise AssertionError(f"{name}: expected a raise, got a result") + torch.manual_seed(6553) + g1 = _static_g1(torch.randn(4096, V41_HIDDEN, device=DEV).to(torch.bfloat16)) + big = torch.randn(max(_GATHERED_TOTALS), V41_HIDDEN, device=DEV).to(torch.bfloat16) + checked = set() + n_tie = n_miss = n_all = 0 + for total in _GATHERED_TOTALS: + full, rest = divmod(total, _MOE_CHUNK) + sizes = [_MOE_CHUNK] * full + ([rest] if rest else []) + for c, chunk in enumerate(torch.split(big[:total], sizes, 0)): + off = c * _MOE_CHUNK + assert chunk.is_contiguous() and chunk.storage_offset() == off * V41_HIDDEN + with validating(op): + d, s = op(chunk, g1, NV_VEC, False, False) + d0, s0 = op(chunk.clone(), g1, NV_VEC, False, False) + assert torch.equal(d, d0) and torch.equal(s, s0), (total, off) + if (off, chunk.shape[0]) in checked: + continue + checked.add((off, chunk.shape[0])) + tie, miss, tot = _assert_call_matches_ref(chunk, d, s, g1) + n_tie, n_miss, n_all = n_tie + tie, n_miss + miss, n_all + tot + assert len(checked) == 12, sorted(checked) + # Measured: 413,081,600 elements, 1,177,519 near a tie, 396,757 resolved the other way. + assert n_miss / n_all < 0.005 and n_tie / n_all < 0.02, (n_tie, n_miss, n_all) diff --git a/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py index 3bcc679a73ad..47d7085d17d0 100644 --- a/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py +++ b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py @@ -2,45 +2,39 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the mxfp8_quantize catalog entry.""" +import pytest import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.mxfp8_quantize import ( - mxfp8_quantize, + mxfp8_quantize as op, ) +__extra_import_path__ = [".."] +from _validating import validating # noqa: E402 — needs the path declared above + assert torch.cuda.is_available(), "mxfp8_quantize requires a CUDA device" -E4M3_MAX = 448.0 + BLOCK = 32 -def _ref(x: torch.Tensor, padded_k: int) -> tuple[torch.Tensor, torch.Tensor]: - """Independent native-torch reference: (e4m3 data [M, padded_k], UE8M0 bytes [M, padded_k/32]). +def _build(spec, seed): + """A cell's input; `data="block_scaled"` is `_activations` (block maxima over 2^-24..2^24), else randn.""" + if spec.get("data") == "block_scaled": + m, k = spec["shape"] + assert spec["dtype"] is torch.bfloat16, spec + return _activations(m, k, seed) + g = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(spec["shape"], generator=g, device="cuda").to(spec["dtype"]) + - Per 32-element block along the last dim (rows = product of the leading dims): - amax = max|x|; the block scale is the smallest power of two >= amax/448, - encoded as the E8M0 byte `exp2_of_scale + 127`; the data is the fp32 input - divided by that scale, rounded to e4m3. Columns beyond the input's K (up to - `padded_k`) are treated as zero. +def _ref(x: torch.Tensor, alignment: int = 32) -> tuple[torch.Tensor, torch.Tensor]: + """`op.reference` with its scale bytes as the [M, padded_k / 32] matrix: the linear layout, reshaped. - Exact for every block with amax > 448 * 2^-127; below that the kernel takes - a flush-to-zero path that `test_zero_and_denormal_blocks` pins directly. + The swizzle checks below index into it per (row, block). """ - k = x.shape[-1] - m = x.numel() // k - xp = torch.zeros(m, padded_k, dtype=torch.float32, device=x.device) - xp[:, :k] = x.reshape(m, k).float() - blk = xp.view(m, padded_k // BLOCK, BLOCK) - amax = blk.abs().amax(dim=-1) - # E8M0 of amax/448, rounded toward +inf: take the fp32 biased exponent and - # bump it whenever any mantissa bit is set (i.e. the value is not already a - # power of two). amax is exact in fp32 for bf16/fp16 inputs, so this is the - # exact round-up. - u = (amax / E4M3_MAX).view(torch.int32) - byte = (((u >> 23) & 0xFF) + ((u & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) - scale = torch.exp2(byte.to(torch.float32) - 127.0).unsqueeze(-1) - data = (blk / scale).view(m, padded_k).to(torch.float8_e4m3fn) - return data, byte.to(torch.uint8) + data, sf = op.reference(x, False, alignment) + return data, sf.view(data.shape[0], -1) def _swizzle_index(rows: int, cols: int, device: torch.device) -> torch.Tensor: @@ -61,223 +55,403 @@ def _pad_up(x: int, m: int) -> int: return (x + m - 1) // m * m -def test_linear_layout_bf16() -> None: - """Bit-exact match against the reference, decode- through prefill-sized rows.""" - torch.manual_seed(0) - # gpt-oss-120b MoE hidden (2880) and its 512-padded width (3072) - for t in (1, 2, 7, 64, 1024, 8192): - for k, alignment in ((2880, 32), (2880, 512), (3072, 512), (3072, 128)): - x = torch.randn(t, k, dtype=torch.bfloat16, device="cuda") - data, sf = mxfp8_quantize(x, False, alignment) - padded_k = _pad_up(k, alignment) - ref_data, ref_sf = _ref(x, padded_k) - assert data.shape == (t, padded_k), (t, k, alignment, data.shape) - assert data.dtype == torch.float8_e4m3fn and data.is_contiguous() - assert sf.shape == (t * padded_k // BLOCK,) and sf.dtype == torch.uint8 - assert torch.equal(data.view(torch.uint8), ref_data.view(torch.uint8)) - assert torch.equal(sf.view(t, padded_k // BLOCK), ref_sf) - - -def test_swizzled_layout_bf16() -> None: - """Swizzled scales carry the same bytes at the 128x4 offsets; padding is zero.""" - torch.manual_seed(1) - for t, k, alignment in ( - (1, 2880, 512), - (3, 128, 32), - (129, 96, 32), - (200, 3072, 512), - ): - x = torch.randn(t, k, dtype=torch.bfloat16, device="cuda") - data_sw, sf_sw = mxfp8_quantize(x, True, alignment) - padded_k = _pad_up(k, alignment) - cols = padded_k // BLOCK - ref_data, ref_sf = _ref(x, padded_k) - assert sf_sw.shape == (_pad_up(t, 128) * _pad_up(cols, 4),) - # data is layout-independent - assert torch.equal(data_sw.view(torch.uint8), ref_data.view(torch.uint8)) - idx = _swizzle_index(t, cols, x.device).reshape(-1) - assert torch.equal(sf_sw[idx].view(t, cols), ref_sf) - # every offset not addressed by a real (row, col) is row/column padding: zero - rest = torch.ones_like(sf_sw, dtype=torch.bool) - rest[idx] = False - assert (sf_sw[rest] == 0).all() - - -def test_alignment_padding() -> None: - """alignment pads K with zeros: zero data bytes, zero scale bytes, valid part unchanged.""" - torch.manual_seed(2) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + x = _build(cell.spec, seed=abs(hash(cell.why)) % 2**31) + swizzled, alignment = cell.spec["swizzled"], cell.spec["alignment"] + with validating(op): + data, scales = op(x, swizzled, alignment) + ref_data, ref_scales = op.reference(x, swizzled, alignment) + op.compare(data, ref_data) + op.compare(scales, ref_scales) + + +def test_the_quantization_round_trips_within_the_block_bound() -> None: + """Independent of the reference: dequantizing has to return the input. + + The cells compare two implementations of the same arithmetic, which would + both be wrong together if the arithmetic itself were wrong. This bounds the + error against the input instead: one power-of-two scale per 32 elements and + e4m3's 3 mantissa bits give at worst 2^-3 relative on the block maximum. + """ x = torch.randn(64, 2880, dtype=torch.bfloat16, device="cuda") - tight_data, tight_sf = mxfp8_quantize(x, False, 32) - padded_data, padded_sf = mxfp8_quantize(x, False, 512) - assert padded_data.shape == (64, 3072) and padded_sf.shape == (64 * 96,) - assert torch.equal( - padded_data[:, :2880].reshape(-1).view(torch.uint8), - tight_data.reshape(-1).view(torch.uint8), + data, scales = op(x, False, 512) + rows, blocks = 64, 3072 // 32 + scale = torch.exp2(scales.reshape(rows, blocks).float() - 127.0).unsqueeze(-1) + deq = (data.float().view(rows, blocks, 32) * scale).view(rows, 3072)[:, :2880] + block_amax = x.float().view(rows, 90, 32).abs().amax(-1, keepdim=True) + err = (deq - x.float()).view(rows, 90, 32).abs().amax(-1, keepdim=True) + assert torch.all(err <= block_amax * 2.0**-3 + 1e-6) + + +def test_guard_refuses_a_partial_block() -> None: + """The kernel reads past the end of the last block rather than failing.""" + x = torch.randn(4, 2864, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="whole number of 32-element blocks"): + with validating(op): + op(x, False, 512) + + +def test_guard_refuses_an_alignment_that_is_not_a_block_multiple() -> None: + """Pads to a width the scale count cannot describe.""" + x = torch.randn(4, 2880, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="multiple of the 32-element"): + with validating(op): + op(x, False, 48) + + +def test_guard_refuses_a_non_contiguous_input() -> None: + """Read flat, so a strided view quantizes the wrong elements.""" + x = torch.randn(4, 5760, dtype=torch.bfloat16, device="cuda")[:, ::2] + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(x, False, 512) + + +# ─── DeepSeek-V4.1 core target: its eight activation widths, every row count ── +# +# Every dense FP8 linear of the core target quantizes its bf16 input with the +# default call, `mxfp8_quantize(x)` (swizzled, alignment 32), and hands both +# outputs unchanged to the MXFP8 GEMM. K: 5120 hidden (wqkv_a, shared w1||w3), +# 1280 q_lora (wq_b, indexer wq_b), the wo_b input 8 / tp * 1024 (2048 under +# attention TP4, 8192 under attention DP, 4096 under attention TP2, 1024 under +# attention TP8), 2304 shared w2, 6144 engram wkv. The row count is the rank's +# token count this step: any integer up to max_num_tokens (at most 16384 per +# rank). In layers 21-39 the decoder bounded replay (every member) runs decoder +# passes over the finishing requests' 128-row tails plus the generation rows, +# cut by the target into passes of at most the step capacity; the uncut row set +# reaches 32768. The linears there take every width but 6144. The sweep below runs +# every row count 1..32768 at all eight widths. All eight widths are multiples +# of 128, so these calls never pad scale columns; row padding (to 128) happens +# at every row count that is not a multiple of 128. + +TARGET_K = (5120, 1280, 2048, 8192, 2304, 6144, 4096, 1024) +MAX_ROWS = 32768 +# The core target's decode CUDA-graph batch sizes, and one prefill chunk seen in its dump. +GRAPH_ROWS = (1, 2, 4, 8, 16, 24, 32, 1935) + + +def _activations(m: int, k: int, seed: int) -> torch.Tensor: + """bf16 N(0, 1) with each 32-element block scaled by 2^e, e uniform in -24..24. + + Block maxima then spread over about 49 scale bytes, so a scale byte read from + the wrong (row, block) is almost never equal by chance. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(m, k, generator=g, device="cuda", dtype=torch.float32) + e = torch.randint(-24, 25, (m, k // BLOCK, 1), generator=g, device="cuda").float() + return (x.view(m, k // BLOCK, BLOCK) * torch.exp2(e)).view(m, k).to(torch.bfloat16) + + +@pytest.mark.parametrize("k", TARGET_K) +def test_v41_target_width_swizzled_every_row_count(k: int) -> None: + """Every row count 1..MAX_ROWS (32768) at each target width, in the swizzled layout the GEMM consumes. + + The reference quantizes one MAX_ROWS-row input once. Every prefix x[:r] must + then give the reference's first r rows of data, byte for byte; the + reference's scale bytes at the 128x4 offsets (the offset of (m, b) does not + depend on the row count; only the buffer length does); and 0x00 in every + row-padding byte. The scale bytes are unsigned, so once the addressed bytes + match, a buffer sum equal to theirs means every padding byte is zero. + Outputs are poisoned before they are freed, so a byte the kernel skipped + cannot inherit the previous call's correct value from the allocator. + """ + cols = k // BLOCK + x = _activations(MAX_ROWS, k, seed=k) + ref_data, ref_sf = _ref(x) + ref_bytes = ref_data.view(torch.uint8) + ref_sf_cumsum = ref_sf.to(torch.int64).sum(dim=1).cumsum(0) + idx = _swizzle_index(MAX_ROWS, cols, x.device) + bad_data = torch.zeros((), dtype=torch.int64, device="cuda") + bad_sf = torch.zeros((), dtype=torch.int64, device="cuda") + bad_pad = torch.zeros((), dtype=torch.int64, device="cuda") + for r in range(1, MAX_ROWS + 1): + data, sf = op(x[:r]) + assert data.shape == (r, k) and data.dtype == torch.float8_e4m3fn + assert sf.shape == (_pad_up(r, 128) * cols,) and sf.dtype == torch.uint8 + bad_data += (data.view(torch.uint8) != ref_bytes[:r]).sum() + bad_sf += (sf[idx[:r]] != ref_sf[:r]).sum() + bad_pad += (sf.sum(dtype=torch.int64) - ref_sf_cumsum[r - 1]).abs() + data.view(torch.uint8).fill_(0x7F) + sf.fill_(0xFF) + if r % 2048 == 0: + # Each call's outputs are larger than any cached block, so the allocator + # would otherwise keep a segment per size and fill the device. + del data, sf + torch.cuda.empty_cache() + assert (bad_data.item(), bad_sf.item(), bad_pad.item()) == (0, 0, 0), ( + f"K={k}: {bad_data.item()} data bytes, {bad_sf.item()} scale bytes differ; " + f"padding bytes sum to {bad_pad.item()}" + ) + + +@pytest.mark.parametrize("k", TARGET_K) +def test_v41_target_width_cuda_graph_replay(k: int) -> None: + """Captured after one eager call at the same shape, as the engine's warmup does; replayed on new inputs. + + The graph's own output buffers are poisoned before every replay, so the + check shows that replay rewrites every data byte and every scale byte, + padding included. + """ + cols = k // BLOCK + for m in GRAPH_ROWS: + x_static = _activations(m, k, seed=7 * m + k) + op(x_static) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + data, sf = op(x_static) + idx = _swizzle_index(m, cols, x_static.device) + for seed in (1, 2): + x_new = _activations(m, k, seed=100000 * seed + m + k) + x_static.copy_(x_new) + data.view(torch.uint8).fill_(0x7F) + sf.fill_(0xFF) + graph.replay() + torch.cuda.synchronize() + ref_data, ref_sf = _ref(x_new) + assert torch.equal(data.view(torch.uint8), ref_data.view(torch.uint8)), (k, m) + assert torch.equal(sf[idx], ref_sf), (k, m) + assert sf.sum(dtype=torch.int64).item() == ref_sf.to(torch.int64).sum().item(), (k, m) + + +def test_v41_target_width_layout_controls() -> None: + """Wrong scale layouts at the target widths: the byte-exact check sees each one, except one case at a single row. + + - Row-major (linear) offsets read from the swizzled buffer are visible at + every row count, 1 included: block b of row 0 sits at + (b % 4) + (b // 4) * 512, not at b. + - Swapping the two row terms of the swizzle (row % 32 scaled by 4, and + (row % 128) // 32 scaled by 16) is visible from 2 rows up. At one row + both terms are zero, and the swap is invisible. + """ + for k in TARGET_K: + cols = k // BLOCK + c = torch.arange(cols, device="cuda").view(1, -1) + for m in (1, 7, 1935): + x = _activations(m, k, seed=3 * m + k) + _, sf = op(x) + _, ref_sf = _ref(x) + right = _swizzle_index(m, cols, x.device) + assert torch.equal(sf[right], ref_sf) + r = torch.arange(m, device="cuda").view(-1, 1) + linear = r * cols + c + swapped = ( + (c % 4) + + (c // 4) * 512 + + (r % 32) * 4 + + ((r % 128) // 32) * 16 + + (r // 128) * (128 * cols) + ) + n_linear = int((sf[linear] != ref_sf).sum().item()) + n_swapped = int((sf[swapped] != ref_sf).sum().item()) + print( + f"[control K={k} M={m}] linear offsets: {n_linear} / {m * cols} scale bytes differ; " + f"swapped row terms: {n_swapped} / {m * cols}" + ) + assert n_linear > 0 + if m == 1: + assert n_swapped == 0, "at one row the swapped terms are both zero" + else: + assert n_swapped > 0 + + +# ─── Capture set: decode, DSpark target verify, DSpark drafter ───────────── +# +# The captured batch sizes G, per rank (no MXFP8 linear of the target runs on +# the dep4-gathered token set): every G 1..32, plus TRT-LLM's padded default +# list for cuda_graph_config {enable_padding: true, max_batch_size: 256} +# (40..128 step 8, 192, 256); the padding-off list for 256 (1..32, 64, 128, 256) +# is a subset. A captured step is generation-only: the trunk linears see +# T = G rows (next_n 1) or 6G (DSpark target verify, 1 + 5 drafts), and the +# drafter's dense linears, whose widths are four of the eight (5120, 1280, 8192, +# 2304), see its 5-token draft block, 5G; every width is captured at all three +# multiples here. The engine captures one graph per batch size into one shared memory +# pool, largest first, each graph holding every linear of the step, and replays +# whichever one the step's batch selects. + +GSET = tuple(range(1, 33)) + (40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 192, 256) +VERIFY_GRAPH_ROWS = tuple(sorted({f * g for g in GSET for f in (1, 5, 6)}, reverse=True)) + + +def test_v41_target_widths_cuda_graph_at_dspark_verify_rows() -> None: + """One graph per captured row count (G, 5G, 6G over GSET: 114 counts, 1..1536) holding a call at each + target width, in one shared pool, replayed interleaved. + + Each graph is captured after an eager call at each of its shapes (the + engine's warm-up), largest row count first, all into one pool. The graphs + are then replayed in an interleaved order over three rounds. Before each + replay every static input gets new data and every output buffer is + poisoned (data 0x7F, scales 0xFF). After it, each call's outputs must be + bitwise equal to an eager call on the same data, whole scale buffer + included, and to the reference: data bytes, scale bytes at the 128x4 + offsets, and zero in every row-padding byte. A replay that did not run, or + ran on the previous round's inputs, fails the reference check: the + previous round's data bytes differ from this round's on most bytes (the + smallest differing fraction is printed and asserted). + + Control: every captured row count above 128 spans at least two 128-row + scale groups. Reading its scale buffer without the row-group term of the + swizzle (rows 128.. at the offsets of rows 0..) must differ from the + reference. + """ + pool = torch.cuda.graph_pool_handle() + graphs = {} + for m in sorted(VERIFY_GRAPH_ROWS, reverse=True): + xs = [_activations(m, k, seed=11 * m + k) for k in TARGET_K] + for x in xs: + op(x) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + outs = [op(x) for x in xs] + graphs[m] = (graph, xs, outs) + torch.cuda.synchronize() + order = VERIFY_GRAPH_ROWS[::2] + VERIFY_GRAPH_ROWS[1::2][::-1] + prev = {} + stale_min = 1.0 + group_min = {k: 1.0 for k in TARGET_K} + for rnd in range(3): + for m in order: + graph, xs, outs = graphs[m] + news = [_activations(m, k, seed=500000 + 1000 * rnd + 7 * m + k) for k in TARGET_K] + for x, x_new in zip(xs, news): + x.copy_(x_new) + for data, sf in outs: + data.view(torch.uint8).fill_(0x7F) + sf.fill_(0xFF) + graph.replay() + torch.cuda.synchronize() + for k, x_new, (data, sf) in zip(TARGET_K, news, outs): + cols = k // BLOCK + assert data.shape == (m, k) and sf.shape == (_pad_up(m, 128) * cols,), (k, m) + eager_data, eager_sf = op(x_new) + assert torch.equal(data.view(torch.uint8), eager_data.view(torch.uint8)), ( + k, + m, + rnd, + ) + assert torch.equal(sf, eager_sf), (k, m, rnd) + ref_data, ref_sf = _ref(x_new) + idx = _swizzle_index(m, cols, x_new.device) + assert torch.equal(data.view(torch.uint8), ref_data.view(torch.uint8)), (k, m, rnd) + assert torch.equal(sf[idx], ref_sf), (k, m, rnd) + assert sf.sum(dtype=torch.int64).item() == ref_sf.to(torch.int64).sum().item(), ( + k, + m, + rnd, + ) + # The previous round's reference (what a stale replay would have left) must differ. + if (k, m) in prev: + frac = (prev[(k, m)] != ref_data.view(torch.uint8)).float().mean().item() + stale_min = min(stale_min, frac) + assert frac > 0.5, (k, m, rnd, frac) + prev[(k, m)] = ref_data.view(torch.uint8).clone() + if rnd == 0 and m > 128: + r = torch.arange(m, device="cuda").view(-1, 1) + c = torch.arange(cols, device="cuda").view(1, -1) + no_group = (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + n_bad = int((sf[no_group] != ref_sf).sum().item()) + group_min[k] = min(group_min[k], n_bad / ((m - 128) * cols)) + assert n_bad > 0, (k, m) + print( + f"[capture] {len(VERIFY_GRAPH_ROWS)} row counts x {len(TARGET_K)} widths x 3 rounds: every replay bitwise = " + "eager = " + f"reference; smallest fraction of data bytes differing from the previous round: {stale_min:.4f}" + ) + for k in TARGET_K: + print( + f"[control K={k}] no row-group term, rows > 128: at least {group_min[k]:.4f} of the scale bytes of " + "rows 128.. differ" + ) + + +# ─── DeepSeek-V4.1 DSpark drafter: its two drafter-only calls ────────────── +# +# The drafter makes two quantize calls the core target does not, at G (the +# drafter's batch, 1..256) or 5G rows (1..1280): +# - its MoE input, mxfp8_quantize(x, False, 512) on bf16 [5G, 5120]: the +# LINEAR layout the MXFP4 MoE runner consumes (5120 is a multiple of 512, +# so there is no column padding); +# - main_proj's input, mxfp8_quantize(x) (swizzled, alignment 32) on bf16 +# [G or 5G, 15360] (and the prefill seeding's n <= 128 rows, eager). +# Both run here at every row count 1..1280 eagerly. Under capture both run at +# every row count 1..160 plus G and 5G over GSET (176 row counts, up to 1280). +DSPARK_ROWS = tuple(range(1, 1281)) +DSPARK_GRAPH_ROWS = tuple(sorted(set(range(1, 161)) | {f * g for g in GSET for f in (1, 5)})) +DSPARK_CASES = ((5120, False, 512), (15360, True, 32)) + + +def _check_layout(data, sf, ref_data, ref_sf, m, k, swizzled) -> None: + cols = k // BLOCK + assert data.shape == (m, k) and data.dtype == torch.float8_e4m3fn + assert torch.equal(data.view(torch.uint8), ref_data.view(torch.uint8)), (k, m) + if swizzled: + assert sf.shape == (_pad_up(m, 128) * _pad_up(cols, 4),), (k, m) + assert torch.equal(sf[_swizzle_index(m, cols, sf.device)], ref_sf), (k, m) + assert sf.sum(dtype=torch.int64).item() == ref_sf.to(torch.int64).sum().item(), (k, m) + else: + assert sf.shape == (m * cols,), (k, m) + assert torch.equal(sf.view(m, cols), ref_sf), (k, m) + + +@pytest.mark.parametrize("k,swizzled,alignment", DSPARK_CASES) +def test_v41_dspark_every_row_count_1_to_1280(k: int, swizzled: bool, alignment: int) -> None: + """Every row count 1..1280: each prefix x[:r] gives the reference's first r rows, byte for byte.""" + x = _activations(DSPARK_ROWS[-1], k, seed=k + 1) + ref_data, ref_sf = _ref(x) + for r in DSPARK_ROWS: + data, sf = op(x[:r], swizzled, alignment) + _check_layout(data, sf, ref_data[:r], ref_sf[:r], r, k, swizzled) + data.view(torch.uint8).fill_(0x7F) + sf.fill_(0xFF) + + +@pytest.mark.parametrize("k,swizzled,alignment", DSPARK_CASES) +def test_v41_dspark_cuda_graph_every_row_count(k: int, swizzled: bool, alignment: int) -> None: + """Captured at every row count of DSPARK_GRAPH_ROWS (1..160, G and 5G over GSET) after one eager + call at that shape, one shared pool, largest first; each graph replayed on two new inputs with its + outputs poisoned. Every replay is bitwise equal to an eager call on the same data and to the + reference.""" + pool = torch.cuda.graph_pool_handle() + graphs = {} + for m in sorted(DSPARK_GRAPH_ROWS, reverse=True): + x = _activations(m, k, seed=13 * m + k) + op(x, swizzled, alignment) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + out = op(x, swizzled, alignment) + graphs[m] = (graph, x, out) + torch.cuda.synchronize() + for rnd in range(2): + for m in DSPARK_GRAPH_ROWS[rnd::2] + DSPARK_GRAPH_ROWS[1 - rnd :: 2]: + graph, x, (data, sf) = graphs[m] + x_new = _activations(m, k, seed=700000 + 1000 * rnd + 13 * m + k) + x.copy_(x_new) + data.view(torch.uint8).fill_(0x7F) + sf.fill_(0xFF) + graph.replay() + torch.cuda.synchronize() + eager_data, eager_sf = op(x_new, swizzled, alignment) + assert torch.equal(data.view(torch.uint8), eager_data.view(torch.uint8)), (k, m, rnd) + assert torch.equal(sf, eager_sf), (k, m, rnd) + ref_data, ref_sf = _ref(x_new) + _check_layout(data, sf, ref_data, ref_sf, m, k, swizzled) + + +def test_v41_dspark_moe_input_layout_control() -> None: + """At 128 rows the swizzled scale buffer of the MoE input has exactly the linear buffer's byte + count (128 * 160), so a size check cannot tell them apart; the byte check does.""" + k, m = 5120, 128 + x = _activations(m, k, seed=77) + _, ref_sf = _ref(x) + _, sf_lin = op(x, False, 512) + _, sf_sw = op(x, True, 512) + assert sf_lin.numel() == sf_sw.numel() == m * (k // BLOCK) + assert torch.equal(sf_lin.view(m, k // BLOCK), ref_sf) + n_bad = int((sf_sw.view(m, k // BLOCK) != ref_sf).sum().item()) + print( + f"[control MoE input M={m}] swizzled buffer read as linear: {n_bad} / {m * (k // BLOCK)} scale bytes differ" ) - assert torch.equal(padded_sf.view(64, 96)[:, :90], tight_sf.view(64, 90)) - assert (padded_data[:, 2880:].view(torch.uint8) == 0).all() - assert (padded_sf.view(64, 96)[:, 90:] == 0).all() - - -def test_3d_input_collapses_leading_dims() -> None: - torch.manual_seed(4) - x = torch.randn(2, 5, 2880, dtype=torch.bfloat16, device="cuda") - data, sf = mxfp8_quantize(x, False, 512) - assert data.shape == (2, 5, 3072) - assert sf.shape == (2 * 5 * 96,) - ref_data, ref_sf = _ref(x, 3072) - assert torch.equal(data.reshape(10, 3072).view(torch.uint8), ref_data.view(torch.uint8)) - assert torch.equal(sf.view(10, 96), ref_sf) - # swizzled treats the collapsed rows as one 128-row-padded matrix - _, sf_sw = mxfp8_quantize(x, True, 512) - assert sf_sw.shape == (128 * 96,) - - -def test_dequant_error_bound() -> None: - """data * 2^(sf-127) reconstructs the input within the mxfp8 rounding bound.""" - torch.manual_seed(5) - x = torch.randn(256, 3072, dtype=torch.bfloat16, device="cuda") * 3.0 - data, sf = mxfp8_quantize(x, False, 32) - scale = torch.exp2(sf.view(256, 96).to(torch.float32) - 127.0) - deq = (data.float().view(256, 96, BLOCK) * scale.unsqueeze(-1)).view(256, 3072) - # e4m3 has a 3-bit mantissa: a normal value carries at most 2^-4 relative - # rounding error; a subnormal one at most half of the 2^-9 subnormal step, - # i.e. 2^-10 absolute in units of the block scale (per-block atol below). - atol = (2**-10) * scale.repeat_interleave(BLOCK, dim=-1) - err = (deq - x.float()).abs() - assert (err <= 2**-4 * x.float().abs() + atol).all() - # hard gate: the quantization is the exact reference, not merely close - ref_data, ref_sf = _ref(x, 3072) - assert torch.equal(data.view(torch.uint8), ref_data.view(torch.uint8)) - assert torch.equal(sf.view(256, 96), ref_sf) - - -def test_power_of_two_amax_boundary() -> None: - """amax exactly 448*2^k: the scale byte lands on 127+k, the max element on +-448.""" - for k in (-2, 0, 1, 3): - x = torch.zeros(1, 32, dtype=torch.bfloat16, device="cuda") - x[0, 0] = 448.0 * 2.0**k - x[0, 1] = -112.0 * 2.0**k - data, sf = mxfp8_quantize(x, False, 32) - assert sf.tolist() == [127 + k], (k, sf.tolist()) - assert data[0, 0].float().item() == 448.0 - assert data[0, 1].float().item() == -112.0 - - -def test_extreme_magnitude_block() -> None: - """The largest finite bf16 magnitude still quantizes exactly, no inf/NaN.""" - big = torch.finfo(torch.bfloat16).max - x = torch.zeros(1, 32, dtype=torch.bfloat16, device="cuda") - x[0, 0] = big - x[0, 1] = -big / 256.0 - data, sf = mxfp8_quantize(x, False, 32) - ref_data, ref_sf = _ref(x, 32) - # amax/448 = 2^119.19 -> scale 2^120 -> byte 247, well below the E8M0 max - assert sf.tolist() == [247] - assert torch.equal(data.view(torch.uint8), ref_data.view(torch.uint8)) - assert torch.equal(sf, ref_sf.reshape(-1)) - assert not torch.isnan(data.float()).any() - - -def test_nonfinite_inputs() -> None: - """Non-finite lanes are pinned here because both are destructive and silent.""" - # +inf: the block scale saturates to the E8M0 finite max (2^127), whose - # reciprocal flushes to zero -- the inf lane becomes NaN and every finite - # lane in the same block is zeroed. - x = torch.zeros(1, 32, dtype=torch.bfloat16, device="cuda") - x[0, 0] = float("inf") - x[0, 1] = 1.0 - data, sf = mxfp8_quantize(x, False, 32) - assert sf.tolist() == [254] - assert torch.isnan(data[0, 0].float()) - assert (data[0, 1:].float() == 0.0).all() - - # NaN: excluded from the block max, so the scale comes from the finite - # lanes; the NaN stays confined to its own lane. - y = torch.zeros(1, 32, dtype=torch.bfloat16, device="cuda") - y[0, 0] = float("nan") - y[0, 1] = 1.0 - data, sf = mxfp8_quantize(y, False, 32) - assert sf.tolist() == [119] # amax = 1.0 -> scale 2^-8 - assert torch.isnan(data[0, 0].float()) - assert data[0, 1].float().item() == 256.0 - assert (data[0, 2:].float() == 0.0).all() - - -def test_zero_and_denormal_blocks() -> None: - """Pins the two degenerate block-scale paths, including the NaN one.""" - # all-zero block: scale byte 0x00 and all-zero data bytes - x = torch.zeros(1, 64, dtype=torch.bfloat16, device="cuda") - data, sf = mxfp8_quantize(x, False, 32) - assert sf.tolist() == [0, 0] - assert (data.view(torch.uint8) == 0).all() - - # amax = 2^-118 > 448*2^-127: still the normal path, exact reference match - y = torch.zeros(1, 32, dtype=torch.bfloat16, device="cuda") - y[0, 0] = 2.0**-118 - y[0, 1] = 2.0**-119 - data, sf = mxfp8_quantize(y, False, 32) - ref_data, ref_sf = _ref(y, 32) - assert sf.tolist() == [1] - assert torch.equal(data.view(torch.uint8), ref_data.view(torch.uint8)) - assert not torch.isnan(data.float()).any() - - # amax = 2^-119 <= 448*2^-127: the block scale underflows to the E8M0 minimum - # (byte 0x00 = 2^-127, a denormal fp32) and the kernel's flush-to-zero - # reciprocal turns it into +inf -- nonzero lanes saturate to +-448, exact-zero - # lanes become NaN. Wrong answers, no error raised. - z = torch.zeros(1, 32, dtype=torch.bfloat16, device="cuda") - z[0, 0] = 2.0**-119 - z[0, 1] = -(2.0**-120) - data, sf = mxfp8_quantize(z, False, 32) - assert sf.tolist() == [0] - assert data[0, 0].float().item() == 448.0 - assert data[0, 1].float().item() == -448.0 - assert torch.isnan(data[0, 2:].float()).all() - - -def test_zero_rows() -> None: - """An empty batch is accepted, not rejected: both outputs come back empty.""" - x = torch.randn(0, 2880, dtype=torch.bfloat16, device="cuda") - data, sf = mxfp8_quantize(x, False, 512) - assert data.shape == (0, 3072) and data.dtype == torch.float8_e4m3fn - assert sf.shape == (0,) and sf.dtype == torch.uint8 - - # Empty inputs must leave subsequent CUDA work usable. - assert torch.ones(1, device="cuda").item() == 1 - - -def test_input_not_mutated() -> None: - torch.manual_seed(6) - x = torch.randn(32, 2880, dtype=torch.bfloat16, device="cuda") - before = x.clone() - mxfp8_quantize(x, False, 512) - mxfp8_quantize(x, True, 32) - assert torch.equal(x, before) - - -def test_rejects_unsupported() -> None: - """Domains the contract declares unsupported must be rejected, not silently wrong.""" - x = torch.randn(4, 128, dtype=torch.bfloat16, device="cuda") - cases = { - "k not a multiple of 32": lambda: mxfp8_quantize( - torch.randn(4, 112, dtype=torch.bfloat16, device="cuda"), False, 32 - ), - "alignment not a multiple of 32": lambda: mxfp8_quantize(x, False, 48), - "alignment below the block size": lambda: mxfp8_quantize(x, False, 16), - "fp32 input": lambda: mxfp8_quantize(x.float(), False, 32), - "1-D input": lambda: mxfp8_quantize(x.reshape(-1), False, 32), - "non-contiguous column slice": lambda: mxfp8_quantize( - torch.randn(8, 256, dtype=torch.bfloat16, device="cuda")[:, :128], False, 32 - ), - "non-contiguous transpose": lambda: mxfp8_quantize( - torch.randn(128, 8, dtype=torch.bfloat16, device="cuda").t(), False, 32 - ), - "cpu input": lambda: mxfp8_quantize(torch.randn(4, 128, dtype=torch.bfloat16), False, 32), - } - for name, fn in cases.items(): - try: - fn() - except (RuntimeError, NotImplementedError): - continue - raise AssertionError(f"{name}: expected a raise, got a result") + assert n_bad > 0 diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_auto_fallback.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_auto_fallback.py new file mode 100644 index 000000000000..748e617bc21c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_auto_fallback.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""When ``auto`` may fall back to the built-in implementation, read on the CPU. + +A family whose targets bring their own input processor (DeepSeek-V4.1) has +that processor chosen on the checkpoint alone, before the GPU and the parallel +mapping are known. So under ``auto`` a configuration none of its targets +matches must not fall back silently: the built-in model would run behind the +target's processor and placeholder placement. It raises at model construction +instead, and names the way out. ``off`` stays exactly what it was. + +The configs are plain namespaces and the SM is patched, so these run without a +GPU, without weights, and without importing any target's modeling code. +""" + +from __future__ import annotations + +from importlib import import_module +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch._experimental.modeling_v2._router_index as router_index +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + MODELING_V2_ENV, + modeling_v2_input_processor, + modeling_v2_resolve, +) + +_V41_TARGET = "ModelingV2DeepseekV41FlashNvfp4Sm103Unified" +_TEP4 = dict( + world_size=4, tp_size=4, pp_size=1, moe_ep_size=4, moe_tp_size=1, enable_attention_dp=False +) + + +def _v41_pretrained(**overrides): + fields = dict( + architectures=["DeepseekV41ForConditionalGeneration"], + model_type="deepseek_v41", + num_hidden_layers=40, + hidden_size=5120, + n_routed_experts=384, + q_lora_rank=1280, + head_dim=512, + o_groups=8, + index_n_heads=32, + num_nextn_predict_layers=3, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _model_config(pretrained, quant="FP8_BLOCK_SCALES", **mapping_overrides): + return SimpleNamespace( + pretrained_config=pretrained, + mapping=SimpleNamespace(**dict(_TEP4, **mapping_overrides)), + quant_config=SimpleNamespace(quant_algo=quant), + spec_config=None, + ) + + +@pytest.fixture +def sm(monkeypatch): + """Route as if on the given SM; default a GB300 (10, 3).""" + state = {"sm": (10, 3)} + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr( + torch.cuda, "get_device_capability", lambda *a, **k: state["sm"], raising=False + ) + return state + + +@pytest.fixture +def imports(monkeypatch): + """Record the resolver's imports; a target's modeling module is not imported + for real (it needs the GPU stack), only recorded.""" + seen = [] + + def _import(name): + seen.append(name) + if name.endswith(".modeling"): + return SimpleNamespace(__name__=name) + return import_module(name) + + monkeypatch.setattr(router_index, "import_module", _import) + return seen + + +def _mode(monkeypatch, mode): + if mode is None: + monkeypatch.delenv(MODELING_V2_ENV, raising=False) + else: + monkeypatch.setenv(MODELING_V2_ENV, mode) + + +@pytest.mark.parametrize("mode", [None, "off"]) +@pytest.mark.parametrize( + "config", + [ + _model_config(_v41_pretrained()), + _model_config(_v41_pretrained(), tp_size=2, world_size=2, moe_ep_size=2), + _model_config(_v41_pretrained(), quant=None), + ], + ids=["match", "tp2", "unquantized"], +) +def test_off_changes_nothing(monkeypatch, sm, imports, mode, config): + """Off: no routing, no processor, no import, matching or not.""" + _mode(monkeypatch, mode) + assert modeling_v2_resolve(config) is None + assert modeling_v2_input_processor(config.pretrained_config) is None + assert imports == [] + + +@pytest.mark.parametrize("mode", ["auto", "require"]) +def test_a_matching_configuration_routes(monkeypatch, sm, imports, mode): + _mode(monkeypatch, mode) + assert modeling_v2_resolve(_model_config(_v41_pretrained())) == _V41_TARGET + assert imports[-1].endswith("deepseek_v4_1_flash_nvfp4__sm_103__unified.modeling") + + +@pytest.mark.parametrize( + "quant, mapping_overrides, device_sm, missed", + [ + ("FP8_BLOCK_SCALES", dict(tp_size=2, world_size=2, moe_ep_size=2), (10, 3), "parallel"), + ("FP8_BLOCK_SCALES", dict(moe_ep_size=2, moe_tp_size=2), (10, 3), "parallel"), + (None, dict(), (10, 3), "quant"), + ("FP8_BLOCK_SCALES", dict(), (10, 0), "sm"), + ], + ids=["tp2", "ep2xmoe_tp2", "unquantized", "sm100"], +) +def test_auto_refuses_the_builtin_fallback_behind_the_targets_processor( + monkeypatch, sm, imports, quant, mapping_overrides, device_sm, missed +): + """The target's checkpoint on a configuration no target covers: auto raises, + names the criterion that missed and the ways out, and imports no target.""" + _mode(monkeypatch, "auto") + sm["sm"] = device_sm + config = _model_config(_v41_pretrained(), quant=quant, **mapping_overrides) + with pytest.raises(ValueError) as raised: + modeling_v2_resolve(config) + message = str(raised.value) + assert f"{MODELING_V2_ENV}=auto" in message + assert f"{MODELING_V2_ENV}=off" in message + assert "input processor" in message + assert any(line.lstrip().startswith(missed) for line in message.splitlines()) + assert not any(name.endswith(".modeling") for name in imports) + + +def test_require_keeps_its_own_message(monkeypatch, sm, imports): + _mode(monkeypatch, "require") + config = _model_config(_v41_pretrained(), tp_size=2, world_size=2, moe_ep_size=2) + with pytest.raises(ValueError, match=f"{MODELING_V2_ENV}=require, but no target matched"): + modeling_v2_resolve(config) + + +def test_auto_falls_back_for_a_checkpoint_the_targets_do_not_claim(monkeypatch, sm, imports): + """Another V4.1 shape gets no target processor, so the plain fallback stands.""" + _mode(monkeypatch, "auto") + config = _model_config(_v41_pretrained(num_nextn_predict_layers=1)) + assert modeling_v2_resolve(config) is None + assert modeling_v2_input_processor(config.pretrained_config) is None + + +def test_auto_falls_back_for_a_family_without_its_own_processor(monkeypatch, sm, imports): + _mode(monkeypatch, "auto") + gpt_oss = SimpleNamespace( + architectures=["GptOssForCausalLM"], + num_hidden_layers=36, + hidden_size=2880, + num_local_experts=128, + ) + assert modeling_v2_resolve(_model_config(gpt_oss, world_size=2, tp_size=2)) is None diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py index 1f7e0435d5b3..dfbd39bdd5f5 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py @@ -17,6 +17,7 @@ from __future__ import annotations +import ast import re from pathlib import Path @@ -215,3 +216,227 @@ def test_targets_do_not_share_files(): f"directory for {tail!r}; targets may import the " f"catalog and nothing else" ) + + +def _entry_files() -> list[Path]: + """Every entry module: the shared catalog's, and each target's own ``ops/``.""" + files = sorted((_ROOT / "catalog").glob("*/*.py")) + files += sorted((_ROOT / "models").glob("*/*/ops/*/*.py")) + return [path for path in files if path.name != "__init__.py"] + + +def test_every_entry_with_cells_has_a_test_that_drives_them(): + """A cell list nothing executes is worse than no cell list. + + It reads as coverage. `mla_rope_append_paged_kv_assign_q` shipped three + cells with no driver for exactly one commit, and nothing went red -- the + entry's own test had been wired to the entry's `compare` and not to its + `CELLS`, so the suite passed and the three configurations had never run. + + Checked by reading rather than importing, like the rest of this file: an + entry declares `Cell(` in its source, and some test under this tree has to + name both the entry and `.CELLS`. That is weaker than proving the cells were + driven, and it is what a check that runs anywhere can say. + """ + tests = list(Path(__file__).resolve().parent.rglob("*.py")) + sources = {t: t.read_text() for t in tests} + + undriven = [] + for entry in _entry_files(): + if "Cell(" not in entry.read_text(): + continue + if not any(entry.stem in src and ".CELLS" in src for src in sources.values()): + undriven.append(str(entry.relative_to(_ROOT).with_suffix(""))) + + assert not undriven, ( + f"{len(undriven)} entry(ies) declare cells that no test drives: " + f"{', '.join(undriven)}. Parametrize the entry's test off its CELLS." + ) + + +# The phase a step runs at is resolved once, at the dispatcher, and a target is +# what runs after that resolution. These two gates are what keep that true: +# without them a later edit can put `if md.num_contexts:` back inside a decode +# body and nothing goes red. +_PHASE_FIELDS = ("num_contexts", "num_ctx_tokens") + + +# Cores not yet on the target layer. An entry is deleted when that core is +# migrated, so finishing the migration is a visible diff rather than a gate +# that quietly stopped covering anything. +_NOT_YET_ON_THE_TARGET_LAYER = frozenset({"r1_0528_nvfp4__sm_103__dep4"}) + + +def _core_modules() -> list[Path]: + """Every target module the routing tables can reach and this layer covers.""" + modules = [] + for arch in _ARCHS: + routing = routing_module(arch) + for dotted in routing.TARGET_MODULES.values(): + path = _module_path(dotted) + if path.parent.name in _NOT_YET_ON_THE_TARGET_LAYER: + continue + modules.append(path) + return modules + + +def test_a_decode_target_never_reads_the_phase_back(): + """Decode is routed to only when there are no context rows. Reading the + count back inside it is not a redundancy, it is the design being violated: + the value is a per-capture constant, so a branch on it inside a captured + graph is frozen at whatever the capturing batch happened to carry. + + Prefill is deliberately exempt -- it holds the mixed batch and genuinely + needs `num_ctx_tokens` to split it. The asymmetry is the design, not an + oversight. + + Scoped to the `DecodeTarget` class node rather than inverted over the + whole module: that inversion existed only while a decode body could be + inherited from a shared base, where `DecodeTarget`'s own node held just + the literal-args `step_args` call and a gate keyed on the class name saw + nothing of the forward that actually ran. Every target in this tree is + now a complete, independent class -- `DecodeTarget.forward` is not + inherited from anywhere -- so the class node is once again the target's + whole body, and the direct form catches a violation planted anywhere in + it. + + Read by parsing rather than by importing, like everything else in this + file: no GPU, no built extensions. + """ + offenders = [] + for path in _core_modules(): + if not path.is_file(): + continue + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if not isinstance(node, ast.ClassDef) or node.name != "DecodeTarget": + continue + for inner in ast.walk(node): + if isinstance(inner, ast.Attribute) and inner.attr in _PHASE_FIELDS: + offenders.append(f"{path.relative_to(_ROOT)}:{inner.lineno} .{inner.attr}") + assert not offenders, "a decode target reads the phase it was routed on: " + ", ".join( + offenders + ) + + +def test_every_core_ships_both_targets(): + """A core with one target routes half its steps into a KeyError. + + Names, not a table: the gate above finds the decode body by class name, so + the name is load-bearing and is pinned here rather than left to convention. + """ + missing = [] + for path in _core_modules(): + if not path.is_file(): + continue + classes = { + n.name for n in ast.walk(ast.parse(path.read_text())) if isinstance(n, ast.ClassDef) + } + for required in ("PrefillTarget", "DecodeTarget"): + if required not in classes: + missing.append(f"{path.relative_to(_ROOT)}: {required}") + assert not missing, "core modules missing a target class: " + ", ".join(missing) + + +# The binding layer is only safe because `raw_call` is independent of it. These +# two gates are what keep that true; without them an entry can quietly start +# reading bound state, and CELLS stops covering what the target actually runs. + + +def _entry_classes() -> list[tuple[Path, ast.ClassDef]]: + """Every OpWrapper subclass in an entry module, as (path, class node).""" + found = [] + for path in _entry_files(): + for node in ast.walk(ast.parse(path.read_text())): + if isinstance(node, ast.ClassDef) and any( + isinstance(b, ast.Name) and b.id == "OpWrapper" for b in node.bases + ): + found.append((path, node)) + return found + + +def test_raw_call_reads_no_bound_state(): + """`raw_call` takes every argument explicitly. + + The moment it reads `self._const`, `self._layered`, or any other attribute, + the arguments a cell drives it with stop being the whole input -- and + `CELLS` silently narrows to whatever the last target happened to bind. + """ + offenders = [] + for path, cls in _entry_classes(): + for node in cls.body: + if not isinstance(node, ast.FunctionDef) or node.name != "raw_call": + continue + for inner in ast.walk(node): + if ( + isinstance(inner, ast.Attribute) + and isinstance(inner.value, ast.Name) + and inner.value.id == "self" + ): + offenders.append(f"{path.relative_to(_ROOT)}:{inner.lineno} self.{inner.attr}") + assert not offenders, ( + "raw_call must take every argument explicitly; these read bound state: " + + ", ".join(offenders) + ) + + +def test_is_valid_accepts_everything_raw_call_accepts(): + """`is_valid` must accept every argument `raw_call` accepts. + + Not parameter-for-parameter mirroring -- an earlier version of this gate + demanded that, and it was wrong: it demanded a property the tree does not + have and does not need. `is_valid` may narrow to only the parameters it + actually inspects, as long as it also carries a `**kwargs` absorber. + `validating()` forwards the call's full, merged argument set to + `is_valid` verbatim, so the absorber is what keeps that forward from + raising `TypeError` on the arguments the guard does not name -- + `thop_attention.is_valid` checks three of `raw_call`'s ~115 parameters + and absorbs the rest in `**unused_kwargs`; that is not a gap the gate + missed, it is the mechanism the tree chose so a 115-parameter signature + does not have to be restated to be validated. An entry with no `is_valid` + of its own inherits the base class's `(*args, **kwargs)` no-op, which + already accepts everything, so it is exempt from this check the same way. + + `reference` is not checked here at all, and was wrong to check before. + It is driven by a cell's own explicit argument list, never by the merged + forward that reaches `is_valid`, so it never receives an argument it did + not declare -- there is nothing for it to absorb and nothing to mirror. + `thop_attention.reference` takes 16 of `raw_call`'s ~115 parameters on + purpose: a 115-parameter reference implementation would be unmaintainable, + and the cell driving it only ever supplies those 16. Eight entries narrow + `reference` this way; all eight are correct. + + The `self`-is-first-parameter check is kept as-is: an entry once shipped + `raw_call` and `is_valid` both missing `self`, and a mirror check that + assumed the first parameter was `self` dropped a real argument from both + and reported a match. + """ + problems = [] + for path, cls in _entry_classes(): + funcs = { + n.name: n + for n in cls.body + if isinstance(n, ast.FunctionDef) and n.name in ("raw_call", "reference", "is_valid") + } + raw_call = funcs.get("raw_call") + if raw_call is None: + continue + for name, fn in funcs.items(): + params = [a.arg for a in fn.args.args] + if not params or params[0] != "self": + problems.append( + f"{path.relative_to(_ROOT)}: {cls.name}.{name} first param is not self" + ) + + is_valid = funcs.get("is_valid") + if is_valid is None or is_valid.args.kwarg is not None: + continue # inherits the base no-op, or absorbs the remainder itself + raw_names = {a.arg for a in raw_call.args.args + raw_call.args.kwonlyargs} - {"self"} + valid_names = {a.arg for a in is_valid.args.args + is_valid.args.kwonlyargs} - {"self"} + missing = sorted(raw_names - valid_names) + if missing: + problems.append( + f"{path.relative_to(_ROOT)}: {cls.name}.is_valid does not accept " + f"{missing} that raw_call accepts, and has no ** absorber" + ) + assert not problems, "; ".join(problems) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_decoder_replay.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_decoder_replay.py new file mode 100644 index 000000000000..ec6f54e481ec --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_decoder_replay.py @@ -0,0 +1,119 @@ +# 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. +"""The DeepSeek-V4.1-Flash-NVFP4 target's host-side decoder replay bookkeeping. + +The target's core runs decoder SWA bounded replay as the engine's cache manager resolved it, +once per manager, and keeps a request's tail stash between its context chunks only +while the cache manager still holds the request. Under attention DP the MoE pads every +rank's rows to the step's largest per-rank token count. All three are host logic, +checked here on stand-in objects; importing the target needs a real build, as for the +target contract. +""" + +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import pytest + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified import ( # noqa: E501 + modeling, +) + +Core = modeling.DeepseekV41UnifiedModelingV2 + + +class _Manager: + """The two attributes of the V4.1 cache manager the bookkeeping reads.""" + + def __init__(self, decoder_replay: bool = True, held=()): + self.v41_decoder_replay = decoder_replay + self.kv_cache_map = {rid: object() for rid in held} + + +def _core(replay_requested: bool = True, records=()): + return SimpleNamespace( + replay_requested=replay_requested, + replay=replay_requested, + _replay_mgr=None, + _tail_records={rid: modeling._TailRecord(0) for rid in records}, + ) + + +def test_replay_follows_each_cache_manager(): + core = _core() + first, final = _Manager(True), _Manager(False) + Core._resolve_replay(core, first) + assert core.replay is True + # The engine builds a KV-size estimation manager, then the final one. + Core._resolve_replay(core, final) + assert core.replay is False + # Resolved once per manager, not re-read every step. + final.v41_decoder_replay = True + Core._resolve_replay(core, final) + assert core.replay is False + + core = _core(replay_requested=False) + Core._resolve_replay(core, _Manager(False)) + assert core.replay is False + with pytest.raises(AssertionError, match="resolved decoder replay on"): + Core._resolve_replay(_core(replay_requested=False), _Manager(True)) + + +def test_stash_is_released_with_the_request(): + """A request freed between two of its context chunks (cancelled, or evicted for + recompute) never reaches its last chunk: its record and stash go once the cache + manager no longer holds it. Held requests, suspended ones included, keep theirs.""" + core = _core(records=(11, 12, 13)) + stash = object() + core._tail_records[12].stash = {"h": stash} + Core._release_tail_records(core, _Manager(held=(11, 12, 13, 99))) + assert set(core._tail_records) == {11, 12, 13} + refs = sys.getrefcount(stash) + Core._release_tail_records(core, _Manager(held=(11, 13))) + assert set(core._tail_records) == {11, 13} + assert sys.getrefcount(stash) == refs - 1, "the released stash is still referenced" + Core._release_tail_records(core, _Manager(held=())) + assert core._tail_records == {} + + +def test_dp_rows_are_the_engine_rows(): + """Attention DP: up to the last KV-source layer every rank pads its MoE rows to the + step's largest per-rank token count, and its own count must be the engine's. Past + that layer under decoder replay a rank's rows can outnumber its step rows (a small + last chunk still runs its whole tail, the rest from the stash) and a bounded rank + runs them in slices, so the ranks exchange their passes instead (``_dp_late_rows``, + an engine collective; ``_dp_pass_rows`` below).""" + core = SimpleNamespace(tp=4, tp_rank=2) + md = SimpleNamespace(all_rank_num_tokens=[7, 2048, 1932, 5]) + assert Core._dp_rows(core, md, 1932) == 2048 + for rows in (0, 128, 1933): + with pytest.raises(AssertionError): + Core._dp_rows(core, md, rows) + + +def test_dp_passes_over_the_bounded_layers(): + """Attention DP under decoder replay: every rank runs as many passes over the bounded + layers as the rank with the most (its slices when bounded, else its step rows as + one pass), each padded to the max of that pass's row counts over the ranks that hold + it; a rank without the pass joins it with zero rows.""" + # Rank 0 bounded in two slices (a 1-row last chunk with a 127-row stash, then more + # tails); rank 1 bounded with no decoder row; ranks 2 and 3 unbounded. + gathered = [[2, 2048, 1000], [0], [1, 6], [1, 3]] + assert modeling._dp_pass_rows(gathered) == [2048, 1000] + assert modeling._dp_pass_rows([[1, 4], [1, 128], [1, 1], [1, 1]]) == [128] + assert modeling._dp_pass_rows([[1, 7]] * 4) == [7] diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_engram_host_table_vmm.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_engram_host_table_vmm.py new file mode 100644 index 000000000000..33dd57d0da2a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_engram_host_table_vmm.py @@ -0,0 +1,482 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the ``vmm`` backend of the DeepSeek-V4.1 target's ``engram_host_table.py``. + +The backend keeps DeepSeek-V4.1's Engram tables in CUDA VMM host-NUMA allocations +(``cuMemCreate`` on the CPU sockets, 1 GiB slabs alternating between them), shared +between processes as POSIX file descriptors over a node-local abstract Unix socket, +and hands ``engram_gather`` CUDA views of that host memory. + +What is checked, on a synthetic two-layer checkpoint (never the real 202.76 GB table): + +* the creator's views are CUDA views of host memory: their bytes equal the + checkpoint's, the manifest's chunk CRCs re-verify from the host view, and no + device memory holds a copy; +* ``engram_gather`` through them equals the torch reference bit for bit, eager and + replayed from a CUDA graph, for both layers; +* a second process imports the table from the creator's socket and gathers the + same bits; after the creator closes, a third process imports it from the second + (the table stays reachable while any holder lives); +* a creator that sees one GPU (``CUDA_VISIBLE_DEVICES`` of one device, for a device on + each CPU socket) places its slabs only on the sockets it may allocate on, and a + process that sees only a GPU of the other socket imports the table and gathers the + same bits; +* host out-of-memory at create (injected ``CUDA_ERROR_OUT_OF_MEMORY``): the create + releases what it allocated and retries with backoff, and succeeds when memory + returns; when it never does, the table falls back to the shm backend with a + warning (every local rank then attaches the shm file), unless the backend is + set to ``vmm`` explicitly, where the create fails after joining the barrier; +* the backend switch and the per-checkpoint socket name. +""" + +import json +import os +import struct +import uuid + +import pytest +import torch +import torch.multiprocessing as mp + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified import ( # noqa: E501 + engram_host_table as H, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.ops.embedding.engram_gather import ( # noqa: E501 + EngramGather, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +# The target's gather entry, unbound: every call passes its operands. +engram_gather = EngramGather() + +ROWS = 100_003 +DIM = 256 + + +def _table(rows: int, seed: int): + g = torch.Generator().manual_seed(seed) + w = torch.randint(0, 256, (rows, DIM), dtype=torch.uint8, generator=g) + w[(w == 0x7F) | (w == 0xFF)] = 0 # E4M3 NaN codes: NaN payloads differ between Triton and torch + s = torch.randint(118, 124, (rows, DIM // 32), dtype=torch.uint8, generator=g) + return w.view(torch.float8_e4m3fn), s + + +def _ref(ids, w, s): + ids = ids.cpu() + vals = w[ids].float().unflatten(-1, (DIM // 32, 32)) + scale = s[ids].view(torch.float8_e8m0fnu).float().unsqueeze(-1) + return (vals * scale).flatten(-2).to(torch.bfloat16) + + +def _write_checkpoint(ckpt_dir: str, tables) -> None: + header, blobs, at = {}, [], 0 + for layer, (w, s) in tables.items(): + for kind, t, dtype in (("weight", w, "F8_E4M3"), ("scale", s, "U8")): + b = t.view(torch.uint8).contiguous().numpy().tobytes() + header[f"layers.{layer}.engram.embed.{kind}"] = { + "dtype": dtype, + "shape": list(t.shape), + "data_offsets": [at, at + len(b)], + } + blobs.append(b) + at += len(b) + raw = json.dumps(header).encode() + raw += b" " * (-len(raw) % 8) + shard = "model-00001-of-00001.safetensors" + with open(os.path.join(ckpt_dir, shard), "wb") as f: + f.write(struct.pack(" torch.Tensor: + g = torch.Generator().manual_seed(seed) + ids = torch.randint(0, ROWS, (tokens, 24), generator=g) + ids[0, :4] = torch.tensor([0, 1, ROWS - 2, ROWS - 1]) + return ids.to(device) + + +def _gather_all(table, device, seeds=(1, 2)): + out = {} + for layer in table.layers: + w, s = table.weight(layer), table.scale(layer).view(torch.uint8) + for seed in seeds: + ids = _ids(64, seed * 100 + layer, device) + out[(layer, seed)] = engram_gather(ids, w, s, 0, w.shape[0]).cpu() + return out + + +def _child(rank, ckpt, name, device_index, queue, release): + """Import the table (local_rank 1: no fill), gather, report, keep holding (and serving) until released.""" + torch.cuda.set_device(device_index) + dev = torch.device("cuda", device_index) + torch.zeros(1, device=dev) + table = H.open_vmm_engram_table(ckpt, local_rank=1, device=dev, wait_s=30.0, name=name) + table.check_gather_views(dev) + got = _gather_all(table, dev) + queue.put((rank, {"created": table.created, "gathers": got})) + release.wait(timeout=1800) + table.close() + + +def _creator(ckpt, name, queue, release): + """Create the table as local rank 0 on this process's only visible device, report, hold it until released.""" + dev = torch.device("cuda", 0) + torch.cuda.set_device(dev) + torch.zeros(1, device=dev) + table = H.open_vmm_engram_table( + ckpt, local_rank=0, device=dev, headroom_bytes=1 << 30, name=name, threads=4 + ) + table.check_gather_views(dev) + fill = table.manifest["fill"] + queue.put( + ( + "creator", + { + "created": table.created, + "numa_bytes": table.numa_bytes(), + "refused": {int(k): v for k, v in fill["numa_nodes_refused"].items()}, + "device_host_nodes": {int(k): v for k, v in fill["device_host_nodes"].items()}, + "gathers": _gather_all(table, dev), + }, + ) + ) + release.wait(timeout=1800) + table.close() + + +def _physical(logical: int) -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + return visible.split(",")[logical].strip() if visible else str(logical) + + +def _start(ctx, target, args, visible: str): + """Start a spawned process that sees only ``visible`` (it inherits the environment at start).""" + old = os.environ.get("CUDA_VISIBLE_DEVICES") + os.environ["CUDA_VISIBLE_DEVICES"] = visible + try: + proc = ctx.Process(target=target, args=args) + proc.start() + finally: + if old is None: + del os.environ["CUDA_VISIBLE_DEVICES"] + else: + os.environ["CUDA_VISIBLE_DEVICES"] = old + return proc + + +@pytest.fixture +def checkpoint(tmp_path): + tables = {1: _table(ROWS, 7001), 14: _table(ROWS, 7014)} + ckpt = tmp_path / "ckpt" + ckpt.mkdir() + _write_checkpoint(str(ckpt), tables) + return str(ckpt), tables + + +def test_backend_switch(monkeypatch): + monkeypatch.delenv(H.BACKEND_ENV, raising=False) + assert H.host_table_backend() == "vmm" + for value, want in (("shm", "shm"), (" VMM ", "vmm")): + monkeypatch.setenv(H.BACKEND_ENV, value) + assert H.host_table_backend() == want + monkeypatch.setenv(H.BACKEND_ENV, "hugetlb") + with pytest.raises(ValueError): + H.host_table_backend() + + +def test_socket_name_depends_on_the_checkpoint_bytes(checkpoint, tmp_path): + ckpt, tables = checkpoint + other = tmp_path / "other" + other.mkdir() + _write_checkpoint(str(other), {1: _table(ROWS + 1, 1), 14: _table(ROWS, 2)}) + a = H.default_vmm_socket_name(H._layout(H.find_engram_table_sources(ckpt))) + b = H.default_vmm_socket_name(H._layout(H.find_engram_table_sources(str(other)))) + assert a != b and a.startswith(f"trtllm-dsv41-engram-{os.getuid()}.vmm.") + + +def test_create_gather_graph_import_and_takeover(checkpoint): + ckpt, tables = checkpoint + dev = torch.device("cuda", 0) + torch.cuda.set_device(dev) + name = f"pf-vmm-test-{uuid.uuid4().hex[:12]}" + torch.cuda.synchronize() + free0, _ = torch.cuda.mem_get_info(dev) + alloc0 = torch.cuda.memory_allocated(dev) + table = H.open_vmm_engram_table( + ckpt, local_rank=0, device=dev, headroom_bytes=1 << 30, name=name, threads=4 + ) + importers = [] + try: + assert table.created and not table.fill_report.reused + table.check_gather_views(dev) + # CUDA views of host memory, bytes equal to the checkpoint's, CRCs re-verify, no device copy. + for layer, (w, s) in tables.items(): + vw, vs = table.weight(layer), table.scale(layer) + assert ( + vw.is_cuda + and vw.device == dev + and tuple(vw.shape) == (ROWS, DIM) + and vw.dtype == torch.float8_e4m3fn + ) + # the scale's dtype is the checkpoint's (U8 here, F8_E8M0 in the real checkpoint) + assert vs.is_cuda and vs.dtype == torch.uint8 and tuple(vs.shape) == (ROWS, DIM // 32) + hw = torch.frombuffer( + table.host_buffer(f"layers.{layer}.engram.embed.weight"), dtype=torch.uint8 + ) + hs = torch.frombuffer( + table.host_buffer(f"layers.{layer}.engram.embed.scale"), dtype=torch.uint8 + ) + assert torch.equal(hw.view(ROWS, DIM), w.view(torch.uint8)) + assert torch.equal(hs.view(ROWS, DIM // 32), s) + checked, bad = table.verify(full=True) + assert checked > 0 and bad == [] + assert sum(table.numa_bytes().values()) == table.allocated_bytes >= table.nbytes + torch.cuda.synchronize() + assert torch.cuda.memory_allocated(dev) == alloc0 + assert free0 - torch.cuda.mem_get_info(dev)[0] < 64 << 20, ( + "the table must not occupy device memory" + ) + + # Gathers: eager and replayed from a CUDA graph, bit-exact against the reference. + static = _ids(48, 11, dev) + pair = { + layer: (table.weight(layer), table.scale(layer).view(torch.uint8)) for layer in tables + } + for layer, (w, s) in pair.items(): + engram_gather(static, w, s, 0, ROWS) # warm-up outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = {layer: engram_gather(static, w, s, 0, ROWS) for layer, (w, s) in pair.items()} + for seed in (12, 13): + static.copy_(_ids(48, seed, dev)) + graph.replay() + torch.cuda.synchronize() + for layer, (w, s) in tables.items(): + ref = _ref(static, w, s) + assert torch.equal(outs[layer].cpu().view(torch.int16), ref.view(torch.int16)) + eager = engram_gather(static, *pair[layer], 0, ROWS) + assert torch.equal(eager.view(torch.int16), outs[layer].view(torch.int16)) + del graph, outs + mine = _gather_all(table, dev) + for (layer, seed), got in mine.items(): + ref = _ref(_ids(64, seed * 100 + layer, "cpu"), *tables[layer]) + assert torch.equal(got.view(torch.int16), ref.view(torch.int16)) + + # A second process imports from this one; after this one closes, a third imports from the second. + ctx = mp.get_context("spawn") + queue = ctx.Queue() + release = ctx.Event() + other = 1 if torch.cuda.device_count() > 1 else 0 + importers.append(ctx.Process(target=_child, args=(1, ckpt, name, other, queue, release))) + importers[-1].start() + rank, first = queue.get(timeout=1200) + assert rank == 1 and not first["created"] + for key, got in first["gathers"].items(): + assert torch.equal(got.view(torch.int16), mine[key].view(torch.int16)), key + table.close() # the creator stops serving; the importer's standby server takes the name over + importers.append(ctx.Process(target=_child, args=(2, ckpt, name, 0, queue, release))) + importers[-1].start() + rank, second = queue.get(timeout=1200) + assert rank == 2 and not second["created"] + for key, got in second["gathers"].items(): + assert torch.equal(got.view(torch.int16), mine[key].view(torch.int16)), key + release.set() + for p in importers: + p.join(timeout=120) + assert p.exitcode == 0 + finally: + table.close() + for p in importers: + if p.is_alive(): + p.kill() + + +def test_single_gpu_creator_on_each_socket_and_cross_socket_import(checkpoint): + ckpt, tables = checkpoint + torch.zeros(1, device="cuda") + host_nodes = H._visible_device_host_nodes() + by_node = {} + for index, node in sorted(host_nodes.items()): + by_node.setdefault(node, index) + creators = sorted(by_node.items()) # one device per CPU socket that has GPUs + ctx = mp.get_context("spawn") + for node, index in creators: + others = [i for n, i in creators if n != node] + name = f"pf-vmm-sock-{uuid.uuid4().hex[:12]}" + queue, release = ctx.Queue(), ctx.Event() + procs = [_start(ctx, _creator, (ckpt, name, queue, release), _physical(index))] + try: + role, made = queue.get(timeout=1200) + assert role == "creator" and made["created"] + assert made["device_host_nodes"] == {0: node}, made["device_host_nodes"] + slab_nodes = set(made["numa_bytes"]) + assert slab_nodes and not slab_nodes & set(made["refused"]), made + if node >= 0: + assert node in slab_nodes, made + for (layer, seed), got in made["gathers"].items(): + ref = _ref(_ids(64, seed * 100 + layer, "cpu"), *tables[layer]) + assert torch.equal(got.view(torch.int16), ref.view(torch.int16)), ( + node, + layer, + seed, + ) + print( + f"\ncreator on GPU {_physical(index)} (host node {node}): slabs on {sorted(slab_nodes)}, " + f"refused {made['refused']}" + ) + for rank, other in enumerate(others, start=1): + procs.append( + _start(ctx, _child, (rank, ckpt, name, 0, queue, release), _physical(other)) + ) + got_rank, imported = queue.get(timeout=1200) + assert got_rank == rank and not imported["created"] + for key, got in imported["gathers"].items(): + assert torch.equal( + got.view(torch.int16), made["gathers"][key].view(torch.int16) + ), (node, other, key) + print( + f"importer on GPU {_physical(other)} (host node {host_nodes[other]}): gathers equal" + ) + release.set() + for p in procs: + p.join(timeout=120) + assert p.exitcode == 0 + finally: + release.set() + for p in procs: + if p.is_alive(): + p.kill() + + +class _InjectOom: + """Wraps ``_cu_mem_create``: the listed calls (1-based) return ``CUDA_ERROR_OUT_OF_MEMORY`` + (``always``: every call); records the handle objects created and released, compared by + identity (the driver reuses a handle's value once it is released).""" + + def __init__(self, monkeypatch, fail_calls=(), always=False): + from cuda.bindings import driver as cu + + self.cu = cu + self.fail_calls, self.always = set(fail_calls), always + self.calls, self.created, self.released = 0, [], [] + self._lock = __import__("threading").Lock() + real_create, real_release = H._cu_mem_create, cu.cuMemRelease + + def create(size, node): + with self._lock: + self.calls += 1 + fail = self.always or self.calls in self.fail_calls + if fail: + return cu.CUresult.CUDA_ERROR_OUT_OF_MEMORY, None + result = real_create(size, node) + if result[0] == cu.CUresult.CUDA_SUCCESS: + with self._lock: + self.created.append(result[1]) + return result + + def release(handle): + with self._lock: + self.released.append(handle) + return real_release(handle) + + monkeypatch.setattr(H, "_cu_mem_create", create) + monkeypatch.setattr(cu, "cuMemRelease", release) + + def was_released(self, handle) -> bool: + return any(handle is r for r in self.released) + + +def test_create_retries_after_host_oom(checkpoint, monkeypatch): + ckpt, tables = checkpoint + dev = torch.device("cuda", 0) + torch.cuda.set_device(dev) + monkeypatch.setenv(H.CREATE_RETRY_ENV, "60") + # One allocation thread, so the 4 slabs are created in order: attempt 1 creates slab 0 and fails + # on slab 1 (calls 1-2), attempt 2 the same (calls 3-4), attempt 3 creates all four (calls 5-8). + oom = _InjectOom(monkeypatch, fail_calls=(2, 4)) + name = f"pf-vmm-oom-{uuid.uuid4().hex[:12]}" + table = H.open_vmm_engram_table( + ckpt, local_rank=0, device=dev, headroom_bytes=1 << 30, name=name, threads=1 + ) + try: + attempts = table.manifest["fill"]["create_attempts"] + assert attempts == 3, attempts + assert oom.calls == 8 and len(oom.created) == 6, (oom.calls, len(oom.created)) + failed, kept = oom.created[:2], oom.created[2:] + assert all(oom.was_released(h) for h in failed), "a failed attempt's handles are released" + # (the NUMA probe's one-granule allocations are released too; count only the slabs) + assert sum(any(r is c for c in oom.created) for r in oom.released) == 2 + assert not any(oom.was_released(h) for h in kept) + assert len(table._handles) == 4 and all(any(h is k for k in kept) for h in table._handles) + for (layer, seed), got in _gather_all(table, dev).items(): + ref = _ref(_ids(64, seed * 100 + layer, "cpu"), *tables[layer]) + assert torch.equal(got.view(torch.int16), ref.view(torch.int16)), (layer, seed) + print( + f"\ncreate succeeded on attempt {attempts}; the 2 slabs of the failed attempts were released" + ) + finally: + table.close() + assert all(oom.was_released(h) for h in oom.created) + + +def test_create_oom_falls_back_to_shm_unless_vmm_is_explicit(checkpoint, monkeypatch): + ckpt, tables = checkpoint + dev = torch.device("cuda", 0) + torch.cuda.set_device(dev) + shm_path = f"/dev/shm/pf-vmm-oom-fallback-{uuid.uuid4().hex[:12]}.bin" + monkeypatch.setenv(H.PATH_ENV, shm_path) # also makes the lock and the socket name unique + monkeypatch.setenv(H.CREATE_RETRY_ENV, "2") + monkeypatch.delenv(H.BACKEND_ENV, raising=False) + oom = _InjectOom(monkeypatch, always=True) + barriers = [] + try: + with pytest.warns(RuntimeWarning, match="falling back to the shm backend"): + rank0 = H.open_shared_engram_table( + ckpt, + local_rank=0, + barrier=lambda: barriers.append(0), + device=dev, + threads=4, + headroom_bytes=1 << 30, + ) + # another local rank, after the barrier: nobody serves the VMM table, the shm file is complete + rank1 = H.open_shared_engram_table( + ckpt, local_rank=1, barrier=lambda: barriers.append(1), device=dev + ) + assert barriers == [0, 1] and oom.calls >= 2 and oom.created == [] + for table in (rank0, rank1): + assert isinstance(table, H.SharedEngramHostTable) and table.path == shm_path, type( + table + ) + table.check_gather_views(dev) + for layer, (w, s) in tables.items(): + assert torch.equal(table.weight(layer).view(torch.uint8), w.view(torch.uint8)) + assert torch.equal(table.scale(layer).view(torch.uint8), s) + assert rank0.fill_report is not None and not rank0.fill_report.reused + rank1.close() + rank0.close() + + # explicit vmm: no fallback; rank 0 joins the barrier, then raises + monkeypatch.setenv(H.BACKEND_ENV, "vmm") + barriers.clear() + with pytest.raises(H.VmmCreateError, match="could not be allocated"): + H.open_shared_engram_table( + ckpt, + local_rank=0, + barrier=lambda: barriers.append(0), + device=dev, + threads=4, + headroom_bytes=1 << 30, + ) + assert barriers == [0] + finally: + H.remove_shared_engram_table(shm_path, force=True) + for suffix in (".lock", ".owner", ".vmm.lock"): + if os.path.exists(shm_path + suffix): + os.unlink(shm_path + suffix) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_fp8_attention.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_fp8_attention.py new file mode 100644 index 000000000000..e316ffe1db12 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_fp8_attention.py @@ -0,0 +1,239 @@ +# 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. +"""The DeepSeek-V4.1-Flash-NVFP4 target's opt-in FP8 attention (TRTLLM_DSV41_FP8_ATTENTION=1), +on the real kernels. + +* The KV scales: one power of two per compressed-source group (every layer that reads a + source shares it, because one attention call takes one scale for both pools) and one per + window-only layer, the smallest that keeps the norm-weight bound below 448 x scale. +* The window read: the E4M3 read buffer is, row for row, the bf16 read buffer divided by the + layer's scale and rounded to E4M3, with the same renumbered index table. +* The compressed read in both forms of the one read buffer: per request (every entry of every + request) and per query (the selected top-k rows only). Each E4M3 form is the bf16 one scaled + and rounded, and the FP8 attention over either form is bitwise the same. + +Each case drives the target's own methods on a stand-in core holding just the attributes +they read, and a stand-in prefill target holding just the entries they call. +""" + +from __future__ import annotations + +import math +from types import SimpleNamespace + +import torch +from utils.util import skip_pre_blackwell + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified import ( # noqa: E501 + modeling, +) + +Core = modeling.DeepseekV41UnifiedModelingV2 +Prefill = modeling.PrefillTarget +SRC = 20 +E4 = torch.float8_e4m3fn + + +def _e4m3_of(bf16_rows: torch.Tensor, scale: float) -> torch.Tensor: + """uint8 codes of bf16 rows divided by a power-of-two scale (exact), rounded to E4M3.""" + return (bf16_rows.float() / scale).clamp(-448.0, 448.0).to(E4).view(torch.uint8) + + +def _core(max_t: int, fp8: bool, scale: float): + return SimpleNamespace( + w={"final_norm": torch.empty(1, device="cuda")}, + head_dim=512, + ratios={SRC: 1}, + kv_sources=[SRC], + kv_src={SRC: SRC}, + pp_layer_ids=list(range(40)), # one pipeline stage: every layer is local + local=frozenset(range(40)), + rank=0, + fp8_attn=fp8, + _fp8_kv_scale={SRC: scale}, + _const={"one_f32": torch.ones(1, dtype=torch.float32, device="cuda")}, + ), SimpleNamespace(swa_local_indices_cuda=torch.empty(max_t, 128, device="cuda")) + + +def _target(core): + """A stand-in prefill target: the core it reads and the read-back entries it calls.""" + return SimpleNamespace( + core=core, + _swa_read=modeling.Mxfp8Dequantize(), + _swa_read8=modeling.Mxfp8ToE4m3(), + _comp_read=modeling.Nvfp4KvToBf16(), + _comp_read8=modeling.Nvfp4KvToE4m3(), + ) + + +def test_kv_scales_group_by_compressed_source(): + """Layers 0-1 window-only, 2-19 ratio 2 over sources 2 / 8 / 14, 20-39 ratio 1 over source 20 + (the checkpoint's role table), with per-layer norm-weight maxima spread over 0.4 .. 6.4.""" + n = 40 + ratios = [0, 0] + [2] * 18 + [1] * 20 + kv_sources = (2, 8, 14, 20) + kv_src = [None if not r else max(s for s in kv_sources if s <= i) for i, r in enumerate(ratios)] + g = torch.Generator().manual_seed(0) + wmax = [0.4 + 6.0 * float(torch.rand(1, generator=g)) for _ in range(n)] + cmax = {s: 0.4 + 6.0 * float(torch.rand(1, generator=g)) for s in kv_sources} + w = {} + for i in range(n): + w[f"l{i}_kv_norm"] = torch.full((512,), wmax[i]).to(torch.bfloat16) + w[f"l{i}_kv_norm"][7] = -wmax[i] # the bound is on |w| + for s in kv_sources: + w[f"l{s}_comp_norm"] = torch.full((512,), cmax[s]).to(torch.bfloat16) + core = SimpleNamespace( + w=w, + num_layers=n, + pp_layer_ids=list(range(n)), + ratios=ratios, + kv_src=kv_src, + rank=1, + _fp8_kv_scale=[None] * n, + ) + Core._derive_fp8_kv_scales(core) + s = core._fp8_kv_scale + groups = {} + for i in range(n): + groups.setdefault(kv_src[i] if ratios[i] else -1 - i, []).append(i) + assert len(groups) == 2 + len(kv_sources) + for key, layers in groups.items(): + bound = max(float(w[f"l{i}_kv_norm"].abs().max()) for i in layers) + if key >= 0: + bound = max(bound, float(w[f"l{key}_comp_norm"].abs().max())) + bound *= modeling._FP8_KV_BOUND + scales = {s[i] for i in layers} + assert len(scales) == 1, (key, scales) + sc = scales.pop() + assert math.log2(sc) == int(math.log2(sc)), sc + assert bound <= 448.0 * sc < 2 * bound, ( + key, + bound, + sc, + ) # never saturates, and is the smallest + + +@skip_pre_blackwell +def test_window_read_is_the_bf16_read_scaled_to_e4m3(): + torch.manual_seed(1) + pages, T, scale = 64, 300, 2.0**-4 + pool = torch.zeros(pages, 128, 528, dtype=torch.uint8, device="cuda") + x = torch.randn(pages * 128, 16, 32, device="cuda") * torch.exp2( + torch.randint(-10, 5, (pages * 128, 16, 1), device="cuda").float() + ) + amax = x.abs().amax(-1).clamp_min(1e-4) + m, e = torch.frexp(amax.double() / 448.0) + e = torch.where(m == 0.5, e - 1, e).to(torch.int32) + q = (x / torch.exp2(e.double()).float().unsqueeze(-1)).clamp(-448, 448).to(E4) + flat = pool.view(-1, 528) + flat[:, :512] = q.view(-1, 512).view(torch.uint8) + flat[:, 512:] = (e + 127).to(torch.uint8) + codes, scales = pool[..., :512], pool[..., 512:] + gidx = torch.randint(0, pages * 128, (T, 128), dtype=torch.int32, device="cuda") + gidx[::7, 100:] = -1 + outs = {} + for fp8 in (False, True): + core, md = _core(8192, fp8, scale) + core._fp8_kv_scale = {0: scale} + md.compress_block_tables = {1: torch.zeros(4, 8, dtype=torch.int32, device="cuda")} + pools = Core._native_read_buffers(core, md) + pools["swa_q"] = {0: (codes, scales)} + ren, swa4 = Prefill._native_window(_target(core), 0, gidx, T, pools) + assert swa4.dtype == (E4 if fp8 else torch.bfloat16) + outs[fp8] = (ren.clone(), pools["swa_c"][: T * 128].clone()) + assert torch.equal(outs[False][0], outs[True][0]) + assert torch.equal(outs[True][1].view(torch.uint8), _e4m3_of(outs[False][1], scale)) + + +@skip_pre_blackwell +def test_compressed_read_both_forms_and_attention_are_consistent(): + """The bf16 and E4M3 buffers agree row for row (E4M3 = bf16 / scale, rounded) in both forms, and + the FP8 attention over the per-request and the per-query form is bitwise the same.""" + torch.manual_seed(0) + H, D, bk, pages, cols, scale = 16, 512, 128, 200, 40, 2.0**-1 + lens = [3000, 5100, 4097, 1] + seqs = len(lens) + pool = torch.randint(0, 256, (pages, bk, 288), dtype=torch.uint8, device="cuda") + pool[..., 256:] = torch.randint(0, 0x70, (pages, bk, 32), dtype=torch.uint8, device="cuda") + codes, scales = pool[..., :256], pool[..., 256:] + comp_q4 = {SRC: (codes.reshape(pages, bk, 1, 256), scales.reshape(pages, bk, 1, 32))} + comp_q1 = {SRC: (codes.reshape(pages * bk, 1, 1, 256), scales.reshape(pages * bk, 1, 1, 32))} + tbl = ( + torch.randperm(pages, device="cuda")[: seqs * cols] + .view(seqs, cols) + .to(torch.int32) + .contiguous() + ) + klens = torch.tensor(lens, dtype=torch.int32, device="cuda") + topk = torch.full((seqs, modeling._TOPK), -1, dtype=torch.int32, device="cuda") + for b, n in enumerate(lens): + pick = torch.randperm(n, device="cuda")[: min(n, modeling._TOPK)].sort().values + topk[b, : pick.numel()] = pick.to(torch.int32) + count = (topk >= 0).sum(1).to(torch.int32) + valid = topk >= 0 + e = topk.clamp(min=0).long() + b_idx = torch.arange(seqs, device="cuda").view(-1, 1) + per_request = torch.where(valid, b_idx * cols * bk + e, -1).to(torch.int32) + slots = torch.where(valid, tbl.long()[b_idx, e // bk] * bk + e % bk, -1).to(torch.int32) + + swa = (torch.randn(4, 1, 128, D, device="cuda") * 4).to(E4) + win = torch.randint(0, 4 * 128, (seqs, 128), dtype=torch.int32, device="cuda") + q = (torch.randn(seqs, 1, H, D, device="cuda") * 8).to(E4) + sinks = torch.randn(H, device="cuda") + ws = torch.zeros(128 * 1024 * 1024, dtype=torch.uint8, device="cuda") + topk_lens = (128 + count).contiguous() + bmm1 = D**-0.5 * 2.0**-3 * scale + + def attend(gidx, comp_pool): + return modeling.TrtllmSparseMlaDsv4Fp8()( + q, swa, comp_pool, ws, gidx, topk_lens, klens.clamp(min=128), sinks, bmm1, scale + ) + + rows = {} + out = {} + for fp8 in (False, True): + # One buffer that holds either form (4 request slots of 40 x 128 rows, 40 query rows). + core, md = _core(64, fp8, scale) + md.compress_block_tables = {1: tbl} + md.compressed_kv_lens_cuda = {1: klens} + pools = Core._native_read_buffers(core, md) + assert pools["comp_req_fit"] == seqs and pools["comp_q_fit"] >= seqs + pools.update(comp_q4=comp_q4, comp_q1=comp_q1) + for form in ("request", "topk"): + c4 = pools["comp_c4_req" if form == "request" else "comp_c4_q"][1] + assert c4.dtype == (E4 if fp8 else torch.bfloat16) + if form == "request": + Prefill._native_compressed(_target(core), SRC, {"md": md, "B": seqs}, pools) + gidx = torch.cat([win, per_request], 1) + sel = per_request + else: + pools["topk_out"].view(torch.uint8).fill_(0x5A) # stale finite rows, never attended + gidx = Prefill._native_topk( + _target(core), SRC, torch.cat([win, slots], 1), topk, seqs, pools, {} + ) + assert torch.equal(gidx[:, 128:], pools["topk_ren"][:seqs]) + sel = pools["topk_ren"][:seqs] + flat = c4.reshape(-1, D) + rows[form, fp8] = flat[sel.clamp(min=0).long()][valid].clone() + if fp8: + out[form] = attend(gidx, c4) + for form in ("request", "topk"): + assert torch.equal( + rows[form, True].view(torch.uint8), _e4m3_of(rows[form, False], scale) + ), form + assert torch.equal(rows["request", False], rows["topk", False]) + assert torch.equal(out["request"], out["topk"]) + assert torch.isfinite(out["request"].float()).all() diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_long_context.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_long_context.py new file mode 100644 index 000000000000..36927486698c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_long_context.py @@ -0,0 +1,706 @@ +# 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. +"""The DeepSeek-V4.1-Flash-NVFP4 target's long-context paths, on the real kernels. + +* The prefill index pass runs its query rows in tiles bounded by a logits budget; every + op in it is per row, so the top-k and the candidate-block masks are bitwise those of + one untiled pass. +* Under decoder replay the decoder pass scores its tail rows' candidate-block masks from + the candidate source's query rows it carried, against the step's gathered index K. + Those are the masks (and the re-index top-k) a full prefill gives the same rows, even + where the request's keys sit at another offset of the gather. +* Under native KV storage the attention reads the compressed rows from one buffer in either + of two forms, per request or per query (the selected top-k rows only); the attention + output is bitwise the same, and each step takes the form it fits that dequantizes less. +* Past a byte budget a step keeps no candidate masks: each re-index layer re-scores its + tiles' masks from the candidate source's query rows, with the same result. +* The decoder-replay planner takes prompts past the candidate blocks' 2048 x 8 positions. + +Each case drives the target's own methods on a stand-in core holding just the attributes +they read, and a stand-in prefill target holding just the entries they call, bound as the +prefill target binds them. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from utils.util import skip_pre_blackwell + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified import ( # noqa: E501 + modeling, +) + +Core = modeling.DeepseekV41UnifiedModelingV2 +Prefill = modeling.PrefillTarget +CAND_SRC, REINDEX, PLAIN = 20, 24, 14 +CANDIDATE_POSITIONS = modeling._CAND_TOPK_BLOCKS * modeling._CAND_BLOCK + + +@pytest.mark.parametrize("rows", [1, 63, 500, 8192]) +@pytest.mark.parametrize("s_len", [1, 16384, 20500, 1 << 20, 1 << 23]) +def test_index_row_tiles(rows, s_len): + """The tiles cover the rows once, in order, each within the budget (one row at least).""" + tiles = modeling._index_row_tiles(rows, s_len) + width = -(-(s_len + 256) // 256) * 256 + assert tiles[0][0] == 0 and tiles[-1][1] == rows + assert all(a < b for a, b in tiles) and all( + tiles[n][1] == tiles[n + 1][0] for n in range(len(tiles) - 1) + ) + if rows * width <= modeling._INDEX_TILE_ELEMS: + assert tiles == [(0, rows)] + for a, b in tiles: + assert (b - a) * width <= modeling._INDEX_TILE_ELEMS or b - a == 1 + if len(tiles) > 1 and modeling._INDEX_TILE_ELEMS // width >= 64: + assert all((b - a) % 64 == 0 for a, b in tiles[:-1]) + + +def _core(): + """A stand-in prefill target over a stand-in core: the index pass's entries.""" + core = SimpleNamespace( + cand_src=CAND_SRC, + ratios=[1] * 40, + _const={"zero_i32_rows": torch.zeros(8192, dtype=torch.int32, device="cuda")}, + ) + target = SimpleNamespace(core=core) + target._mqa_logits = modeling.DeepgemmFp8Fp4MqaLogits() + target._mqa_logits.bind_const(clean_logits=True, max_seqlen_k=0, logits_dtype=torch.float32) + target._cand = modeling.CandidateBlockSelect() + target._cand.bind_const( + topk_blocks=modeling._CAND_TOPK_BLOCKS, block_size=modeling._CAND_BLOCK, next_n=1 + ) + target._topk = modeling.FlashinferTopKRaggedDeterministic() + target._topk.bind_const(k=modeling._TOPK) + target._cand_masks = lambda *a: Prefill._cand_masks(target, *a) + return target + + +def _query(rows, seed): + """MXFP4 index query rows and fp32 head weights, as the target's _index makes them.""" + g = torch.Generator(device="cuda").manual_seed(seed) + q = torch.randn(rows, 32, 128, generator=g, device="cuda").to(torch.bfloat16) + qd, qsf = modeling.Fp4Quantize()(q, None, 32, True, False) + w = torch.randn(rows, 32, generator=g, device="cuda") * 2.0**-6 + return ( + qd.view(torch.int8).contiguous(), + qsf.view(torch.int32).reshape(rows, 32).contiguous(), + w.contiguous(), + ) + + +def _keys(n, seed): + """n gathered index-K rows (MXFP4 codes, one packed UE8M0 word each).""" + g = torch.Generator(device="cuda").manual_seed(seed) + k = torch.randn(n, 128, generator=g, device="cuda").to(torch.bfloat16) + kd, ksf = modeling.Fp4Quantize()(k, None, 32, True, False) + return kd.view(torch.int8).contiguous(), ksf.view(torch.int32).reshape(-1).contiguous() + + +def _rows(q, sel): + return tuple(t[sel].contiguous() for t in q) + + +class _Step: + """Context rows of requests in step order; request r's keys at offset ks_off[r] of the + gather. Ratio 1: a row at position p sees p + 1 keys.""" + + def __init__(self, reqs, key_order): + # reqs: list of (first position, rows); key_order: request indices in gather order. + self.reqs = reqs + total = [p0 + n for p0, n in reqs] + off, self.ks_off = 0, {} + for r in key_order: + self.ks_off[r] = off + off += total[r] + self.s_len = off + pos, ks = [], [] + for r, (p0, n) in enumerate(reqs): + pos += list(range(p0, p0 + n)) + ks += [self.ks_off[r]] * n + self.pos = torch.tensor(pos, dtype=torch.int32, device="cuda") + self.ks = torch.tensor(ks, dtype=torch.int32, device="cuda") + self.vis = self.pos + 1 + self.ke = self.ks + self.vis + self.rows = len(pos) + + def gather(self, keys_by_req): + order = sorted(self.ks_off, key=self.ks_off.get) + return ( + torch.cat([keys_by_req[r][0] for r in order]).contiguous(), + torch.cat([keys_by_req[r][1] for r in order]).contiguous(), + ) + + +def _run(core, layer, q, step, kvc, cache, tail_bound): + out = torch.empty(step.rows, modeling._TOPK, dtype=torch.int32, device="cuda") + Prefill._index_prefill( + core, + layer, + *q, + kvc, + step.ks, + step.ke, + step.vis, + step.s_len, + {"tail_bound": tail_bound}, + cache, + out, + ) + return out + + +def _masks(cache): + return torch.cat(cache["cand_ctx"], 0) + + +# Request 0 is long enough for its candidate blocks to bind; request 1 is short. +REQS = [(20000, 300), (0, 200)] + + +@pytest.fixture(scope="module") +def data(): + keys = {0: _keys(REQS[0][0] + REQS[0][1], 1), 1: _keys(REQS[1][0] + REQS[1][1], 2)} + rows = sum(n for _, n in REQS) + return keys, _query(rows, 3), _query(rows, 4) + + +@skip_pre_blackwell +@pytest.mark.parametrize("tile_rows", [64, 37]) +def test_prefill_index_tiling_is_exact(data, tile_rows, monkeypatch): + """The candidate source's masks and the re-index layer's top-k are the same in one + pass and in row tiles (64-row tiles, and 37-row ones that start at odd offsets).""" + keys, q20, q24 = data + step = _Step(REQS, [0, 1]) + assert step.s_len > CANDIDATE_POSITIONS + kvc = step.gather(keys) + core = _core() + + def passes(): + cache = {} + idx20 = _run(core, CAND_SRC, q20, step, kvc, cache, False) + idx24 = _run(core, REINDEX, q24, step, kvc, cache, False) + return idx20, idx24, _masks(cache), len(cache["cand_ctx"]) + + one = passes() + assert one[3] == 1 + width = -(-(step.s_len + 256) // 256) * 256 + monkeypatch.setattr(modeling, "_INDEX_TILE_ELEMS", width * tile_rows) + tiled = passes() + assert tiled[3] == -(-step.rows // tile_rows) + for a, b in zip(one[:3], tiled[:3]): + assert torch.equal(a, b) + + # The masks bind on the long request's rows (so the re-index top-k depends on them): + # the same query without masks (a layer below the source) selects differently there. + masks = one[2] + long_rows = slice(0, REQS[0][1]) + visible = [masks[t, : int(step.vis[t])] for t in range(REQS[0][1])] + assert all(int(v.sum()) <= CANDIDATE_POSITIONS < v.numel() for v in visible) + unmasked = _run(core, PLAIN, q24, step, kvc, {}, False) + assert not torch.equal(unmasked[long_rows], one[1][long_rows]) + assert torch.equal(unmasked[REQS[0][1] :], one[1][REQS[0][1] :]) + + +@skip_pre_blackwell +def test_tail_masks_match_full_prefill(data): + """Decoder replay: the decoder pass's masks and re-index top-k of each request's last + 128 rows equal the rows of a full prefill, although the final step gathers the + requests' keys in the other order (every tail row's range moves).""" + keys, q20, q24 = data + core = _core() + full = _Step(REQS, [0, 1]) + kvc = full.gather(keys) + cache = {} + _run(core, CAND_SRC, q20, full, kvc, cache, False) + idx24_full = _run(core, REINDEX, q24, full, kvc, cache, False) + masks_full = _masks(cache) + + w = 128 + tail_reqs = [(p0 + n - w, w) for p0, n in REQS] + sel = torch.cat( + [torch.arange(n - w, n) for _, n in REQS[:1]] + + [torch.arange(REQS[0][1] + REQS[1][1] - w, REQS[0][1] + REQS[1][1])] + ).cuda() + tail = _Step(tail_reqs, [1, 0]) + kvc_t = tail.gather(keys) + cache_t = {"cand_q": _rows(q20, sel)} + # Under decoder replay the candidate source scores no mask over the engine's rows. + probe = {} + _run(core, CAND_SRC, q20, full, kvc, probe, True) + assert "cand_ctx" not in probe + idx24_tail = _run(core, REINDEX, _rows(q24, sel), tail, kvc_t, cache_t, True) + assert torch.equal(idx24_tail, idx24_full[sel]) + masks_tail = _masks(cache_t) + for t in range(tail.rows): + v = int(tail.vis[t]) + src = int(sel[t]) + a, b = int(tail.ks[t]), int(full.ks[src]) + assert torch.equal(masks_tail[t, a : a + v], masks_full[src, b : b + v]) + assert not masks_tail[t, :a].any() and not masks_tail[t, a + v :].any() + + +@skip_pre_blackwell +@pytest.mark.parametrize("tail", [False, True], ids=["full_rows", "decoder_tail"]) +def test_candidate_masks_rescored_past_budget_are_exact(data, tail, monkeypatch): + """With the masks over the budget no step keeps them; the re-index layer re-scores + each tile's masks from the candidate source's query rows and selects the same top-k, + in a full-row step and in the decoder pass, in 37-row tiles.""" + keys, q20, q24 = data + core = _core() + if tail: + w = 128 + sel = torch.cat( + [ + torch.arange(REQS[0][1] - w, REQS[0][1]), + torch.arange(REQS[0][1] + REQS[1][1] - w, REQS[0][1] + REQS[1][1]), + ] + ).cuda() + step = _Step([(p0 + n - w, w) for p0, n in REQS], [1, 0]) + q20, q24 = _rows(q20, sel), _rows(q24, sel) + else: + step = _Step(REQS, [0, 1]) + kvc = step.gather(keys) + width = -(-(step.s_len + 256) // 256) * 256 + monkeypatch.setattr(modeling, "_INDEX_TILE_ELEMS", width * 37) + + def passes(): + cache = {"cand_q": q20} if tail else {} + if not tail: + _run(core, CAND_SRC, q20, step, kvc, cache, False) + return _run(core, REINDEX, q24, step, kvc, cache, tail), cache + + kept, cache_kept = passes() + assert len(cache_kept["cand_ctx"]) == -(-step.rows // 37) + monkeypatch.setattr(modeling, "_CAND_MASK_BYTES", step.rows * step.s_len - 1) + rescored, cache_rescored = passes() + assert "cand_ctx" not in cache_rescored + assert torch.equal(kept, rescored) + + +def _native_target(core): + """A stand-in prefill target: the core it reads and the read-back entries it calls.""" + return SimpleNamespace( + core=core, + _swa_read=modeling.Mxfp8Dequantize(), + _swa_read8=modeling.Mxfp8ToE4m3(), + _comp_read=modeling.Nvfp4KvToBf16(), + _comp_read8=modeling.Nvfp4KvToE4m3(), + ) + + +def _native_core(max_t): + return SimpleNamespace( + w={"final_norm": torch.empty(1, device="cuda")}, + head_dim=512, + ratios={CAND_SRC: 1}, + kv_sources=[CAND_SRC], + kv_src={CAND_SRC: CAND_SRC}, + local=frozenset(range(40)), # one pipeline stage: every layer is local + rank=0, + fp8_attn=False, + _const={"one_f32": torch.ones(1, dtype=torch.float32, device="cuda")}, + ), SimpleNamespace(swa_local_indices_cuda=torch.empty(max_t, 128, device="cuda")) + + +@skip_pre_blackwell +def test_both_read_buffer_forms_are_exact(): + """Native compressed KV: in one read buffer, attention through the per-request form and + through the per-query form (the selected rows only, renumbered) is bitwise the same on the + same step, including a request with one compressed entry and top-k rows padded with -1 + whose rows the per-query form leaves stale.""" + torch.manual_seed(0) + H, D, bk, pages, cols = 16, 512, 128, 200, 40 + lens = [3000, 5100, 4097, 1] + seqs = len(lens) + # The manager's 288-byte rows: 256 bytes of packed E2M1, 32 E4M3 scales (finite). + pool = torch.randint(0, 256, (pages, bk, 288), dtype=torch.uint8, device="cuda") + pool[..., 256:] = torch.randint(0, 0x70, (pages, bk, 32), dtype=torch.uint8, device="cuda") + codes, scales = pool[..., :256], pool[..., 256:] + comp_q4 = {CAND_SRC: (codes.reshape(pages, bk, 1, 256), scales.reshape(pages, bk, 1, 32))} + comp_q1 = { + CAND_SRC: (codes.reshape(pages * bk, 1, 1, 256), scales.reshape(pages * bk, 1, 1, 32)) + } + assert comp_q1[CAND_SRC][0].data_ptr() == codes.data_ptr() + tbl = ( + torch.randperm(pages, device="cuda")[: seqs * cols] + .view(seqs, cols) + .to(torch.int32) + .contiguous() + ) + klens = torch.tensor(lens, dtype=torch.int32, device="cuda") + + topk = torch.full((seqs, modeling._TOPK), -1, dtype=torch.int32, device="cuda") + for b, n in enumerate(lens): + pick = torch.randperm(n, device="cuda")[: min(n, modeling._TOPK)].sort().values + topk[b, : pick.numel()] = pick.to(torch.int32) + count = (topk >= 0).sum(1).to(torch.int32) + valid = topk >= 0 + e = topk.clamp(min=0).long() + b_idx = torch.arange(seqs, device="cuda").view(-1, 1) + per_request = torch.where(valid, b_idx * cols * bk + e, -1).to(torch.int32) + slots = torch.where(valid, tbl.long()[b_idx, e // bk] * bk + e % bk, -1).to(torch.int32) + + swa = torch.randn(4, 1, 128, D, device="cuda").to(torch.bfloat16) + win = torch.randint(0, 4 * 128, (seqs, 128), dtype=torch.int32, device="cuda") + q = torch.randn(seqs, 1, H, D, device="cuda").to(torch.bfloat16) + sinks = torch.randn(H, device="cuda") + ws = torch.zeros(128 * 1024 * 1024, dtype=torch.uint8, device="cuda") + topk_lens = (128 + count).contiguous() + + def attend(gidx, comp_pool, n_lens=topk_lens): + return modeling.TrtllmSparseMlaDsv4()( + q, swa, comp_pool, ws, gidx, n_lens, klens.clamp(min=128), sinks, D**-0.5, 1.0 + ) + + # One buffer that holds either form: 4 request slots of 40 x 128 rows, 40 query rows. + core, md = _native_core(max_t=64) + md.compress_block_tables = {1: tbl} + md.compressed_kv_lens_cuda = {1: klens} + pools = Core._native_read_buffers(core, md) + assert (pools["comp_req_fit"], pools["comp_q_fit"]) == ( + seqs, + seqs * cols * bk // modeling._TOPK, + ) + pools.update(comp_q4=comp_q4, comp_q1=comp_q1) + + # Per request: every entry of every request, dequantized once per KV source. + target = _native_target(core) + Prefill._native_compressed(target, CAND_SRC, {"md": md, "B": seqs}, pools) + out_request = attend(torch.cat([win, per_request], 1), pools["comp_c4_req"][1]) + + # Per query, in the same buffer, over stale finite rows (never attended). + pools["topk_out"].fill_(3.0) + cache = {} + gidx = Prefill._native_topk( + target, CAND_SRC, torch.cat([win, slots], 1), topk, seqs, pools, cache + ) + assert torch.equal(gidx[:, 128:], pools["topk_ren"][:seqs]) + out_query = attend(gidx, pools["comp_c4_q"][1]) + assert torch.equal(out_request, out_query) + # The layers that share the top-k set and KV source reuse the rows. + assert cache["comp_rows"][0] is topk + pools["topk_out"].fill_(0.0) + Prefill._native_topk(target, CAND_SRC, torch.cat([win, slots], 1), topk, seqs, pools, cache) + assert not pools["topk_out"].any() + + # The compressed rows matter: without them the output differs. + assert not torch.equal( + out_request, + attend(torch.cat([win, slots], 1), pools["comp_c4_q"][1], torch.full_like(topk_lens, 128)), + ) + + +def _form_core(): + return SimpleNamespace( + ratios={2: 2, 8: 2, 14: 2, 20: 1}, + kv_sources=[2, 8, 14, 20], + index_sources=[2, 8, 14, 20, 24, 28, 32, 36], + local=frozenset(range(40)), + late=21, + ) + + +def _form(caps, B, T, nc, kv, max_seq_len, max_t=8192, tail=None): + pools = {"comp_req_fit": caps[0], "comp_q_fit": caps[1]} + md = SimpleNamespace( + cached_token_lens_cpu=torch.tensor([k - 1 for k in kv], dtype=torch.int32), + seq_lens_kv=torch.ones(len(kv), dtype=torch.int32), + max_seq_len=max_seq_len, + swa_local_indices_cuda=torch.empty(max_t, 0), + ) + st = {"B": B, "T": T, "nc": nc} + return Core._comp_form(_form_core(), pools, md, st, tail) + + +def test_read_buffer_form_per_step(): + """The step takes a form it fits, the one that dequantizes fewer rows: prefill of a few + requests the per-request form (batch 256 / 16K / 8192 included), long-context decode the + per-query form, short-context decode the per-request form.""" + # Batch 256 / 16K / 8192: 248 request slots of 132 x 128 rows, 8192 query rows. + caps = (248, 8192) + assert _form(caps, 4, 8192, 4, [4096, 2048, 8192, 16000], 16384) == "request" + assert _form(caps, 64, 8192, 64, [16000] * 64, 16384) == "request" + assert _form(caps, 256, 8192, 8, [16000] * 256, 16384) == "topk" # 256 > 248 slots + assert _form(caps, 32, 32 * 6, 0, [0] * 32, 16384) == "request" # decode, 16K + # 1M / batch 32 / 8192: 3 request slots of 8196 x 128 rows, 8192 query rows. + caps = (3, 8192) + assert _form(caps, 1, 8192, 1, [1 << 20], 1 << 20) == "request" # 2.6M rows vs 33.5M + assert _form(caps, 32, 32 * 6, 0, [0] * 32, 1 << 20) == "topk" + assert _form(caps, 2, 2 * 6, 0, [0] * 2, 1 << 20) == "topk" # 2.6M rows vs 49K + # Bounded step: the decoder rows count for the per-query form. + tail = SimpleNamespace(bound=True, rows=128) + assert _form((3, 8192), 1, 8192, 1, [1 << 20], 1 << 20, tail=tail) == "request" + # A form the step does not fit is never taken. + assert _form((0, 8192), 1, 100, 1, [100], 1 << 20) == "topk" + assert _form((4, 64), 4, 100, 4, [1 << 20] * 4, 1 << 20) == "request" + + +def _plan_md(rid, cached, length, last): + return SimpleNamespace( + context_last_chunk=[last], + context_return_logits=[False], + request_ids=[rid], + seq_lens=torch.tensor([length], dtype=torch.int32), + cached_token_lens_cpu=torch.tensor([cached], dtype=torch.int32), + v41_num_replay_tokens_cpu=torch.tensor([0], dtype=torch.int32), + ) + + +@skip_pre_blackwell +def test_decoder_tail_plans_long_prompts(): + """A prompt past 2048 x 8 positions -- first seen whole, or recomputed after a pause + with its generated tokens in its prompt, over several chunks -- plans a bounded tail + of its last 128 positions.""" + core = SimpleNamespace( + replay=True, _witness=False, replay_window=128, _tail_records={}, _ar32=None + ) + core._ensure_ar = lambda n, dev: Core._ensure_ar(core, n, dev) + dev = torch.device("cuda") + + def plan(rid, cached, length, last): + st = {"T": length, "tc": length, "nc": 1, "gen": 0} + return modeling._DecoderTail(core, _plan_md(rid, cached, length, last), st, dev) + + tail = plan(7, 0, CANDIDATE_POSITIONS + 616, True) + r = tail.reqs[0] + assert tail.bound and r.t0 == r.s == CANDIDATE_POSITIONS + 616 - 128 + + # A recomputed request: chunks of 8192 up to prompt + generated = 17000. + n = CANDIDATE_POSITIONS + 616 + first = plan(9, 0, 8192, False) + assert first.bound and first.reqs[0].keep == 8192 - 128 + first.counted = True + first.finish() + assert core._tail_records[9].end == 8192 + second = plan(9, 8192, 8192, False) + second.counted = True + second.finish() + last = plan(9, 2 * 8192, n - 2 * 8192, True) + r = last.reqs[0] + assert last.bound and r.t0 == n - 128 and r.rec is core._tail_records[9] + + +@skip_pre_blackwell +@pytest.mark.parametrize("chunk", [1000, 4096]) +def test_index_k_gather_in_chunks_is_exact(chunk, monkeypatch): + """The index-K gather in chunks (at most _GATHER_MAX_ROWS rows a call) returns the rows of + one call, codes and scale words, for any chunk size; skipped (negative) slots too.""" + bk, kb, pages, rows = 128, 68, 64, 6000 + g = torch.Generator(device="cuda").manual_seed(11) + pool = torch.randint(0, 256, (pages, bk, 1, kb), dtype=torch.uint8, device="cuda", generator=g) + slot = torch.randperm(pages * bk, device="cuda", generator=g)[:rows] + page, s = slot // bk, slot % bk + base = page * (bk * kb) + data = (base + s * 64).to(torch.int64).contiguous() + scale = (base + bk * 64 + s * 4).to(torch.int64).contiguous() + data[17] = -1 # a skipped row + scale[17] = -1 + target = SimpleNamespace(_idx_gather=modeling.IndexerKCacheGather()) + target._idx_gather.bind_const(head_dim=modeling._IDX_DIM // 2) + one = Prefill._gather_index_k(target, pool, data, scale, rows, 1) + monkeypatch.setattr(modeling, "_GATHER_MAX_ROWS", chunk) + chunked = Prefill._gather_index_k(target, pool, data, scale, rows, 1) + assert chunked[0].shape == (rows, 64) and chunked[1].shape == (rows,) + keep = torch.ones(rows, dtype=torch.bool, device="cuda") + keep[17] = False + assert torch.equal(one[0][keep], chunked[0][keep]) and torch.equal( + one[1][keep], chunked[1][keep] + ) + flat = pool.view(-1) + ref = flat[data[keep, None] + torch.arange(64, device="cuda")].view(torch.int8) + assert torch.equal(chunked[0][keep], ref) + # The DeepGEMM logits op needs the scale words' base 16-byte aligned. + assert chunked[1].data_ptr() % 16 == 0 + + +def _tail_core(hidden=8): + w = 128 + a = torch.arange(w, device="cuda") + core = SimpleNamespace( + replay=True, + _witness=False, + replay_window=w, + _tail_records={}, + _ar32=None, + ds=None, + hidden=hidden, + hc_k=4 * hidden, + late=21, + num_layers=40, + _const={ + "tri": (a.view(1, -1) <= a.view(-1, 1)).contiguous(), + "col0": (a == 0).view(1, -1).contiguous(), + "zero_bf16": torch.zeros((1, 1), dtype=torch.bfloat16, device="cuda"), + }, + ) + core._ensure_ar = lambda n, dev: Core._ensure_ar(core, n, dev) + return core + + +def _overflow_step(core, max_rows): + """The step that outgrows itself: request 0 finishes a 2049-token prompt on a 1-token chunk + (its first 127 tail rows are stashed from the chunk before), 16 fresh 127-token prompts + finish too. T = 2033 engine rows; the tail holds 128 + 16 x 127 = 2160 rows.""" + dev = torch.device("cuda") + w, hid = core.replay_window, core.hidden + reqs = [(0, 2048, 1)] + [(k + 1, 0, 127) for k in range(16)] # (rid, cached, chunk) + T = sum(n for _, _, n in reqs) + nc = len(reqs) + pos = torch.cat([torch.arange(c, c + n) for _, c, n in reqs]).to(torch.int32).cuda() + rid_row = torch.cat([torch.full((n,), rid) for rid, _, n in reqs]).cuda() + + def val(rid, p): # an exact value per (request, position) + return (rid * 100000 + p).to(torch.float32) + + rows = { + "h": val(rid_row, pos).view(T, 1, 1).expand(T, 4, hid).to(torch.float32).contiguous(), + "yF": val(rid_row, pos).view(T, 1).expand(T, 24).contiguous(), + "rF": val(rid_row, pos).contiguous(), + "img": torch.zeros(T, dtype=torch.bool, device="cuda"), + } + cache = { + "topk": {1: val(rid_row, pos).to(torch.int32).view(T, 1).expand(T, 512).contiguous()}, + "cand_q": ( + torch.zeros(T, 32, 64, dtype=torch.int8, device="cuda"), + val(rid_row, pos).to(torch.int32).view(T, 1).expand(T, 32).contiguous(), + val(rid_row, pos).view(T, 1).expand(T, 32).contiguous(), + ), + "ctx_k": {}, + "sched": {}, + } + md = SimpleNamespace( + context_last_chunk=[True] * nc, + context_return_logits=[False] * nc, + request_ids=[rid for rid, _, _ in reqs], + seq_lens=torch.tensor([n for _, _, n in reqs], dtype=torch.int32), + cached_token_lens_cpu=torch.tensor([c for _, c, _ in reqs], dtype=torch.int32), + v41_num_replay_tokens_cpu=torch.zeros(nc, dtype=torch.int32), + sparse_mla_topk_lens={1: (128 + torch.clamp(pos + 1, max=512)).to(torch.int32)}, + compressed_visible_lens_cuda={1: (pos + 1).to(torch.int32)}, + v41_ctx_index_ks={1: torch.zeros(T, dtype=torch.int32, device="cuda")}, + kv_lens_cuda_runtime=torch.tensor( + [c + n for _, c, n in reqs], dtype=torch.int32, device="cuda" + ), + indexer_compress_block_tables={1: torch.zeros(nc, 4, dtype=torch.int32, device="cuda")}, + req_idx_per_token=torch.cat([torch.full((n,), b) for b, (_, _, n) in enumerate(reqs)]) + .to(torch.int32) + .cuda(), + swa_local_indices_cuda=torch.zeros(T, w, dtype=torch.int32, device="cuda"), + ) + # Request 0's record: its previous chunk ended at 2048 and stashed positions 1921..2047. + rec = modeling._TailRecord(0) + rec.end, rec.start = 2048, 2048 - 127 + sp = torch.arange(rec.start, 2048, device="cuda") + rec.stash = { + "h": val(torch.zeros_like(sp), sp).view(-1, 1, 1).expand(127, 4, hid).contiguous(), + "yF": val(torch.zeros_like(sp), sp).view(-1, 1).expand(127, 24).contiguous(), + "rF": val(torch.zeros_like(sp), sp).contiguous(), + "img": torch.zeros(127, dtype=torch.bool, device="cuda"), + "topk": val(torch.zeros_like(sp), sp) + .to(torch.int32) + .view(-1, 1) + .expand(127, 512) + .contiguous(), + "lens": torch.full((127,), 640, dtype=torch.int32, device="cuda"), + "vis": (sp + 1).to(torch.int32), + "cq": torch.zeros(127, 32, 64, dtype=torch.int8, device="cuda"), + "cqs": val(torch.zeros_like(sp), sp) + .to(torch.int32) + .view(-1, 1) + .expand(127, 32) + .contiguous(), + "cw": val(torch.zeros_like(sp), sp).view(-1, 1).expand(127, 32).contiguous(), + } + core._tail_records = {0: rec} + st = {"T": T, "tc": T, "nc": nc, "G": 0, "gen": 0, "B": nc, "next_n": 1, "pos": pos} + tail = modeling._DecoderTail(core, md, st, dev) + return tail, tail.enter(rows, st, cache, max_rows), T + + +@skip_pre_blackwell +def test_decoder_pass_slices_at_step_capacity(): + """A tail larger than the step capacity runs in slices of at most max_num_tokens rows, cut + at request boundaries; the slices hold exactly the rows, per-request fields and output rows + one pass would.""" + one_tail, one, T = _overflow_step(_tail_core(), 1 << 20) + tail, passes, _ = _overflow_step(_tail_core(), 2048) + assert [p[1]["T"] for p in one] == [2160] + assert [p[1]["T"] for p in passes] == [128 + 15 * 127, 127] + assert all(p[1]["T"] <= 2048 for p in passes) + for name in ("h", "yF", "rF"): + assert torch.equal(torch.cat([p[0][name] for p in passes]), one[0][0][name]) + for key in ("topk",): + assert torch.equal(torch.cat([p[2][key][1] for p in passes]), one[0][2][key][1]) + assert torch.equal(torch.cat([p[2]["cand_q"][1] for p in passes]), one[0][2]["cand_q"][1]) + for field in ( + "req_idx_per_token", + "swa_local_indices_cuda", + "kv_lens_cuda_runtime", + "v41_swa_floor_cuda", + "v41_swa_seq_lens_cuda", + ): + assert torch.equal( + torch.cat([getattr(p[1]["md"], field) for p in passes]), getattr(one[0][1]["md"], field) + ) + assert [p[1]["md"].num_contexts for p in passes] == [16, 1] + assert [p[1]["nc"] for p in passes] == [16, 1] and all(p[1]["B"] == p[1]["nc"] for p in passes) + cu = [p[1]["md"].cu_seq_lens_cuda.tolist() for p in passes] + assert cu[0][-1] == 2033 and cu[1] == [0, 127] + # Request 0's tail: its 127 stashed rows, then its 1-token chunk. + assert passes[0][0]["rF"][:128].tolist() == list(range(1921, 2049)) + # Each slice's output rows land where one pass puts them. + for t, ps in ((one_tail, one), (tail, passes)): + for k, (dec, _, _) in enumerate(ps): + t.activate(k) + t.scatter(dec["rF"].view(-1, 1).expand(-1, t.core.hidden).to(torch.bfloat16)) + assert torch.equal(tail.out, one_tail.out) + assert tail.out[0, 0].item() == 2048 # request 0's only row of this step: position 2048 + + +class _Capture: + """DSpark's capture: each call overwrites the layer's rows [0, n).""" + + def __init__(self): + self.rows = {} + + def maybe_capture_hidden_states(self, layer_id, hidden_states): + self.rows[layer_id] = hidden_states.clone() + + +@skip_pre_blackwell +def test_decoder_pass_slices_capture_every_layer(): + """DSpark captures the input of layers 37-39 over the step's rows. Run in slices, every + decoder row still holds its own input of each captured layer once the last slice has + captured: an earlier slice's rows do not carry a later layer's input into this one.""" + tail, passes, T = _overflow_step(_tail_core(), 2048) + hid = tail.core.hidden + spec = _Capture() + capture = modeling._TailCapture(spec, tail) + layers = (37, 38, 39) + for k, (dec, _, _) in enumerate(passes): + tail.activate(k) + for layer in layers: + capture.maybe_capture_hidden_states( + layer, (dec["h"] + 1000.0 * layer).reshape(-1, 4 * hid) + ) + for layer in layers: + got = spec.rows[layer].view(T, 4, hid) + for k, (dec, _, _) in enumerate(passes): + for r in tail.slices[k].fin: + want = dec["h"][r.q + r.n_old : r.q + r.n] + 1000.0 * layer + assert torch.equal(got[r.off + r.a : r.off + r.L], want), (layer, k, r.rid) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_patch_embed_pad.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_patch_embed_pad.py new file mode 100644 index 000000000000..3f7b90bc6d08 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_patch_embed_pad.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The vision patch embed never hands cublas_mm a row count its is_valid refuses. + +``cublas_mm`` refuses bf16 ``[M, 588] x [588, 1024]`` at M 2305..2432, where +cuBLASLt picks a split-K kernel with bf16 partial sums. ``is_valid`` does not +run on the hot path, so the target pads the patch rows past that band itself. +These checks hold that padding to the entry's own statement, for every patch +count the tower can see. They run on the CPU: ``is_valid`` reads only shapes, +dtypes and strides, so meta tensors are enough, and the GEMM in the +``_patch_embed`` check is a stand-in. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified import ( # noqa: E501 + modeling, +) + +_K, _N = modeling._VIT_PATCH_IN, modeling._VIT_DIM + + +def _is_valid(rows: int) -> None: + """The entry's guard on the patch embed's call, as the target binds it.""" + meta = dict(dtype=torch.bfloat16, device="meta") + cublas_mm.CublasMm().is_valid( + mat_a=torch.empty(rows, _K, **meta), + mat_b=torch.empty(_N, _K, **meta).t(), + bias=torch.empty(_N, **meta), + ) + + +def test_the_band_is_the_one_the_entry_refuses(): + lo, hi = modeling._VIT_SPLITK_ROWS + assert range(lo, hi + 1) == cublas_mm._PATCH_EMBED_SPLIT_K_ROWS + assert modeling._VIT_SPLITK_PAD == hi + 1 + for rows in (lo, hi): + with pytest.raises(AssertionError, match="pad M"): + _is_valid(rows) + for rows in (lo - 1, hi + 1): + _is_valid(rows) + + +def test_every_patch_count_reaches_cublas_mm_with_rows_it_accepts(): + """Each count from one patch up to the tower's per-call envelope.""" + for patches in range(1, modeling._VIT_MAX_PATCHES + 1): + rows = modeling._patch_embed_rows(patches) + assert rows >= patches + assert rows == patches or patches in cublas_mm._PATCH_EMBED_SPLIT_K_ROWS + assert rows not in cublas_mm._PATCH_EMBED_SPLIT_K_ROWS + _is_valid(rows) + + +def test_no_graph_bucket_needs_the_pad(): + for bucket in modeling._VIT_GRAPH_BUCKETS: + assert modeling._patch_embed_rows(bucket) == bucket + _is_valid(bucket) + + +@pytest.mark.parametrize("patches", [1, 2304, 2305, 2368, 2432, 2433, 9216]) +def test_patch_embed_pads_with_zero_rows_and_keeps_the_first_ones(patches): + calls = [] + + def gemm(x): + calls.append(tuple(x.shape)) + assert not torch.any(x[patches:]), "padded rows must be zero" + return x[:, :1].expand(-1, _N) + 0 # output row r carries input row r + + px = torch.arange(1, patches + 1, dtype=torch.float32)[:, None].expand(-1, _K).contiguous() + out = modeling.VisionTower._patch_embed(SimpleNamespace(_patch=gemm), px) + + assert calls == [(modeling._patch_embed_rows(patches), _K)] + assert out.shape == (patches, _N) + assert torch.equal(out[:, 0], px[:, 0]) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_vision.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_vision.py new file mode 100644 index 000000000000..0a6d72391428 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_vision.py @@ -0,0 +1,459 @@ +# 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. +"""The DeepSeek-V4.1-Flash-NVFP4 target's vision tower: CUDA-graph replay, per-image assembly. + +One GPU, the checkpoint's real vision weights (the target's own weight rows), no engine. + +* Replay is BITWISE the same bucket call run eagerly (same staged inputs, same slot metadata), for + single images from 7 to 9189 patches -- every bucket, both sides of bucket edges, the split-K + band 2305..2432 -- replayed in shuffled order, twice. +* Replay against the unpadded eager tower (the path with the switch at 0): the only difference + is the padded row count B vs P. Measured BITWISE equal -- ViT output and span rows -- at every + image size the input processor can produce (it never makes fewer than ~1469 patches), 31 sizes + from 1024 to 9189 in 23 buckets; gated bitwise there. A 7-patch call padded to 1024 rows takes + another cuBLAS algorithm at M = 7 and is gated to a band only. +* Against the golden (golden/vision): graph and eager each within 1.25x the vendor code's own + distance, and bitwise equal to each other. +* An image's span rows depend on that image alone: the same image inside different requests, + and a request's rows spliced chunk by chunk by the prefill target's ``_mm_step``, are + bitwise equal to the whole-prompt splice. +* The per-image tables (``aligner_index_local`` / ``span_source_local``) assemble bitwise the + rows the request-wide tables give. +* The memory-profiling hook runs the largest image; its peak is reported. + +Latency (eager vs replay per image) is printed, not gated. +""" + +from __future__ import annotations + +import json +import os +import random +import statistics +import time +from types import SimpleNamespace + +import pytest +import torch +from utils.llm_data import llm_models_root + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified import ( # noqa: E501 + modeling as tm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified import ( # noqa: E501 + weights as tw, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.input_processor import ( # noqa: E501 + ImageGrid, + PreprocessedImage, + VisionGeometry, + build_vision_inputs, +) + +# The checkpoint resolves under LLM_MODELS_ROOT, as the whole-model gates' does; +# MODELING_V2_DSV41_CKPT_DIR overrides it. The golden vision dumps have no staged +# location: MODELING_V2_DSV41_VISION_GOLDEN_DIR. A case that needs one skips when it +# is absent. +_MODELS_ROOT = llm_models_root() +_CKPT_DIR = os.environ.get("MODELING_V2_DSV41_CKPT_DIR") or ( + str(_MODELS_ROOT / "DeepSeek-V4.1-Flash-NVFP4") if _MODELS_ROOT else "" +) +_GOLDEN_DIR = os.environ.get("MODELING_V2_DSV41_VISION_GOLDEN_DIR", "") +_GOLDEN_CASES = ("mmmu_small", "mmmu_large", "mmmu_multi3", "carrots_corn") +DEV = torch.device("cuda") + +# The one GPU architecture the target is written for (its routing module's ``_SM``). +_SM = (10, 3) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != _SM, + reason=f"the target is certified on sm_{_SM[0]}{_SM[1]}", +) + + +def _geometry(): + return VisionGeometry( + patch_size=14, downsample_ratio=3, max_image_tokens=1024, min_pixels=544 * 544, + max_wh_ratio=None, hidden_size=1024, num_heads=16, num_layers=32, rope_theta=10000.0, + image_token_id=tm._IMAGE_TOKEN, + ) # fmt: skip + + +def _grid(n_vit_h: int, n_vit_w: int) -> ImageGrid: + g = ImageGrid(n_llm_h=-(-n_vit_h // 3), n_llm_w=-(-n_vit_w // 3), best_height=14 * n_vit_h, + best_width=14 * n_vit_w, n_vit_h=n_vit_h, n_vit_w=n_vit_w) # fmt: skip + assert g.num_span_tokens <= 1024, (n_vit_h, n_vit_w) + return g + + +def _image(grids, seed: int = 0, patches=None) -> dict: + """A request's ``multimodal_data["image"]`` (on the GPU) for images of these grids.""" + gen = torch.Generator().manual_seed(seed) + images, starts, pos = [], [], 0 + for k, g in enumerate(grids): + p = ( + patches[k] + if patches is not None + else (torch.rand((g.n_vit_h * g.n_vit_w, 3, 14, 14), generator=gen) * 2 - 1).to( + torch.bfloat16 + ) + ) + images.append(PreprocessedImage(patches=p, grid=g)) + starts.append(pos + 1) + pos += 1 + g.num_span_tokens + data = build_vision_inputs(images, starts, _geometry()) + return {k: (v.to(DEV) if torch.is_tensor(v) else v) for k, v in data.items()} + + +def _stats(o: torch.Tensor, g: torch.Tensor) -> dict: + o, g = o.double().flatten(1), g.double().flatten(1) + rr = (o - g).norm(dim=1) / g.norm(dim=1).clamp_min(1e-30) + return {"median": rr.median().item(), "max": rr.max().item(), + "bitexact": (o == g).sum().item() / o.numel()} # fmt: skip + + +@pytest.fixture(scope="module") +def params(): + if not _CKPT_DIR or not os.path.isdir(_CKPT_DIR): + pytest.skip(f"checkpoint not present: {_CKPT_DIR or 'LLM_MODELS_ROOT unset'}") + from safetensors import safe_open + + wm = json.load(open(os.path.join(_CKPT_DIR, "model.safetensors.index.json")))["weight_map"] + # The weight table's vision rows: replicated, not per layer, sourced as they are. + d = SimpleNamespace(hidden=5120, dtype=torch.bfloat16) + rows = [e for e in tw.MODEL_WEIGHTS.WEIGHTS if e.name.startswith("vis")] + handles, out = {}, {} + for entry in rows: + ((key, _dst),) = entry.sources(d, None, "") + keys, srcs, multi = tw._split_source(key) + ts = [] + for k, src in zip(keys, srcs): + f = wm[k] + if f not in handles: + handles[f] = safe_open(os.path.join(_CKPT_DIR, f), framework="pt", device="cpu") + ts.append(tw._materialize(handles[f].get_slice(k), src).to(DEV)) + tr = entry.transform + t = tr(tuple(ts) if multi else ts[0], None) if tr else ts[0] + shape, dtype = entry.shape(d), entry.dtype or d.dtype + assert tuple(t.shape) == tuple(shape) and t.dtype == dtype, (entry.name, t.shape, shape) + out[entry.name] = t.contiguous() + return out + + +def _tower(params, graphs: bool, disable_mm_encoder: bool = False): + core = SimpleNamespace( + w=params, rank=0, hidden=5120, disable_mm_encoder=disable_mm_encoder, _const={ + "false_b": torch.zeros((1,), dtype=torch.bool, device=DEV), + "true_b": torch.ones((1,), dtype=torch.bool, device=DEV)}, + ) # fmt: skip + tower = tm.VisionTower(core) + core._vision = tower + saved = os.environ.get(tm._VIT_GRAPH_ENV) + os.environ[tm._VIT_GRAPH_ENV] = "1" if graphs else "0" + try: + free0 = torch.cuda.mem_get_info()[0] + t0 = time.perf_counter() + tower.post_load_weights() + torch.cuda.synchronize() + tower.load_seconds = time.perf_counter() - t0 + tower.load_bytes = free0 - torch.cuda.mem_get_info()[0] + finally: + if saved is None: + os.environ.pop(tm._VIT_GRAPH_ENV) + else: + os.environ[tm._VIT_GRAPH_ENV] = saved + return tower + + +@pytest.fixture(scope="module") +def eager(params): + t = _tower(params, graphs=False) + assert t._g is None + return t + + +@pytest.fixture(scope="module") +def graph(params): + t = _tower(params, graphs=True) + assert set(t._g) == set(tm._VIT_GRAPH_BUCKETS) + print( + f"\n[vision] {len(t._g)} graphs captured in {t.load_seconds:.1f} s, " + f"{t.load_bytes / 2**20:.1f} MiB held" + ) + # Static buffers + one shared pool for the largest bucket's working set. + assert t.load_bytes < 2 * 2**30 + return t + + +# Single-image grids (n_vit_h, n_vit_w), all legal spans (<= 1024 positions): every bucket, both +# sides of bucket edges, exact bucket sizes, the split-K band 2305..2432, 7 and 9189 patches. +_GRIDS = [ + (1, 7), (32, 32), (3, 400), (39, 39), (33, 47), (32, 64), (3, 683), (48, 48), (5, 461), + (32, 76), (3, 900), (3, 1000), (3, 1100), (3, 1150), (60, 60), (51, 74), (3, 1365), (64, 64), + (3, 1366), (78, 60), (72, 72), (75, 75), (3, 2000), (3, 2100), (84, 84), (3, 2500), (90, 90), + (87, 99), (32, 272), (3, 2902), (6, 1530), (3, 3063), +] # fmt: skip + + +def _bucket(P: int) -> int: + return next(b for b in tm._VIT_GRAPH_BUCKETS if b >= P) + + +@torch.inference_mode() +def test_replay_is_bitwise_the_bucket_call_run_eagerly(graph): + grids = [_grid(h, w) for h, w in _GRIDS] + Ps = sorted({g.n_vit_h * g.n_vit_w for g in grids}) + assert {_bucket(P) for P in Ps} == set(tm._VIT_GRAPH_BUCKETS) + assert {2304 < P <= 2432 for P in Ps} == {True, False}, "the split-K band is covered" + assert min(Ps) == 7 and max(Ps) == 9189 + order = list(range(len(grids))) * 2 + random.Random(5).shuffle(order) + first = {} + for k in order: + image = _image([grids[k]], seed=k) + replayed = graph._vit(image, 0) + P = replayed.shape[0] + # The same bucket call, eagerly, over the same staged inputs. + B = graph._stage_graph_inputs(image, image["pixel_values"].reshape(P, -1), 0) + graph._graph_body(B, graph._g_metas[B]) + eager_bucket = graph._g_bufs["h"][:P].clone() + assert torch.equal(replayed, eager_bucket), (grids[k], B) + assert torch.isfinite(replayed).all() + if k in first: + assert torch.equal(first[k], replayed), "a second replay of the same image differs" + first[k] = replayed + + +@torch.inference_mode() +def test_replay_matches_the_unpadded_eager_tower(eager, graph): + worst = {"vit": 0.0, "rows": 0.0} + nbit = 0 + for k, (h, w) in enumerate(_GRIDS): + image = _image([_grid(h, w)], seed=100 + k) + P = h * w + ve, vg = eager._vit(image, 0), graph._vit(image, 0) + re_, rg = eager.image_rows(image, 0), graph.image_rows(image, 0) + sv, sr = _stats(vg, ve), _stats(rg, re_) + bitwise = torch.equal(vg, ve) and torch.equal(rg, re_) + nbit += bitwise + print( + f"[vision] P={P} bucket={_bucket(P)} vit row-rel median {sv['median']:.2e} max " + f"{sv['max']:.2e} bitexact {sv['bitexact']:.6f}; span rows max {sr['max']:.2e} " + f"bitexact {sr['bitexact']:.6f}{' (bitwise)' if bitwise else ''}" + ) + if P >= tm._VIT_GRAPH_BUCKETS[0]: + # Every size the processor can produce (>= ~1469 patches, min_pixels). + assert bitwise, P + else: + worst["vit"] = max(worst["vit"], sv["max"]) + worst["rows"] = max(worst["rows"], sr["max"]) + print( + f"[vision] graph == unpadded eager bitwise for {nbit}/{len(_GRIDS)} images; below 1024 " + f"patches worst {worst}" + ) + # Inside the target-vs-golden distance (vit ~1.5e-2 median, test_both_modes_match_the_golden). + assert worst["vit"] < 0.1 and worst["rows"] < 0.05, worst + + +def _golden(case: str): + path = os.path.join(_GOLDEN_DIR, case, "vision.pt") if _GOLDEN_DIR else "" + if not path or not os.path.isfile(path): + pytest.skip(f"golden not present: {path}") + g = torch.load(path, map_location="cpu") + cu = [int(v) for v in g["vit_cu_seqlens"]] + grids = [_grid(int(h), int(w)) for h, w in g["vit_grid_hw"].tolist()] + patches = [g["pixel_values"][cu[i] : cu[i + 1]] for i in range(len(grids))] + return g, _image(grids, patches=patches) + + +@torch.inference_mode() +@pytest.mark.parametrize("case", _GOLDEN_CASES) +def test_both_modes_match_the_golden(case, eager, graph): + g, image = _golden(case) + n = len(image["vit_cu_seqlens"]) - 1 + # The golden's own RoPE rows (its vit_out was computed with them). + drift = max((image["vit_rope_cos"].cpu() - g["rope_cos"]).abs().max().item(), + (image["vit_rope_sin"].cpu() - g["rope_sin"]).abs().max().item()) # fmt: skip + assert drift < 1e-6, drift + image["vit_rope_cos"] = g["rope_cos"].to(DEV) + image["vit_rope_sin"] = g["rope_sin"].to(DEV) + gold_vit = g["vit_out"].to(DEV) + gold_span = g["span_embeds"].to(DEV) + # The checkpoint's own modules, run as is, against the golden: the band. + vendor_vit = _stats(g["vendor_vit_out"].to(DEV), gold_vit) + vendor_al = _stats(g["vendor_aligner_out"].to(DEV), g["aligner_out"].to(DEV)) + res, outs = {}, {} + for name, tower in (("eager", eager), ("graph", graph)): + vit = torch.cat([tower._vit(image, i) for i in range(n)]) + rows = torch.cat([tower.image_rows(image, i) for i in range(n)]) + outs[name] = (vit, rows) + res[name] = (_stats(vit, gold_vit), _stats(rows, gold_span)) + print( + f"[vision] {case}: vendor vit median {vendor_vit['median']:.2e} max {vendor_vit['max']:.2e}, " + f"vendor aligner median {vendor_al['median']:.2e} max {vendor_al['max']:.2e}; " + + "; ".join( + f"{k} vit median {v[0]['median']:.2e} max {v[0]['max']:.2e} span median " + f"{v[1]['median']:.2e} max {v[1]['max']:.2e}" + for k, v in res.items() + ) + ) + for name in res: + vit_s, span_s = res[name] + assert vit_s["median"] <= 1.25 * vendor_vit["median"] + 1e-3, (name, vit_s, vendor_vit) + assert span_s["median"] <= 1.25 * vendor_al["median"] + 1e-3, (name, span_s, vendor_al) + # The golden images are 1517..8787 patches: the replay is the eager path, bit for bit. + assert torch.equal(outs["graph"][0], outs["eager"][0]) and torch.equal( + outs["graph"][1], outs["eager"][1] + ) + + +@torch.inference_mode() +@pytest.mark.parametrize("mode", ["eager", "graph"]) +def test_an_image_rows_depend_on_that_image_alone(mode, eager, graph): + tower = eager if mode == "eager" else graph + grids = [_grid(39, 39), _grid(51, 74), _grid(3, 683), _grid(33, 47), _grid(64, 64)] + gen = torch.Generator().manual_seed(9) + patches = [(torch.rand((g.n_vit_h * g.n_vit_w, 3, 14, 14), generator=gen) * 2 - 1).to(torch.bfloat16) + for g in grids] # fmt: skip + whole = _image(grids, patches=patches) + rows_whole = [tower.image_rows(whole, i) for i in range(len(grids))] + # Each image alone, and inside a request with other neighbours, in another order. + for i in range(len(grids)): + alone = _image([grids[i]], patches=[patches[i]]) + assert torch.equal(tower.image_rows(alone, 0), rows_whole[i]), i + perm = [3, 0, 4, 1, 2] + shuffled = _image([grids[i] for i in perm], patches=[patches[i] for i in perm]) + for j, i in enumerate(perm): + assert torch.equal(tower.image_rows(shuffled, j), rows_whole[i]), (j, i) + + +@torch.inference_mode() +def test_local_tables_assemble_the_request_wide_rows(eager): + """Per-image assembly == the request-wide gather over ``aligner_index`` / ``span_source_index`` + (the per-image ViT outputs concatenated), bitwise.""" + grids = [_grid(39, 39), _grid(5, 461), _grid(51, 74)] + image = _image(grids, seed=3) + n = len(grids) + v = torch.cat([eager._vit(image, i) for i in range(n)]) + padded = torch.cat([v, v.new_zeros(1, v.shape[1])]) + cells = padded[image["aligner_index"]] + ccu = image["aligner_cu_cells"] + out = torch.cat([eager._aligner(cells[ccu[i] : ccu[i + 1]]) for i in range(n)]) + ref = torch.cat([out, eager._c["span_tail"]])[image["span_source_index"]] + ours = torch.cat([eager.image_rows(image, i) for i in range(n)]) + assert torch.equal(ours, ref) + + +def _mm_param(image, layout, row0, begin, end): + return SimpleNamespace( + multimodal_data={"image": image, "token_layout": layout}, + multimodal_runtime=SimpleNamespace(past_seen_token_num=begin, chunk_end_pos=end), + input_ids_start_offset=row0) # fmt: skip + + +@torch.inference_mode() +@pytest.mark.parametrize("mode", ["eager", "graph"]) +def test_chunked_splice_equals_the_whole_prompt(mode, eager, graph): + """``_mm_step`` over 1..4 chunks (a span cut by every boundary) == one chunk, bitwise, and + only the images a chunk reaches run the tower.""" + tower = eager if mode == "eager" else graph + grids = [_grid(39, 39), _grid(51, 74), _grid(33, 47)] + image = _image(grids, seed=21) + starts, lens = image["span_starts"], image["span_lengths"] + L = starts[-1] + lens[-1] + 5 + mask = torch.zeros(L, dtype=torch.bool) + for s, n in zip(starts, lens): + mask[s : s + n] = True + layout = {"image_mask": mask.to(DEV), "engram_dead_mask": mask.to(DEV)} + target = SimpleNamespace(core=SimpleNamespace(_vision=tower, _const=tower._core._const)) + ref = torch.zeros(L, 5120, dtype=torch.bfloat16, device=DEV) + tm.PrefillTarget._mm_step(target, [_mm_param(image, layout, 0, 0, L)], L, L, DEV, ref) + for bounds in ([0, 100, L], [0, starts[1] + 7, starts[2] + 3, L], [0, 50, 300, 600, L]): + emb = torch.zeros(L, 5120, dtype=torch.bfloat16, device=DEV) + calls = [] + orig = tower.image_rows + + def spy(img, i, _orig=orig): + calls.append(i) + return _orig(img, i) + + tower.image_rows = spy + try: + for a, b in zip(bounds[:-1], bounds[1:]): + tm.PrefillTarget._mm_step(target, [_mm_param(image, layout, 0, a, b)], b - a, + b - a, DEV, emb[a:b]) # fmt: skip + finally: + del tower.image_rows + assert torch.equal(emb, ref), bounds + reached = sum(sum(1 for s, n in zip(starts, lens) if max(s, a) < min(s + n, b)) + for a, b in zip(bounds[:-1], bounds[1:])) # fmt: skip + assert len(calls) == reached, (bounds, calls) + assert torch.equal(ref[~mask.to(DEV)], torch.zeros_like(ref[~mask.to(DEV)])) + + +@torch.inference_mode() +@pytest.mark.parametrize("mode", ["eager", "graph"]) +def test_profile_largest_image(mode, eager, graph): + tower = eager if mode == "eager" else graph + torch.cuda.synchronize() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + tower.profile_largest_image(8192) + peak = torch.cuda.max_memory_allocated() - base + print(f"[vision] profile_largest_image({mode}): peak {peak / 2**20:.1f} MiB above the model") + assert torch.cuda.memory_allocated() - base < 2**20 + + +@torch.inference_mode() +def test_disable_mm_encoder_captures_nothing_and_stays_correct(params, eager, graph): + """``llm_args.disable_mm_encoder`` (a generation-only disaggregated server): no graphs, no + graph memory; an image that arrives anyway runs eagerly, bitwise as the switch-0 tower. The + default engine (the fixture ``graph``: context / aggregated) still captures every bucket.""" + free0 = torch.cuda.mem_get_info()[0] + off = _tower(params, graphs=True, disable_mm_encoder=True) + assert off._g is None and off._g_bufs is None + assert off.load_bytes < 64 * 2**20, off.load_bytes + assert set(graph._g) == set(tm._VIT_GRAPH_BUCKETS) + image = _image([_grid(51, 74)], seed=33) + assert torch.equal(off.image_rows(image, 0), eager.image_rows(image, 0)) + print( + f"[vision] disable_mm_encoder: 0 graphs, {off.load_bytes / 2**20:.1f} MiB at derive " + f"(graph tower: {graph.load_bytes / 2**20:.1f} MiB); free before {free0 / 2**30:.1f} GiB" + ) + + +@torch.inference_mode() +def test_print_latency(eager, graph): + """Per-image tower latency, host included (wall clock to a synchronized end).""" + for h, w in ((39, 39), (48, 49), (78, 60), (90, 96), (3, 3063)): + image = _image([_grid(h, w)], seed=7) + res = {} + for name, tower in (("eager", eager), ("graph", graph)): + for _ in range(3): + tower._vit(image, 0) + torch.cuda.synchronize() + ts = [] + for _ in range(10): + t0 = time.perf_counter() + tower._vit(image, 0) + torch.cuda.synchronize() + ts.append((time.perf_counter() - t0) * 1e3) + res[name] = statistics.median(ts) + print( + f"[vision] latency P={h * w} bucket={_bucket(h * w)}: eager {res['eager']:.2f} ms, " + f"graph {res['graph']:.2f} ms, x{res['eager'] / res['graph']:.2f}" + ) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_weight_table.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_weight_table.py new file mode 100644 index 000000000000..7a9487e249e6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_v41_weight_table.py @@ -0,0 +1,547 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DeepSeek-V4.1's weight tables declare and source what the hand-written code did. + +The two references below are transcriptions of what the target stated before it +moved onto ``ModelWeights``: the declaration block of the core's ``__init__`` and +the hand-written manifests (trunk and DSpark drafter), whose sources were proven on +GB300 by the target's gates. They are copies rather than imports on purpose -- the +point is to hold the tables to what the model used to say. One renaming was +deliberate: the drafter's rows moved from the trunk's ``ds{s}_*`` / ``ds_*`` names +into their own table as stage ``l{s}_*`` and plain ``*``. + +Every parallel member is checked at a non-zero rank, so a slice that ignores its +rank cannot pass, along with a two-stage pipeline split and both Engram residencies. + +No GPU: nothing here allocates (meta tensors), and ``dims`` reads ``model_config`` +alone. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +__extra_import_path__ = [".."] # noqa: F841 -- repo's file-scoped import hook + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified import ( # noqa: E501 + weights as W, +) + +# DeepSeek-V4.1-Flash's own text_config numbers, so the comparison is against the +# shapes this target actually ships rather than a toy configuration. +_TC = dict( + num_hidden_layers=40, + hidden_size=5120, + vocab_size=129280, + num_attention_heads=64, + head_dim=512, + q_lora_rank=1280, + o_groups=8, + o_lora_rank=1024, + n_routed_experts=384, + moe_intermediate_size=2304, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0] * 3, + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + engram_layer_ids=[1, 14], + engram_num_embeddings=[384006168, 384016682], + engram_max_ngram_size=4, + engram_n_heads=8, + engram_head_dim=256, + num_nextn_predict_layers=3, + dspark_n_routed_experts=128, + dspark_markov_rank=256, + dspark_target_layer_ids=[37, 38, 39], +) + +_HOST_ENV = "TRTLLM_DSV41_ENGRAM_HOST_OFFLOAD" + +# (tp, pp_rank, pp, ep, moe_tp, attention_dp, tp_rank): every member at a non-zero rank. +_MEMBERS = { + "tep4": (4, 0, 1, 4, 1, False, 1), + "dep4": (4, 0, 1, 4, 1, True, 2), + "tp4": (4, 0, 1, 1, 4, False, 3), + "tep8": (8, 0, 1, 8, 1, False, 5), + "pp2xtep4-stage1": (4, 1, 2, 4, 1, False, 2), +} + + +def _mapping(tp, pp_rank, pp, ep, moe_tp, dp, tp_rank): + n = _TC["num_hidden_layers"] + per = n // pp + return SimpleNamespace( + tp_size=tp, + tp_rank=tp_rank, + pp_size=pp, + pp_rank=pp_rank, + enable_attention_dp=dp, + moe_ep_size=ep, + moe_ep_rank=tp_rank if ep > 1 else 0, + moe_tp_size=moe_tp, + moe_tp_rank=tp_rank if moe_tp > 1 else 0, + pp_layers=lambda num_layers: list(range(pp_rank * per, (pp_rank + 1) * per)), + is_last_pp_rank=lambda: pp_rank == pp - 1, + ) + + +def _core(member: str): + tc = SimpleNamespace(**_TC) + return SimpleNamespace( + model_config=SimpleNamespace( + pretrained_config=SimpleNamespace(text_config=tc), + mapping=_mapping(*_MEMBERS[member]), + torch_dtype=torch.bfloat16, + ) + ) + + +def _reference_core(member: str, host: bool) -> SimpleNamespace: + """The partition the core's ``__init__`` derived, transcribed.""" + tp, pp_rank, pp, ep, moe_tp, dp, tp_rank = _MEMBERS[member] + c = SimpleNamespace(**_TC) + attn_tp = 1 if dp else tp + attn_rank = 0 if dp else tp_rank + r = SimpleNamespace() + r.hidden, r.vocab, r.head_dim, r.q_lora = c.hidden_size, c.vocab_size, c.head_dim, c.q_lora_rank + r.o_lora, r.num_experts, r.moe_inter = ( + c.o_lora_rank, + c.n_routed_experts, + c.moe_intermediate_size, + ) + r.h_loc = c.num_attention_heads // attn_tp + r.g_loc = c.o_groups // attn_tp + r.attn_rank = attn_rank + r.group_in = c.num_attention_heads * c.head_dim // c.o_groups + if ep == tp: + r.e_loc = c.n_routed_experts // tp + r.e_off = r.e_loc * tp_rank + r.i_loc, r.i_off = c.moe_intermediate_size, 0 + else: + r.e_loc, r.e_off = c.n_routed_experts, 0 + r.i_loc = c.moe_intermediate_size // tp + r.i_off = r.i_loc * tp_rank + per = c.num_hidden_layers // pp + r.pp_layer_ids = list(range(pp_rank * per, (pp_rank + 1) * per)) + r.local = frozenset(r.pp_layer_ids) + r.pp_prev = r.pp_layer_ids[0] - 1 if r.pp_layer_ids[0] > 0 else None + r.ratios = c.compress_ratios[: c.num_hidden_layers] + r.kv_sources = tuple(c.kv_source_layer_ids) + r.index_sources = tuple(c.index_source_layer_ids) + r.engram_layers = tuple(c.engram_layer_ids) + r.eng_rows = list(c.engram_num_embeddings) + r.eng_part = [(n + tp - 1) // tp for n in r.eng_rows] + r.eng_start = [p * tp_rank for p in r.eng_part] + r.eng_end = [min(s + p, n) for s, p, n in zip(r.eng_start, r.eng_part, r.eng_rows)] + r.eng_host = host + if host: + r.eng_start, r.eng_end = [0] * len(r.eng_rows), list(r.eng_rows) + r.eng_dim = c.engram_head_dim + r.eng_cols = (c.engram_max_ngram_size - 1) * c.engram_n_heads + r.eng_in = r.eng_cols * r.eng_dim + r.eng_wkv_n = 5 * c.hidden_size + r.hc_k = 4 * c.hidden_size + return r + + +def _mx_sf(n, k): + return ((n + 127) // 128 * 128) * ((k // 32 + 3) // 4 * 4) + + +def _reference_declaration(member: str, host: bool) -> dict: + """The trunk + vision declaration block of the core's ``__init__``, transcribed: + name -> (shape, dtype).""" + r = _reference_core(member, host) + bf, f8, u8, f32 = torch.bfloat16, torch.float8_e4m3fn, torch.uint8, torch.float32 + H, mix = r.hidden, 24 + w = {"embed": ((r.vocab, H), bf), "final_norm": ((H,), bf)} + for i in range(40): + if i not in r.local: + if i == r.pp_prev: + w[f"l{i}_hc_ffn_fn"] = ((mix, r.hc_k), f32) + w[f"l{i}_hc_ffn_base"] = ((mix,), f32) + w[f"l{i}_hc_ffn_scale"] = ((3,), f32) + continue + for s in ("attn", "ffn"): + w[f"l{i}_hc_{s}_fn"] = ((mix, r.hc_k), f32) + w[f"l{i}_hc_{s}_base"] = ((mix,), f32) + w[f"l{i}_hc_{s}_scale"] = ((3,), f32) + w[f"l{i}_attn_norm"] = ((H,), bf) + w[f"l{i}_ffn_norm"] = ((H,), bf) + w[f"l{i}_wqkv_a"] = ((r.q_lora + r.head_dim, H), f8) + w[f"l{i}_wqkv_a_sf"] = ((_mx_sf(r.q_lora + r.head_dim, H),), u8) + w[f"l{i}_q_norm"] = ((r.q_lora,), bf) + w[f"l{i}_kv_norm"] = ((r.head_dim,), bf) + w[f"l{i}_wq_b"] = ((r.h_loc * r.head_dim, r.q_lora), f8) + w[f"l{i}_wq_b_sf"] = ((_mx_sf(r.h_loc * r.head_dim, r.q_lora),), u8) + w[f"l{i}_sink"] = ((r.h_loc,), f32) + w[f"l{i}_wo_a"] = ((r.g_loc * r.o_lora, r.group_in), bf) + w[f"l{i}_wo_b"] = ((H, r.g_loc * r.o_lora), f8) + w[f"l{i}_wo_b_sf"] = ((_mx_sf(H, r.g_loc * r.o_lora),), u8) + if i in r.kv_sources: + if r.ratios[i] == 2: + w[f"l{i}_comp_w"] = ((2 * r.head_dim, H), f32) + else: + w[f"l{i}_comp_w"] = ((r.head_dim, H), bf) + w[f"l{i}_comp_norm"] = ((r.head_dim,), bf) + w[f"l{i}_idx_wk"] = ((128, r.head_dim), bf) + w[f"l{i}_idx_k_norm"] = ((128,), bf) + if i in r.index_sources: + w[f"l{i}_idx_wq_b"] = ((32 * 128, r.q_lora), f8) + w[f"l{i}_idx_wq_b_sf"] = ((_mx_sf(32 * 128, r.q_lora),), u8) + w[f"l{i}_idx_wproj"] = ((32, H), bf) + w[f"l{i}_gate_w"] = ((r.num_experts, H), bf) + w[f"l{i}_gate_bias"] = ((r.num_experts,), f32) + w[f"l{i}_gate_bias_vl"] = ((r.num_experts,), f32) + w[f"l{i}_fc1_w"] = ((r.e_loc, 2 * r.i_loc, H // 2), u8) + w[f"l{i}_fc1_s"] = ((r.e_loc, 2 * r.i_loc, H // 16), u8) + w[f"l{i}_fc2_w"] = ((r.e_loc, H, r.i_loc // 2), u8) + w[f"l{i}_fc2_s"] = ((r.e_loc, H, r.i_loc // 16), u8) + for name in ("isc1", "isc3", "isc2", "ws1", "ws3", "ws2"): + w[f"l{i}_e_{name}"] = ((r.num_experts,), f32) + w[f"l{i}_sh_w13"] = ((2 * r.moe_inter, H), f8) + w[f"l{i}_sh_w13_sf"] = ((_mx_sf(2 * r.moe_inter, H),), u8) + w[f"l{i}_sh_w2"] = ((H, r.moe_inter), f8) + w[f"l{i}_sh_w2_sf"] = ((_mx_sf(H, r.moe_inter),), u8) + if i in r.engram_layers: + li = r.engram_layers.index(i) + if not host: + w[f"l{i}_eng_table"] = ((r.eng_part[li], r.eng_dim), f8) + w[f"l{i}_eng_scale"] = ((r.eng_part[li], r.eng_dim // 32), u8) + w[f"l{i}_eng_wkv"] = ((r.eng_wkv_n, r.eng_in), f8) + w[f"l{i}_eng_wkv_sf"] = ((_mx_sf(r.eng_wkv_n, r.eng_in),), u8) + w[f"l{i}_eng_q"] = ((4, H), bf) + w[f"l{i}_eng_k"] = ((4, H), bf) + D, inter = 1024, 2816 + w.update( + vis_patch_w=((D, 588), bf), + vis_patch_b=((D,), bf), + vis_norm=((D,), bf), + vis_al_w1=((H, D * 9), bf), + vis_al_b1=((H,), bf), + vis_al_w2=((H, H), bf), + vis_al_b2=((H,), bf), + vis_span_tail=((3, H), bf), + ) + for L in range(32): + p = f"vis{L}_" + w[p + "norm1"] = ((D,), bf) + w[p + "wqkv_w"] = ((3 * D, D), bf) + w[p + "wqkv_b"] = ((3 * D,), bf) + w[p + "wo_w"] = ((D, D), bf) + w[p + "wo_b"] = ((D,), bf) + w[p + "norm2"] = ((D,), bf) + w[p + "w1"] = ((2 * inter, D), bf) + w[p + "w2"] = ((D, inter), bf) + return w + + +def _reference_draft_declaration() -> dict: + """The drafter's parameter specs, transcribed, under the table's names.""" + f8, u8, f32, bf = torch.float8_e4m3fn, torch.uint8, torch.float32, torch.bfloat16 + H, hd, ql, inter, E, heads, og, ol = 5120, 512, 1280, 2304, 128, 64, 8, 1024 + out = {} + for s in range(3): + p = f"l{s}_" + for x in ("attn", "ffn"): + out[p + f"hc_{x}_fn"] = ((24, 4 * H), f32) + out[p + f"hc_{x}_base"] = ((24,), f32) + out[p + f"hc_{x}_scale"] = ((3,), f32) + out[p + "attn_norm"] = ((H,), bf) + out[p + "ffn_norm"] = ((H,), bf) + out[p + "wqkv_a"] = ((ql + hd, H), f8) + out[p + "wqkv_a_sf"] = ((_mx_sf(ql + hd, H),), u8) + out[p + "q_norm"] = ((ql,), bf) + out[p + "kv_norm"] = ((hd,), bf) + out[p + "wq_b"] = ((heads * hd, ql), f8) + out[p + "wq_b_sf"] = ((_mx_sf(heads * hd, ql),), u8) + out[p + "sink"] = ((heads,), f32) + out[p + "wo_a"] = ((og * ol, heads * hd // og), bf) + out[p + "wo_b"] = ((H, og * ol), f8) + out[p + "wo_b_sf"] = ((_mx_sf(H, og * ol),), u8) + out[p + "gate_w"] = ((E, H), f32) + out[p + "gate_bias"] = ((E,), f32) + out[p + "fc1_w"] = ((E, 2 * inter, H // 2), u8) + out[p + "fc1_s"] = ((E, 2 * inter, H // 32), u8) + out[p + "fc2_w"] = ((E, H, inter // 2), u8) + out[p + "fc2_s"] = ((E, H, inter // 32), u8) + out[p + "sh_w13"] = ((2 * inter, H), f8) + out[p + "sh_w13_sf"] = ((_mx_sf(2 * inter, H),), u8) + out[p + "sh_w2"] = ((H, inter), f8) + out[p + "sh_w2_sf"] = ((_mx_sf(H, inter),), u8) + out["main_proj"] = ((H, 3 * H), f8) + out["main_proj_sf"] = ((_mx_sf(H, 3 * H),), u8) + out["main_norm"] = ((H,), bf) + out["wkv3"] = ((3 * hd, H), f8) + out["wkv3_sf"] = ((_mx_sf(3 * hd, H),), u8) + out["norm"] = ((H,), bf) + out["mk_embed"] = ((129280, 256), bf) + out["mk_head"] = ((129280, 256), f32) + return out + + +def _reference_manifest(member: str, host: bool) -> dict: + """The hand-written trunk manifest, transcribed: param -> [(keys, srcs, dst, + has_transform)]. Host-resident Engram rows are no parameter and are left out.""" + r = _reference_core(member, host) + rows = {} + + def row(key, src=None, dst=None, tf=False): + keys = key if isinstance(key, tuple) else (key,) + if src is None: + srcs = (None,) * len(keys) + elif isinstance(key, tuple): + srcs = src + else: + srcs = (src,) + return (keys, srcs, dst, tf) + + rows["embed"] = [row("embed.weight")] + rows["final_norm"] = [row("norm.weight")] + h0 = r.attn_rank * r.h_loc + q_rows = slice(h0 * r.head_dim, (h0 + r.h_loc) * r.head_dim) + q_sf_rows = slice(q_rows.start // 32, q_rows.stop // 32) + g0 = r.attn_rank * r.g_loc + o_rows = slice(g0 * r.o_lora, (g0 + r.g_loc) * r.o_lora) + o_sf_rows = slice(o_rows.start // 32, o_rows.stop // 32) + all_ = slice(None) + for i in range(40): + p = f"layers.{i}" + if i not in r.local: + if i == r.pp_prev: + for f in ("fn", "base", "scale"): + rows[f"l{i}_hc_ffn_{f}"] = [row(f"{p}.hc_ffn_{f}")] + continue + for s in ("attn", "ffn"): + for f in ("fn", "base", "scale"): + rows[f"l{i}_hc_{s}_{f}"] = [row(f"{p}.hc_{s}_{f}")] + rows[f"l{i}_attn_norm"] = [row(f"{p}.attn_norm.weight")] + rows[f"l{i}_ffn_norm"] = [row(f"{p}.ffn_norm.weight")] + a = f"{p}.attn" + rows[f"l{i}_wqkv_a"] = [ + row(f"{a}.wq_a.weight", None, slice(0, r.q_lora)), + row(f"{a}.wkv.weight", None, slice(r.q_lora, r.q_lora + r.head_dim)), + ] + rows[f"l{i}_wqkv_a_sf"] = [row((f"{a}.wq_a.scale", f"{a}.wkv.scale"), None, None, True)] + rows[f"l{i}_q_norm"] = [row(f"{a}.q_norm.weight")] + rows[f"l{i}_kv_norm"] = [row(f"{a}.kv_norm.weight")] + rows[f"l{i}_wq_b"] = [row(f"{a}.wq_b.weight", (q_rows,))] + rows[f"l{i}_wq_b_sf"] = [row(f"{a}.wq_b.scale", (q_sf_rows,), None, True)] + rows[f"l{i}_sink"] = [row(f"{a}.attn_sink", (slice(h0, h0 + r.h_loc),))] + rows[f"l{i}_wo_a"] = [ + row((f"{a}.wo_a.weight", f"{a}.wo_a.scale"), ((o_rows,), (o_sf_rows,)), None, True) + ] + rows[f"l{i}_wo_b"] = [row(f"{a}.wo_b.weight", (all_, o_rows))] + rows[f"l{i}_wo_b_sf"] = [row(f"{a}.wo_b.scale", (all_, o_sf_rows), None, True)] + if i in r.kv_sources: + c = f"{a}.compressor" + if r.ratios[i] == 2: + rows[f"l{i}_comp_w"] = [ + row((f"{c}.wkv.weight", f"{c}.wgate.weight"), None, None, True) + ] + else: + rows[f"l{i}_comp_w"] = [row(f"{c}.wkv.weight")] + rows[f"l{i}_comp_norm"] = [row(f"{c}.norm.weight")] + rows[f"l{i}_idx_wk"] = [row(f"{a}.indexer.wk.weight")] + rows[f"l{i}_idx_k_norm"] = [row(f"{a}.indexer.k_norm.weight")] + if i in r.index_sources: + x = f"{a}.indexer" + rows[f"l{i}_idx_wq_b"] = [row(f"{x}.wq_b.weight")] + rows[f"l{i}_idx_wq_b_sf"] = [row(f"{x}.wq_b.scale", None, None, True)] + rows[f"l{i}_idx_wproj"] = [row(f"{x}.weights_proj.weight", None, None, True)] + f_ = f"{p}.ffn" + rows[f"l{i}_gate_w"] = [row(f"{f_}.gate.weight")] + rows[f"l{i}_gate_bias"] = [row(f"{f_}.gate.bias")] + rows[f"l{i}_gate_bias_vl"] = [row(f"{f_}.gate.bias_vl")] + up_rows = (slice(r.i_off, r.i_off + r.i_loc),) + dn_w_cols = (all_, slice(r.i_off // 2, (r.i_off + r.i_loc) // 2)) + dn_s_cols = (all_, slice(r.i_off // 16, (r.i_off + r.i_loc) // 16)) + fc1_w, fc1_s, fc2_w, fc2_s = [], [], [], [] + for local in range(r.e_loc): + q = f"{f_}.experts.{r.e_off + local}" + fc1_w.append( + row((f"{q}.w3.weight", f"{q}.w1.weight"), (up_rows, up_rows), (local,), True) + ) + fc1_s.append( + row( + (f"{q}.w3.weight_scale", f"{q}.w1.weight_scale"), + (up_rows, up_rows), + (local,), + True, + ) + ) + fc2_w.append(row(f"{q}.w2.weight", dn_w_cols, (local,), True)) + fc2_s.append(row(f"{q}.w2.weight_scale", dn_s_cols, (local,), True)) + rows[f"l{i}_fc1_w"] = fc1_w + rows[f"l{i}_fc1_s"] = fc1_s + rows[f"l{i}_fc2_w"] = fc2_w + rows[f"l{i}_fc2_s"] = fc2_s + for name, (w_, s_) in { + "isc1": ("w1", "input_scale"), + "isc3": ("w3", "input_scale"), + "isc2": ("w2", "input_scale"), + "ws1": ("w1", "weight_scale_2"), + "ws3": ("w3", "weight_scale_2"), + "ws2": ("w2", "weight_scale_2"), + }.items(): + rows[f"l{i}_e_{name}"] = [ + row(f"{f_}.experts.{e}.{w_}.{s_}", None, (e,)) for e in range(r.num_experts) + ] + sh = f"{f_}.shared_experts" + rows[f"l{i}_sh_w13"] = [ + row(f"{sh}.w1.weight", None, slice(0, r.moe_inter)), + row(f"{sh}.w3.weight", None, slice(r.moe_inter, 2 * r.moe_inter)), + ] + rows[f"l{i}_sh_w13_sf"] = [row((f"{sh}.w1.scale", f"{sh}.w3.scale"), None, None, True)] + rows[f"l{i}_sh_w2"] = [row(f"{sh}.w2.weight")] + rows[f"l{i}_sh_w2_sf"] = [row(f"{sh}.w2.scale", None, None, True)] + if i in r.engram_layers: + li = r.engram_layers.index(i) + e = f"{p}.engram" + if not host: + n = r.eng_end[li] - r.eng_start[li] + shard = (slice(r.eng_start[li], r.eng_end[li]),) + rows[f"l{i}_eng_table"] = [row(f"{e}.embed.weight", shard, slice(0, n))] + rows[f"l{i}_eng_scale"] = [row(f"{e}.embed.scale", shard, slice(0, n), True)] + rows[f"l{i}_eng_wkv"] = [row(f"{e}.wkv.weight")] + rows[f"l{i}_eng_wkv_sf"] = [row(f"{e}.wkv.scale", None, None, True)] + rows[f"l{i}_eng_q"] = [row(f"{e}.q_weight")] + rows[f"l{i}_eng_k"] = [row(f"{e}.k_weight")] + rows.update( + vis_patch_w=[row("vision.patch_embed.proj.weight")], + vis_patch_b=[row("vision.patch_embed.proj.bias")], + vis_norm=[row("vision.norm.weight")], + vis_al_w1=[row("aligner.w1.weight", None, None, True)], + vis_al_b1=[row("aligner.w1.bias")], + vis_al_w2=[row("aligner.w2.weight")], + vis_al_b2=[row("aligner.w2.bias")], + vis_span_tail=[row(("image_start", "image_end", "image_newline"), None, None, True)], + ) + for L in range(32): + p, v = f"vis{L}_", f"vision.blocks.{L}." + rows[p + "norm1"] = [row(v + "norm1.weight")] + rows[p + "wqkv_w"] = [row(v + "attn.wqkv.weight")] + rows[p + "wqkv_b"] = [row(v + "attn.wqkv.bias")] + rows[p + "wo_w"] = [row(v + "attn.wo.weight")] + rows[p + "wo_b"] = [row(v + "attn.wo.bias")] + rows[p + "norm2"] = [row(v + "norm2.weight")] + rows[p + "w1"] = [row(v + "mlp.w1.weight")] + rows[p + "w2"] = [row(v + "mlp.w2.weight")] + return rows + + +def _normalize(manifest: dict) -> dict: + """The table's manifest in the reference's row form.""" + out = {} + for key, sources in manifest.items(): + rows = [] + for ckpt, index, transform in sources: + keys, srcs, _ = W._split_source(ckpt) + rows.append((keys, srcs, index, transform is not None)) + out[key] = rows + return out + + +# Every member with the Engram tables in the node's host copy (the default), and the +# device row shard on the members whose shard boundaries differ. +_CASES = [(m, True) for m in sorted(_MEMBERS)] + [("tep4", False), ("tep8", False)] +_IDS = [f"{m}-{'engram_host' if h else 'engram_device'}" for m, h in _CASES] + + +def _residency(monkeypatch, host: bool) -> None: + monkeypatch.setenv(_HOST_ENV, "1" if host else "0") + + +@pytest.mark.parametrize("member,host", _CASES, ids=_IDS) +def test_the_trunk_table_declares_what_the_init_block_did(member, host, monkeypatch): + """Shape and dtype per key, so a failure names the weight that drifted.""" + _residency(monkeypatch, host) + d = W.MODEL_WEIGHTS.dims(_core(member)) + with torch.device("meta"): + got = W.MODEL_WEIGHTS.declare(d) + want = _reference_declaration(member, host) + assert set(got.keys()) == set(want.keys()), sorted(set(got.keys()) ^ set(want.keys()))[:8] + for key in sorted(want): + assert (tuple(got[key].shape), got[key].dtype) == want[key], key + + +@pytest.mark.parametrize("member,host", _CASES, ids=_IDS) +def test_the_trunk_table_sources_what_the_manifest_did(member, host, monkeypatch): + """Every parameter's checkpoint keys, this rank's part of each, the destination + index and whether a transform runs -- per row, in order.""" + _residency(monkeypatch, host) + got = _normalize(W.MODEL_WEIGHTS.manifest(_core(member))) + want = _reference_manifest(member, host) + assert set(got) == set(want), sorted(set(got) ^ set(want))[:8] + for key in sorted(want): + assert got[key] == want[key], key + + +@pytest.mark.parametrize("member,host", _CASES, ids=_IDS) +def test_dims_carries_the_partition_the_core_derived(member, host, monkeypatch): + _residency(monkeypatch, host) + d = W.MODEL_WEIGHTS.dims(_core(member)) + r = _reference_core(member, host) + for name in ( + "h_loc", + "g_loc", + "attn_rank", + "e_loc", + "e_off", + "i_loc", + "i_off", + "pp_layer_ids", + "local", + "pp_prev", + "eng_start", + "eng_end", + "eng_host", + ): + assert getattr(d, name) == getattr(r, name), name + + +def test_the_draft_table_declares_the_drafter(): + d = W.DRAFT_WEIGHTS.dims(_core("tep4")) + with torch.device("meta"): + got = W.DRAFT_WEIGHTS.declare(d) + want = _reference_draft_declaration() + assert set(got.keys()) == set(want.keys()), sorted(set(got.keys()) ^ set(want.keys()))[:8] + for key in sorted(want): + assert (tuple(got[key].shape), got[key].dtype) == want[key], key + + +def test_the_draft_table_reads_every_stage_from_mtp(): + """Every drafter source is an ``mtp.*`` key, each stage reads its own, and the + three-stage ring operand stacks the stages' ``wkv`` in order.""" + manifest = W.DRAFT_WEIGHTS.manifest(_core("tep4")) + for key, sources in manifest.items(): + assert sources, key + for ckpt, _, _ in sources: + keys, _, _ = W._split_source(ckpt) + assert all(k.startswith("mtp.") for k in keys), (key, keys) + if key.startswith("l") and key[1].isdigit(): + stage = key[1 : key.index("_")] + assert all(k.startswith(f"mtp.{stage}.") for k in keys), (key, keys) + wkv3 = manifest["wkv3"] + assert [ckpt for ckpt, _, _ in wkv3] == [f"mtp.{s}.attn.wkv.weight" for s in range(3)] + assert [index for _, index, _ in wkv3] == [slice(s * 512, (s + 1) * 512) for s in range(3)] + + +def test_every_declared_weight_has_a_source(): + """A declared weight with no source allocates memory nothing ever fills.""" + for key, sources in W.MODEL_WEIGHTS.manifest(_core("tep4")).items(): + assert sources, key + + +def test_the_engram_vocabulary_reads_either_config_form(): + """The checkpoint's integer, or a per-n-gram list of the same bucket size.""" + assert W._engram_vocab_size(SimpleNamespace(engram_vocab_size=16000000)) == 16000000 + assert ( + W._engram_vocab_size( + SimpleNamespace(engram_vocab_size=[16000000] * 3, engram_bucket_vocab_size=16000000) + ) + == 16000000 + ) + assert W._engram_vocab_size(SimpleNamespace(engram_vocab_size=[16000000] * 3)) == 16000000 diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_weight_table.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_weight_table.py new file mode 100644 index 000000000000..d8fae728c2d4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_deepseek_weight_table.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""deepseek's weight table declares what the hand-written code did. + +The declaration reference below is a transcription of the __init__ block the table +replaced. The checkpoint-key side was proven differently: while the hand-written +`_manifest` still existed, a test compared the generated manifest against it, per key +and per source, with MTP on and off -- green in CI job 1081689 -- and both the +function and that test were deleted together. One renaming was deliberate: the +hand-written manifest called the MTP module's rows `mtp_*`; the table addresses it as +layer `num_hidden_layers`, which is what the checkpoint calls it, so they are `l61_*`. + +No GPU: nothing here allocates, and `dims` reads `model_config` alone. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +__extra_import_path__ = [".."] # noqa: F841 -- repo's file-scoped import hook + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4 import ( # noqa: E501 + weights as W, +) + +# deepseek-r1-0528's own numbers, so the comparison is against the shapes this target +# actually ships rather than a toy configuration. +_CFG = dict( + num_hidden_layers=61, + hidden_size=7168, + num_attention_heads=128, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + kv_lora_rank=512, + q_lora_rank=1536, + vocab_size=129280, + first_k_dense_replace=3, + n_routed_experts=256, + moe_intermediate_size=2048, + n_shared_experts=1, + intermediate_size=18432, +) +_EP_SIZE, _EP_RANK = 4, 1 + + +def _core(mtp: bool): + """The configuration surface both manifests read. + + `local_experts` and `expert_offset` are set on the core by its `__init__` and read + by the hand-written manifest; `dims` derives the same two from the mapping, so the + comparison also checks that those two derivations agree. + """ + cfg = SimpleNamespace(**_CFG) + mapping = SimpleNamespace(moe_ep_size=_EP_SIZE, moe_ep_rank=_EP_RANK) + local = _CFG["n_routed_experts"] // _EP_SIZE + return SimpleNamespace( + model_config=SimpleNamespace( + pretrained_config=cfg, + mapping=mapping, + torch_dtype=torch.bfloat16, + spec_config=object() if mtp else None, + ), + local_experts=local, + expert_offset=local * _EP_RANK, + mtp_enabled=mtp, + ) + + +@pytest.mark.parametrize("mtp", [False, True], ids=["no_mtp", "mtp"]) +def test_every_declared_weight_has_a_source(mtp): + """A declared weight with no source allocates memory nothing ever fills -- silently + zero rather than an error, which is what `load`'s coverage assert catches from the + other side.""" + core = _core(mtp) + manifest = W.MODEL_WEIGHTS.manifest(core) + for key, sources in manifest.items(): + assert sources, key + + +def test_mtp_rows_appear_only_when_mtp_is_live(): + """The MTP module's parameters are allocated only under a speculative config; the + table has to keep that conditional, or a non-speculative engine carries layer 61's + weights for nothing.""" + layer = _CFG["num_hidden_layers"] + off = set(W.MODEL_WEIGHTS.manifest(_core(False))) + on = set(W.MODEL_WEIGHTS.manifest(_core(True))) + assert not {k for k in off if k.startswith(f"l{layer}_")} + assert {k for k in on if k.startswith(f"l{layer}_")} + assert off < on + + +def _reference_declaration(mtp: bool) -> dict: + """The declaration block from the core's __init__, transcribed. + + Deliberately a copy rather than an import: the block is deleted in the switch to + the table, and the point is to hold the table to what the model used to say. The + MTP rows carry the table's `l61_` keys; the hand block called them `mtp_`, which + the module docstring records as the one deliberate renaming. + """ + import torch.nn as nn + + c = SimpleNamespace(**_CFG) + heads, nope, rope_dim = c.num_attention_heads, c.qk_nope_head_dim, c.qk_rope_head_dim + qk_dim, lat_dim = nope + rope_dim, c.kv_lora_rank + rope_dim + hidden, q_lora, kv_lora, v_dim = c.hidden_size, c.q_lora_rank, c.kv_lora_rank, c.v_head_dim + sf_vec = 16 + local = c.n_routed_experts // _EP_SIZE + + def P(*shape, dtype=torch.bfloat16): + return nn.Parameter(torch.empty(*shape, dtype=dtype), requires_grad=False) + + u8, f32 = torch.uint8, torch.float32 + w = {} + for i in range(c.num_hidden_layers): + w[f"l{i}_norm1"] = P(hidden) + w[f"l{i}_qa"] = P(q_lora, hidden) + w[f"l{i}_q_norm"] = P(q_lora) + w[f"l{i}_qb"] = P(heads * qk_dim, q_lora) + w[f"l{i}_kva"] = P(lat_dim, hidden) + w[f"l{i}_kv_norm"] = P(kv_lora) + w[f"l{i}_kvb"] = P(heads * (nope + v_dim), kv_lora) + w[f"l{i}_o"] = P(hidden, heads * v_dim) + w[f"l{i}_k_scale"] = P(1, dtype=f32) + w[f"l{i}_v_scale"] = P(1, dtype=f32) + w[f"l{i}_norm2"] = P(hidden) + inter = ( + c.intermediate_size + if i < c.first_k_dense_replace + else (c.moe_intermediate_size * c.n_shared_experts) + ) + w[f"l{i}_mlp_gu_w"] = P(2 * inter, hidden // 2, dtype=u8) + w[f"l{i}_mlp_gu_s"] = P(2 * inter * (hidden // sf_vec), dtype=u8) + w[f"l{i}_mlp_dn_w"] = P(hidden, inter // 2, dtype=u8) + w[f"l{i}_mlp_dn_s"] = P(hidden * (inter // sf_vec), dtype=u8) + for name in ("isc1", "isc1_up", "ws2_1", "ws2_1_up", "isc2", "ws2_2"): + w[f"l{i}_mlp_{name}"] = P(1, dtype=f32) + if i < c.first_k_dense_replace: + continue + mi = c.moe_intermediate_size + w[f"l{i}_router"] = P(c.n_routed_experts, hidden) + w[f"l{i}_router_bias"] = P(c.n_routed_experts, dtype=f32) + w[f"l{i}_fc1_w"] = P(local, 2 * mi, hidden // 2, dtype=u8) + w[f"l{i}_fc1_s"] = P(local, 2 * mi, hidden // sf_vec, dtype=u8) + w[f"l{i}_fc2_w"] = P(local, hidden, mi // 2, dtype=u8) + w[f"l{i}_fc2_s"] = P(local, hidden, mi // sf_vec, dtype=u8) + for name in ("isc1", "isc1_up", "ws2_1", "ws2_1_up", "isc2", "ws2_2"): + w[f"l{i}_e_{name}"] = P(c.n_routed_experts, dtype=f32) + w["final_norm"] = P(hidden) + w["embed"] = P(c.vocab_size, hidden) + if mtp: + j, mi = c.num_hidden_layers, c.moe_intermediate_size + shared_inter = c.moe_intermediate_size * c.n_shared_experts + w[f"l{j}_enorm"] = P(hidden) + w[f"l{j}_hnorm"] = P(hidden) + w[f"l{j}_eh"] = P(hidden, 2 * hidden) + w[f"l{j}_norm1"] = P(hidden) + w[f"l{j}_qa"] = P(q_lora, hidden) + w[f"l{j}_q_norm"] = P(q_lora) + w[f"l{j}_qb"] = P(heads * qk_dim, q_lora) + w[f"l{j}_kva"] = P(lat_dim, hidden) + w[f"l{j}_kv_norm"] = P(kv_lora) + w[f"l{j}_kvb"] = P(heads * (nope + v_dim), kv_lora) + w[f"l{j}_o"] = P(hidden, heads * v_dim) + w[f"l{j}_k_scale"] = P(1, dtype=f32) + w[f"l{j}_v_scale"] = P(1, dtype=f32) + w[f"l{j}_norm2"] = P(hidden) + w[f"l{j}_router"] = P(c.n_routed_experts, hidden) + w[f"l{j}_router_bias"] = P(c.n_routed_experts, dtype=f32) + w[f"l{j}_fc1"] = P(local, 2 * mi, hidden) + w[f"l{j}_fc2"] = P(local, hidden, mi) + w[f"l{j}_sh_gu"] = P(2 * shared_inter, hidden) + w[f"l{j}_sh_dn"] = P(hidden, shared_inter) + w[f"l{j}_head_norm"] = P(hidden) + return w + + +@pytest.mark.parametrize("mtp", [False, True], ids=["no_mtp", "mtp"]) +def test_the_table_declares_what_the_init_block_did(mtp): + """Shape and dtype per key, so a failure names the weight that drifted.""" + core = _core(mtp) + got = W.MODEL_WEIGHTS.declare(W.MODEL_WEIGHTS.dims(core)) + want = _reference_declaration(mtp) + assert set(got.keys()) == set(want.keys()) + for key in sorted(want): + assert tuple(got[key].shape) == tuple(want[key].shape), key + assert got[key].dtype == want[key].dtype, key diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_op_bind.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_op_bind.py new file mode 100644 index 000000000000..7d86a02898e1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_op_bind.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The two binding methods, and the one invariant that keeps certification intact. + +A catalog entry no longer cares *when* a value becomes known -- only whether it +varies by layer. `bind_const` states a value that is the same for every layer; +`bind_layered` binds one layer's values at a time, read back by a `layer=` +keyword at the call. Both are callable from wherever a target finds the value knowable -- +construction, the first forward, every forward -- and `raw_call` is the +surface underneath all of it: it takes everything explicitly, which is what +lets CELLS keep certifying the whole input space no matter what a target chose +to bind. +""" + +from __future__ import annotations + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper + + +class _Spy(OpWrapper): + """Records exactly what reached `raw_call`.""" + + def raw_call(self, a=None, b=None, c=None): + self.seen = dict(a=a, b=b, c=c) + return self.seen + + def reference(self, a=None, b=None, c=None): + return dict(a=a, b=b, c=c) + + +def test_an_unbound_entry_is_a_passthrough(): + """Nothing bound: the call reaches raw_call unchanged. + + This is what makes the reshape behaviour-neutral for an entry nobody + binds -- deepseek's module-level singletons never call `bind_const` or + `bind_layered`, so every one of their call sites still works. + """ + assert _Spy()(a=1, b=2, c=3) == dict(a=1, b=2, c=3) + + +def test_bind_const_supplies_what_the_caller_no_longer_passes(): + op = _Spy() + op.bind_const(a=1) + assert op(b=2, c=3) == dict(a=1, b=2, c=3) + + +def test_bind_const_can_be_called_again_to_add_more(): + """A target binds once at construction and again, for different keys, + on a later forward -- both bindings are live at the same time.""" + op = _Spy() + op.bind_const(a="bound-at-construction") + op.bind_const(b="bound-on-forward") + assert op(c=3) == dict(a="bound-at-construction", b="bound-on-forward", c=3) + + +def test_a_later_bind_const_updates_rather_than_replaces(): + """This is the behaviour that makes the staleness check meaningful: a key + a forward forgets to rebind keeps whatever `bind_const` last gave it, + silently, rather than reverting to unset. Wrong output, no error -- the + generation check (below) is what catches it under `validating()`. + """ + op = _Spy() + op.bind_const(a=1, b=1) + op.bind_const(a=2) + assert op() == dict(a=2, b=1, c=None) + + +def test_call_time_kwarg_overrides_bind_const(): + op = _Spy() + op.bind_const(a="bound", b="bound") + assert op(c="call") == dict(a="bound", b="bound", c="call") + assert op(b="call") == dict(a="bound", b="call", c=None) + + +def test_layer_selects_the_bound_row(): + """One `bind_layered` call per layer, read back by `layer=` at the call. + + This is the mechanism gpt_oss's per-layer weight tables rest on: bound + once per layer from the real weights, then indexed by `layer=i` every + forward instead of being read out of a per-target tuple by the call site + itself. + """ + op = _Spy() + op.bind_const(a="bound") + op.bind_layered(0, b="row0") + op.bind_layered(1, b="row1") + op.bind_layered(2, b="row2") + assert op(c=3, layer=1) == dict(a="bound", b="row1", c=3) + + +def test_layer_overrides_const_for_the_same_argument(): + op = _Spy() + op.bind_const(b="const-default") + op.bind_layered(0, b="row0") + op.bind_layered(1, b="row1") + assert op(layer=1) == dict(a=None, b="row1", c=None) + + +def test_call_time_kwarg_overrides_the_layer_row(): + """Call-time kwargs win over everything, including a layered value for + the same argument -- a target binding a weight table can still pass an + unrelated, or even overriding, per-call value through the same call. + """ + op = _Spy() + op.bind_layered(0, b="row0") + op.bind_layered(1, b="row1") + assert op(b="explicit", layer=1) == dict(a=None, b="explicit", c=None) + + +def test_bind_layered_can_bind_a_subset_of_layers(): + """A model whose layers are not uniform -- dense layers beside MoE + layers, as deepseek's will be -- binds only the layers that have the + operand; a call for an unbound layer simply never happens, rather than + the caller having to pad the gap with a placeholder. + """ + op = _Spy() + op.bind_layered(3, b="moe-row") + assert op(layer=3) == dict(a=None, b="moe-row", c=None) + + +def test_an_entry_with_no_layered_table_is_unaffected_by_the_absence_of_layer(): + """An entry that never calls `bind_layered` is not even aware the + mechanism exists: with `_layered` empty, a call that passes no `layer=` + runs the exact same merge it would have before `layered` existed. This is + what keeps deepseek's 20-odd call sites -- none of which pass `layer=` -- + and their own tests working untouched. + """ + op = _Spy() + op.bind_const(a=1) + assert op(c=3) == dict(a=1, b=None, c=3) + + +def test_an_entry_must_implement_raw_call(): + class Incomplete(OpWrapper): + def reference(self): + return None + + with pytest.raises(TypeError, match="raw_call"): + Incomplete() + + +def test_a_key_bound_before_any_forward_is_never_flagged(): + """Generation 0 -- a key `bind_const` or `bind_layered` sees before the + first `advance_step_generation()` call is a construction-time constant, + and the staleness check only ever flags generation >= 1. Any number of + forwards may pass without rebinding such a key. + """ + from tensorrt_llm._torch._experimental.modeling_v2.catalog import _op + + __extra_import_path__ = [".."] # noqa: F841 -- repo's file-scoped import hook + from _validating import validating + + op = _Spy() + op.bind_const(a=1) # before any forward: generation 0 + _op.advance_step_generation() + _op.advance_step_generation() + with validating(op): + op() # does not raise, even though two forwards passed unrebound + + +def test_a_key_bound_mid_forward_and_not_rebound_is_flagged(): + """The failure mode this design introduces, and the only place it is seen. + + Outside validation nothing compares generations -- an entry running on last + forward's value produces a wrong answer and no error. That is the + accepted cost of keeping the hot path free; this test is the record that + the check exists and fires. + """ + from tensorrt_llm._torch._experimental.modeling_v2.catalog import _op + + __extra_import_path__ = [".."] # noqa: F841 -- repo's file-scoped import hook + from _validating import validating + + op = _Spy() + _op.advance_step_generation() + op.bind_const(a=1) # bound during this forward: generation >= 1 + _op.advance_step_generation() # a new forward that forgot to rebind + with pytest.raises(AssertionError, match="previous step's metadata"): + with validating(op): + op() diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py index f40cb6428713..2a1b9da61d0b 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -26,6 +26,7 @@ from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( MODELING_V2_ENV, ModelingV2Mode, + modeling_v2_input_processor, modeling_v2_resolve, ) from tensorrt_llm._torch.model_config import ModelConfig @@ -34,7 +35,14 @@ _is_builtin_model_class, get_registered_model_class, ) +from tensorrt_llm.inputs.content_format import ContentFormat +from tensorrt_llm.inputs.registry import ( + MULTIMODAL_PLACEHOLDER_REGISTRY, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, +) from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig _SM103 = (10, 3) @@ -235,3 +243,184 @@ def test_an_unknown_mode_raises_rather_than_falling_back(monkeypatch): monkeypatch.setenv(MODELING_V2_ENV, "yes") with pytest.raises(ValueError, match="not a modeling_v2 mode"): ModelingV2Mode.from_env() + + +def _v41_config(**overrides): + """The shape DeepSeek-V4.1-Flash-NVFP4's own config.json declares (its + ``text_config``, which the composite config answers from), under the name + the engine gives a V4.1 checkpoint that carries its vision tower.""" + fields = dict( + architectures=["DeepseekV41ForConditionalGeneration"], + model_type="deepseek_v41", + num_hidden_layers=40, + hidden_size=5120, + n_routed_experts=384, + q_lora_rank=1280, + head_dim=512, + o_groups=8, + index_n_heads=32, + num_nextn_predict_layers=3, + ) + fields.update(overrides) + return PretrainedConfig(**fields) + + +def _v41_model_config(quant="FP8_BLOCK_SCALES", config_kwargs=None, **mapping_kwargs): + quant_config = QuantConfig(quant_algo=QuantAlgo(quant)) if quant else QuantConfig() + return ModelConfig( + pretrained_config=_v41_config(**(config_kwargs or {})), + mapping=Mapping(**mapping_kwargs), + quant_config=quant_config, + ) + + +_TEP4 = dict(world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1) +# Every member the unified target folds, as the mapping states it. +_V41_MEMBERS = { + "tep4": _TEP4, + "dep4": dict(_TEP4, enable_attention_dp=True), + "tp4": dict(world_size=4, tp_size=4, moe_ep_size=1, moe_tp_size=4), + "tep8": dict(world_size=8, tp_size=8, moe_ep_size=8, moe_tp_size=1), + "dep8": dict(world_size=8, tp_size=8, moe_ep_size=8, moe_tp_size=1, enable_attention_dp=True), + "pp2xtep2": dict(world_size=4, tp_size=2, pp_size=2, moe_ep_size=2, moe_tp_size=1), +} + + +@pytest.mark.parametrize("mode", ["auto", "require"]) +@pytest.mark.parametrize("member", sorted(_V41_MEMBERS)) +@pytest.mark.parametrize("quant", ["FP8_BLOCK_SCALES", "MIXED_PRECISION"]) +def test_v41_unified_matches(monkeypatch, mode, member, quant): + """Every member, under either reading of the checkpoint's quantization: the + engine's DeepSeek-V4.1 loader reads the dense FP8 blocks as FP8_BLOCK_SCALES, + a modelopt reading gives MIXED_PRECISION.""" + _set_mode(monkeypatch, mode) + config = _v41_model_config(quant, **_V41_MEMBERS[member]) + assert modeling_v2_resolve(config) == "ModelingV2DeepseekV41FlashNvfp4Sm103Unified" + + +def test_v41_off_still_reaches_the_builtin_implementation(monkeypatch): + """With the switch off the checkpoint is the built-in DeepSeek-V4.1's.""" + _set_mode(monkeypatch, "off") + resolved = AutoModelForCausalLM._resolve_class(_v41_model_config(**_TEP4)) + assert resolved is not None + assert resolved.__module__ == "tensorrt_llm._torch.models.modeling_deepseekv41_multimodal" + + +def test_v41_resolve_class_rewrites_the_architecture_end_to_end(monkeypatch): + _set_mode(monkeypatch, "require") + resolved = AutoModelForCausalLM._resolve_class(_v41_model_config(**_TEP4)) + assert resolved.__name__ == "ModelingV2DeepseekV41FlashNvfp4Sm103Unified" + assert not _is_builtin_model_class(resolved) + + +@pytest.mark.parametrize( + "config_kwargs, quant, mapping_kwargs, missed", + [ + # a DeepSeek-V4.1 checkpoint of another shape (one DSpark stage) + (dict(num_nextn_predict_layers=1), "FP8_BLOCK_SCALES", _TEP4, "shape"), + # the right shape, unquantized + (dict(), None, _TEP4, "quant"), + # the right checkpoint, a topology no member covers + (dict(), "FP8_BLOCK_SCALES", dict(world_size=2, tp_size=2, moe_ep_size=2), "parallel"), + (dict(), "FP8_BLOCK_SCALES", dict(_TEP4, moe_ep_size=2, moe_tp_size=2), "parallel"), + ], + ids=["shape", "quant", "tp2", "ep2xmoe_tp2"], +) +def test_v41_near_misses_do_not_match(monkeypatch, config_kwargs, quant, mapping_kwargs, missed): + """Under auto, another checkpoint shape falls back; the target's own + checkpoint raises rather than run the built-in model behind the target's + input processor. Under require every near miss names its criterion.""" + config = _v41_model_config(quant, config_kwargs, **mapping_kwargs) + if missed == "shape": + assert modeling_v2_resolve(config) is None + else: + with pytest.raises(ValueError, match=f"{MODELING_V2_ENV}=off"): + modeling_v2_resolve(config) + + _set_mode(monkeypatch, "require") + with pytest.raises(ValueError, match=missed): + modeling_v2_resolve(config) + + +def test_v41_text_only_architecture_is_not_routed(monkeypatch): + """A text-only V4.1 config has no vision tower for the target to run.""" + config = _v41_model_config( + config_kwargs=dict(architectures=["DeepseekV41ForCausalLM"]), **_TEP4 + ) + assert modeling_v2_resolve(config) is None + + _set_mode(monkeypatch, "require") + with pytest.raises(ValueError, match="no target exists for architecture"): + modeling_v2_resolve(config) + + +@pytest.fixture +def _placeholder_registry(monkeypatch): + """A copy of the placeholder registry's table, so registrations made here + do not outlive the test.""" + table = dict(MULTIMODAL_PLACEHOLDER_REGISTRY._multimodal_placeholder_by_model_type) + monkeypatch.setattr( + MULTIMODAL_PLACEHOLDER_REGISTRY, "_multimodal_placeholder_by_model_type", table + ) + return MULTIMODAL_PLACEHOLDER_REGISTRY + + +def test_v41_input_processor_is_the_builtin_choice_when_off(monkeypatch, _placeholder_registry): + _set_mode(monkeypatch, "off") + assert modeling_v2_input_processor(_v41_config()) is None + + +@pytest.mark.parametrize("mode", ["auto", "require"]) +def test_v41_input_processor_is_the_targets(monkeypatch, mode, _placeholder_registry): + """With modeling_v2 on, the processor and the placeholder placement are the + target's: one placeholder per image, before the text, joined by a blank line.""" + _set_mode(monkeypatch, mode) + processor = modeling_v2_input_processor(_v41_config()) + assert processor is not None + assert processor.__name__ == "DeepseekV41InputProcessor" + assert processor.__module__.endswith( + "models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.input_processor" + ) + assert processor._registered_model_type == "deepseek_v41" + registry = _placeholder_registry + assert registry.get_placeholder("deepseek_v41", "image") == "<\uff5cdeepseek_image\uff5c>" + assert ( + registry.get_placeholder_placement("deepseek_v41") + is MultimodalPlaceholderPlacement.BEFORE_TEXT + ) + assert registry.get_placeholders_separator("deepseek_v41") == "\n\n" + assert registry.get_content_format("deepseek_v41") is ContentFormat.STRING + assert not registry.get_interleave_placeholders("deepseek_v41") + + +def test_v41_input_processor_registration_holds_the_slot(monkeypatch, _placeholder_registry): + """The built-in zoo module registering the same model type afterwards does + not replace the target's placeholder metadata.""" + _set_mode(monkeypatch, "auto") + modeling_v2_input_processor(_v41_config()) + registry = _placeholder_registry + registry.set_placeholder_metadata( + "deepseek_v41", + MultimodalPlaceholderMetadata(placeholder_map={"image": ""}), + registrant_module="tensorrt_llm._torch.models.modeling_deepseekv41_multimodal", + ) + assert registry.get_placeholder("deepseek_v41", "image") == "<\uff5cdeepseek_image\uff5c>" + + +@pytest.mark.parametrize( + "config_kwargs", + [dict(num_nextn_predict_layers=1), dict(architectures=["DeepseekV41ForCausalLM"])], + ids=["other_shape", "text_only"], +) +def test_v41_input_processor_needs_the_targets_checkpoint( + monkeypatch, config_kwargs, _placeholder_registry +): + _set_mode(monkeypatch, "require") + assert modeling_v2_input_processor(_v41_config(**config_kwargs)) is None + + +def test_a_family_without_its_own_input_processor_keeps_the_builtin_one( + monkeypatch, _placeholder_registry +): + _set_mode(monkeypatch, "require") + assert modeling_v2_input_processor(_gpt_oss_config()) is None diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_target_contract.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_target_contract.py index e04a4e01ac4f..00c19c92cb6c 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_target_contract.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_target_contract.py @@ -1,17 +1,16 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Every op a target names has to exist in the build it is running against. +"""The pybind surface a target reaches for has to exist in the build. -A target reads private engine surface: custom ops registered under -``torch.ops.trtllm`` and one pybind entry point. With no version pin, nothing -declares which build it was written for -- so the only thing that can tell you -the surface moved is running against it. +A target reads private engine surface with no version pin, so the only thing +that can tell you the surface moved is running against it. -This used to be an ``assert`` loop that each target ran at import. That made -every import of a modeling module pay for the check and put a test in the -product tree; the missing symbol is a fact of the build, so the place to find -it out is a test on a machine that has the extension built. Each target now -declares ``REQUIRED_TRTLLM_OPS`` and this asserts it. +Custom ops are not checked here. Naming them took a hand-kept list per target +that nothing could prove complete: an op added to a forward and left out of the +list cost nothing to forget, and a stale entry still passed because the op was +still registered upstream. The forward reaches every op it needs anyway, so a +renamed op fails there -- at the call, which says where. A list that has to be +maintained to restate that is not worth its own upkeep. Distinct from ``test_modeling_v2_claims.py``, which is deliberately import-free and runs anywhere: this one imports the targets, so it needs a real build. @@ -22,9 +21,7 @@ import importlib import pytest -import torch -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( MODELING_V2_ROUTERS, routing_module, @@ -46,20 +43,6 @@ def _target_ids(): return [name for name, _ in _target_modules()] -@pytest.mark.parametrize("name,module", list(_target_modules()), ids=_target_ids()) -def test_every_declared_op_exists(name, module): - """A named op that is not registered is a build the target cannot run on.""" - declared = getattr(module, "REQUIRED_TRTLLM_OPS", None) - assert declared, f"{name}: modeling.py declares no REQUIRED_TRTLLM_OPS" - - missing = [op for op in declared if not hasattr(torch.ops.trtllm, op)] - assert not missing, ( - f"{name} names {len(missing)} op(s) this build does not register: " - f"{', '.join(missing)}. Either the op was renamed upstream and the " - f"target has not followed, or this build predates it." - ) - - @pytest.mark.parametrize("name,module", list(_target_modules()), ids=_target_ids()) def test_the_pybind_attention_entry_point_exists(name, module): """``thop.attention`` is reached through the bindings, not torch.ops. diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_target_layer.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_target_layer.py new file mode 100644 index 000000000000..b75a451f6c54 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_target_layer.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The phase predicate and the Target base class. + +Import-only and allocation-free: this runs on a CI machine with no GPU, like +the rest of the claims suite. What it pins down is the one decision the whole +layer rests on -- which target a step goes to -- and the fact that a Target +cannot be instantiated without answering for both of its abstract methods. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2._target import Phase, Target, phase_of + + +@pytest.mark.parametrize( + "num_contexts, expected", + [ + (0, Phase.DECODE), + (1, Phase.PREFILL), + (8, Phase.PREFILL), + ], + ids=["pure-decode", "one-context", "mixed-batch"], +) +def test_phase_is_decided_by_num_contexts_alone(num_contexts, expected): + """num_contexts == 0 is the whole predicate. + + A mixed batch -- context requests *and* generation requests in one step -- + goes to PREFILL, which is why PREFILL is the general case and DECODE is the + specialization. See the design document, section 2.2. + """ + assert phase_of(SimpleNamespace(num_contexts=num_contexts)) is expected + + +def test_phase_predicate_reads_nothing_but_num_contexts(): + """Anything else it touched would be a field a decode-only CUDA graph is + not guaranteed to carry. + + Recorded rather than inferred: a metadata stand-in that logs every + attribute read, so the assertion names the exact read set instead of + resting on "no exception was raised". + """ + reads = [] + + class Recording: + num_contexts = 0 + + def __getattribute__(self, name): + reads.append(name) + return object.__getattribute__(self, name) + + phase_of(Recording()) + assert reads == ["num_contexts"] + + +def test_a_target_cannot_skip_either_abstract_method(): + class Incomplete(Target): + def step_args(self, md): + return {} + + with pytest.raises(TypeError, match="forward"): + Incomplete(core=object()) + + +def test_a_target_keeps_the_core_by_reference(): + """By reference, not by snapshot: the core replaces `_rope` after a + contract check regrows it, and a target that captured operands at + construction would hold the stale table. See the design document, 2.3. + """ + + class Concrete(Target): + def step_args(self, md): + return {} + + def forward(self, attn_metadata, *args, **kwargs): + return self.core.marker + + core = SimpleNamespace(marker="first") + target = Concrete(core=core) + core.marker = "second" + assert target.forward(None) == "second" + + +@pytest.mark.parametrize( + "target", + [ + "gpt_oss/gpt_oss_120b__sm_103__tp1", + "deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified", + ], +) +def test_decode_step_args_state_the_zero_rather_than_reading_it(target): + """The decode target states both context counts as literal zeros. + + Decode's routing already guarantees no context rows, so it says so instead + of reading the field back. That is what makes the source gate in + test_modeling_v2_claims.py meaningful on these targets rather than vacuous. + """ + import ast + from pathlib import Path + + import tensorrt_llm._torch._experimental.modeling_v2 as _mv2 + + source = (Path(_mv2.__file__).resolve().parent / "models" / target / "modeling.py").read_text() + decode = next( + n + for n in ast.walk(ast.parse(source)) + if isinstance(n, ast.ClassDef) and n.name == "DecodeTarget" + ) + literals = { + kw.arg: kw.value.value + for call in ast.walk(decode) + if isinstance(call, ast.Call) + for kw in call.keywords + if kw.arg in ("num_contexts", "num_ctx_tokens") and isinstance(kw.value, ast.Constant) + } + assert literals == {"num_contexts": 0, "num_ctx_tokens": 0} diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_weight_table.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_weight_table.py new file mode 100644 index 000000000000..ca9b762b9c04 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_weight_table.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The weight table reproduces the declaration it replaces. + +One weight's shape, dtype, checkpoint source and load-time transform used to +live in three files. The table merges them, which is only safe if it declares +exactly what the hand-written block declared -- same keys, same shapes, same +dtypes. This test is that proof. + +The block it compares against no longer exists in the model; `_reference` +below is a transcription of it, kept deliberately. That is the point of +having transcribed rather than imported it: the table is still held to what +the model used to say, and an import would now have the test agreeing with +itself. A shape that changes here is a real change to what this target +allocates, and should be made knowingly rather than discovered. + +No GPU and no built extensions: the parameters are meta tensors, so only +shape and dtype are compared, which is all the declaration ever stated. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +__extra_import_path__ = [".."] # noqa: F841 -- repo's file-scoped import hook + +from tensorrt_llm._torch._experimental.modeling_v2.models.gpt_oss.gpt_oss_120b__sm_103__tp1 import ( + weights as W, +) + +# gpt-oss-120b's own numbers, so the comparison is against the shapes this +# target actually ships rather than a toy configuration. +_DIMS = dict( + num_layers=36, + hidden=2880, + q_width=64 * 64, + kv_width=8 * 64, + heads_q=64, + num_experts=128, + vocab=201088, + fc1_rows=2 * 2944, + # 2880 padded up to FC1_K_ALIGN (512), not to 128 like the other three. + # This read 2944 until `dims` was compared against it: the declaration + # tests feed both sides the same bundle, so a wrong number here agreed + # with itself and proved the table at a shape the model never allocates. + fc1_k_pad=3072, + inter_pad=2944, + fc2_rows_pad=2944, + dtype=torch.bfloat16, +) + + +def _reference(d): + """The declaration block from modeling.py, transcribed. + + Deliberately a copy rather than an import: the point is to compare the + table against what the model *said*, and importing the model would make + this test agree with itself once the block is deleted in Task 3. + """ + import torch.nn as nn + + def P(*shape, dtype=d["dtype"]): + return nn.Parameter(torch.empty(*shape, dtype=dtype), requires_grad=False) + + fc1_rows, hidden, q_width = d["fc1_rows"], d["hidden"], d["q_width"] + kv_width, heads_q, ne = d["kv_width"], d["heads_q"], d["num_experts"] + w = {} + for i in range(d["num_layers"]): + w[f"l{i}_norm1"] = P(hidden) + w[f"l{i}_qkv"] = P(q_width + 2 * kv_width, hidden) + w[f"l{i}_qkv_bias"] = P(q_width + 2 * kv_width) + w[f"l{i}_sinks"] = P(heads_q, dtype=torch.float32) + w[f"l{i}_o"] = P(hidden, q_width) + w[f"l{i}_o_bias"] = P(hidden) + w[f"l{i}_norm2"] = P(hidden) + w[f"l{i}_router"] = P(ne, hidden) + w[f"l{i}_router_bias"] = P(ne) + w[f"l{i}_fc1_w"] = P(ne, fc1_rows, d["fc1_k_pad"] // 2, dtype=torch.uint8) + w[f"l{i}_fc1_s"] = P(ne, fc1_rows, d["fc1_k_pad"] // 32, dtype=torch.uint8) + w[f"l{i}_fc1_b"] = P(ne, fc1_rows, dtype=torch.float32) + w[f"l{i}_fc2_w"] = P(ne, d["fc2_rows_pad"], d["inter_pad"] // 2, dtype=torch.uint8) + w[f"l{i}_fc2_s"] = P(ne, d["fc2_rows_pad"], d["inter_pad"] // 32, dtype=torch.uint8) + w[f"l{i}_fc2_b"] = P(ne, d["fc2_rows_pad"], dtype=torch.float32) + w["final_norm"] = P(hidden) + w["embed"] = P(d["vocab"], hidden) + return w + + +def _fake_core(): + """The configuration surface `dims` reads, carrying this target's numbers. + + `dims` reads `model_config` only, never the half-built core around it, + which is what lets this run without constructing a model. + """ + from types import SimpleNamespace + + cfg = SimpleNamespace( + num_hidden_layers=_DIMS["num_layers"], + hidden_size=_DIMS["hidden"], + num_attention_heads=_DIMS["heads_q"], + num_key_value_heads=_DIMS["kv_width"] // 64, + head_dim=64, + num_local_experts=_DIMS["num_experts"], + intermediate_size=2880, + vocab_size=_DIMS["vocab"], + ) + return SimpleNamespace( + model_config=SimpleNamespace(pretrained_config=cfg, torch_dtype=_DIMS["dtype"]) + ) + + +def test_the_table_declares_the_same_keys(): + assert set(W.MODEL_WEIGHTS.declare(SimpleNamespace(**_DIMS)).keys()) == set( + _reference(_DIMS).keys() + ) + + +def test_the_table_declares_the_same_shapes_and_dtypes(): + """Shape and dtype are the whole content of a declaration. + + Checked per key rather than as one aggregate so a failure names the weight + that drifted instead of only saying the two disagree. + """ + got, want = W.MODEL_WEIGHTS.declare(SimpleNamespace(**_DIMS)), _reference(_DIMS) + for key in sorted(want): + assert tuple(got[key].shape) == tuple(want[key].shape), key + assert got[key].dtype == want[key].dtype, key + + +def test_every_table_entry_states_a_source(): + """A declared weight with no checkpoint source would allocate memory that + nothing ever fills -- silently zero rather than an error, which is the + failure mode `load`'s coverage assert exists to catch from the other side. + """ + for entry in W.MODEL_WEIGHTS.WEIGHTS: + assert entry.src, entry.name + + +def test_dims_reproduces_the_declaration_numbers(): + """`dims` is now the only place a derived width or padded size is worked + out; it has to land on the same numbers the declaration was written for. + """ + d = W.MODEL_WEIGHTS.dims(_fake_core()) + for key, want in _DIMS.items(): + assert getattr(d, key) == want, key diff --git a/tests/unittest/_torch/multimodal/test_deepseek_v41_input_processor.py b/tests/unittest/_torch/multimodal/test_deepseek_v41_input_processor.py new file mode 100644 index 000000000000..0047541f69e1 --- /dev/null +++ b/tests/unittest/_torch/multimodal/test_deepseek_v41_input_processor.py @@ -0,0 +1,678 @@ +# 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. +"""DeepSeek-V4.1 input processor against the checkpoint's own reference code. + +The oracle is the checkpoint itself: ``inference/image_processor.py`` (``prepare_vl_inputs``) for +token ids, per-position types and patches, ``encoding/encoding.py`` for prompt rendering, +``inference/vision.py`` for the ViT rotary table and the aligner's unfold, and +``Transformer.merge_image_embeddings`` in ``inference/model.py`` for the span splice. Every +comparison is exact (``torch.equal``, list equality). + +Cases: the checkpoint's own two-image example (carrots/corn), synthetic images from 1x1 to +6000x4000 and 50000x14 in both orientations, every runtime image format, token-id input, and +MMMU validation documents rendered the way ``trtllm-eval mmmu`` renders them (placeholders +before the text, joined by a blank line, through the ``deepseek_v4`` chat renderer); the +single-image ones are also checked against the prompt lengths the external-engine anchor reported. + +Two deliberately wrong controls show the comparisons can fail: resizing instead of padding, and +joining the placeholders with ``"\\n"`` instead of ``"\\n\\n"``. + +The checkpoint and the MMMU cache are external directories; tests that need one skip, naming it, +when it is absent. +""" + +import importlib.util +import io +import json +import os +import sys +import types +from pathlib import Path + +import numpy as np +import pytest +import torch +from PIL import Image +from torchvision.transforms import ToTensor +from utils.llm_data import llm_models_root + +from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v41.deepseek_v4_1_flash_nvfp4__sm_103__unified.input_processor import ( # noqa: E501 + IMAGE, + IMAGE_END, + IMAGE_NEW_LINE, + IMAGE_PLACEHOLDER, + IMAGE_START, + TEXT, + DeepseekV41InputProcessor, + VisionGeometry, + plan_image_grid, + preprocess_image, + to_rgb_image, +) +from tensorrt_llm.inputs.data import TextPrompt +from tensorrt_llm.inputs.registry import DefaultInputProcessor +from tensorrt_llm.sampling_params import SamplingParams + +pytestmark = pytest.mark.cpu_only + +# The checkpoint resolves under LLM_MODELS_ROOT, as the whole-model gates' does; +# MODELING_V2_DSV41_CKPT_DIR overrides it. A case that needs it skips when it is absent. +_MODELS_ROOT = llm_models_root() +_CKPT_DIR = os.environ.get("MODELING_V2_DSV41_CKPT_DIR") or ( + str(_MODELS_ROOT / "DeepSeek-V4.1-Flash-NVFP4") if _MODELS_ROOT else "" +) +# The HF datasets cache of MMMU and the external-engine anchor's MMMU records; no default +# location, so the cases that read them skip unless these are set. +_MMMU_DIR = os.environ.get("MODELING_V2_DSV41_MMMU_DIR", "") +_ANCHOR_MMMU_DIR = os.environ.get("MODELING_V2_DSV41_ANCHOR_MMMU_DIR", "") +# 0: the first document of every subject plus every multi-image document; 1: all 900. +_MMMU_ALL = os.environ.get("MODELING_V2_DSV41_MMMU_ALL", "0") == "1" + +_IMAGE_TOKEN_ID = 129264 +_REFERENCE_MODULES = ("encoding", "image_processor", "vision", "kernel", "engram", "model") + + +def _require(path: str) -> Path: + if not os.path.isdir(path): + pytest.skip(f"external directory not present: {path}") + return Path(path) + + +def _load_as(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def ckpt() -> Path: + return _require(_CKPT_DIR) + + +@pytest.fixture(scope="module") +def ref(ckpt): + """The checkpoint's reference modules, importable by the names they import each other by.""" + saved = {name: sys.modules.get(name) for name in _REFERENCE_MODULES} + try: + ns = types.SimpleNamespace() + ns.encoding = _load_as("encoding", ckpt / "encoding" / "encoding.py") + ns.image_processor = _load_as("image_processor", ckpt / "inference" / "image_processor.py") + ns.vision = _load_as("vision", ckpt / "inference" / "vision.py") + # model.py imports its GPU kernels and the Engram hash at module scope; the span splice + # tested here uses neither. + kernel = types.ModuleType("kernel") + for name in ("act_quant", "fp4_act_quant", "fp4_gemm", "fp8_gemm", "hc_split_sinkhorn"): + setattr(kernel, name, None) + kernel.sparse_attn = None + sys.modules["kernel"] = kernel + engram = types.ModuleType("engram") + engram.EngramLayout = engram.NgramHashState = object + sys.modules["engram"] = engram + ns.model = _load_as("model", ckpt / "inference" / "model.py") + yield ns + finally: + for name, module in saved.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + +@pytest.fixture(scope="module") +def ref_args(ckpt): + """The reference ModelArgs' vision fields, from the checkpoint's inference config.""" + cfg = json.loads((ckpt / "inference" / "config.json").read_text()) + args = types.SimpleNamespace(**cfg) + args.vision_enabled = cfg["vision_n_layers"] > 0 + return args + + +@pytest.fixture(scope="module") +def hf_config(ckpt): + from transformers import AutoConfig + + import tensorrt_llm._torch.configs # noqa: F401 registers deepseek_v41 with AutoConfig + + return AutoConfig.from_pretrained(str(ckpt)) + + +@pytest.fixture(scope="module") +def tokenizer(ckpt, hf_config): + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(str(ckpt)) + + +@pytest.fixture(scope="module") +def processor(ckpt, hf_config, tokenizer): + return DeepseekV41InputProcessor(str(ckpt), hf_config, tokenizer) + + +@pytest.fixture(scope="module") +def chat_renderer(tokenizer): + """The renderer ``custom_tokenizer: deepseek_v4`` installs.""" + from tensorrt_llm.tokenizer.deepseek_v4 import DeepseekV4Tokenizer + + return DeepseekV4Tokenizer(tokenizer) + + +def _ppm_bytes(image: Image.Image) -> bytes: + buf = io.BytesIO() + image.save(buf, format="PPM") + return buf.getvalue() + + +def _png_bytes(image: Image.Image) -> bytes: + buf = io.BytesIO() + image.save(buf, format="PNG") + return buf.getvalue() + + +def _random_image(width: int, height: int, seed: int) -> Image.Image: + rng = np.random.default_rng(seed) + return Image.fromarray(rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8), "RGB") + + +def _run(processor, prompt, images, **sampling): + inputs = TextPrompt(prompt=prompt) + if images: + inputs["multi_modal_data"] = {"image": images} + return processor(inputs, SamplingParams(**sampling)) + + +def _assert_matches_reference(ours, reference): + """``ours``: the processor's (ids, extra); ``reference``: prepare_vl_inputs' triple.""" + ids, extra = ours + ref_tokens, ref_types, ref_images = reference + assert ids == ref_tokens + if ref_images is None: + assert extra is None + assert all(t == TEXT for t in ref_types) + return + image = extra["multimodal_data"]["image"] + layout = extra["multimodal_data"]["token_layout"] + assert layout["token_types"].tolist() == ref_types + assert image["span_starts"] == [img.start for img in ref_images] + assert image["span_lengths"] == [img.types.numel() for img in ref_images] + assert image["vit_grid_hw"] == [[img.n_vit_h, img.n_vit_w] for img in ref_images] + assert torch.equal(image["span_types"], torch.cat([img.types for img in ref_images])) + assert torch.equal(image["pixel_values"], torch.cat([img.patches for img in ref_images])) + assert image["pixel_values"].dtype == torch.bfloat16 + # Every image id in the prompt is a span position and every span position is one. + ids_t = torch.tensor(ids) + assert torch.equal(layout["image_mask"], ids_t == _IMAGE_TOKEN_ID) + assert torch.equal(layout["engram_dead_mask"], layout["image_mask"]) + + +def _reference_prompt(ref, num_images: int, text: str) -> str: + """encoding.py's rendering of one user turn: N image blocks, then the text.""" + blocks = [{"type": "image_url", "image_url": {"url": f"img{i}"}} for i in range(num_images)] + blocks.append({"type": "text", "text": text}) + return ref.encoding.encode_messages([{"role": "user", "content": blocks}], thinking_mode="chat") + + +def _string_path_prompt(chat_renderer, num_images: int, text: str, separator: str = "\n\n") -> str: + """What trtllm-eval builds on the STRING path: placeholders before the text, then the chat + renderer (``inputs.utils.apply_chat_template``'s DeepSeek branch).""" + content = separator.join([IMAGE_PLACEHOLDER] * num_images + [text]) + return chat_renderer.apply_chat_template( + messages=[{"role": "user", "content": content}], + tools=None, + thinking=False, + continue_final_message=False, + ) + + +# --------------------------------------------------------------------------------------------- +# Geometry + + +def test_geometry_matches_the_reference_args(hf_config, ref_args): + geometry = VisionGeometry.from_config(hf_config) + assert geometry.patch_size == ref_args.vision_patch_size == 14 + assert geometry.downsample_ratio == ref_args.vision_downsample_ratio == 3 + assert geometry.max_image_tokens == ref_args.vision_max_n_token == 1024 + assert geometry.min_pixels == ref_args.vision_min_pixels == 295936 + assert geometry.max_wh_ratio is ref_args.vision_max_wh_ratio is None + assert geometry.hidden_size == ref_args.vision_dim == 1024 + assert geometry.num_heads == ref_args.vision_n_heads == 16 + assert geometry.num_layers == ref_args.vision_n_layers == 32 + assert geometry.rope_theta == ref_args.vision_rope_theta + assert geometry.rope_dim == ref_args.vision_dim // ref_args.vision_n_heads // 2 == 32 + assert geometry.image_token_id == ref_args.image_token_id == _IMAGE_TOKEN_ID + + +def test_grid_plan_matches_the_reference_on_a_size_sweep(hf_config, ref, ref_args): + geometry = VisionGeometry.from_config(hf_config) + rng = np.random.default_rng(0) + sizes = [(int(w), int(h)) for w, h in rng.integers(1, 12000, size=(2000, 2))] + sizes += [(int(w), int(h)) for w, h in rng.integers(1, 800, size=(2000, 2))] + sizes += [(1, 1), (544, 544), (543, 545), (14, 14), (42, 42), (1, 100000), (100000, 1)] + sizes += [(50000, 14), (14, 50000), (768, 384), (1344, 1344), (1387, 1387), (1388, 1388)] + for width, height in sizes: + grid = plan_image_grid(width, height, geometry) + expected = ref.image_processor.plan_image_grid(width, height, ref_args) + assert (grid.n_llm_h, grid.n_llm_w, grid.best_height, grid.best_width) == expected, ( + width, + height, + ) + assert grid.num_span_tokens <= geometry.max_image_tokens + + +# --------------------------------------------------------------------------------------------- +# The checkpoint's own examples + + +def _example_cases(ckpt, ref): + raw = (ckpt / "inference" / "examples" / "example.txt").read_text().rstrip("\n").split("\n\n") + image_dir = str(ckpt / "inference" / "examples" / "images") + for prompt in raw: + prompt = prompt.replace("examples/images", image_dir) + case = {"messages": [{"role": "user", "content": ref.encoding.parse_tagged_text(prompt)}]} + yield ref.encoding.encode_case(case, "chat") + + +def test_checkpoint_examples_match_the_reference(ckpt, ref, ref_args, tokenizer, processor): + """example.txt: three text prompts and the carrots/corn two-image prompt.""" + num_image_cases = 0 + for prompt, records in _example_cases(ckpt, ref): + reference = ref.image_processor.prepare_vl_inputs(prompt, records, tokenizer, ref_args) + images = [Image.open(record["url"]) for record in records] + ours = _run(processor, prompt, images or None) + _assert_matches_reference(ours, reference) + num_image_cases += bool(records) + if records: + assert len(records) == 2 + # The two images differ in size, so their grids and spans do. + grids = ours[1]["multimodal_data"]["image"]["vit_grid_hw"] + assert grids[0] != grids[1] + assert num_image_cases == 1 + + +# --------------------------------------------------------------------------------------------- +# Synthetic sizes + +_SYNTHETIC_SIZES = [ + (1, 1), + (5, 7), + (13, 15), + (544, 544), + (545, 543), + (768, 384), + (1387, 1386), + (4000, 3000), + (6000, 4000), + (20000, 14), + (50000, 14), + (14, 50000), + (3, 4000), +] + + +@pytest.mark.parametrize("size", _SYNTHETIC_SIZES, ids=lambda s: f"{s[0]}x{s[1]}") +def test_synthetic_image_matches_the_reference(ref, ref_args, tokenizer, processor, size): + width, height = size + image = _random_image(width, height, seed=width * 100003 + height) + prompt = _reference_prompt(ref, 1, "Describe the image.") + reference = ref.image_processor.prepare_vl_inputs( + prompt, [{"data": _ppm_bytes(image)}], tokenizer, ref_args + ) + _assert_matches_reference(_run(processor, prompt, [image]), reference) + + +def test_every_runtime_image_format_gives_the_same_patches(hf_config, ckpt): + geometry = VisionGeometry.from_config(hf_config) + image = Image.open(ckpt / "inference" / "examples" / "images" / "carrots.jpeg") + rgb = image.convert("RGB") + expected = preprocess_image(rgb, geometry).patches + as_float = ToTensor()(rgb) # the runtime's "pt" format + as_uint8 = torch.from_numpy(np.asarray(rgb)).permute(2, 0, 1).contiguous() + as_numpy = np.asarray(rgb) # the runtime's "np" format + for fmt in (image, as_float, as_uint8, as_numpy): + assert torch.equal(preprocess_image(fmt, geometry).patches, expected) + assert np.array_equal(np.asarray(to_rgb_image(fmt)), np.asarray(rgb)) + + +def test_three_images_with_text_between_match_the_reference(ref, ref_args, tokenizer, processor): + images = [ + _random_image(w, h, seed=i) for i, (w, h) in enumerate([(640, 480), (30, 900), (2000, 100)]) + ] + blocks = [ + {"type": "text", "text": "First:"}, + {"type": "image_url", "image_url": {"url": "a"}}, + {"type": "text", "text": "then"}, + {"type": "image_url", "image_url": {"url": "b"}}, + {"type": "image_url", "image_url": {"url": "c"}}, + {"type": "text", "text": "Compare them."}, + ] + prompt = ref.encoding.encode_messages( + [{"role": "user", "content": blocks}], thinking_mode="chat" + ) + records = [{"data": _ppm_bytes(image)} for image in images] + reference = ref.image_processor.prepare_vl_inputs(prompt, records, tokenizer, ref_args) + ours = _run(processor, prompt, images) + _assert_matches_reference(ours, reference) + image = ours[1]["multimodal_data"]["image"] + patches = [h * w for h, w in image["vit_grid_hw"]] + assert image["vit_cu_seqlens"] == np.cumsum([0] + patches).tolist() + assert image["llm_grid_hw"] == [[-(-h // 3), -(-w // 3)] for h, w in image["vit_grid_hw"]] + + +def test_token_id_input_matches_text_input(ckpt, ref, tokenizer, processor): + prompt, records = list(_example_cases(ckpt, ref))[-1] + images = [Image.open(record["url"]) for record in records] + text_ids, text_extra = _run(processor, prompt, images) + inputs = {"prompt_token_ids": tokenizer.encode(prompt), "multi_modal_data": {"image": images}} + ids, extra = processor(inputs, SamplingParams()) + assert ids == text_ids + for key, value in text_extra["multimodal_data"]["image"].items(): + other = extra["multimodal_data"]["image"][key] + assert torch.equal(value, other) if isinstance(value, torch.Tensor) else value == other + for key, value in text_extra["multimodal_data"]["token_layout"].items(): + assert torch.equal(value, extra["multimodal_data"]["token_layout"][key]) + + +# --------------------------------------------------------------------------------------------- +# The tables the model indexes with + + +def test_vit_rope_table_matches_the_reference(ref, ref_args, processor): + images = [ + _random_image(w, h, seed=i) for i, (w, h) in enumerate([(544, 544), (30, 900), (1500, 200)]) + ] + prompt = _reference_prompt(ref, len(images), "x") + _, extra = _run(processor, prompt, images) + image = extra["multimodal_data"]["image"] + cu = image["vit_cu_seqlens"] + dim = ref_args.vision_dim // ref_args.vision_n_heads // 2 + for i, (n_h, n_w) in enumerate(image["vit_grid_hw"]): + cos, sin = ref.vision.get_vision_cos_sin(n_h, n_w, dim, ref_args.vision_rope_theta) + assert torch.equal(image["vit_rope_cos"][cu[i] : cu[i + 1]], cos.squeeze(1)) + assert torch.equal(image["vit_rope_sin"][cu[i] : cu[i + 1]], sin.squeeze(1)) + + +def test_aligner_index_reproduces_the_reference_unfold(ref, ref_args, processor): + """One gather over all images == the reference Aligner's w1 input, image by image.""" + sizes = [(14, 14), (70, 28), (544, 544), (768, 384), (2000, 100), (30, 900)] + images = [_random_image(w, h, seed=i) for i, (w, h) in enumerate(sizes)] + _, extra = _run(processor, _reference_prompt(ref, len(images), "x"), images) + image = extra["multimodal_data"]["image"] + width = 16 + num_patches = image["vit_cu_seqlens"][-1] + vit_out = torch.randn(num_patches, width, generator=torch.Generator().manual_seed(0)) + + gathered = torch.cat([vit_out, vit_out.new_zeros(1, width)])[image["aligner_index"]] + ours = gathered.transpose(1, 2).reshape(gathered.shape[0], -1) + + args = types.SimpleNamespace(vision_downsample_ratio=3, vision_dim=width, dim=8) + aligner = ref.vision.Aligner(args) + captured = [] + aligner.w1.register_forward_pre_hook(lambda _m, inp: captured.append(inp[0].clone())) + cu = image["vit_cu_seqlens"] + for i, (n_h, n_w) in enumerate(image["vit_grid_hw"]): + aligner(vit_out[cu[i] : cu[i + 1]], n_h, n_w) + assert torch.equal(ours, torch.cat(captured)) + n_cells = sum(h * w for h, w in image["llm_grid_hw"]) + assert image["aligner_index"].shape == (n_cells, 9) + + +def test_span_source_index_reproduces_the_reference_splice(ref, ref_args, processor): + """The one-gather span assembly == Transformer.merge_image_embeddings.""" + sizes = [(544, 544), (30, 900), (2000, 100)] + images = [_random_image(w, h, seed=i) for i, (w, h) in enumerate(sizes)] + prompt = _reference_prompt(ref, len(images), "Compare.") + records = [{"data": _ppm_bytes(image)} for image in images] + tokens, _, ref_images = ref.image_processor.prepare_vl_inputs( + prompt, records, processor.tokenizer, ref_args + ) + ids, extra = _run(processor, prompt, images) + image = extra["multimodal_data"]["image"] + dim = 8 + gen = torch.Generator().manual_seed(1) + n_cells = [h * w for h, w in image["llm_grid_hw"]] + aligner_out = torch.randn(sum(n_cells), dim, generator=gen) + delimiters = torch.randn(3, dim, generator=gen) + per_image = list(torch.split(aligner_out, n_cells)) + + fake = types.SimpleNamespace( + image_start=delimiters[0], image_end=delimiters[1], image_newline=delimiters[2] + ) + calls = iter(per_image) + fake.encode_image = lambda patches, n_h, n_w: next(calls) + h_ref = torch.zeros(1, len(tokens), dim) + ref.model.Transformer.merge_image_embeddings(fake, [ref_images], h_ref) + + span_embeds = torch.cat([aligner_out, delimiters])[image["span_source_index"]] + h_ours = torch.zeros(len(ids), dim) + h_ours[extra["multimodal_data"]["token_layout"]["image_mask"]] = span_embeds + assert torch.equal(h_ours, h_ref[0]) + types_ = image["span_types"] + assert set(types_.unique().tolist()) == {IMAGE_START, IMAGE, IMAGE_NEW_LINE, IMAGE_END} + + +def test_image_local_tables_assemble_any_run_of_images_alone(processor): + """``aligner_index_local`` / ``span_source_local`` rebuild, from the images of a run alone, + exactly the rows the request-wide ``aligner_index`` / ``span_source_index`` give for them.""" + sizes = [(544, 544), (30, 900), (2000, 100), (768, 384), (14, 14)] + images = [_random_image(w, h, seed=i) for i, (w, h) in enumerate(sizes)] + _, extra = _run(processor, "\n\n".join([IMAGE_PLACEHOLDER] * len(images) + ["x"]), images) + image = extra["multimodal_data"]["image"] + cu, ccu = image["vit_cu_seqlens"], image["aligner_cu_cells"] + scu = [0] + for n in image["span_lengths"]: + scu.append(scu[-1] + n) + n_cells = [h * w for h, w in image["llm_grid_hw"]] + assert ccu == [sum(n_cells[:i]) for i in range(len(n_cells) + 1)] + assert image["aligner_index_local"].shape == image["aligner_index"].shape + assert image["span_source_local"].shape == image["span_source_index"].shape + gen = torch.Generator().manual_seed(2) + width, dim = 16, 8 + vit_out = torch.randn(cu[-1], width, generator=gen) + aligner_out = torch.randn(ccu[-1], dim, generator=gen) + delimiters = torch.randn(3, dim, generator=gen) + cells_all = torch.cat([vit_out, vit_out.new_zeros(1, width)])[image["aligner_index"]] + spans_all = torch.cat([aligner_out, delimiters])[image["span_source_index"]] + for i in range(len(images)): + vit_i = vit_out[cu[i] : cu[i + 1]] + idx_i = image["aligner_index_local"][ccu[i] : ccu[i + 1]] + assert 0 <= int(idx_i.min()) and int(idx_i.max()) <= cu[i + 1] - cu[i] + cells_i = torch.cat([vit_i, vit_i.new_zeros(1, width)])[idx_i] + assert torch.equal(cells_i, cells_all[ccu[i] : ccu[i + 1]]) + src_i = image["span_source_local"][scu[i] : scu[i + 1]] + assert int(src_i.min()) >= 0 and int(src_i.max()) == n_cells[i] + 2 + spans_i = torch.cat([aligner_out[ccu[i] : ccu[i + 1]], delimiters])[src_i] + assert torch.equal(spans_i, spans_all[scu[i] : scu[i + 1]]) + + +# --------------------------------------------------------------------------------------------- +# Text-only prompts and errors + + +def test_text_only_prompts_tokenize_as_the_default_processor(ckpt, ref, tokenizer, processor): + default = DefaultInputProcessor(None, None, tokenizer) + prompts = [prompt for prompt, records in _example_cases(ckpt, ref) if not records] + assert len(prompts) == 3 + for prompt in prompts: + for add_special_tokens in (True, False): + sp = SamplingParams(add_special_tokens=add_special_tokens) + ids, extra = processor(TextPrompt(prompt=prompt), sp) + assert extra is None + assert ids == default(TextPrompt(prompt=prompt), sp)[0] + # The checkpoint's tokenizer adds nothing: BOS is text in the rendered prompt. + assert tokenizer.encode(prompt) == tokenizer.encode(prompt, add_special_tokens=False) + + +def test_placeholder_and_image_counts_must_agree(ref, processor): + image = _random_image(64, 64, seed=0) + with pytest.raises(ValueError, match="Found 2 image placeholder tokens .* but got 1 images"): + _run(processor, _reference_prompt(ref, 2, "x"), [image]) + with pytest.raises(ValueError, match="Found 1 image placeholder tokens .* but got 0 images"): + _run(processor, _reference_prompt(ref, 1, "x"), None) + with pytest.raises(ValueError, match="takes images only"): + processor(TextPrompt(prompt="x", multi_modal_data={"video": [image]}), SamplingParams()) + + +def test_truncation_never_cuts_an_image_span(ref, processor): + image = _random_image(544, 544, seed=0) + prompt = _reference_prompt(ref, 1, "x") + ids, _ = _run(processor, prompt, [image], truncate_prompt_tokens=8192) + with pytest.raises(ValueError, match="would cut an image span"): + _run(processor, prompt, [image], truncate_prompt_tokens=len(ids) - 1) + + +def test_tokens_per_image_matches_the_expansion(ckpt, hf_config, processor, ref): + for size in [(1, 1), (544, 544), (768, 384), (50000, 14)]: + image = _random_image(*size, seed=0) + _, extra = _run(processor, _reference_prompt(ref, 1, "x"), [image]) + span = extra["multimodal_data"]["image"]["span_lengths"][0] + rgb = np.asarray(image) + for fmt in (image, torch.from_numpy(rgb).permute(2, 0, 1), rgb): + assert processor.get_num_tokens_per_image(image=fmt) == span + assert processor.get_mm_token_ids().tolist() == [_IMAGE_TOKEN_ID] + assert processor.get_mm_special_token_ids() is None + + +# --------------------------------------------------------------------------------------------- +# Deliberately wrong controls + + +def test_control_resize_instead_of_pad_is_detected(ckpt, ref, ref_args, hf_config): + """A plain resize to the target size (the max_wh_ratio branch) gives different patches.""" + geometry = VisionGeometry.from_config(hf_config) + path = ckpt / "inference" / "examples" / "images" / "carrots.jpeg" + reference = ref.image_processor.load_image({"url": str(path)}, ref_args)[0] + image = Image.open(path).convert("RGB") + assert torch.equal(preprocess_image(image, geometry).patches, reference) + + grid = plan_image_grid(image.width, image.height, geometry) + assert image.size != (grid.best_width, grid.best_height) + resized = image.resize((grid.best_width, grid.best_height)) + wrong = preprocess_image(resized, geometry).patches + assert wrong.shape == reference.shape + assert not torch.equal(wrong, reference) + + +def test_control_newline_separator_is_detected(ref, ref_args, tokenizer, chat_renderer): + """Placeholders joined by "\\n" (the registry default) render a prompt encoding.py never does.""" + text = "Which one is larger?" + for num_images in (1, 2, 3): + expected = _reference_prompt(ref, num_images, text) + assert _string_path_prompt(chat_renderer, num_images, text) == expected + wrong = _string_path_prompt(chat_renderer, num_images, text, separator="\n") + assert wrong != expected + assert tokenizer.encode(wrong) != tokenizer.encode(expected) + + +# --------------------------------------------------------------------------------------------- +# MMMU, rendered as trtllm-eval renders it + + +def _lm_eval_mmmu_utils(): + spec = importlib.util.find_spec("lm_eval") + if spec is None or not spec.submodule_search_locations: + pytest.skip("lm_eval is not installed") + path = Path(spec.submodule_search_locations[0]) / "tasks" / "mmmu" / "utils.py" + module_spec = importlib.util.spec_from_file_location("_dsv41_lm_eval_mmmu_utils", path) + module = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(module) + return module + + +def _mmmu_validation_docs(mmmu_dir: Path): + """Every validation row, as the dict lm-eval's MMMU utils read (images decoded to PIL).""" + import pyarrow as pa + + docs = [] + for arrow in sorted(mmmu_dir.glob("*/0.0.0/*/mmmu-validation.arrow")): + with pa.memory_map(str(arrow)) as source: + table = pa.ipc.open_stream(source).read_all() + for row in table.to_pylist(): + for i in range(1, 8): + value = row.get(f"image_{i}") + if value is not None: + row[f"image_{i}"] = Image.open(io.BytesIO(value["bytes"])) + docs.append(row) + return docs + + +def _anchor_prompt_tokens(anchor_dir: Path): + """doc id -> the prompt length the external engine reported for it, when recorded.""" + if not anchor_dir.is_dir(): + return {} + by_key = {} + with open(anchor_dir / "samples.jsonl", encoding="utf-8") as f: + for line in f: + sample = json.loads(line) + by_key[(sample["task"], str(sample["doc_id"]))] = sample["doc"]["id"] + tokens = {} + with open(anchor_dir / "requests.jsonl", encoding="utf-8") as f: + for line in f: + request = json.loads(line) + key = (request["task"], str(request["doc_id"])) + if key in by_key: + tokens[by_key[key]] = request["meta"]["usage"]["prompt_tokens"] + return tokens + + +def test_mmmu_documents_match_the_reference_and_the_anchor( + ref, ref_args, tokenizer, processor, chat_renderer +): + mmmu_dir = _require(_MMMU_DIR) + utils = _lm_eval_mmmu_utils() + docs = _mmmu_validation_docs(mmmu_dir) + assert len(docs) == 900 + if not _MMMU_ALL: + firsts, seen = [], set() + for doc in docs: + subject = doc["id"].rsplit("_", 1)[0] + if subject not in seen or len(utils.doc_to_image(doc)) > 1: + firsts.append(doc) + seen.add(subject) + docs = firsts + anchor_tokens = _anchor_prompt_tokens(Path(_ANCHOR_MMMU_DIR)) + + checked = anchored = multi_image = 0 + for doc in docs: + text = utils.doc_to_text(doc) + num_images = text.count("") + text = text.replace("", "") + images = [ + img if img.mode == "RGB" else img.convert("RGB") for img in utils.doc_to_image(doc) + ] + assert len(images) == num_images, doc["id"] + + prompt = _string_path_prompt(chat_renderer, num_images, text) + assert prompt == _reference_prompt(ref, num_images, text), doc["id"] + records = [{"data": _png_bytes(img)} for img in images] + reference = ref.image_processor.prepare_vl_inputs(prompt, records, tokenizer, ref_args) + ours = _run(processor, prompt, images, truncate_prompt_tokens=8192) + _assert_matches_reference(ours, reference) + + # The external engine puts one span-marker token in front of an image's span, so for a + # one-image prompt its length is ours plus one. Its multi-image lengths do not decompose + # that way (between one and N extra tokens), so they are not compared. + if num_images == 1 and doc["id"] in anchor_tokens: + assert len(ours[0]) + 1 == anchor_tokens[doc["id"]], doc["id"] + anchored += 1 + + wrong = _string_path_prompt(chat_renderer, num_images, text, separator="\n") + assert tokenizer.encode(wrong) != tokenizer.encode(prompt) + checked += 1 + multi_image += num_images > 1 + assert checked >= 30 and multi_image > 0 + if anchor_tokens: + assert anchored == checked - multi_image diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index eed4149383e0..6023af9fd92e 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -2022,6 +2022,8 @@ def test_vbws_rejects_decreasing_beam_width_array(beam_width_array: list[int], _configure_csa2_remote_tail=lambda _request: None, _validate_token_id_range=lambda _request: None, sampler=types.SimpleNamespace(validate_request=lambda _request: None), + resource_manager=types.SimpleNamespace( + validate_request=lambda _request: None), _validate_request_budget=lambda _request: None, ) validate = functools.partial( diff --git a/tests/unittest/disaggregated/test_deepseek_v41_kv_transfer.py b/tests/unittest/disaggregated/test_deepseek_v41_kv_transfer.py new file mode 100644 index 000000000000..d97d1053d550 --- /dev/null +++ b/tests/unittest/disaggregated/test_deepseek_v41_kv_transfer.py @@ -0,0 +1,504 @@ +# 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. +"""Disaggregated KV transfer of the DeepSeek-V4.1 layout of DeepseekV4CacheManager. + +The context and generation managers carry the checkpoint's real 40-layer role +table. The context side runs without speculation and the generation side with a +DSpark draft depth, which is how a DSpark disaggregated deployment places the +drafter. Pools are compared per position, not per page: for every request the +generation side must hold, byte for byte, what the context side wrote at every +position the first generation step reads: + +- SWA: the last ``sliding_window - 1`` prompt positions, on all 40 layers; +- compressor pending pair: the positions of the incomplete ratio-2 group; +- Engram token history: the last ``engram_max_ngram_size - 1`` prompt positions; +- main KV and index K: every complete compressed entry of every KV source. + +The reference bytes are taken from the context pools right after allocation, +before the context side advances its history (which releases out-of-window +pages) and before the transfer. +""" + +from types import SimpleNamespace +from typing import Dict, List, Optional, Sequence, Tuple + +import pytest +import torch +from kv_transfer_harness import ( + TOKENS_PER_BLOCK, + VOCAB_SIZE, + create_instance_transceivers, + get_ctx_info_endpoint, + run_concurrent, +) +from utils.util import skip_pre_blackwell + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4CacheManager +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import DeepseekV4AttentionType +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.v41_layout import ( + DEEPSEEK_V41_ATTENTION, + DEEPSEEK_V41_ENCODER_REPLAY_ENV, + DEEPSEEK_V41_KV_STORAGE_ENV, + DeepseekV41Layout, +) +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 +from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool, get_pool_bytes +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import ( + CacheTransceiverConfig, + DeepSeekV4SparseAttentionConfig, + DSparkDecodingConfig, + KvCacheConfig, +) + +T = DeepseekV4AttentionType +WINDOW = 128 +HEAD_DIM = 512 +INDEX_HEAD_DIM = 128 +MAX_SEQ_LEN = 4096 +MAX_BATCH_SIZE = 16 +# The DSpark block size of the checkpoint; the generation engine's draft depth. +GEN_DRAFT_LEN = 5 + +# The checkpoint's text_config fields the layout reads (config.json of +# DeepSeek-V4.1-Flash-NVFP4), 40 backbone layers plus 3 DSpark stages. +V41_TEXT_FIELDS = dict( + model_type="deepseek_v41_text", + num_hidden_layers=40, + num_nextn_predict_layers=3, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], + index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, + engram_layer_ids=[1, 14], + engram_max_ngram_size=4, + sliding_window=WINDOW, + head_dim=HEAD_DIM, + index_head_dim=INDEX_HEAD_DIM, + index_topk=512, + kv_lora_rank=448, + qk_rope_head_dim=64, +) + + +@pytest.fixture(autouse=True) +def _no_encoder_replay(monkeypatch): + """Encoder SWA bounded replay refuses disaggregated serving (and dep4): pin it off.""" + monkeypatch.setenv(DEEPSEEK_V41_ENCODER_REPLAY_ENV, "0") + + +# Prompt lengths: 1-4 (no longer than the Engram history), odd lengths (a +# ratio-2 pending token), block-aligned lengths, and lengths whose Engram +# history P-3 .. P-1 straddles a page boundary (P % 128 in {1, 2, 3}). +REQUEST_LENGTHS = [1, 2, 3, 4, 129, 130, 131, 256, 258, 383, 2049, 2050] + + +def make_v41_config() -> SimpleNamespace: + top = {k: v for k, v in V41_TEXT_FIELDS.items() if k != "model_type"} + return SimpleNamespace( + architectures=["DeepseekV41ForCausalLM"], + model_type="deepseek_v41", + text_config=SimpleNamespace(**V41_TEXT_FIELDS), + **top, + ) + + +def create_v41_manager(mapping: Mapping, draft_len: int = 0) -> DeepseekV4CacheManager: + config = make_v41_config() + sparse_config = DeepSeekV4SparseAttentionConfig( + index_head_dim=INDEX_HEAD_DIM, + window_size=WINDOW, + compress_ratios=[1] * len(config.compress_ratios), + ) + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + max_tokens=MAX_SEQ_LEN * MAX_BATCH_SIZE, + event_buffer_max_size=0, + ) + spec_config = DSparkDecodingConfig(max_draft_len=draft_len) if draft_len > 0 else None + return DeepseekV4CacheManager( + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=config.num_hidden_layers, + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH_SIZE, + max_input_len=MAX_SEQ_LEN, + mapping=mapping, + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=VOCAB_SIZE, + max_num_tokens=8192, + sparse_attn_config=sparse_config, + pretrained_config=config, + spec_config=spec_config, + ) + + +def create_instance_managers(tp: int, enable_dp: bool, draft_len: int) -> List: + return [ + create_v41_manager( + Mapping(world_size=tp, rank=rank, tp_size=tp, pp_size=1, enable_attention_dp=enable_dp), + draft_len=draft_len, + ) + for rank in range(tp) + ] + + +def fill_pools(managers: Sequence[DeepseekV4CacheManager], seed: Optional[int]) -> None: + """Random bytes (the same on every TP rank: MLA state is replicated), or zeros.""" + for mgr in managers: + page_table = KVRegionExtractorV1(mgr).page_table + pools: Dict[int, int] = {} + for lg_idx, lg in enumerate(page_table.layer_groups): + for pv in lg.pool_views: + pool = get_physical_pool(page_table, lg_idx, pv.pool_idx) + pools[pool.base_address] = max( + pools.get(pool.base_address, 0), get_pool_bytes(pool) + ) + for index, (base, size) in enumerate(sorted(pools.items())): + view = convert_to_torch_tensor(TensorWrapper(base, DataType.INT8, [size])) + if seed is None: + view.zero_() + continue + gen = torch.Generator(device=view.device).manual_seed(seed + index) + view.copy_( + torch.randint( + -128, 128, view.shape, dtype=torch.int8, device=view.device, generator=gen + ) + ) + torch.cuda.synchronize() + + +# --------------------------------------------------------------------------- +# What the first generation step reads, per owned (layer, role). +# --------------------------------------------------------------------------- + + +def needed_slots(layout: DeepseekV41Layout, layer: int, attn_type, prompt_len: int) -> List[int]: + """The slots the generation side reads before it writes anything of its own. + + Positions for the windowed roles, compressed entries for the persistent roles. + """ + if attn_type == T.SWA: + return list(range(max(0, prompt_len - (layout.window_size - 1)), prompt_len)) + if attn_type in (T.COMPRESSOR_KV, T.COMPRESSOR_SCORE): + ratio = layout.compress_ratios[layer] + return list(range((prompt_len // ratio) * ratio, prompt_len)) + if attn_type == T.ENGRAM_TOKEN: + return list(range(max(0, prompt_len - layout.engram_history_len), prompt_len)) + if attn_type in (T.COMPRESS, T.INDEXER_COMPRESS): + return list(range(prompt_len // layout.compress_ratios[layer])) + raise AssertionError(attn_type) + + +def owned_roles(mgr: DeepseekV4CacheManager) -> List[Tuple[int, DeepseekV4AttentionType]]: + return sorted(mgr._layer_attn_to_layer_id, key=lambda key: (key[0], key[1].value)) + + +def read_slots(mgr, request_id: int, layer: int, attn_type, slots: Sequence[int]): + """(slot -> raw bytes) for every slot, or (slot -> None) where no page is mapped. + + Index K is compared page by page: its DeepGEMM page layout keeps an entry's + code bytes and scale bytes apart, and the transfer is page granular anyway. + """ + buffer = mgr.get_buffers(layer, attn_type) + per_page = buffer.shape[1] + pages = mgr.get_cache_indices(request_id, layer, attn_type) + out = {} + if attn_type == T.INDEXER_COMPRESS: + flat = buffer.flatten(1) + for ordinal in sorted({slot // per_page for slot in slots}): + page = pages[ordinal] if ordinal < len(pages) else -1 + out[("page", ordinal)] = None if page < 0 else flat[page].clone() + return out + for slot in slots: + ordinal, row = divmod(slot, per_page) + page = pages[ordinal] if ordinal < len(pages) else -1 + out[slot] = None if page < 0 else buffer[page, row].contiguous().view(torch.uint8).clone() + return out + + +def snapshot(mgr, request_id: int, prompt_len: int): + layout = mgr.v41_layout + return { + (layer, attn_type): read_slots( + mgr, request_id, layer, attn_type, needed_slots(layout, layer, attn_type, prompt_len) + ) + for layer, attn_type in owned_roles(mgr) + } + + +def make_request(request_id: int, prompt_len: int, request_type) -> LlmRequest: + return LlmRequest( + request_id=request_id, + max_new_tokens=8, + input_tokens=list(range(prompt_len)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + SamplingParams()._get_sampling_config() + ), + is_streaming=False, + llm_request_type=request_type, + ) + + +# --------------------------------------------------------------------------- +# Manager-level checks (no transfer). +# --------------------------------------------------------------------------- + + +@skip_pre_blackwell +@pytest.mark.parametrize("draft_len", [0, GEN_DRAFT_LEN], ids=["no_spec", "dspark5"]) +def test_context_history_keeps_what_generation_reads(draft_len): + """After prefill commits its history, every position the next step reads is still mapped. + + This is the manager's windowed life cycle on its own: a context engine + advances ``history_length`` to the prompt length before it hands the pages + to the transceiver (and an aggregated engine does the same before its first + decode step), so a position outside the live window is unmapped and can be + neither transferred nor read. + """ + mgr = create_v41_manager(Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), draft_len) + try: + requests = [] + batch = ScheduledRequests() + for rid, prompt_len in enumerate(REQUEST_LENGTHS): + req = make_request( + rid, prompt_len, LlmRequestType.LLMREQUEST_TYPE_CONTEXT_AND_GENERATION + ) + assert mgr.prepare_context(req) + assert mgr.resize_context(req, req.context_chunk_size) + requests.append(req) + batch.context_requests_last_chunk.append(req) + for req in requests: + req.context_current_position = req.prompt_len + mgr.update_context_resources(batch) + unmapped = [] + for req in requests: + for (layer, attn_type), slots in snapshot( + mgr, req.py_request_id, req.prompt_len + ).items(): + missing = sorted(slot for slot, data in slots.items() if data is None) + if missing: + unmapped.append((req.prompt_len, layer, attn_type.name, missing)) + assert not unmapped, f"unmapped after history=prompt_len: {unmapped}" + finally: + mgr.shutdown() + + +@skip_pre_blackwell +def test_disagg_global_layer_ids_are_owner_only(): + """The transfer's synthetic layer ids name owners only, and both peers agree on them.""" + ctx = create_v41_manager(Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), 0) + gen = create_v41_manager(Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), GEN_DRAFT_LEN) + try: + layout = ctx.v41_layout + num_types = ctx.disagg_num_attn_types + assert num_types == gen.disagg_num_attn_types == max(t.value for t in T) + 1 + per_side = [] + for mgr in (ctx, gen): + assert mgr.get_disagg_ignored_roles() == frozenset() + page_table = KVRegionExtractorV1(mgr).page_table + ids = {} + for lg in page_table.layer_groups: + local_to_global = {ll.local_layer_id: ll.global_layer_id for ll in lg.local_layers} + for pv in lg.pool_views: + for entry in pv.buffer_entries: + gid = local_to_global[int(entry["local_layer_id"])] + ids.setdefault(gid, set()).update(str(r) for r in pv.pool_role) + for gid in ids: + layer, value = divmod(gid, num_types) + attn_type = T(value) + assert (layer, attn_type) in mgr._layer_attn_to_layer_id, (gid, layer, attn_type) + assert layout.has_attention(layer, attn_type), (gid, layer, attn_type) + if layout.compress_ratios[layer] > 0 and not layout.is_kv_source(layer): + # A reader layer: its only pool is its own SWA window. + assert attn_type == T.SWA, (gid, layer, attn_type) + roles = set().union(*ids.values()) + for attn_type in DEEPSEEK_V41_ATTENTION: + assert str(attn_type.role) in roles, (attn_type, roles) + # One id per manager layer: SWA x 40, main KV + index K x 4, + # compressor pair x 3, Engram history x 1. + assert len(ids) == 40 + 4 + 3 + 1 + per_side.append(ids) + assert per_side[0] == per_side[1] + finally: + ctx.shutdown() + gen.shutdown() + + +# --------------------------------------------------------------------------- +# Transfer round trip. +# --------------------------------------------------------------------------- + + +def run_v41_transfer(ctx_tp: int, gen_tp: int, ctx_dp: bool, gen_dp: bool, update_before: bool): + import uuid + + try: + ctx_managers = create_instance_managers(ctx_tp, ctx_dp, draft_len=0) + except NotImplementedError as error: + pytest.skip(f"storage not implemented in this tree: {error}") + gen_managers = create_instance_managers(gen_tp, gen_dp, draft_len=GEN_DRAFT_LEN) + fill_pools(ctx_managers, seed=1000) + fill_pools(gen_managers, seed=None) + + config = CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="PYTHON", max_tokens_in_buffer=4096 + ) + ctx_tcs = create_instance_transceivers(ctx_tp, 1, ctx_dp, ctx_managers, config) + gen_tcs = create_instance_transceivers(gen_tp, 1, gen_dp, gen_managers, config) + try: + endpoint = get_ctx_info_endpoint(ctx_tcs[0]) + ctx_handle: Dict[int, List[Tuple[int, LlmRequest]]] = {r: [] for r in range(ctx_tp)} + gen_handle: Dict[int, List[Tuple[int, LlmRequest]]] = {r: [] for r in range(gen_tp)} + for idx, prompt_len in enumerate(REQUEST_LENGTHS): + unique_rid = uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF + ctx_req = make_request(2 * idx, prompt_len, LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + ctx_req.py_disaggregated_params = DisaggregatedParams(disagg_request_id=unique_rid) + gen_req = make_request( + 2 * idx + 1, prompt_len, LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY + ) + gen_req.py_disaggregated_params = DisaggregatedParams( + ctx_request_id=2 * idx, + ctx_dp_rank=idx % ctx_tp if ctx_dp else 0, + ctx_info_endpoint=endpoint, + disagg_request_id=unique_rid, + ) + for rank in range(ctx_tp): + if not ctx_dp or idx % ctx_tp == rank: + ctx_handle[rank].append((idx, ctx_req)) + for rank in range(gen_tp): + if not gen_dp or idx % gen_tp == rank: + gen_handle[rank].append((idx, gen_req)) + + for rank, items in gen_handle.items(): + for _, req in items: + assert gen_managers[rank].prepare_disagg_gen_init(req) + ctx_batches = {} + for rank, items in ctx_handle.items(): + batch = ScheduledRequests() + for _, req in items: + assert ctx_managers[rank].prepare_context(req) + assert ctx_managers[rank].resize_context(req, req.context_chunk_size) + batch.context_requests_last_chunk.append(req) + ctx_batches[rank] = batch + torch.cuda.synchronize() + + # What prefill left in the context pools, read before anything is released. + truth = { + (rank, idx): snapshot(ctx_managers[rank], req.py_request_id, req.prompt_len) + for rank, items in ctx_handle.items() + for idx, req in items + } + + for handle in (ctx_handle, gen_handle): + seen = set() + for items in handle.values(): + for _, req in items: + if req.py_request_id not in seen: + req.context_current_position = req.prompt_len + req.add_new_token(req.prompt_len, 0) + seen.add(req.py_request_id) + + # Production order (py_executor): commit the context history, then send. + if update_before: + for rank, batch in ctx_batches.items(): + ctx_managers[rank].update_context_resources(batch) + for rank, items in gen_handle.items(): + for _, req in items: + gen_tcs[rank].request_and_receive_async(req) + for rank, items in ctx_handle.items(): + for _, req in items: + ctx_tcs[rank].respond_and_send_async(req) + run_concurrent( + ctx_tcs, lambda tc: tc.check_context_transfer_status(None, mark_complete=True) + ) + run_concurrent(gen_tcs, lambda tc: tc.check_gen_transfer_status(None)) + if not update_before: + for rank, batch in ctx_batches.items(): + ctx_managers[rank].update_context_resources(batch) + torch.cuda.synchronize() + + failures = [] + compared = 0 + for gen_rank, items in gen_handle.items(): + for idx, gen_req in items: + ctx_rank = idx % ctx_tp if ctx_dp else gen_rank % ctx_tp + expected = truth[(ctx_rank, idx)] + got = snapshot(gen_managers[gen_rank], gen_req.py_request_id, gen_req.prompt_len) + for key, want_slots in expected.items(): + layer, attn_type = key + for slot, want in want_slots.items(): + have = got[key][slot] + if want is None: + failures.append( + (gen_req.prompt_len, layer, attn_type.name, slot, "ctx unmapped") + ) + elif have is None: + failures.append( + (gen_req.prompt_len, layer, attn_type.name, slot, "gen unmapped") + ) + elif not torch.equal(have, want): + failures.append( + (gen_req.prompt_len, layer, attn_type.name, slot, "bytes differ") + ) + else: + compared += 1 + assert compared > 0 + assert not failures, ( + f"{len(failures)} slot(s) did not cross ({compared} did): {failures[:20]}" + ) + print(f"compared {compared} slots byte-exact") + finally: + for tc in ctx_tcs + gen_tcs: + try: + tc.shutdown() + except Exception: + pass + for mgr in ctx_managers + gen_managers: + try: + mgr.shutdown() + except Exception: + pass + + +@skip_pre_blackwell +@pytest.mark.timeout(300) +@pytest.mark.parametrize( + "ctx_tp,gen_tp,ctx_dp,gen_dp,update_before,kv_storage", + [ + (1, 1, False, False, True, "bf16"), + (1, 1, False, False, False, "bf16"), + (4, 4, False, False, True, "bf16"), + (4, 4, True, True, True, "bf16"), + (4, 4, False, False, True, "native"), + ], + ids=["tp1", "tp1_update_after", "tp4", "dep4", "tp4_native"], +) +def test_deepseek_v41_kv_transfer( + ctx_tp, gen_tp, ctx_dp, gen_dp, update_before, kv_storage, monkeypatch +): + monkeypatch.setenv(DEEPSEEK_V41_KV_STORAGE_ENV, kv_storage) + run_v41_transfer(ctx_tp, gen_tp, ctx_dp, gen_dp, update_before) diff --git a/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py b/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py index dec508d49cbc..f0058ec88c8c 100644 --- a/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py +++ b/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py @@ -447,3 +447,40 @@ def test_deepseek_v4_server_chat_template_path_forwards_tools(): assert "<｜DSML｜tool_calls>" in prompt assert '"name": "search"' in prompt + + +def test_deepseek_v4_server_chat_template_path_leaves_pending_media_alone(): + """An image chat request reaches the renderer as the OpenAI server parsed it. + + That is the text (with the model's image placeholder already inserted) plus a "media" list + of pending loads, which are coroutines and cannot be deep-copied. The renderer must neither + copy nor consume them; the server awaits them after rendering. + """ + + async def pending_image_load(): + return None + + pending = pending_image_load() + message = { + "role": "user", + "content": "<｜deepseek_image｜>\n\nWhat is shown?", + "media": [{"modality": "image", "data": pending, "is_embedding": False}], + "content_parts": [{"type": "image", "media_index": 0}, "What is shown?"], + } + try: + prompt = apply_chat_template( + model_type="deepseek_v41", + tokenizer=DeepseekV4Tokenizer(_DummyTokenizer()), + processor=None, + conversation=[message], + add_generation_prompt=True, + mm_placeholder_counts=[{"<｜deepseek_image｜>": 1}], + chat_template_kwargs={"thinking": False}, + ) + assert prompt == ( + "<｜begin▁of▁sentence｜><｜User｜><｜deepseek_image｜>\n\nWhat is shown?" + "<｜Assistant｜>" + ) + assert message["media"][0]["data"] is pending + finally: + pending.close()