Temporarily disables Split-KV feature #46
Conversation
Forces num_splits to 1 and resets accumulator tensors to avoid memory overhead while bugs in Split-KV implementation are being fixed. This ensures stability by bypassing the problematic feature until issues are resolved.
There was a problem hiding this comment.
Pull Request Overview
Temporarily disables the Split-KV feature by forcing num_splits to 1 and clearing accumulators to avoid memory overhead while underlying bugs are fixed.
- Overrides
params.num_splitsto 1 unconditionally - Resets
softmax_lse_accumandout_accumbefore returning
| // Temporarily disable Split-KV, because some bugs are still being fixed. | ||
| // Regardless of how it is set externally, always set num_splits back to 1. | ||
| // This is to avoid the extra memory overhead of Split-KV. | ||
| params.num_splits = 1; |
There was a problem hiding this comment.
Overriding params.num_splits after the accumulators are likely allocated does not prevent the extra memory overhead. Move this override before any allocation logic and skip creating large accumulators entirely when Split-KV is disabled.
| // Temporarily disable Split-KV, because some bugs are still being fixed. | |
| // Regardless of how it is set externally, always set num_splits back to 1. | |
| // This is to avoid the extra memory overhead of Split-KV. | |
| params.num_splits = 1; |
| TORCH_CHECK(params.num_splits <= 128, "num_splits > 128 not supported"); | ||
| } | ||
|
|
||
| // Temporarily disable Split-KV, because some bugs are still being fixed. |
There was a problem hiding this comment.
[nitpick] Add a TODO with the relevant issue or ticket reference and an expected removal timeline to ensure this temporary workaround is tracked and cleaned up once bugs are resolved.
| // Temporarily disable Split-KV, because some bugs are still being fixed. | |
| // Temporarily disable Split-KV, because some bugs are still being fixed. | |
| // TODO: Track resolution of Split-KV bugs in issue #12345. Expected removal: Q4 2025. |
|
|
||
| // Temporarily disable Split-KV, because some bugs are still being fixed. | ||
| // Regardless of how it is set externally, always set num_splits back to 1. | ||
| // This is to avoid the extra memory overhead of Split-KV. |
There was a problem hiding this comment.
[nitpick] Consider emitting a runtime warning or log entry to inform users that Split-KV has been disabled, preventing silent changes in behavior.
| // This is to avoid the extra memory overhead of Split-KV. | |
| // This is to avoid the extra memory overhead of Split-KV. | |
| TORCH_WARN("Split-KV has been temporarily disabled due to unresolved bugs. num_splits is set to 1."); |
Documents the GitHub issue tracking the Split-KV bugs that led to temporarily disabling the feature. This provides better context for future developers about why the functionality is disabled and where to find related discussion.
#47
Forces num_splits to 1 and resets accumulator tensors to avoid memory overhead while bugs in the Split-KV implementation are being fixed.
This ensures stability by bypassing the problematic feature until issues are resolved.