Add GDR CP controls and legacy kernel override#4779
Open
lzhangzz wants to merge 2 commits into
Open
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds explicit runtime controls for TurboMind GDR context-parallel behavior and legacy kernel dispatch, and propagates the selected configuration through planning/metadata and tests.
Changes:
- Replace
cp_modewithcp_level(off/exact/all) across the benchmark harness, Python bindings, and kernel planning/metadata. - Add process-cached parsing of
TM_GDR_FORCE_LEGACYand make kernel-family selection explicit via an architecture key in the GDR registry lookup. - Adjust kernel planning/execution details to respect CP level (e.g., warmup gating) and broaden pre-SM90 kernel compilation across configured CUDA targets.
Reviewed changes
Copilot reviewed 27 out of 27 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/turbomind/linear_attn/turbomind_gated_delta_rule.py | Propagates cp_level through TurboMind benchmark task/plan execution. |
| tests/turbomind/linear_attn/test_gated_delta_rule.py | Updates tests for cp_level defaults/choices and CP plan metadata assertions. |
| tests/turbomind/linear_attn/reference.py | Updates reference benchmark row metadata to use cp_level. |
| tests/turbomind/linear_attn/flashqla_gated_delta_rule.py | Switches FlashQLa benchmark integration to cp_level and removes fixed chunk-size enforcement. |
| tests/turbomind/linear_attn/fla_gated_delta_rule.py | Updates FLA benchmark row metadata to use cp_level. |
| tests/turbomind/linear_attn/benchmark.py | Renames CP knobs to cp_level, updates CLI choices/defaults, and simplifies CP request validation flow. |
| src/turbomind/models/llama/GatedDeltaNetLayer.h | Stores selected GDR CP level as a layer member. |
| src/turbomind/models/llama/GatedDeltaNetLayer.cc | Parses TM_GDR_CP_LEVEL and propagates CP level into GDR operation planning; logs selected level. |
| src/turbomind/kernels/linear_attn/registry.h | Extends registry lookup API to include explicit architecture selection. |
| src/turbomind/kernels/linear_attn/registry.cu | Filters registered kernels by architecture key during lookup. |
| src/turbomind/kernels/linear_attn/python_bind.cpp | Exposes cp_level in plan dicts and parses Python cp_level strings into ContextParallelLevel. |
| src/turbomind/kernels/linear_attn/kernel/sm_90/prepare_h.h | Gates warmup behavior based on CP level (allow/disallow warmup). |
| src/turbomind/kernels/linear_attn/kernel/sm_90/plan.cc | Plans CP based on cp_level (off vs enabled), threads CP level into disabled plan metadata. |
| src/turbomind/kernels/linear_attn/kernel/sm_90/entry.cu | Removes redundant arch checks from per-kernel Match (arch now selected by registry). |
| src/turbomind/kernels/linear_attn/kernel/sm_120/prepare_h.h | Mirrors SM90 warmup gating changes for SM120. |
| src/turbomind/kernels/linear_attn/kernel/sm_120/plan.cc | Mirrors SM90 CP planning changes for SM120. |
| src/turbomind/kernels/linear_attn/kernel/sm_120/entry.cu | Removes redundant arch checks from per-kernel Match. |
| src/turbomind/kernels/linear_attn/kernel/pre_sm90/recurrent.cu | Removes arch-guarded compilation blocks to allow building legacy kernels for all configured targets. |
| src/turbomind/kernels/linear_attn/kernel/pre_sm90/plan.cc | Threads cp_level into the disabled CP plan metadata for legacy kernels. |
| src/turbomind/kernels/linear_attn/kernel/pre_sm90/internal.h | Updates legacy planning function signature to accept Operation. |
| src/turbomind/kernels/linear_attn/kernel/pre_sm90/entry.cu | Updates legacy kernel entry to pass Operation into planning; removes redundant arch checks. |
| src/turbomind/kernels/linear_attn/kernel/pre_sm90/chunked.cu | Removes arch-guarded compilation blocks to allow building legacy kernels for all configured targets. |
| src/turbomind/kernels/linear_attn/kernel/plan.h | Updates disabled CP plan builder signature to accept ContextParallelLevel. |
| src/turbomind/kernels/linear_attn/kernel/plan.cc | Stores cp_level into disabled CP plan metadata. |
| src/turbomind/kernels/linear_attn/kernel/CMakeLists.txt | Builds pre-SM90 kernels for all configured CUDA architectures. |
| src/turbomind/kernels/linear_attn/delta_rule.h | Introduces ContextParallelLevel and threads it through Operation/ContextParallelPlan. |
| src/turbomind/kernels/linear_attn/delta_rule.cu | Adds cached parsing of TM_GDR_FORCE_LEGACY and explicit architecture selection for registry lookup. |
Comments suppressed due to low confidence (1)
tests/turbomind/linear_attn/flashqla_gated_delta_rule.py:149
- When
cp_patternis explicitly set, FlashQLa currently requirescp_level == 'all', which unnecessarily rejectscp_level == 'exact'once exact is accepted as a CP-enabled level. The gating here should only reject the CP-disabled case (off).
if request.cp_pattern != 'auto':
if request.cp_level != 'all':
raise ValueError('cp_not_selected')
inputs = apply_cp_pattern(
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+75
to
79
| if cp_level == 'all': | ||
| run_initial_state, run_cu_seqlens, cp_seq_map, raw_cu_seqlens = imports[ | ||
| 'intra_card_cp_preprocess' | ||
| ]( | ||
| k=k, |
Comment on lines
112
to
+116
| def _cp_pattern_segment_tokens(inputs: InputTensors) -> int: | ||
| if inputs.q.shape[0] != 1: | ||
| raise ValueError('cp_not_selected') | ||
|
|
||
| raw_q_offsets = inputs.offsets | ||
| if raw_q_offsets is None: | ||
| raw_q_offsets = torch.tensor( | ||
| (0, inputs.q.shape[1]), | ||
| dtype=torch.int32, | ||
| device=inputs.q.device, | ||
| ) | ||
| use_cp, cp_q_offsets, sequence_starts, _, _, _ = _imports()['calc_cp_seqs']( | ||
| raw_q_offsets, | ||
| FLASHQLA_CHUNK_SIZE, | ||
| inputs.v.shape[2], | ||
| imports = _imports() | ||
| a = imports['kkt_solve']( | ||
| k=inputs.k, | ||
| b=inputs.beta, |
Comment on lines
+180
to
+191
| ContextParallelLevel ParseContextParallelLevel(const std::string& level) | ||
| { | ||
| if (mode == "auto") { | ||
| return ContextParallelMode::kAuto; | ||
| if (level == "off") { | ||
| return ContextParallelLevel::kOff; | ||
| } | ||
| if (level == "exact") { | ||
| return ContextParallelLevel::kExact; | ||
| } | ||
| if (mode == "off") { | ||
| return ContextParallelMode::kOff; | ||
| if (level == "all") { | ||
| return ContextParallelLevel::kAll; | ||
| } | ||
| throw py::value_error("cp_mode must be one of: auto, off"); | ||
| throw py::value_error("cp_level must be one of: off, exact, all"); |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
TurboMind GDR needs explicit controls for context-parallel execution and a way to force the legacy implementation on newer GPUs for debugging, validation, and performance comparison.
Changes
ContextParallelLevelwith three supported levels:off: disable context parallelism.exact: allow exact CP only, without warmup CP.all: allow both exact and warmup CP.cp_levelvaluesoff,exact, andall.TM_GDR_CP_LEVELvalues0,1, and2.TM_GDR_FORCE_LEGACY:0: use native architecture dispatch.1: force the pre-SM90 kernel family.TM_GDR_FORCE_LEGACYonce per process.Validation
TM_GDR_CP_LEVEL=0TM_GDR_CP_LEVEL=1TM_GDR_CP_LEVEL=2TM_GDR_FORCE_LEGACY=1